import sys
import numpy as np
import pandas as pd
from flask import render_template, request
from functions import SigFigs, Piringer, WilkeChang, SheetRelease, SheetRates, RatePlot
from functions import Piecewise, PowerLaw
from qrf.functions import QRF_Apply, QRF_Ceramic
from . import blueprint
from polymers import Polymers, Polymers3
from ChemID import *
import rdkit
from rdkit.Chem import AllChem as Chem

# get additional physical properties, options are: logp, rho, mp
# MP is looked up lazily only when it can affect ceramic classification.
get_properties = []
# show additional physical properties
show_properties = False
# output additional info for physical properties
debug = False
# flag for testing new polymer categories
#use_new = True

# load polymer data including Ap values
polymers, categories, params = Polymers3()

@blueprint.get("/exposure3_module/COU")
def COU():
    return render_template(
        "exposure3_COU.html",
        tool_title="CHRIS - Bulk Chemicals (v3)",
        nav_links=[
            ("Tool", "/exposure3"),
            ("CoU", "/exposure3_module/COU"),
            ("About", "/exposure3_module/RST"),
        ],
        footer_text="Office of Science and Engineering Laboratories · ",
        contact_email="david.saylor@fda.hhs.gov",
    )
@blueprint.get("/exposure3_module/RST")
def RST():
    return render_template("exposure3_RST.html",
        tool_title="CHRIS - Bulk Chemicals (v3)",
        nav_links=[
            ("Tool", "/exposure3"),
            ("CoU", "/exposure3_module/COU"),
            ("About", "/exposure3_module/RST"),
        ],
        footer_text="Office of Science and Engineering Laboratories · ",
        contact_email="david.saylor@fda.hhs.gov",
    )

@blueprint.route('/exposure3', methods=['GET'])
def exposure():
    return render_template('exposure3_index.html', polymers=polymers,
        tool_title="CHRIS - Bulk Chemicals (v3)",
        nav_links=[
            ("Tool", "/exposure3"),
            ("CoU", "/exposure3_module/COU"),
            ("About", "/exposure3_module/RST"),
        ],
        footer_text="Office of Science and Engineering Laboratories · ",
        contact_email="david.saylor@fda.hhs.gov",
    )


