print('Starter opp...')
from models.kempers01 import Kempers01
from pyctp import saftvrmie, saftvrqmie
import pandas as pd
import numpy as np
from numpy import sqrt
import matplotlib.pyplot as plt
from scipy.constants import Boltzmann, Avogadro
from scipy.optimize import fsolve
import file_handling as fh
import os
from datetime import datetime
import warnings
print('Nu kjør vi!!')

df = pd.read_csv('data/mie_simulation.csv')
mie_df = pd.read_excel('models/mie.xlsx')

sig_Ar = np.array(mie_df.loc[mie_df['comp'] == 'AR']['sigma'].iloc[0]) * 1e-10
eps_Ar = np.array(mie_df.loc[mie_df['comp'] == 'AR']['epsilon'].iloc[0]) * Boltzmann

eps = eps_Ar / Boltzmann

M_Ar = 39.948

reduced_T = 0.85
reduced_rho = 0.81

# eps1 and sig1 are the ratio of sigma and epsilon to sigma_Ar and epsilon_Ar
T = lambda eps1 : reduced_T * eps_Ar * np.sqrt(eps1) / Boltzmann
rho = lambda sig1 : reduced_rho / (0.5 * (sig1 * sig_Ar + sig_Ar))**3

df_m = df.loc[df['m1'] != 1.0]
df_s = df.loc[df['sigma1'] != 1.00]
df_e = df.loc[df['epsilon1'] != 1.00]

ROOT_PATH = os.path.dirname(os.path.abspath(__file__))
DATA_PATH = ROOT_PATH + '/data/mie_data'

