from flask import Flask, render_template, request, redirect

from .chemprop_lime import *
from .chemprop_path_explain import *
from .explainer_utils import *

from chemprop.utils import load_checkpoint

from matplotlib.pyplot import cm
import os
import hashlib


def mkdir(path):
    if not os.path.isdir(path):
        os.makedirs(path)
        return True
    return False


cmap = cm.get_cmap("bone_r")
color = lambda x: tuple(map(lambda x: int(x*255), cmap(x)[:3]))
explainers = list(enumerate(["PathExplain", "LIME"]))
exps_ids = ["path_explain", "lime"]

tmp_path = "/tmp/chemexp_data/"

app = Flask(__name__)


class GUIServer:
    def __init__(self):
        self.models_dir = None
        self.checkpoints = None
        self.models = None
        self.model_names = None
        
        # creating temporary folder (for generated figures)
        mkdir(tmp_path)
        self.cd = "/".join(__file__.split("/")[:-1])
        if not os.path.islink(self.cd + "/static/tmp"):
            os.symlink(tmp_path, self.cd + "/static/tmp")
        
    
    def set_models_dir(self, models_dir):
        """
        models_dir: folder containing checkpoints generated by chemprop (.pt files)
        """
        models_dir += "/"
        self.models_dir = models_dir
        self.checkpoints = {
            file.rstrip(".pt"): load_checkpoint(models_dir+file)
            for file in os.listdir(models_dir)
            if file.endswith(".pt") and os.path.isfile(models_dir+file)
        }
        self.models = [
            {
                "name": name,
                "lime": LIMEExplainer(self.checkpoints[name]),
                "path_explain": PathExplainer(self.checkpoints[name])
            }
            for name in self.checkpoints
        ]
        self.model_names = list(enumerate([m["name"] for m in self.models]))

    
    def run(self, **kwargs):
        """
        **kwargs: arguments for the Flask server
        """
        app.run(**kwargs)


@app.route("/")
def index():
    return render_template("index.html", models=server.model_names)


@app.route("/display", methods=["POST"])
def display():
    mol = request.form["smiles"]
    try:
        model = server.models[int(request.form["model"])]["lime"]
    except:
        model = server.models[0]["lime"]
    # image generation
    mkdir(tmp_path + "mols")
    mol_md5 = hashlib.md5(mol.encode()).hexdigest()
    mol_to_png(mol, tmp_path + "mols/" + mol_md5 + ".png", use_color=True)
    # model prediction
    try:
        model._set_molecule(mol)
        predictions = model.predict_proba([model.input])[0]
        n_atoms = model.mol.n_atoms
        n_bonds = model.mol.n_bonds
    except:
        predictions = [0]
        n_atoms = 0
        n_bonds = 0
    return render_template("display.html", mol_md5=mol_md5,
                           explainers=explainers,
                           predictions=predictions,
                           predcolors=[color(x) for x in predictions],
                           n_predictions=len(predictions),
                           n_atoms=n_atoms,
                           n_bonds=n_bonds)


@app.route("/explain", methods=["POST"])
def explain():
    classNb = int(request.form["class"])
    mol = request.form["smiles"]
    try:
        explainer = exps_ids[int(request.form["exp"])]
        model = server.models[int(request.form["model"])]
    except:
        explainer = exps_ids[0]
        model = server.models[0]
    print(mol)
    # create an explanation and save it in a json file,
    # or load it if it already exists
    base_dir = f"{server.cd}/static/{model['name']}/{explainer}"
    mkdir(base_dir)
    mol_hash = hashlib.md5(mol.encode()).hexdigest()
    filename = f"{base_dir}/{mol_hash}_{classNb}.json"
    if not os.path.isfile(filename):
        if explainer == "lime":
            exp = model["lime"].explain_molecule(mol, classNb)
        else:
            exp = model["path_explain"].explain_molecule(mol, classNb)
        exp_to_json(exp, filename)
    else:
        exp = json_to_exp(filename)
    # generate images
    static_path = f"static/tmp/{model['name']}/{explainer}/{mol_hash}_{classNb}"
    mkdir(f"{server.cd}/{static_path}")
    exp_to_png(exp, f"{server.cd}/{static_path}/avg.png")
    for i in range(8):
        atoms_exp_to_png(exp, i, f"{server.cd}/{static_path}/exp_atom_{i}.png")
    for i in range(4):
        bonds_exp_to_png(exp, i, f"{server.cd}/{static_path}/exp_bond_{i}.png")
    bonds_max = abs(exp.f_bonds).max() if len(exp.f_bonds) > 0 else 0
    atoms_max = abs(exp.f_atoms).max() if len(exp.f_atoms) > 0 else 1
    bound = round(max(bonds_max, atoms_max), 3)

    return render_template("results.html",
                           mol=mol,
                           path=static_path,
                           bound=bound,
                           classNb=classNb,
                           explainer=explainer)


@app.route("/interactions", methods=["POST"])
def interactions():
    classNb = int(request.form["class"])
    mol = request.form["smiles"]
    try:
        feature1 = int(request.form["feature1"])
        feature2 = int(request.form["feature2"])
        model = server.models[int(request.form["model"])]
    except:
        feature1 = 0
        feature2 = 5
        model = server.models[0]
    print(mol)
    # create an explanation and save it in a json file,
    # or load it if it already exists
    mkdir(f"{server.cd}/static/{model['name']}/path_explain/interactions")
    mol_hash = hashlib.md5(mol.encode()).hexdigest()
    filename = f"{server.cd}/static/{model['name']}/path_explain/interactions/{mol_hash}_{classNb}_{feature1}_{feature2}.json"
    if not os.path.isfile(filename):
        exp = model["path_explain"].explain_interactions(mol, feature1, feature2, classNb)
        exp_to_json(exp, filename)
    else:
        exp = json_to_exp(filename)
    # generate images
    static_path = "static/tmp/path_explain" + f"/{mol_hash}_{model['name']}_{classNb}"
    mkdir(f"{server.cd}/{static_path}")
    img_name = f"{static_path}/interactions_{classNb}_{feature1}_{feature2}.png"
    atoms_exp_to_png(exp, 0, f"{server.cd}/{img_name}")
    atoms_max = abs(exp.f_atoms).max() if len(exp.f_atoms) > 0 else 1
    bound = round(atoms_max, 3)

    return render_template("interactions.html",
                           path=img_name,
                           bound=bound,
                           classnb=classNb)


server = GUIServer()
