import pandas as pd
import gradio as gr
from pyterrier_doc2query import Doc2Query

MODEL = 'macavaney/doc2query-t5-base-msmarco'

doc2query = Doc2Query(MODEL, append=True, num_samples=5)

def df2code(df):
  rows = []
  for row in df.itertuples(index=False):
    rows.append(f'  {dict(row._asdict())},')
  rows = '\n'.join(rows)
  return f'''pd.DataFrame([
{rows}
])'''

def predict(input, model, append, num_samples):
  assert model == MODEL
  doc2query.append = append
  doc2query.num_samples = num_samples
  code = f'''
**Code:**

```python
import pandas as pd
from pyterrier_doc2query import Doc2Query
doc2query = Doc2Query({repr(model)}, append={append}, num_samples={num_samples})
doc2query({df2code(input)})
```
'''
  return (doc2query(input), code)

example_inp = pd.DataFrame([
  {'docno': '0', 'text': 'The presence of communication amid scientific minds was equally important to the success of the Manhattan Project as scientific intellect was. The only cloud hanging over the impressive achievement of the atomic researchers and engineers is what their success truly meant; hundreds of thousands of innocent lives obliterated.'},
  {'docno': '86', 'text': 'Usually, you can feel the pain reverberating from the upper portion of the left side of your abdomen towards the left side of your ribcage. Irritation on the Spleen – There is a chance that your spleen has already ruptured because of various reasons and this can cause some pains on the left rib cage.'},
  {'docno': '985', 'text': 'Continue on Hollins Ferry Road to Patapsco Avenue. Make a right onto Patapsco Avenue for approximately 2.5 miles. The courthouse is at the corner of Patapsco Avenue and 7th Street. The commissioner\'s office is on the first (ground) floor.'}
])

example_out = predict(example_inp, MODEL, doc2query.append, doc2query.num_samples)

gr.Interface(
    predict,
    inputs=[gr.Dataframe(
      headers=["docno", "text"],
      datatype=["str", "str"],
      col_count=(2, "fixed"),
      row_count=1,
      wrap=True,
      label='Pipeline Input',
      value=example_inp,
    ), gr.Dropdown(
      choices=[MODEL],
      value=MODEL,
      label='Model',
      interactive=False,
    ), gr.Checkbox(
      value=doc2query.append,
      label="Append",
    ), gr.Slider(
      minimum=1,
      maximum=10,
      value=doc2query.num_samples,
      step=1.,
      label='# Queries'
    )],
    outputs=[gr.Dataframe(
      headers=["docno", "text", "querygen"],
      datatype=["str", "str", "str"],
      col_count=3,
      row_count=1,
      wrap=True,
      label='Pipeline Output',
      value=example_out[0],
    ), gr.Markdown(value=example_out[1])],
    title="🐕 PyTerrier: Doc2Query",
    description=open('README.md', 'rt').read().split('\n---\n')[-1],
    allow_flagging='never',
    css="table.font-mono td, table.font-mono th { white-space: pre-line; font-size: 11px; line-height: 16px; } table.font-mono td input { width: 95%; } th .cursor-pointer {display: none;} th .min-h-\[2\.3rem\] {min-height: inherit;}",
).launch(share=False)