from math import sqrt,exp
from CoolProp.CoolProp import Props
import numpy as np
import matplotlib.pyplot as plt
from scipy.odr import *
from math import log
E_K = {'REFPROP-Ammonia':386,
       'REFPROP-Argon':143.2
       }
SIGMA = {'REFPROP-Ammonia':0.2957,
         'REFPROP-Argon':0.335
         }

E_K['REFPROP-Propane']=263.88
SIGMA['REFPROP-Propane']=0.49748
E_K['REFPROP-R32']=289.65
SIGMA['REFPROP-R32']=0.4098
E_K['REFPROP-R245fa'] = 329.72
SIGMA['REFPROP-R245fa'] = 0.5529


def viscosity_dilute(fluid,T,e_k,sigma):
    """
    T in [K], e_k in [K], sigma in [nm]
    viscosity returned is in [Pa-s]
    """
    
    Tstar = T/e_k
    molemass = Props(fluid,'molemass')
    if fluid == 'Propane' or fluid == 'REFPROP-Propane':
        a = [0.25104574,-0.47271238,0,0.060836515,0]
        theta_star = exp(a[0]*pow(log(Tstar),0)+a[1]*pow(log(Tstar),1)+a[3]*pow(log(Tstar),3));
        eta_star = 0.021357*sqrt(molemass*T)/(pow(sigma,2)*theta_star)/1e6;
        return eta_star
    
    # From Neufeld, 1972, Journal of Chemical Physics - checked coefficients
    OMEGA_2_2 = 1.16145*pow(Tstar,-0.14874)+ 0.52487*exp(-0.77320*Tstar)+2.16178*exp(-2.43787*Tstar)
    # Using the leading constant from McLinden, 2000 since the leading term from Huber 2003 gives crazy values
    eta_star = 26.692e-3*sqrt(molemass*T)/(pow(sigma,2)*OMEGA_2_2)/1e6
    return eta_star

def viscosity_linear(fluid, T, rho, e_k, sigma):
    """
    Implements the method of Vogel 1998 (Propane) for the linear part
    """
    N_A=6.02214129e23
    molemass = Props(fluid,'molemass')
    Tstar = T/e_k
    b= [-19.572881,219.73999,-1015.3226,2471.01251,-3375.1717,2491.6597,-787.26086,14.085455,-0.34664158]
    s = sum([b[i]*pow(Tstar,-0.25*i) for i in range(7)])
    
    B_eta_star = s+b[7]*pow(Tstar,-2.5)+b[8]*pow(Tstar,-5.5) #//[no units]
    B_eta = N_A*pow(sigma/1e9,3)*B_eta_star #[m3/mol]
    return viscosity_dilute(fluid,T,e_k,sigma)*B_eta*rho/molemass*1000

from PDSim.misc.datatypes import Collector
RHO = Collector()
TT = Collector()
DELTA = Collector()
TAU = Collector()
VV = Collector()
VV0 = Collector()
VV1 = Collector()
VVH = Collector()

fluid = 'REFPROP-R32'
Tc = Props(fluid,'Tcrit')
rhoc = Props(fluid,'rhocrit')
for T in np.linspace(290,Props(fluid,'Tcrit')-0.1,100):
    rhoV = Props('D','T',T,'Q',1,fluid)
    rhoL = Props('D','T',T,'Q',0,fluid)
    rhomax = Props('D','T',Props(fluid,'Tmin'),'Q',0,fluid)
    for rho in list(np.linspace(rhoL,rhomax,100)):#+list(np.linspace(rhoV,0.0001,100)):
    #for rho in list(np.linspace(rhoV,0.0001,100)):
        mu_0 = viscosity_dilute(fluid,T,E_K[fluid],SIGMA[fluid])
        mu_1 = viscosity_linear(fluid,T,rho,E_K[fluid],SIGMA[fluid])
        mu = Props('V','T',T,'D',rho,fluid)
        VV << mu
        VV0 << mu_0
        VV1 << mu_1
        VVH << mu-mu_0-mu_1
        TT << T
        RHO << rho
        DELTA << rho/rhoc
        TAU << Tc/T
        
def f_RHS(E, DELTA_TAU, VV):
    k = 0
    sum = 0
    DELTA = DELTA_TAU[0,:]
    TAU = DELTA_TAU[1,:]
    for i in range(2,5):
        for j in range(3):
            sum += E[k]*DELTA**i/TAU**j
            k += 1
#    f1,f2,f3,g1,g2 = E[k],E[k+1],E[k+2],E[k+3],E[k+4]
#    DELTA0 = g1*(1+g2*np.sqrt(TAU))
#    sum += (f1+f2/TAU+f3/TAU/TAU)*(DELTA/(DELTA0-DELTA)-DELTA/DELTA0)
    print np.mean(np.abs(((sum/VV-1)*100))),'%'
    return sum

log_muH = np.log(VVH.v().T)

x = np.c_[DELTA.v().T, TAU.v().T].T
y = VVH.v()

linear = Model(f_RHS, extra_args = (y,) )
mydata = Data(x, y)
myodr = ODR(mydata, linear, beta0=np.array([0.1]*17),)
myoutput = myodr.run()
E = myoutput.beta
print E

#plt.plot(TT.vec, y,'b.',TT.vec, f_RHS(E, x, y),'r.')
#plt.show()
#plt.plot()
plt.plot(y.T,f_RHS(E, x, y))
plt.show()