import pickle
import random
import numpy as np
import pandas as pd
from faerun import Faerun, host
from rdkit.Chem import AllChem
from rdkit.Chem import rdMolDescriptors as Descriptors
from sklearn.decomposition import PCA


def get_fingerprint(in_smiles):
    results = []
    smiles = []
    for s in in_smiles:
        mol = AllChem.MolFromSmiles(s)
        if mol is not None:
            if mol.GetNumAtoms() > 100:
                continue
            mqn = Descriptors.MQNs_(mol)
            results.append(np.array(mqn))
            smiles.append(s)

    return results, smiles


def load():
    drugbank = []
    with open("drugbank.smi") as f:
        for line in f.readlines():
            drugbank.append(line.split()[0].strip())

    mqns, smiles = get_fingerprint(drugbank)
    mqns = np.array(mqns)

    pca = PCA(n_components=2)
    result = pca.fit_transform(mqns)
    return result, mqns, smiles


def main():
    coords, mqns, smiles = load()

    data = {"x": [], "y": [], "c": [], "s": [], "labels": smiles}

    for i, _ in enumerate(coords):
        data["x"].append(coords[i][0])
        data["y"].append(coords[i][1])
        data["c"].append(mqns[i][0])
        data["s"].append(random.randint(1, 5))

    df = pd.DataFrame.from_dict(data)

    faerun = Faerun()
    faerun.add_scatter(
        "drugbank",
        df,
        shader="smoothCircle",
        point_scale=2.0,
        colormap="jet",
        has_legend=True,
        categorical=False,
    )

    with open("example.faerun", "wb+") as f:
        pickle.dump(faerun.create_python_data(), f, protocol=pickle.HIGHEST_PROTOCOL)

    info = (
        "#Welcome to Fearun",
        "This is a small demonstration visualising drugs stored in [Drugbank](http://drugbank.ca) using coordinates generated by a PCA projection."
        "",
        "By the way, this info was generated using markdown, this means that you can easily:",
        "- Add lists",
        "- Build tables",
        "- Insert images and links",
        "- Add code examples",
        "- ...",
    )

    host(
        "example.faerun",
        label_type="smiles",
        legend=True,
        theme="light",
        info="\n".join(info),
    )


if __name__ == "__main__":
    main()
