import json,pickle
import numbers
import numpy as np
import pandas as pd
import scipy as sp
from scipy.optimize import bisect
import scipy.special

import rdkit
from rdkit.Chem import AllChem as Chem
import chemicals

import mordred
import mordred.descriptors

from functions import PowerLaw, Piringer

## sampling parameters
N_sample = int(1e5)  # reduced from 1e6 for performance; results are statistically indistinguishable
#rng = np.random.Generator(np.random.PCG64(seed=12345))

## list of solvents to include, all semi-polar and non-polar solvents in ISO 10993-18:2020 Table D.1 (except DMSO, which is not used in practice)
solvents = ['acetonitrile','methanol','acetone','ethanol','tetrahydrofuran','propanol','isopropanol','dichloromethane','toluene','cyclohexane','heptane','hexane']
#solvents = ['acetonitrile','methanol','acetone','ethanol','50% ethanol','tetrahydrofuran','propanol','isopropanol','dichloromethane','toluene','cyclohexane','heptane','hexane']

## c distribution parameters
T_cut = 20
MW_cut = 20

use_new = True
T_cut_new = 0.5

#### read data files
# CHRIS parameter distributions
if not use_new:
    param_dists = {}
    with open('quantity_module/data/param_distribution_37.json','r') as fp:
        param_dists[37] = json.load(fp)
    with open('quantity_module/data/param_distribution_50.json','r') as fp:
        param_dists[50] = json.load(fp)
else:
    with open('quantity_module/data/param_distribution_allT.json','r') as fp:
        param_dists = json.load(fp)
# other stuff
df_visc = pd.read_parquet('quantity_module/data/solvent-viscosity.parquet')
df_desc = pd.read_parquet('quantity_module/data/data-descriptors-mordred-numconfs51.parquet', columns=['Solute_InChIKey', 'Vabc','VMcGowan'])
if not use_new:
    ## clean data
    df_final_37 = pd.read_parquet('quantity_module/data/db-D-interp-37-clean.parquet')
    df_final_50 = pd.read_parquet('quantity_module/data/db-D-interp-50-clean.parquet')
    # convert all T to K
    df_final_37['T'] = df_final_37['T'] + 273.15
    df_final_37['Polymer_Tg'] = df_final_37['Polymer_Tg'] + 273.15
    df_final_37['Polymer_Tm'] = df_final_37['Polymer_Tm'] + 273.15
    df_final_50['T'] = df_final_50['T'] + 273.15
    df_final_50['Polymer_Tg'] = df_final_50['Polymer_Tg'] + 273.15
    df_final_50['Polymer_Tm'] = df_final_50['Polymer_Tm'] + 273.15
    # add volumes
    df_final_37 = pd.merge(df_final_37, df_desc[['Solute_InChIKey', 'Vabc', 'VMcGowan']], how='left', on='Solute_InChIKey', suffixes=('', '_dupe'))
    df_final_50 = pd.merge(df_final_50, df_desc[['Solute_InChIKey', 'Vabc', 'VMcGowan']], how='left', on='Solute_InChIKey', suffixes=('', '_dupe'))
else:
    ## clean data
    df_final = pd.read_parquet('quantity_module/data/db-D-interp-allT-clean.parquet')
    # convert all T to K
    df_final['T'] = df_final['T'] + 273.15
    df_final['Polymer_Tg'] = df_final['Polymer_Tg'] + 273.15
    df_final['Polymer_Tm'] = df_final['Polymer_Tm'] + 273.15
    # add volumes
    df_final = pd.merge(df_final, df_desc[['Solute_InChIKey', 'Vabc', 'VMcGowan']], how='left', on='Solute_InChIKey', suffixes=('', '_dupe'))

