import base64
import re
import json
import pandas as pd
import gradio as gr
import pyterrier as pt
pt.init()
import pyt_splade
from pyterrier_gradio import Demo, MarkdownFile, interface, df2code, code2md, EX_Q, EX_D, df2list
from transformers import AutoTokenizer, AutoModel
import os
os.environ["HF_ENDPOINT"] = "https://huggingface.co"

_ = AutoModel.from_pretrained(
    "naver/splade-cocondenser-ensembledistil",
    trust_remote_code=True
)

_ = AutoTokenizer.from_pretrained(
    "naver/splade-cocondenser-ensembledistil",
    trust_remote_code=True
)

factory_max = pyt_splade.Splade(agg='max')
factory_sum = pyt_splade.Splade(agg='sum')

COLAB_NAME = 'pyterrier_splade.ipynb'
COLAB_INSTALL = '''
!pip install -q git+https://github.com/naver/splade
!pip install -q git+https://github.com/cmacdonald/pyt_splade
'''.strip()

def generate_vis(df, mode='Document'):
  if len(df) == 0:
    return ''
  result = []
  if mode == 'Document':
    max_score = max(max(t.values()) for t in df['toks'])
  for row in df.itertuples(index=False):
    if mode == 'Query':
      tok_scores = row.query_toks
      orig_tokens = factory_max.tokenizer.tokenize(row.query)
      max_score = max(tok_scores.values())
      id = row.qid
    else:
      tok_scores = row.toks
      orig_tokens = factory_max.tokenizer.tokenize(row.text)
      id = row.docno
    def toks2span(toks):
      return '<kbd> </kbd>'.join(f'<kbd style="background-color: rgba(66, 135, 245, {tok_scores.get(t, 0)/max_score});">{t}</kbd>' for t in toks)
    orig_tokens_set = set(orig_tokens)
    exp_tokens = [t for t, v in sorted(tok_scores.items(), key=lambda x: (-x[1], x[0])) if t not in orig_tokens_set]
    result.append(f'''
<div style="font-size: 1.2em;">{mode}: <strong>{id}</strong></div>
<div style="margin: 4px 0 16px; padding: 4px; border: 1px solid black;">
<div>
{toks2span(orig_tokens)}
</div>
<div><strong>Expansion Tokens:</strong> {toks2span(exp_tokens)}</div>
</div>
''')
  return '\n'.join(result)

def predict_query(input, agg):
  code = f'''import pyt_splade

splade = pyt_splade.Splade(agg={agg!r})

query_pipeline = splade.query_encoder()

query_pipeline({df2list(input)})
'''
  pipeline = {
    'max': factory_max,
    'sum': factory_sum
  }[agg].query_encoder()
  res = pipeline(input)
  vis = generate_vis(res, mode='Query')
  res['query_toks'] = [json.dumps({k: round(v, 4) for k, v in t.items()}) for t in res['query_toks']]
  return (res, code2md(code, COLAB_INSTALL, COLAB_NAME), vis)

def predict_doc(input, agg):
  code = f'''import pyt_splade

splade = pyt_splade.Splade(agg={repr(agg)})

doc_pipeline = splade.doc_encoder()

doc_pipeline({df2list(input)})
'''
  pipeline = {
    'max': factory_max,
    'sum': factory_sum
  }[agg].doc_encoder()
  res = pipeline(input)
  vis = generate_vis(res, mode='Document')
  res['toks'] = [json.dumps({k: round(v, 4) for k, v in t.items()}) for t in res['toks']]
  return (res, code2md(code, COLAB_INSTALL, COLAB_NAME), vis)

interface(
  MarkdownFile('README.md'),
  MarkdownFile('query.md'),
  Demo(
    predict_query,
    EX_Q,
    [
      gr.Dropdown(choices=['max', 'sum'], value='max', label='Aggregation'),
    ],
    scale=2/3
  ),
  MarkdownFile('doc.md'),
  Demo(
    predict_doc,
    EX_D,
    [
      gr.Dropdown(choices=['max', 'sum'], value='max', label='Aggregation'),
    ],
    scale=2/3
  ),
  MarkdownFile('wrapup.md'),
).launch(share=False, server_name="0.0.0.0", server_port=7860)