# -*- coding: utf-8 -*-
"""
Author: Philipp Seidl
        ELLIS Unit Linz, LIT AI Lab, Institute for Machine Learning
        Johannes Kepler University Linz
Contact: seidl@ml.jku.at

File contains functions that
"""

from . import model
import torch
import os

MODEL_PATH = 'data/model/'

def smarts2svg(smarts, useSmiles=True, highlightByReactant=True, save_to=''):
    """
    draws smiles of smarts to an SVG and displays it in the Notebook,
    or optinally can be saved to a file `save_to`
    adapted from https://www.kesci.com/mw/project/5c7685191ce0af002b556cc5
    """
    # adapted from https://www.kesci.com/mw/project/5c7685191ce0af002b556cc5
    from rdkit import RDConfig
    from rdkit import Chem
    from rdkit.Chem import Draw, AllChem
    from rdkit.Chem.Draw import rdMolDraw2D
    from rdkit import Geometry
    import matplotlib.pyplot as plt
    import matplotlib.cm as cm
    import matplotlib
    from IPython.display import SVG, display

    rxn = AllChem.ReactionFromSmarts(smarts,useSmiles=useSmiles)
    d = Draw.MolDraw2DSVG(900, 100)

    # rxn = AllChem.ReactionFromSmarts('[CH3:1][C:2](=[O:3])[OH:4].[CH3:5][NH2:6]>CC(O)C.[Pt]>[CH3:1][C:2](=[O:3])[NH:6][CH3:5].[OH2:4]',useSmiles=True)
    colors=[(0.3, 0.7, 0.9),(0.9, 0.7, 0.9),(0.6,0.9,0.3),(0.9,0.9,0.1)]
    try:
        d.DrawReaction(rxn,highlightByReactant=highlightByReactant)
        d.FinishDrawing()

        txt = d.GetDrawingText()
        # self.assertTrue(txt.find("<svg") != -1)
        # self.assertTrue(txt.find("</svg>") != -1)

        svg = d.GetDrawingText()
        svg2 = svg.replace('svg:','')
        svg3 = SVG(svg2)
        display(svg3)

        if save_to!='':
            with open(save_to, 'w') as f_handle:
                f_handle.write(svg3.data)
    except:
        print('Error drawing')

    return svg2

def list_models(model_path=MODEL_PATH):
    """returns a list of loadable models"""
    return dict(enumerate(list(filter(lambda k: str(k)[-3:]=='.pt', os.listdir(model_path)))))

def load_clf(model_fn='', model_path=MODEL_PATH, device='cpu', model_type='mhn'):
    """ returns the model with loaded weights given a filename"""
    import json
    config_fn = '_'.join(model_fn.split('_')[-2:]).split('.pt')[0]
    conf_dict = json.load( open( f"{model_path}{config_fn}_config.json" ) )
    train_conf_dict = json.load( open( f"{model_path}{config_fn}_config.json" ) )

    # specify the config the saved model had
    conf = model.ModelConfig(**conf_dict)
    conf.device = device
    print(conf.__dict__)

    if model_type == 'staticQK':
        clf = model.StaticQK(conf)
    elif model_type == 'mhn':
        clf = model.MHN(conf)
    elif model_type == 'segler':
        clf = model.SeglerBaseline(conf)
    elif model_type == 'fortunato':
        clf = model.SeglerBaseline(conf)
    else:
        raise NotImplementedError('model_type',model_type,'not found')

    # load the model
    PATH = model_path+model_fn
    params = torch.load(PATH, map_location=torch.device('cpu')) #!!!
    clf.load_state_dict(params, strict=False)
    if 'templates+noise' in params.keys():
        print('loading templates+noise')
        clf.templates = params['templates+noise']
        #clf.templates.to(clf.config.device)
    return clf