#### solvent-specific viscosity
# add MW -- already in the spreadsheet
#mws = []
#for solv in df_visc['Solvent_Name']:
#    mw = chemicals.search_chemical(solv).MW
#    mws.append(mw)
#df_visc['MW'] = mws
## selected solvent MWs
#Solvent_MWs = {solv:df_visc.loc[df_visc['Solvent_Name']==solv,'MW'].iloc[0] for solv in solvents}
Solvent_MWs = dict(df_visc[['Solvent_Name','MW']].itertuples(index=False))
#Solvent_Densities = {solv:string2density(solv)[0] for solv in solvents} # can fail sometimes, so these are hardcoded for now
#Solvent_Densities = {'ethanol': 0.79, 'isopropanol': 0.787515, 'acetonitrile': 0.78467, 'toluene': 0.86925, 'dichloromethane': 1.326875, 'hexane': 0.6602}
Solvent_Densities = {'acetonitrile': 0.7847, 'methanol': 0.7955, 'acetone': 0.7893, 'ethanol': 0.7898, '50% ethanol': (0.7898+1)/2, 'tetrahydrofuran': 0.8878, 'propanol': 0.8057, 'isopropanol': 0.787, 'dichloromethane': 1.3269, 'toluene': 0.8673, 'cyclohexane': 0.7793, 'heptane': 0.6808, 'hexane': 0.6599}
# linear relation to estimate Vabc when it fails for a molecule
Vabc = df_desc['Vabc']
Vmcg = df_desc['VMcGowan']
m = ~pd.isna(Vabc)
popt_V = np.polyfit(Vmcg[m], Vabc[m], 1)

# ---- model: Grunberg–Nissan style with polynomial interaction that vanishes at x=0,1 ----
# fitted to data from R. Belda, J. V. Herráez, O. Diez, Rheological study and thermodynamic analysis of the binary system (water/ethanol): Influence of concentration. Physics and Chemistry of Liquids 42, 467-479 (2004).
popt_etoh = np.array([-6.35036532e+00, 1.86507282e+03, -5.30902320e+00, 1.60463200e+03, -1.03040657e+01, 3.05646061e+00, -4.93824317e+00, 4.16274239e+03, -1.18411097e+03, 1.69557649e+03])
def predict_lneta(p, T, x, n_poly=3, interaction_has_T=True):
    Aw, Bw, Ae, Be = p[:4]  # ln(eta_w)=Aw+Bw/T, ln(eta_e)=Ae+Be/T
    ln_eta_w = Aw + Bw / T
    ln_eta_e = Ae + Be / T
    xc = 2.0*x - 1.0  # map wt frac [0,1] -> [-1,1]
    Phi = np.vstack([xc**k for k in range(n_poly)])  # (n_poly, N)
    if interaction_has_T:
        a = p[4:4+n_poly]
        b = p[4+n_poly:4+2*n_poly]
        G = (a @ Phi) + (b @ Phi) / T
    else:
        a = p[4:4+n_poly]
        G = (a @ Phi)
    return x*ln_eta_e + (1-x)*ln_eta_w + x*(1-x)*G

def get_WC(T,solv,V):
    params = df_visc[df_visc['Solvent_Name']==solv].iloc[0]
    if params['Equation'] == '10^A(1/T-1/B)':
        eta = 10**(params['A']*(1/T-1/params['B']))
    elif params['Equation'] == 'A*exp(B/T)':
        eta = params['A']*np.exp(params['B']/T)
    elif params['Equation'] == 'E*exp(A+B/(T/298.15)+C/(T/298.15)^2+D/(T/298.15)^3)':
        eta = params['E']*np.exp(params['A'] + params['B']/(T/298.15) + params['C']/(T/298.15)**2 + params['D']/(T/298.15)**3)
    elif params['Equation'] == 'A*exp(-0.01*B*(T-298.15))':
        eta = params['A']*np.exp(-0.01*params['B']*(T-298.15))
    elif params['Equation'] == 'A+BT/1+CT+DT^2':
        eta = (params['A']+params['B']*T) / (1 + params['C']*T + params['D']*T**2)
    elif params['Equation'] == 'A+B/T+C/T^2+D/T^3':
        eta = params['A'] + params['B']/T + params['C']/T**2 + params['D']/T**3
    elif params['Equation'] == 'A*298.15/T':
        eta = params['A'] * 298.15/T
    elif params['Equation'] == 'A*T+B':
        eta = params['A'] * T + params['A']
    elif params['Equation'] == 'fitted_EtOH':
        # assuming 50% is by volume --> by mass for consistency with fitted model
        eta = np.exp(predict_lneta(popt_etoh, T, 0.5*0.7898/(0.5*0.7898+0.5*1.000), n_poly=3, interaction_has_T=True))
    else:
        eta = np.nan
    D_WC = 7.4e-8*(params['MW']*params['WC_assoc_param'])**0.5*T/eta/V**0.6
    return D_WC, eta, params['MW']