# build the report page for the exposure module
@blueprint.route('/exposure3', methods=['POST'])
def exp_post():

    chemName = request.form['chemName']
    IDtype = request.form['IDtype']

    if debug:
        iupac, cas, smiles, MW, LogP, LogP_origin, rho, rho_origin, mp, mp_origin, molImage, error = ResolveChemical(chemName, IDtype, debug=debug, get_properties=['logp','rho'])
    else:
        LogP_origin, rho_origin = '',''
        iupac, cas, smiles, MW, LogP, rho, mp, molImage, error = ResolveChemical(chemName, IDtype, get_properties=get_properties)

    if error > 0:
        return render_template('chemError.html', chemName=chemName, IDtype=IDtype, error=error,
                               error_message=ERROR_CODES.get(error, 'Unknown chemical identification error'),
                               tool_title="CHRIS - Bulk Chemicals (v3)",
                               nav_links=[
                                   ("Tool", "/exposure3"),
                                   ("CoU", "/exposure3_module/COU"),
                                   ("About", "/exposure3_module/RST"),
                               ],
                               footer_text="Office of Science and Engineering Laboratories · ",
                               contact_email="david.saylor@fda.hhs.gov",
                               )

    # metals/ceramics logic; only look up MP if the structural checks need it
    is_metal, is_ceramic, mp, material_mp_origin, mp_searched = ClassifyMaterial(cas, iupac, smiles)
    if debug:
        LogP_origin = f' ({LogP_origin})'
        rho_origin = f' ({rho_origin})'
        mp_origin = f' ({material_mp_origin})' if mp_searched else ''
    else:
        mp_origin = ''

    #MW = SigFigs(MW, 6)
    if 'logp' not in get_properties:
        LogP = 'Not searched'
    elif LogP is np.nan or LogP is None:
        LogP = 'Not found'
    else:
        LogP = SigFigs(LogP, 4)
    if 'rho' not in get_properties:
        rho = 'Not searched'
    elif rho is np.nan or rho is None:
        rho = 'Not found'
    else:
        rho = SigFigs(rho, 4)
    if not mp_searched:
        mp = 'Not searched'
    elif mp is np.nan or mp is None:
        mp = 'Not found'

    #ceramic = False
    #mol = Chem.MolFromSmiles(smiles)
    #atom_num_list = [a.GetAtomicNum() for a in mol.GetAtoms()]
    #is_metal = set(atom_num_list) <= METAL_ATOM_SET
    if is_metal:
        # if all atoms are metals -> this is a metal
        return render_template('metalError.html', show_properties=show_properties, chemName=chemName, MW=MW,
                               LogP=LogP, rho=rho, mp=mp, iupac=iupac,
                               cas=cas, smiles=smiles, molImage=molImage,
                               LogP_origin=LogP_origin, rho_origin=rho_origin, mp_origin=mp_origin,
                               tool_title="CHRIS - Bulk Chemicals (v3)",
                               nav_links=[
                                   ("Tool", "/exposure3"),
                                   ("CoU", "/exposure3_module/COU"),
                                   ("About", "/exposure3_module/RST"),
                               ],
                               footer_text="Office of Science and Engineering Laboratories · ",
                               contact_email="david.saylor@fda.hhs.gov",
                               )
    #else:
    if is_ceramic:
        MW = 1100.
        # get number of carbon-carbon bonds
        #num_CC_bonds = sum([1 if b.GetBeginAtom().GetAtomicNum() == 6 and b.GetEndAtom().GetAtomicNum() == 6 else 0 for b in mol.GetBonds()])
        #if not num_CC_bonds and (mp is not None) and mp > 700.:
            # if not a metal, no C-C bonds, and mp > 700 (sodium chloride has mp ~ 800), assume ceramic...
            #MW = 1100.
            #ceramic = True

    amount = float(request.form['amount'])
    mass = float(request.form['mass'])
    density = float(request.form['density'])
    vol = mass / density
    area = float(request.form['area'])
    exposure = request.form['exposure']

    polymer = request.form['polymer']
    pIndex = np.argmax(polymers == polymer)

    if exposure != 'limited':
        time = 24.
    else:
        time = float(request.form['exptime'])

    if exposure != 'long-term':
        TTC = 0.12
    else:
        TTC = 0.0015

    units = request.form['units']
    if units == 'µg':
        TTC *= 1000

    polytg = ''
    use_qrf = False
    if polymer == 'Other polymer':
        polytg = request.form['polytg']
        if polytg == '':
            # left blank, use old default (Wilke Chang), which is taken care of by pIndex
            pass
        else:
            polytg = float(polytg)
            use_qrf = True
        #print('b', polytg, type(polytg), file=sys.stderr)

    if use_qrf:
        #print('using qrf', file=sys.stderr)
        method = 'qrf'
        if is_ceramic:
            diff,domain_extrap = QRF_Ceramic(density, polytg, quantiles=[0.03,0.5,0.97], worstcase='hi')
        else:
            diff,domain_extrap = QRF_Apply(density, polytg, smiles, quantiles=[0.03,0.5,0.97])
        diff = diff[2] # upper bound
        if domain_extrap:
            # outside training domain, default to Wilke-Chang
            diff = Piecewise(MW, params[None])
            method = 'qrf/wc'
    else:
        ## use categories
        category = categories[pIndex]
        diff = Piecewise(MW, params[category])
        domain_extrap = False
        if category:
            method = 'category'
        else:
            method = 'wc'

    release = SheetRelease(amount, vol, area, time, diff)

    MOS = TTC / release

    # Generate the rate plot using matplotlib
    #tarray = np.arange(1., 31., 1.)
    #rates = SheetRates(amount, vol, area, tarray, diff)
    #pngImageB64String = RatePlot(tarray, rates)

    # Generate cumulative release at given times
    tlist = [1.]
    if len(request.form.getlist('exposure30d')):
        tlist.append(30.)
    if len(request.form.getlist('exposureuser')):
        tlist.append(float(request.form['exposureuserbox']))
    release_list = [SheetRelease(amount, vol, area, t*24, diff) for t in tlist]
    #print(release_list, file=sys.stderr)
    frac_list = [str(SigFigs(r/amount*100,2))+'%' for r in release_list]
    release_list = [SigFigs(r,2) for r in release_list]
    df_rel = pd.DataFrame(np.transpose([tlist,release_list,frac_list]),columns=['Exposure duration (days)',f'Cumulative exposure ({units})','Cumulative exposure (%)'])
    df_rel['Exposure duration (days)'] = df_rel['Exposure duration (days)'].astype(float)
    release_table = df_rel.sort_values(by='Exposure duration (days)').to_html(index=False)

    release = SigFigs(release, 2)
    MOS = SigFigs(MOS, 2)
    diff = SigFigs(diff, 2)

    return render_template('exposure3_report.html', show_properties=show_properties, polymers=polymers, pIndex=pIndex, release=release, release_table=release_table,
                           area=area, vol=vol, amount=amount, diff=diff, time=time, exposure=exposure, TTC=TTC, units=units,
                           MOS=MOS, chemName=chemName, MW=MW, LogP=LogP, rho=rho, mp=mp, iupac=iupac, cas=cas, smiles=smiles, molImage=molImage,
                           LogP_origin=LogP_origin, rho_origin=rho_origin, mp_origin=mp_origin, ceramic=is_ceramic, methods=[method,polytg,density],
                           tool_title="CHRIS - Bulk Chemicals (v3)",
                           nav_links=[
                               ("Tool", "/exposure3"),
                               ("CoU", "/exposure3_module/COU"),
                               ("About", "/exposure3_module/RST"),
                           ],
                           footer_text="Office of Science and Engineering Laboratories · ",
                           contact_email="david.saylor@fda.hhs.gov",
                           )