class MieRunner:
    def __init__(self, mode='cov', eos=saftvrmie.saftvrmie, eos_name='SAFT-VR-MIE', version=None):
        verify = input('Initialising MieTester with mode : '+mode+', eos : ' + eos_name+'. "Y" to continiue.')
        if verify != 'Y':
            exit(0)

        self.mode = mode
        self.model = Kempers01
        self.model_name = 'Kempers01'
        self.eos = eos
        self.eos_name = eos_name

        if version is None:
            # Make output dir for this benchmark run
            i = 0
            while os.path.isdir(ROOT_PATH + '/output/mie/V' + str(i)):
                i += 1

            self.outdir = ROOT_PATH + '/output/mie/V' + str(i)
            self.outpath = self.outdir + '/' + mode
            os.mkdir(self.outdir)
            os.mkdir(self.outpath)


            # Write a little meta file to the output dir
            with open(self.outdir+'/meta.txt', 'w') as file:
                tags = ['Date', 'Model', 'Mode', 'EoS']
                vals = [datetime.now().strftime("%d/%m/%Y %H:%M:%S"), self.model_name, mode, self.eos_name]
                for tag, val in zip(tags, vals):
                    file.write(tag + ' : ' + str(val) + '\n')

                file.write('\nContains runs for :\n')
        else:
            self.outdir = ROOT_PATH + '/output/mie/V' + str(version)
            self.outpath = self.outdir + '/'+mode
            if os.path.isdir(self.outpath):
                print('This BenchmarkRunner instance may overwrite files in', self.outpath)
                verify = input("'Y' to verify")
                if verify != 'Y':
                    exit(0)

                with open(self.outdir + '/meta.txt', 'a') as file:
                    file.write('\n#########################\n')
                    tags = ['Date', 'Model', 'Mode', 'EoS']
                    vals = [datetime.now().strftime("%d/%m/%Y %H:%M:%S"), self.model_name, mode, self.eos_name]
                    for tag, val in zip(tags, vals):
                        file.write(tag + ' : ' + str(val) + '\n')

                    file.write('\nContains runs for :\n')
            else:
                os.mkdir(self.outpath)


    def update_meta(self, filename):
        with open(self.outdir + '/meta.txt', 'a') as file:
            file.write(filename + '\n')

    def run_sigma(self):
        s_list = df_s['sigma1'].tolist()

        ST_list_p_rep = np.empty_like(s_list, dtype=float)
        ST_list_p_calc = np.empty_like(s_list, dtype=float)
        p_calc_list = np.empty_like(s_list, dtype=float)
        vm_calc_list = np.empty_like(s_list, dtype=float)
        phase_list = np.empty_like(s_list, dtype=int)

        eos = self.eos()
        eos.init('AR,NC6')
        for i in range(len(s_list)):
            s = s_list[i]

            eos.set_pure_fluid_param(1, 1, s * sig_Ar, eps, 6, 12)
            eos.set_pure_fluid_param(2, 1, sig_Ar, eps, 6, 12)
            eos.redefine_critical_parameters(silent=False)

            p_calc_list[i] = eos.pressure_tv(T(1), Avogadro / rho(s), [0.5, 0.5])[0]
            vm_calc_list[i] = eos.specific_volume(T(1), df_s['p'].tolist()[i] * 1e7, [0.5, 0.5], 1)[0]
            phase_list[i] = eos.guess_phase(T(1), df_s['p'].tolist()[i] * 1e7, [0.5, 0.5])
            rep_model =  Kempers01('AR,NC6', eos, x=[0.5, 0.5], temp=T(1), pres=df_s['p'].tolist()[i] * 1e7,
                                        mole_weights=[M_Ar, M_Ar], particle_density=rho(s), sigma=[s * sig_Ar, sig_Ar], epsilon=[eps_Ar, eps_Ar])

            calc_model = Kempers01('AR,NC6', eos, x=[0.5, 0.5], temp=T(1), pres=p_calc_list[i],
                                        mole_weights=[M_Ar, M_Ar], particle_density=rho(s), sigma=[s * sig_Ar, sig_Ar],
                                        epsilon=[eps_Ar, eps_Ar])

            ST_list_p_calc[i] = calc_model.get_soret(self.mode)[0]
            ST_list_p_rep[i] = rep_model.get_soret(self.mode)[0]


        out_df = pd.DataFrame({'s1/s2' : s_list, 'ST_p_calc' : ST_list_p_calc, 'ST_p_reported' : ST_list_p_rep,
                               'p_calc' : p_calc_list, 'vm_calc' : vm_calc_list, 'p_reported' : df_s['p'].tolist(),
                               'vm_reported' : [Avogadro/rho(s) for s in s_list],
                               'phase' : phase_list})

        filename = 'sigma' + '_' + self.eos_name
        savepath = self.outpath + '/' + filename + '.csv'
        out_df.to_csv(savepath)
        self.update_meta(filename)

        print('Saved sigma run to : ', savepath)

    def run_epsilon(self):
        e_list = df_e['epsilon1'].tolist()

        ST_list_p_rep = np.empty_like(e_list, dtype=float)
        ST_list_p_calc = np.empty_like(e_list, dtype=float)
        p_calc_list = np.empty_like(e_list, dtype=float)
        vm_calc_list = np.empty_like(e_list, dtype=float)
        T_calc_list = np.empty_like(e_list, dtype=float)

        eos = self.eos()
        eos.init('AR,NC6')
        for i in range(len(e_list)):
            e = e_list[i]

            eos.set_pure_fluid_param(1, 1, sig_Ar, e * eps, 6, 12)
            eos.set_pure_fluid_param(2, 1, sig_Ar, eps, 6, 12)
            eos.redefine_critical_parameters(silent=False)

            p_calc_list[i] = eos.pressure_tv(T(e), Avogadro / rho(1), [0.5, 0.5])[0]
            vm_calc_list[i] = eos.specific_volume(T(e), df_s['p'].tolist()[i] * 1e7, [0.5, 0.5], 1)[0]
            T_calc_list[i] = fsolve(lambda T: eos.pressure_tv(T, 1500/rho(1), [750/Avogadro, 750/Avogadro])[0] - df_e['p'].tolist()[i] * 1e7, T(e))[0]

            model_p_rep = Kempers01('AR,NC6', eos, x=[0.5, 0.5], temp=T(e), pres=df_e['p'].tolist()[i] * 1e7,
                                        mole_weights=[M_Ar, M_Ar], particle_density=rho(1), sigma=[sig_Ar, sig_Ar], epsilon=[e * eps, eps])

            model_p_calc = Kempers01('AR,NC6', eos, x=[0.5, 0.5], temp=T(e), pres=p_calc_list[i],
                                         mole_weights=[M_Ar, M_Ar], particle_density=rho(1),
                                         sigma=[sig_Ar, sig_Ar],
                                         epsilon=[e * eps, eps])

            ST_list_p_calc[i] = model_p_calc.get_soret(self.mode)[0]
            ST_list_p_rep[i] = model_p_rep.get_soret(self.mode)[0]

        out_df = pd.DataFrame({'e1/e2': e_list, 'ST_p_calc': ST_list_p_calc, 'ST_p_reported': ST_list_p_rep,
                               'p_calc': p_calc_list, 'vm_calc': vm_calc_list, 'p_reported': df_e['p'].tolist(),
                               'vm_reported': [Avogadro / rho(1) for e in e_list], 'T_calc' : T_calc_list,
                                'T_reported' : [T(e) for e in e_list]})


        filename = 'epsilon_'+self.eos_name
        savepath = self.outpath + '/' + filename + '.csv'
        out_df.to_csv(savepath)
        self.update_meta(filename)

        print('Saved epsilon run to : ', savepath)

    def run_mass(self):
        eos = self.eos()
        eos.init('AR,NC6')

        eos.set_pure_fluid_param(1, 1, sig_Ar, eps, 6, 12)
        eos.set_pure_fluid_param(2, 1, sig_Ar, eps, 6, 12)
        eos.redefine_critical_parameters(silent=False)

        pres = eos.pressure_tv(T(1), Avogadro / rho(1), [0.5, 0.5])[0]
        vm = np.array([eos.specific_volume(T(1), df_s['p'].tolist()[i] * 1e7, [0.5, 0.5], 1)[0] for i in range(len(df_m))])
        #print('Pressure : ', pres / 1e5, df_m['p']*1e2)
        #print('vm : ', Avogadro / vm, rho(1))

        m_list = df_m['m1'].tolist()

        p_calc_list = [pres for m in m_list]
        vm_calc_list = [vm for p in m_list]

        ST_list_p_rep = np.empty_like(m_list, dtype=float)
        kin_list = np.empty_like(m_list, dtype=float)
        ST_list_p_calc = np.empty_like(m_list, dtype=float)

        for i, m in enumerate(m_list):

            model_p_rep = Kempers01('AR,NC6', eos=eos, pres=df_m['p'][i] * 1e7,
                                  temp=T(1),
                                  particle_density=rho(1),
                                  mole_weights=[m * M_Ar, M_Ar],
                                  sigma=[sig_Ar, sig_Ar],
                                  epsilon=[eps_Ar, eps_Ar])

            model_p_calc = Kempers01('AR,NC6', eos=eos, pres=pres,
                                  temp=T(1),
                                  particle_density=rho(1),
                                  mole_weights=[m * M_Ar, M_Ar],
                                  sigma=[sig_Ar, sig_Ar],
                                  epsilon=[eps_Ar, eps_Ar])


            ST, kin = model_p_rep.get_soret(self.mode, kin=True)
            ST_list_p_rep[i] = ST[0]
            kin_list[i] = kin[0]

            ST_list_p_calc[i] = model_p_calc.get_soret(self.mode)[0]

        filename = 'mass_'+self.eos_name
        out_df = pd.DataFrame({'m1/m2': m_list, 'ST_p_calc': ST_list_p_calc, 'ST_p_reported': ST_list_p_rep,
                               'p_calc': p_calc_list, 'vm_calc': vm_calc_list, 'p_reported': df_m['p'].tolist(),
                               'vm_reported': [Avogadro / rho(1) for m in m_list], 'Kinetic' : kin_list})

        savepath = self.outpath + '/' + filename + '.csv'
        out_df.to_csv(savepath)
        self.update_meta(filename)

        print('Saved mass_mie_test to : ', savepath)

    def test_div(self):

        systems = ['XE,KR', 'KR,AR', 'AR,C1', 'XE,AR', 'KR,C1', 'XE,C1']

        reported = {'XE,KR' : 5.1, 'KR,AR' : 14.4, 'AR,C1' : 9.3, 'XE,AR' : 18.6, 'KR,C1' : 22.3, 'XE,C1' : 23.1}
        params = {'XE' : (3.975, 1.72 * eps_Ar, 131.29), 'AR' : (3.405, eps_Ar, 39.95), 'KR' : (3.633, 1.39 * eps_Ar, 83.80), 'C1' : (3.74, 1.27 * eps_Ar, 16.04)}

        outstr = ''
        outstr += 'Mixture' + ' '*7 + 'Reported' + ' '*7 + 'Predicted' + ' '*7 + 'Kinetic' + ' '*7 + 'Pressure [bar]' + ' '*7 + 'T [K]' + ' '*7 + 'rho [kmol/m3]\n'

        for comps in systems:
            c1, c2 = comps.split(',')

            eos = saftvrmie.saftvrmie()
            eos.init(comps)

            #eos.set_pure_fluid_param(1, 1, params[c1][0]*1e-10, params[c1][1]/Boltzmann, 6, 12)
            #eos.set_pure_fluid_param(2, 1, params[c2][0]*1e-10, params[c2][1]/Boltzmann, 6, 12)
            #eos.redefine_critical_parameters(silent=False)

            e12 = np.sqrt(params[c1][1] * params[c2][1])
            s12 = 0.5 * (params[c1][0] + params[c2][0])

            T = reduced_T * e12 / Boltzmann
            rho = reduced_rho / ((s12*1e-10) ** 3)

            pres = eos.pressure_tv(T, Avogadro/rho, [0.5, 0.5])[0]

            model = Kempers01(comps, eos, x=[0.5, 0.5], pres=pres, temp=T)#,
                              #particle_density=rho)#,
                              #sigma=[params[c1][0], params[c2][0]],
                              #epsilon=[params[c1][1], params[c2][1]],
                              #mole_weights=[params[c1][2], params[c2][2]])

            soret, kin = model.get_soret(self.mode, kin=True)

            outstr += comps + ' '*10 + str(reported[comps]) + ' '*(17 - len(str(reported[comps]))) +\
                    str(round(soret[0] * 1e3, 1)) + ' '*(14 - len(str(round(soret[0] * 1e3, 1)))) +\
                    str(round(kin[0] * 1e3, 1)) + ' '*(17 - len(str(round(kin[0] * 1e3, 1)))) +\
                    str(round(pres/1e5, 1)) + ' '*(17 - len(str(round(pres/1e5, 1))))\
                    + str(round(T, 1)) + ' '*(15 - len(str(round(T, 1))))\
                    + str(round(rho/(1e3 * Avogadro), 1)) + '\n'

        with open(self.outpath + '/div.txt', 'w') as file:
            file.write(outstr)

        print('Saved :')
        print('#'*40)
        print(outstr)
        print('#'*40, '\n')
        print('To', self.outpath+'/div.txt')
        print('\n', '#'*40, sep='')

    def deviation(self):
        ST_list_p_rep = np.empty_like(df['p'], dtype=float)
        ST_list_p_calc = np.empty_like(df['p'], dtype=float)
        p_calc_list = np.empty_like(df['p'], dtype=float)

        eos = self.eos()
        eos.init('AR,NC6')

        for i in range(len(df)):
            e = df['epsilon1'][i]
            s = df['sigma1'][i]
            m = df['m1'][i]
            p_rep = df['p'][i]

            eos.set_pure_fluid_param(1, 1, s * sig_Ar, eps, 6, 12)
            eos.set_pure_fluid_param(2, 1, sig_Ar, e * eps, 6, 12)
            eos.set_sigma_lij(1, 2, sig_Ar)
            eos.set_eps_kij(1, 2, 0)
            eos.set_lr_gammaij(1, 2, 0)
            eos.redefine_critical_parameters(silent=False)

            p_calc = eos.pressure_tv(T(e), Avogadro / rho(s), [0.5, 0.5])[0]

            model_p_rep  = Kempers01('AR,NC6', eos, x=[0.5, 0.5], temp=T(e), pres=p_rep * 1e7,
                                            particle_density=rho(s),
                                            sigma=[s * sig_Ar * 1e10, sig_Ar * 1e10],
                                            epsilon=[eps, e * eps],
                                            mole_weights=[m * M_Ar, M_Ar])

            model_p_calc = Kempers01('AR,NC6', eos, x=[0.5, 0.5], temp=T(e), pres=p_calc,
                                             particle_density=rho(s),
                                             sigma=[s * sig_Ar * 1e10, sig_Ar * 1e10],
                                             epsilon=[eps, e * eps],
                                             mole_weights=[m * M_Ar, M_Ar])

            p_calc_list[i] = p_calc

            ST_list_p_rep[i] = model_p_rep.get_soret(self.mode)[0]
            ST_list_p_calc[i] = model_p_calc.get_soret(self.mode)[0]

        out_df = pd.DataFrame({'ST_reported' : df['ST'].tolist(), 'ST_p_calc': ST_list_p_calc, 'ST_p_reported': ST_list_p_rep,
                               'p_calc': p_calc_list, 'p_reported': df['p'].tolist()})

        out_df.to_csv(self.outpath + '/deviation_'+self.eos_name+'.csv')
        self.update_meta('deviation_'+self.eos_name)

        print('Saved : ', self.outpath+ '/deviation_'+self.eos_name+'.csv')

if __name__ == '__main__':
    runner = MieRunner(eos=saftvrmie.saftvrmie, eos_name='SAFT-VR-MIE', mode='cov', version=7)

    runner.test_div()

    #runner.run_sigma()
    #runner.run_epsilon()
    #runner.run_mass()
    #runner.deviation()

