import numbers
import numpy as np
from flask import render_template, request
from functions import SigFigs, Piringer, WilkeChang, SheetRelease, SheetRates, RatePlot
from functions import Piecewise, PowerLaw
from polymers import Polymers, Polymers3
from ChemID import *
from . import blueprint

import rdkit
from rdkit.Chem import AllChem as Chem

# get additional physical properties, options are: logp, rho, mp
#get_properties = [] # don't get any; this breaks ceramics logic
#get_properties = ['logp','rho','mp'] # get all three
get_properties = ['mp'] # only get mp
# show additional physical properties
show_properties = False
# output additional info for physical properties
debug = False
# flag for testing new polymer categories
use_new = True

#ORGANIC_ATOM_SET = {5, 6, 7, 8, 9, 15, 16, 17, 35, 53}
#METAL_ATOM_SET = set([3,4,11,12,13] + list(range(19,31+1)) + list(range(37,50+1)) + list(range(55,84+1)) + list(range(87,114+1)) + [116])

# load polymer data including Ap values
if not use_new:
    polymers, Ap = Polymers()
else:
    polymers, categories, params = Polymers3()

# load the index page for the exposure module
@blueprint.route('/exposure2', methods=['GET'])
def exposure():
    return render_template('exposure2_index.html', polymers=polymers)


# build the report page for the exposure module
@blueprint.route('/exposure2', 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','mp'])
        LogP_origin, rho_origin, mp_origin = f' ({LogP_origin})', f' ({rho_origin})', f' ({mp_origin})', 
    else:
        LogP_origin, rho_origin, mp_origin = '','',''
        iupac, cas, smiles, MW, LogP, rho, mp, molImage, error = ResolveChemical(chemName, IDtype, get_properties=get_properties)

    if error > 0:
        # TODO output more useful info
        return render_template('exposure2_chemError.html')

    #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 'mp' not in get_properties:
        mp = 'Not searched'
    elif mp is np.nan or mp is None:
        mp = 'Not found'

    # metals/ceramics logic
    if isinstance(mp, numbers.Number):
        is_metal, is_ceramic = CeramicOrMetal(smiles,mp)
    else:
        is_metal, is_ceramic = CeramicOrMetal(smiles,100)
    #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
    #if natoms == 1 and smiles != '[C]':
        # only one atom, except for carbon -> assumed metal
        # return render_template('chemError.html')
        return render_template('exposure2_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)
    #else:
    if is_ceramic:
        MW = 1100.

    if MW < 100. and not use_new:
        return render_template('exposure2_MwError.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)

    amount = float(request.form["amount"])
    mass = float(request.form["mass"])
    density = float(request.form["density"])
    vol = mass / density
    polymer = request.form["polymer"]
    #pIndex = (np.where(polymers == polymer)[0])[0]
    pIndex = np.argmax(polymers == polymer)
    area = float(request.form["area"])
    exposure = request.form["exposure"]

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

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

    assume = np.array((request.form.get("assume1") is not None, request.form.get("assume2") is not None,
                      request.form.get("assume3") is not None, request.form.get("assume4") is not None,
                      request.form.get("assume5") is not None))

    if use_new:
        category = categories[pIndex]
        diff = Piecewise(MW, params[category])
    else:
        if not np.isnan(Ap[pIndex]):
            diff = Piringer(MW, Ap[pIndex])
        else:
            diff = WilkeChang(MW)

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

    MOS = TTC / release

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

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

    return render_template('exposure2_report.html', show_properties=show_properties, polymers=polymers, pIndex=pIndex, release=release,
                           assume=assume, area=area, vol=vol, amount=amount, diff=diff, time=time, exposure=exposure, TTC=TTC,
                           MOS=MOS, chemName=chemName, image=pngImageB64String, 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)