#### add Wilke-Chang
if not use_new:
    ## 50 C
    # estimate Vabc for those with nan values
    m = pd.isna(df_final_50['Vabc'])
    v = np.polyval(popt_V, df_final_50['VMcGowan'][m])
    df_final_50.loc[m, 'Vabc'] = v
    T = df_final_50['T']
    V = df_final_50['Vabc']
    for solv in solvents:
        D_WC, eta, MW_solvent = get_WC(T, solv, V)
        df_final_50[f'eta_{solv}'] = eta
        df_final_50[f'D_WC_{solv}'] = D_WC
        df_final_50[f'MW_solvent_{solv}'] = MW_solvent
    ## 37 C
    # estimate Vabc for those with nan values
    m = pd.isna(df_final_37['Vabc'])
    v = np.polyval(popt_V, df_final_37['VMcGowan'][m])
    df_final_37.loc[m, 'Vabc'] = v
    T = df_final_37['T']
    V = df_final_37['Vabc']
    for solv in solvents:
        D_WC, eta, MW_solvent = get_WC(T, solv, V)
        df_final_37[f'eta_{solv}'] = eta
        df_final_37[f'D_WC_{solv}'] = D_WC
        df_final_37[f'MW_solvent_{solv}'] = MW_solvent
else:
    # estimate Vabc for those with nan values
    m = pd.isna(df_final['Vabc'])
    v = np.polyval(popt_V, df_final['VMcGowan'][m])
    df_final.loc[m, 'Vabc'] = v
    T = df_final['T']
    V = df_final['Vabc']
    for solv in solvents:
        D_WC, eta, MW_solvent = get_WC(T, solv, V)
        df_final[f'eta_{solv}'] = eta
        df_final[f'D_WC_{solv}'] = D_WC
        df_final[f'MW_solvent_{solv}'] = MW_solvent

def get_V(smiles):
    mol = Chem.MolFromSmiles(smiles)
    calc = mordred.Calculator([mordred.descriptors.VdwVolumeABC, mordred.descriptors.McGowanVolume])
    Vabc,Vmcg = list(calc(mol).values())
    if not isinstance(Vabc, numbers.Number):
        Vabc = np.polyval(popt_V, Vmcg)
    return Vabc

#### Vrentas-Duda setup
df_vd_solv = pd.read_parquet('quantity_module/data/vrentas-duda-params-solutes.parquet')
df_vd_solv.drop_duplicates(keep='first', inplace=True, ignore_index=True) # drop exact duplicates
df_vd_poly = pd.read_parquet('quantity_module/data/vrentas-duda-params-polymers.parquet')
df_vd_poly.drop_duplicates(keep='first', inplace=True, ignore_index=True) # drop exact duplicates
df_props = pd.read_parquet('data/db-polymer-properties-and-categories.parquet')
df_vd_poly = pd.merge(df_vd_poly, df_props[['Polymer_Name','Polymer_Tg','Polymer_Tm', 'CHRIS Class', 'New Class']], how='left', on='Polymer_Name')
df_vd_poly['New Class'] = df_vd_poly['New Class'].fillna('none')
df_vd_poly['CHRIS Class'] = df_vd_poly['CHRIS Class'].fillna('none')

## Calculate c
dfs_vd_allT = []
for T in np.arange(100,800,20):
    fV_polyT = (df_vd_poly['K12']*(df_vd_poly['K22-Tg2']+T))
    #fV_polyT[fV_polyT<0.025] = 0.025
    for solvname in set(df_vd_solv['Solute_Name']):
        df_sol = df_vd_solv[df_vd_solv['Solute_Name']==solvname]
        for row in df_sol.iterrows():
            row = row[1]
            fV_sol = (row['K11']*(row['K21-Tg1']+T))
            c_sol = fV_sol / fV_polyT
            df_vd_allT = pd.concat([row]*len(df_vd_poly), axis=1, ignore_index=True).T
            df_vd_allT = pd.concat([df_vd_allT, df_vd_poly], axis=1)
            df_vd_allT['c'] = c_sol
            df_vd_allT['T'] = T
            dfs_vd_allT.append(df_vd_allT)
df_vd_allT = pd.concat(dfs_vd_allT, ignore_index=True)
df_vd_allT['T-Tg'] = df_vd_allT['T']-df_vd_allT['Tg2']
df_vd_allT['T-Tg1'] = df_vd_allT['T']-df_vd_allT['Tg1']

