import sys
import numpy as np
import pandas as pd
from flask import render_template, request
from functions import SigFigs, HtmlNumber, Piringer, WilkeChang, CdfPlot
#from functions import Piecewise, PowerLaw
from qrf.functions import QRF_Apply, QRF_Ceramic
from . import blueprint
from polymers import Polymers, Polymers3
from ChemID import *
from quantity_module.functions 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

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

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

# load the index page for the exposure module
@blueprint.route('/quantity', methods=['GET'])
def exposure():
    return render_template('quantity_index.html', polymers=polymers, solvents=solvents,
                           tool_title="CHRIS - Total quantity",
                           nav_links=[
                               ("Tool", "/quantity"),
                               ("CoU", "/quantity_module/COU"),
                               ("About", "/quantity_module/RST"),
                           ],
                           footer_text="Office of Science and Engineering Laboratories · ",
                           contact_email="david.saylor@fda.hhs.gov",
                           )
    ## NOTE solvents are defined in quantity_functions.py


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

    Polymer_Tg = float(request.form['Polymer_Tg'])  ## NOTE Tg is provided in C
    T = float(request.form['T'])
    Polymer_Tg += 273.15
    T += 273.15
    #rng = np.random.Generator(np.random.PCG64(seed=12345))
    CHRIS_category = None

    if T<Polymer_Tg:
        return render_template('quantity_error.html', message='Unfortunately, the CHRIS-Total Quantity module can only be used for polymers with glass transition temperatures lower than the extraction temperature.',
                               tool_title="CHRIS - Total quantity",
                               nav_links=[
                                   ("Tool", "/quantity"),
                                   ("CoU", "/quantity_module/COU"),
                                   ("About", "/quantity_module/RST"),
                               ],
                               footer_text="Office of Science and Engineering Laboratories · ",
                               contact_email="david.saylor@fda.hhs.gov",
                               )

    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 - Total quantity",
                               nav_links=[
                                   ("Tool", "/quantity"),
                                   ("CoU", "/quantity_module/COU"),
                                   ("About", "/quantity_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'

    if is_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 - Total quantity",
                               nav_links=[
                                   ("Tool", "/quantity"),
                                   ("CoU", "/quantity_module/COU"),
                                   ("About", "/quantity_module/RST"),
                               ],
                               footer_text="Office of Science and Engineering Laboratories · ",
                               contact_email="david.saylor@fda.hhs.gov",
                               )
    if is_ceramic:
        MW = 1100.

    M_expt = float(request.form['amount']) # amount
    units = request.form['units']
    mass = float(request.form['mass'])
    Polymer_Density = float(request.form['density'])
    Polymer_Volume = mass / Polymer_Density # vol
    Surface_Area = float(request.form['area']) # area
    Solvent_Volume = float(request.form['solventvol'])
    Solvent_Name = request.form['solventname']
    Swelling_percent = float(request.form['swelling'])
    Swelling_wtfrac = Swelling_percent/100
    Extraction_Time = float(request.form['time'])
    iterations = float(request.form['iterations'])
    K_expt = float(request.form['K'])
    Solvent_MW = Solvent_MWs[Solvent_Name]
    Solvent_Density = Solvent_Densities[Solvent_Name]
    Solute_MW = MW
    Solute_Vabc = get_V(smiles)

    ## extent of swelling, V_swelled/V_unswelled
    if Swelling_wtfrac > 0:
        Swelling_volfrac = Polymer_Density / (Solvent_Density * (1/Swelling_wtfrac-1) + Polymer_Density)
        Qv = 1+Swelling_volfrac/(1-Swelling_volfrac)
    else:
        Swelling_volfrac = 0
        Qv = 1
    # check that swelling is physically possible (swollen polymer volume is less than polymer+solvent volume)
    if Solvent_Volume < Polymer_Volume*(Qv-1):
        return render_template('quantity_error.html', message=f'Based on the information provided, the amount of swelling appears physically impossible. Please check the values you entered.',
                               tool_title="CHRIS - Total quantity",
                               nav_links=[
                                   ("Tool", "/quantity"),
                                   ("CoU", "/quantity_module/COU"),
                                   ("About", "/quantity_module/RST"),
                               ],
                               footer_text="Office of Science and Engineering Laboratories · ",
                               contact_email="david.saylor@fda.hhs.gov",
                               )
    if units == 'mg':
        mass_units = mass*1e3
    elif units == 'µg':
        mass_units = mass*1e6

    polymer = request.form['polymer']

    pIndex = np.argmax(polymers == polymer)
    # QRF is implemented for 25-75 C
    if polymer == 'Other polymer':
        use_qrf = True
    else:
        use_qrf = False

    if use_qrf:
        quantiles = list(np.linspace(0.0,1.0,201))
        if is_ceramic:
            diff,domain_extrap = QRF_Ceramic(Polymer_Density, Polymer_Tg-273.15, quantiles=quantiles, T=T-273.15, worstcase='hi')
        else:
            diff,domain_extrap = QRF_Apply(Polymer_Density, Polymer_Tg-273.15, smiles, quantiles=quantiles, T=T-273.15)
        if domain_extrap:
            # outside training domain, default to Wilke-Chang
            method = 'qrf/wc'
            D_dist_noswell, D_dist_swell = get_D_dists_new(Swelling_wtfrac, T, Polymer_Tg, Solvent_Name, Solvent_MW, Solute_MW, Solute_Vabc, 'G2', return_DCs=False, N=N_sample)
            M0_pred = get_M_dist(D_dist_swell, M_expt, Polymer_Volume, Surface_Area, Solvent_Volume, Extraction_Time*3600, K_expt=K_expt, Qv=Qv, iterations=iterations)
        else:
            method = 'qrf'
            D_dist_noswell, D_dist_swell = get_D_dists_new(Swelling_wtfrac, T, Polymer_Tg, Solvent_Name, Solvent_MW, Solute_MW, Solute_Vabc, CHRIS_category, return_DCs=False, N=N_sample, input_Ds=diff)
            M0_pred = get_M_dist(D_dist_swell, M_expt, Polymer_Volume, Surface_Area, Solvent_Volume, Extraction_Time*3600, K_expt=K_expt, Qv=Qv, iterations=iterations)
    else:
        ## use categories
        CHRIS_category = categories[pIndex]
        CHRIS_flag = None
        if CHRIS_category is None:
            ## worst-case for a generic polymer --> G2
            CHRIS_flag = 'wc'
            CHRIS_category = 'G2'
        D_dist_noswell, D_dist_swell = get_D_dists_new(Swelling_wtfrac, T, Polymer_Tg, Solvent_Name, Solvent_MW, Solute_MW, Solute_Vabc, CHRIS_category, return_DCs=False, N=N_sample, input_Ds=None)
        M0_pred = get_M_dist(D_dist_swell, M_expt, Polymer_Volume, Surface_Area, Solvent_Volume, Extraction_Time*3600, K_expt=K_expt, Qv=Qv, iterations=iterations)
        if CHRIS_flag is None:
            method = 'category'
        else:
            method = 'wc'
#    if 0:
#        print('DEBUG')
#        print('Swelling_wtfrac, T, Polymer_Tg, Solvent_Name, Solvent_MW, Solute_MW, CHRIS_category')
#        print(Swelling_wtfrac, T, Polymer_Tg, Solvent_Name, Solvent_MW, Solute_MW, CHRIS_category)
#        print(np.nanquantile(D_dist_noswell, [0.05,0.5,0.95]))
#        print(np.nanquantile(D_dist_swell, [0.05,0.5,0.95]))
#        print('M_expt, Polymer_Volume, Surface_Area, Solvent_Volume, Extraction_Time*3600, K_expt, method, Qv, Swelling_volfrac')
#        print(M_expt, Polymer_Volume, Surface_Area, Solvent_Volume, Extraction_Time*3600, K_expt, method, Qv, Swelling_volfrac)
#        print(np.nanquantile(M0_pred, [0.05,0.5,0.95]))

    # Generate the rate plot using matplotlib
    #pngImageB64String = CdfPlot(M0_pred[~np.isnan(M0_pred)], units=units)

    # table of percentiles
    vals = np.nanquantile(M0_pred, [0.05,0.25,0.5,0.75,0.95])
    pcts = np.round(vals/vals[2]*100,1)
    Mvals = M_expt/vals
    Dvals = np.nanquantile(D_dist_swell, [0.05,0.25,0.5,0.75,0.95])
    vals = [HtmlNumber(v,3) for v in vals]
    Dvals = [HtmlNumber(v,3) for v in Dvals]
    Mvals = [HtmlNumber(v,3) for v in Mvals]
    df_table = pd.DataFrame(data=np.array([[r'5<sup>th</sup>', r'25<sup>th</sup>', r'50<sup>th</sup> (median)', r'75<sup>th</sup>', r'95<sup>th</sup>'], Dvals, Mvals, vals, pcts]).T, columns=[r'Percentile', r'\( D \) (cm<sup>2</sup>/s)', r'\( M/M_0 \)', r'\( M_0 \)'+f' ({units})', r'\( M_0 \) (% median)'])
    styles = [dict(selector="th", props=[("padding", "0 0.5em")]), dict(selector="td", props=[("padding", "0 0.5em")]),]
    table = df_table.style.set_table_styles(styles).set_properties(subset=[r'\( D \) (cm<sup>2</sup>/s)', r'\( M/M_0 \)', fr'\( M_0 \) ({units})', r'\( M_0 \) (% median)'], **{'text-align': 'right'}).set_table_attributes('border="2"').hide(axis='index').to_html(index=False, escape=False, justify='center')

    tau = np.nanquantile(D_dist_swell,0.5) * (Extraction_Time*3600) / (Polymer_Volume/Surface_Area)**2 * iterations

    M0_out = SigFigs(np.nanquantile(M0_pred,0.5),6)
    tau_out = SigFigs(tau,6)
    mass_units = SigFigs(mass_units,6)

    return render_template('quantity_report.html', show_properties=show_properties, polymers=polymers, pIndex=pIndex,
                           area=Surface_Area, vol=Polymer_Volume, units=units, M=M_expt, M0=M0_out, time=Extraction_Time, iterations=iterations,
                           solventvol=Solvent_Volume, solventname=Solvent_Name, swelling=Swelling_percent, K=K_expt, T=T, tau=tau_out,
                           chemName=chemName, MW=MW, LogP=LogP, rho=rho, mp=mp, iupac=iupac, cas=cas, smiles=smiles, molImage=molImage, table=table,
                           LogP_origin=LogP_origin, rho_origin=rho_origin, mp_origin=mp_origin, ceramic=is_ceramic, methods=[method,round(Polymer_Tg-273.15),Polymer_Density],
                           mass=mass, mass_units=mass_units, density=Polymer_Density,
                           tool_title="CHRIS - Total quantity",
                           nav_links=[
                               ("Tool", "/quantity"),
                               ("CoU", "/quantity_module/COU"),
                               ("About", "/quantity_module/RST"),
                           ],
                           footer_text="Office of Science and Engineering Laboratories · ",
                           contact_email="david.saylor@fda.hhs.gov",
                           )