import gradio as gr
import pandas as pd


# Goal: Maximize
# Data Source: Shields et al. https://www.nature.com/articles/s41586-021-03213-y
def lookup(
    concentration: float, temperature: float, solvent: str, ligand: str, base: str
):
    data = pd.read_csv("./data.csv", delimiter=",", index_col=None)
    other = pd.DataFrame.from_records(
        {
            "Concentration": concentration,
            "Temp_C": temperature,
            "Ligand": ligand,
            "Solvent": solvent,
            "Base": base,
        },
        index=[0],
    )

    merged_df = pd.merge(
        left=data.reset_index(),
        right=other,
        on=other.columns.tolist(),
        how="right",
        sort=False,
    ).dropna(subset="index")
    idxs_matched = pd.Index(merged_df["index"].values)

    if len(idxs_matched) < len(other):
        raise IndexError(
            "When doing exact lookup some rows could not be found. Try "
            "approximate lookup or check reference data."
        )
    if len(idxs_matched) > len(other):
        raise IndexError(
            "When doing exact lookup some rows in the reference dataframe appear "
            "duplicated. Check reference data."
        )

    return data["yield"][idxs_matched[0]]


solvents = {
    "DMAc": r"CC(N(C)C)=O",
    "Butyornitrile": r"CCCC#N",
    "Butyl Ester": r"CCCCOC(C)=O",
    "p-Xylene": r"CC1=CC=C(C)C=C1",
}
ligands = {
    "BrettPhos": r"CC(C)C1=CC(C(C)C)=C(C(C(C)C)=C1)C2=C(P(C3CCCCC3)C4CCCCC4)C(OC)=CC=C2OC",
    "Di-tert-butylphenylphosphine": r"CC(C)(C)P(C1=CC=CC=C1)C(C)(C)C",
    "(t-Bu)PhCPhos": r"CN(C)C1=CC=CC(N(C)C)=C1C2=CC=CC=C2P(C(C)(C)C)C3=CC=CC=C3",
    "Tricyclohexylphosphine": r"P(C1CCCCC1)(C2CCCCC2)C3CCCCC3",
    "PPh3": r"P(C1=CC=CC=C1)(C2=CC=CC=C2)C3=CC=CC=C3",
    "XPhos": r"CC(C1=C(C2=CC=CC=C2P(C3CCCCC3)C4CCCCC4)C(C(C)C)=CC(C(C)C)=C1)C",
    "P(2-furyl)3": r"P(C1=CC=CO1)(C2=CC=CO2)C3=CC=CO3",
    "Methyldiphenylphosphine": r"CP(C1=CC=CC=C1)C2=CC=CC=C2",
    "1268824-69-6": r"CC(OC1=C(P(C2CCCCC2)C3CCCCC3)C(OC(C)C)=CC=C1)C",
    "JackiePhos": r"FC(F)(F)C1=CC(P(C2=C(C3=C(C(C)C)C=C(C(C)C)C=C3C(C)C)C(OC)=CC=C2OC)C4=CC(C(F)(F)F)=CC(C(F)(F)F)=C4)=CC(C(F)(F)F)=C1",
    "SCHEMBL15068049": r"C[C@]1(O2)O[C@](C[C@]2(C)P3C4=CC=CC=C4)(C)O[C@]3(C)C1",
    "Me2PPh": r"CP(C)C1=CC=CC=C1",
}

bases = {
    "Potassium acetate": r"O=C([O-])C.[K+]",
    "Potassium pivalate": r"O=C([O-])C(C)(C)C.[K+]",
    "Cesium acetate": r"O=C([O-])C.[Cs+]",
    "Cesium pivalate": r"O=C([O-])C(C)(C)C.[Cs+]",
}

defaults = (0.1, 90, "DMAc", "JackiePhos", "Cesium acetate")

iface = gr.Interface(
    fn=lookup,
    inputs=[
        gr.Radio(label="Concentration", choices=[0.057, 0.1, 0.153], value=defaults[0]),
        gr.Radio(label="Temperature", choices=[90, 105, 120], value=defaults[1]),
        gr.Radio(label="Solvent", choices=list(solvents.keys()), value=defaults[2]),
        gr.Radio(label="Ligand", choices=list(ligands.keys()), value=defaults[3]),
        gr.Radio(label="Base", choices=list(bases.keys()), value=defaults[4]),
    ],
    outputs=gr.Number(lookup(*defaults), label="yield value (maximize)"),
)
iface.launch()