def get_c_dist(T,Tg,MW):
    m = (~pd.isna(df_vd_allT['c'])) & (np.abs(df_vd_allT['T-Tg']-max(T_cut,T-Tg))<T_cut) & (np.abs(df_vd_allT['M1']-MW)<MW_cut)
    cs = df_vd_allT.loc[m, 'c']
    cs = np.sort(cs) # ensures same order so results are deterministic
    cs = cs[~np.isnan(cs)]
    cs = cs[cs>0]
    return cs

def get_c_dist_cat(T,CHRIS_category,MW):
    m = (~pd.isna(df_vd_allT['c'])) & (df_vd_allT['T-Tg']>0) & (np.abs(df_vd_allT['T']-T)<T_cut) & (np.abs(df_vd_allT['M1']-MW)<MW_cut) & (df_vd_allT['New Class']==CHRIS_category)
    cs = df_vd_allT.loc[m, 'c']
    cs = np.sort(cs) # ensures same order so results are deterministic
    cs = cs[~np.isnan(cs)]
    cs = cs[cs>0]
    return cs

def get_D_dists(w,T,Polymer_Tg,Solvent_Name,Solvent_MW,Solute_MW,CHRIS_category,N=100000,return_DCs=False,input_Ds=None):
    """
    D_dist_noswell, D_dist_swell, (c, Ds, DWCs, DCs) = get_D_dists(Swelling_wtfrac,T,Polymer_Tg,Solvent_Name,Solvent_MW,Solute_MW,CHRIS_category,return_DCs=True,N=N)
    """
    if not use_new:
        if np.abs(T-310)<2:
            df_final_T = df_final_37
        if np.abs(T-323)<2:
            df_final_T = df_final_50
    else:
        df_final_T = df_final.loc[np.abs(df_final['T']-T)<T_cut_new]
    #Solvent_MW = solvmws[Solvent_Name]
    cs = get_c_dist(T,Polymer_Tg,Solvent_MW)
    if Solute_MW < 50:
        m50 = df_final_T['MW']<=50
    else:
        m50 = df_final_T['MW']>50
    ## within cutoffs, with at least N closest (by sorting, separating at MW = 50)
    m2 = (np.abs(df_final_T['Polymer_Tg']-Polymer_Tg)<T_cut) & (np.abs(df_final_T['MW']-Solute_MW)<MW_cut) & m50
    if m2.sum()<25:
        vT = df_final_T.loc[m50,'Polymer_Tg']-Polymer_Tg; vM = df_final_T.loc[m50,'MW']-Solute_MW; m3 = pd.concat([np.abs(vT), np.abs(vM)], axis=1).sort_values(by=['Polymer_Tg', 'MW']).index[1:26]
        m2 = list(set(m2.index[m2]).union(set(m3)))
    rng = np.random.Generator(np.random.PCG64(seed=12345))
    if return_DCs:
        Ds,DWCs,DCs = rng.choice([df_final_T.loc[m2,'D'], df_final_T.loc[m2,f'D_WC_{Solvent_Name}'], df_final_T.loc[m2,f'D_CHRIS_q50']], N, axis=1)
    else:
        Ds,DWCs = rng.choice([df_final_T.loc[m2,'D'], df_final_T.loc[m2,f'D_WC_{Solvent_Name}']], N, axis=1)
    c = rng.choice(cs, N)
    lnD_D0 = c*w/(1+(c-1)*w) * np.log(DWCs/Ds)
    ## distribution of D_CHRIS
    if not use_new:
        if input_Ds is None:
            if Solute_MW > 50:
                params = param_dists[T-273.15][f'{CHRIS_category}_hi']
            else:
                params = param_dists[T-273.15][f'{CHRIS_category}_lo']
            if params[0] == 'pir':
                A_list = params[1:]
                D_list = [Piringer(Solute_MW, Ai, T) for Ai in A_list]
            else:
                Ball = params[1]
                A_list = params[2:]
                D_list = np.array([PowerLaw(Solute_MW, Ai, Ball) for Ai in A_list])
        else:
            D_list = input_Ds
    else:
        if input_Ds is None:
            if Solute_MW > 50:
                subkey = f'{CHRIS_category}_hi'
            else:
                subkey = f'{CHRIS_category}_lo'
            allparams = [param_dists[Ti][subkey] for Ti in param_dists if T+T_cut_new >= int(Ti)+273.15 >= T-T_cut_new]
            #allparams = [param_dists[Ti][subkey] for Ti in param_dists if T+T_cut_new > int(Ti)+273.15 > T-T_cut_new]
            D_list = []
            for params in allparams:
                if params[0] == 'pir':
                    A_list = params[1:]
                    D_list += [Piringer(Solute_MW, Ai, T) for Ai in A_list]
                else:
                    Ball = params[1]
                    A_list = params[2:]
                    D_list += [PowerLaw(Solute_MW, Ai, Ball) for Ai in A_list]
        else:
            D_list = input_Ds
    D_dist_noswell = rng.choice(D_list, N)
    D_dist_swell = np.exp(np.log(D_dist_noswell)+lnD_D0)
    if return_DCs:
        return D_dist_noswell, D_dist_swell, (c, Ds, DWCs, DCs)
    else:
        return D_dist_noswell, D_dist_swell

