from CoolProp import CoolProp as CP
from PDSim.misc.datatypes import Collector
import numpy as np
import matplotlib.pyplot as plt
from scipy.odr import *
import textwrap

#
#fluid = 'Propane'
#Rfluid = 'REFPROP-propane'
#e_k = 263.88
#sigma = 0.49748
#
fluid = 'DimethylEther'
Rfluid = 'REFPROP-DME'
e_k = 329.72
sigma = 0.5529
molemass = CP.Props(fluid, 'molemass')

Ttriple = CP.Props(fluid, 'Ttriple')
Tcrit = CP.Props(fluid, 'Tcrit')
rhocrit = CP.Props(fluid, 'rhocrit')
n = 6
m = 3
NP = 1
Nb = 0
N = (n - 1) * (m + 1) + 3 + Nb

mu, mu_dilute, RHO, TTT = Collector(), Collector(), Collector(), Collector()

rhomax = CP.Props('D', 'T', Ttriple, 'Q', 0, fluid)
# Build a database of "experimental" data
for T in np.linspace(Ttriple, Tcrit + 30, 400):
    for rho in np.linspace(1e-10, rhomax, 400):
        muval = CP.Props('V', 'T', T, 'D', rho, Rfluid)
        mudilute = CP.viscosity_dilute(fluid, T, rho, e_k, sigma)

        # Want positive value, and single-phase
        if (muval > 0 and T > Tcrit or rho > CP.rhosatL_anc(fluid, T) or rho < CP.rhosatV_anc(fluid, T)):
            mu << muval
            mu_dilute << mudilute
            TTT << T
            RHO << rho

from CoolProp.Plots.Plots import Trho
Trho(fluid)
plt.plot(RHO.vec, TTT.vec, '.')
plt.show()

#tau = np.array(TTT.vec)/Tcrit
tau = np.array(TTT.vec) / Tcrit
delta = np.array(RHO.vec) / rhocrit
Tstar = np.array(TTT.vec) / e_k

# Define the objective function


def OBJECTIVE_fit(c, x):
    tau = x[0, :]
    delta = x[1, :]
    # Unpack the inputs into e matrix and f vector
    e = np.zeros((n + 1, m + 1))

    sum = 0
    k = 0
    for i in range(2, n + 1):
        for j in range(0, m + 1):
            e[i][j] = c[k]
            sum += e[i][j] * delta**i / tau**j
            k += 1

    for o in range(0, NP):
        f1 = c[k + o * 3]
        g1 = c[k + 1 + o * 3]
        g2 = c[k + 2 + o * 3]

    delta_0 = g1 * (1 + g2 * tau**0.5)
    sum += f1 * (delta / (delta_0 - delta) - delta / delta_0)
    return sum + np.array(mu_dilute.vec)


print('starting fit')
XXX = np.r_[np.array(tau, ndmin=2), np.array(delta, ndmin=2)]
mod = Model(OBJECTIVE_fit)
mydata = Data(XXX.copy(), np.array(mu.vec))
beta0 = [1 for _ in range(N)]
myodr = ODR(mydata, mod, beta0=beta0)
myoutput = myodr.run()
myoutput.pprint()
print(myoutput.sum_square)
YFIT = OBJECTIVE_fit(myoutput.beta, XXX)
plt.plot(np.array(mu.vec), YFIT / np.array(mu.vec), 'o')
plt.show()

rel_error = (YFIT) / np.array(mu.vec) - 1
MAE = np.mean(np.abs(rel_error)) * 100
SSE = np.sum(np.power(YFIT - np.array(mu.vec), 2))
print(SSE)


def write_output(c):
    e = np.zeros((n + 1, m + 1))
    k = 0
    edata = ''
    for i in range(2, n + 1):
        erow = ''
        for j in range(0, m + 1):
            e[i][j] = c[k]
            erow += 'e[{i:d}][{j:d}] = {val:0.16g}; '.format(val=e[i][j], i=i, j=j)
            k += 1
        edata += erow + '\n'
    f1 = c[k]
    g1 = c[k + 1]
    g2 = c[k + 2]

    template = textwrap.dedent(
    """
    double {name:s}Class::viscosity_Trho(double T, double rho)
    {{
        // This function was generated by fitting REFPROP ECS data
        // to the functional form of Vogel, 1998 (propane viscosity)
        // The script entitled dev/fit_avoid_ECS.py was used to make this
        // function.  The mean absolute error of the fit is equal to
        // {MAE:g} %

        double delta_0, sum, DELTA_H_eta, e_k, sigma, tau, delta;
        double e[{n:d}+1][{m:d}+1];

        tau = T/reduce.T; //[Opposite to normal definition]
        delta = rho/reduce.rho;

        // Load the coefficients
        double f1 = {f1:0.16g}, g1 = {g1:0.16g}, g2 = {g2:0.16g};
        for (int i=0;i<={n:d};i++){{ for(int j=0;j<={m:d};j++){{ e[i][j]=0.0; }} }}
        {edata:s}
        delta_0=g1*(1+g2*sqrt(tau)); //[no units]
        sum=0;
        for (int i=2;i<={n:d};i++){{
            for (int j=0;j<={m:d};j++){{
                sum += e[i][j]*pow(delta,i)/pow(tau,j);
            }}
        }}
        DELTA_H_eta = sum + f1*(delta/(delta_0-delta)-delta/delta_0); //[Pa-s]

        try{{
            // Get the ECS params for the fluid if it has them
            ECSParams(&e_k,&sigma);
        }}
        catch(NotImplementedError)
        {{
            throw ValueError(format("Your fluid does not implement ECSParams"));
        }}

        return viscosity_dilute(T,e_k,sigma) + DELTA_H_eta;
    }}
    """
    )

    values = dict(f1=f1,
                  g1=g1,
                  g2=g2,
                  n=n,
                  m=m,
                  name=fluid,
                  edata=edata,
                  MAE=MAE)

    print(template.format(**values))


write_output(myoutput.beta)