def get_D_dists_new(w,T,Polymer_Tg,Solvent_Name,Solvent_MW,Solute_MW,Solute_Vabc,CHRIS_category,N=10000,return_DCs=False,input_Ds=None):
    """
    D_dist_noswell, D_dist_swell = get_D_dists_new(w, 323.15, Polymer_Tg, Solvent_Name, Solvent_MW, Solute_MW, None, CHRIS_category, return_DCs=False, N=int(1e5), input_Ds=diff)
    """
    df_final_T = df_final.loc[np.abs(df_final['T']-T)<T_cut_new]
    if (T <= Polymer_Tg) or (input_Ds is not None):
        cs = get_c_dist(T,Polymer_Tg,Solvent_MW)
    else:
        cs = get_c_dist_cat(T,CHRIS_category,Solvent_MW)
    if not len(cs):
        cs = get_c_dist(T,Polymer_Tg,Solvent_MW)
    rng = np.random.Generator(np.random.PCG64(seed=12345))
    c = rng.choice(cs, N)
    if Solute_Vabc is None:
        if Solute_MW < 50:
            m50 = df_final_T['MW']<=50
        else:
            m50 = df_final_T['MW']>50
        ## within cutoffs, with at least N closest (by sorting, separating at MW = 50)
        m2 = (np.abs(df_final_T['Polymer_Tg']-Polymer_Tg)<T_cut) & (np.abs(df_final_T['MW']-Solute_MW)<MW_cut) & m50
        if m2.sum()<25:
            vT = df_final_T.loc[m50,'Polymer_Tg']-Polymer_Tg; vM = df_final_T.loc[m50,'MW']-Solute_MW; m3 = pd.concat([np.abs(vT), np.abs(vM)], axis=1).sort_values(by=['Polymer_Tg', 'MW']).index[1:26]
            m2 = list(set(m2.index[m2]).union(set(m3)))
        if return_DCs:
            Ds,DWCs,DCs = rng.choice([df_final_T.loc[m2,'D'], df_final_T.loc[m2,f'D_WC_{Solvent_Name}'], df_final_T.loc[m2,f'D_CHRIS_q50']], N, axis=1)
        else:
            Ds,DWCs = rng.choice([df_final_T.loc[m2,'D'], df_final_T.loc[m2,f'D_WC_{Solvent_Name}']], N, axis=1)
    else:
        DWCs, eta, MW_solvent = get_WC(T, Solvent_Name, Solute_Vabc)
    ## distribution of D_CHRIS
    if input_Ds is None:
        if Solute_MW > 50:
            subkey = f'{CHRIS_category}_hi'
        else:
            subkey = f'{CHRIS_category}_lo'
        allparams = [param_dists[Ti][subkey] for Ti in param_dists if T+T_cut_new >= int(Ti)+273.15 >= T-T_cut_new]
        D_list = []
        for params in allparams:
            if params[0] == 'pir':
                A_list = params[1:]
                D_list += [Piringer(Solute_MW, Ai, T) for Ai in A_list]
            else:
                Ball = params[1]
                A_list = params[2:]
                D_list += [PowerLaw(Solute_MW, Ai, Ball) for Ai in A_list]
    else:
        D_list = input_Ds
    D_dist_noswell = rng.choice(D_list, N)
    if Solute_Vabc is None:
        lnD_D0 = c*w/(1+(c-1)*w) * np.log(DWCs/Ds)
    else:
        lnD_D0 = c*w/(1+(c-1)*w) * np.log(DWCs/D_dist_noswell)
    D_dist_swell = np.exp(np.log(D_dist_noswell)+lnD_D0)
    if return_DCs:
        return D_dist_noswell, D_dist_swell, (c, Ds, DWCs, DCs)
    else:
        return D_dist_noswell, D_dist_swell

def PlaneSheetAnalytical_old(M0,D,K,PolymerVolume,SurfaceArea,SolventVolume,ExtractionTime,nterms=5,Qv=1):
    L = PolymerVolume/SurfaceArea #effective length scale of the component
    T = D*ExtractionTime/L**2.
    result = (T>0.05) *  PlaneSheetFiniteBathMass_old(M0,D,K,PolymerVolume,SurfaceArea,SolventVolume,ExtractionTime,nterms,Qv=Qv) + \
             (T<=0.05) * PlaneSheetFiniteBathMassApprox_old(M0,D,K,PolymerVolume,SurfaceArea,SolventVolume,ExtractionTime,Qv=Qv)
    return result

def PlaneSheetFiniteBathMass_old(M0,D,K,PolymerVolume,SurfaceArea,SolventVolume,ExtractionTime,nterms,Qv=1):
    ### works with scalar- or vector-valued M0 and D
    L = PolymerVolume/SurfaceArea #effective length scale of the component
    #alpha = SolventVolume/PolymerVolume/K
    alpha = (SolventVolume-PolymerVolume*(Qv-1))/(Qv*PolymerVolume*K)
    Minfty = M0/(1.+1./(alpha))
    eps = 1e-12
    f = lambda x: np.tan(x)+alpha*x
    qn = np.zeros((nterms))
    for j in range(nterms):
        rts = bisect(f,np.pi/2.+j*np.pi+eps, np.pi*(1.+j)-eps)
        qn[j] = rts
    result = 1.
    for j in range(nterms):
        result = result - (2.*alpha*(1.+alpha))*np.exp(-D*qn[j]**2.*ExtractionTime/L**2.)/(1.+alpha+alpha**2.*qn[j]**2.)
    result = Minfty*result
    return result

def PlaneSheetFiniteBathMassApprox_old(M0,D,K,PolymerVolume,SurfaceArea,SolventVolume,ExtractionTime,Qv=1):
    ### works with scalar- or vector-valued M0 and D
    L = PolymerVolume/SurfaceArea #effective length scale of the component
    alpha = (SolventVolume-PolymerVolume*(Qv-1))/(Qv*PolymerVolume*K)
    Minfty = M0/(1.+1./(alpha))
    T = D*ExtractionTime/L**2.
    # if exp will blow up, use asymptotic expansion instead
    if not np.ndim(T):
        if(T/alpha**2.<100.):
            result = (1.+alpha)*(1.-np.exp(T/alpha**2.)*sp.special.erfc(np.sqrt(T)/alpha))
        else:
            result = (1.+alpha)*(1.-alpha/(np.sqrt(np.pi)*np.sqrt(T))+alpha**3./(2.*np.sqrt(np.pi)*(T)**1.5)-3.*alpha**5./(4.*np.sqrt(np.pi)*(T)**2.5))
    else:
        result = np.zeros(len(T))
        m = T/alpha**2.<100.
        result[m] = (1.+alpha)*(1.-np.exp(T[m]/alpha**2.)*sp.special.erfc(np.sqrt(T[m])/alpha))
        m = T/alpha**2.>=100.
        result[m] = (1.+alpha)*(1.-alpha/(np.sqrt(np.pi)*np.sqrt(T[m]))+alpha**3./(2.*np.sqrt(np.pi)*(T[m])**1.5)-3.*alpha**5./(4.*np.sqrt(np.pi)*(T[m])**2.5))
    result = Minfty*result
    return result

def PlaneSheetFiniteBathMass(tau,alpha,nterms=5):
    ### works with scalar- or vector-valued tau
    Minfty = 1.0/(1.+1./(alpha))
    eps = 1e-12
    f = lambda x: np.tan(x)+alpha*x
    qn = np.zeros((nterms))
    for j in range(nterms):
        rts = bisect(f,np.pi/2.+j*np.pi+eps, np.pi*(1.+j)-eps)
        qn[j] = rts
    result = 1.
    for j in range(nterms):
        result = result - (2.*alpha*(1.+alpha))*np.exp(-tau*qn[j]**2.)/(1.+alpha+alpha**2.*qn[j]**2.)
    result = Minfty*result
    return result

def PlaneSheetFiniteBathMassApprox(tau,alpha):
    ### works with scalar- or vector-valued tau
    Minfty = 1/(1.+1./(alpha))
    # if exp will blow up, use asymptotic expansion instead
    if not np.ndim(tau):
        if(tau/alpha**2.<100.):
            result = (1.+alpha)*(1.-np.exp(tau/alpha**2.)*sp.special.erfc(np.sqrt(tau)/alpha))
        else:
            result = (1.+alpha)*(1.-alpha/(np.sqrt(np.pi)*np.sqrt(tau))+alpha**3./(2.*np.sqrt(np.pi)*(tau)**1.5)-3.*alpha**5./(4.*np.sqrt(np.pi)*(tau)**2.5))
    else:
        result = np.zeros(len(tau))
        m = tau/alpha**2.<100.
        result[m] = (1.+alpha)*(1.-np.exp(tau[m]/alpha**2.)*sp.special.erfc(np.sqrt(tau[m])/alpha))
        m = tau/alpha**2.>=100.
        result[m] = (1.+alpha)*(1.-alpha/(np.sqrt(np.pi)*np.sqrt(tau[m]))+alpha**3./(2.*np.sqrt(np.pi)*(tau[m])**1.5)-3.*alpha**5./(4.*np.sqrt(np.pi)*(tau[m])**2.5))
    result = Minfty*result
    return result

def PlaneSheetAnalytical(tau,alpha,nterms=5):
    result = (tau>0.05)  * PlaneSheetFiniteBathMass(tau,alpha,nterms) + \
             (tau<=0.05) * PlaneSheetFiniteBathMassApprox(tau,alpha)
    return result

def Conservative(tau):
    release = (tau <= 0.2) * 2. * np.sqrt(tau / np.pi) + \
              (tau > 0.2)  * (1. - (8. / (np.pi ** 2.)) * np.exp(-tau * np.pi ** 2. / 4.))
    return release

def multiEquilSwell(alpha, Kps, swell, iterations):
    # alpha is initial before swelling
    result = 1.-(Kps*(swell+alpha)/(Kps*(swell+alpha)+swell-1.))*(swell/(swell+alpha))**iterations
    return result

def multiELSwell(tau, alpha, Kps, swell, iterations):
    resultEq = multiEquilSwell(alpha, Kps, swell, iterations)
    resultKinetic = Conservative(tau*iterations)
    result = (resultEq < resultKinetic) * resultEq + (resultEq >= resultKinetic) * resultKinetic
    return result

def Extraction(tau, alpha, Kps, swell, iterations):
    tauSpec = tau*iterations
    alphaSpec = (alpha*iterations/swell)-(swell-1)/swell/Kps # for a single iteration
    result = (tauSpec > 1.) * multiELSwell(tau, alpha, Kps, swell, iterations) + (tauSpec <= 1.) * PlaneSheetAnalytical(tauSpec,alphaSpec)
    return result

def get_M_dist(D_dist, M_expt, PolymerVolume, SurfaceArea, SolventVolume, ExtractionTime, K_expt=10, Qv=1, iterations=1):
    L = PolymerVolume/SurfaceArea 
    tau = D_dist*ExtractionTime/L**2. 
    if 1:
        alpha = (SolventVolume)/(PolymerVolume*K_expt) # without swelling; swelling is dealt with downstream
        M_M0 = Extraction(tau, alpha, Kps=K_expt, swell=Qv, iterations=iterations)
    if 0:
        alpha = (SolventVolume-PolymerVolume*(Qv-1))/(Qv*PolymerVolume*K_expt) # with swelling
        M_M0 = PlaneSheetAnalytical(tau, alpha, nterms=5) 
        #M_M0 = PlaneSheetAnalytical_old(1.0, D_dist, K_expt, PolymerVolume, SurfaceArea, SolventVolume, ExtractionTime, nterms=5, Qv=Qv) # much faster and indistinguishable distribution from Mfunc_film
    M0_pred = M_expt/M_M0
    return M0_pred