#!/usr/bin/env python
import numpy as np
from scipy.interpolate import interp1d
import collections
import matplotlib.pyplot as plt
##set values to dictionary
def set_dic(a,param_d):
  param_d['V0'] = a[0]
  param_d['K0'] = a[1]
  param_d['K0p'] = a[2]
  param_d['theta0'] = np.exp(a[3])
  param_d['gamma0'] = np.exp(a[4])
  param_d['q'] = a[5]
  return param_d

# def get_dic(a):
#   param_d['V0'] = a[0]
#   param_d['K0'] = a[1]
#   param_d['K0p'] = a[2]
#   param_d['theta0'] = np.exp(a[3])
#   param_d['gamma0'] = np.exp(a[4])
#   param_d['q'] = a[5]

#   return param_d


def set_const(param_d):
  param_d['Natom'] = np.nan
  param_d['kB'] = 13.806488*10**(-24) #J per K
  param_d['P_conv_Fac'] = 160.217657*(1/1.6021764e-19) #GPa in 1 J/Ang^3
  param_d['C_DP'] = 3*param_d['kB']*param_d['Natom']*param_d['P_conv_Fac'] #3nR

def set_material_const(Natom, param_d): #Natom, kB, P_conv_fac, C_DP
  param_d['Natom'] = Natom
  param_d['C_DP'] = 3*param_d['kB']*param_d['Natom']*param_d['P_conv_Fac'] #3nR
  # print "C_DP", param_d['C_DP']
  



def set_constshifted(param_d):
  param_d['Natom'] = 4
  param_d['kB'] = 8.6173324e-5 #eV per K
  param_d['P_conv_Fac'] = 160.217657#GPa in 1 eV/Ang^3
  param_d['C_DP'] = 3*param_d['kB'] *param_d['Natom']#Dulong-Petit limit for Cv

def set_material_constshifted(Natom,param_d): #Natom, kB, P_conv_fac, C_DP
  param_d['Natom'] = Natom
  param_d['C_DP'] = 3*param_d['kB'] *param_d['Natom']#Dulong-Petit limit for Cv

def debye_fun(x):
    """
    Return debye integral value

    - calculation done using interpolation in a lookup table
    - interpolation done in log-space where behavior is close to linear
    - linear extrapolation is implemented manually
    """

    if(np.isscalar(x)):
        assert x >= 0, 'x values must be greater than zero.'
    else:
        #np.absolute(x)
        assert all(x >= 0), 'x values must be greater than zero.'
    # Lookup table
    # interpolate in log space where behavior is nearly linear
    debyex = np.array([0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0,
                       1.1, 1.2, 1.3, 1.4, 1.6, 1.8, 2.0, 2.2, 2.4, 2.6, 2.8,
                       3.0, 3.2, 3.4, 3.6, 3.8, 4.0, 4.2, 4.4, 4.6, 4.8, 5.0,
                       5.5, 6.0, 6.5, 7.0, 7.5, 8.0, 8.5, 9.0, 9.5, 10.0])
    debyelogf = np.array([ 0.0, -0.03770187, -0.07580279, -0.11429475,
                          -0.15316866, -0.19241674, -0.2320279 , -0.27199378,
                          -0.31230405, -0.35294619, -0.39390815, -0.43518026,
                          -0.47674953, -0.51860413, -0.56072866, -0.64573892,
                          -0.73167389, -0.81841793, -0.90586032, -0.99388207,
                          -1.08236598, -1.17119911, -1.26026101, -1.34944183,
                          -1.43863241, -1.52771969, -1.61660856, -1.70519469,
                          -1.79338479, -1.88108917, -1.96822938, -2.05471771,
                          -2.14049175, -2.35134476, -2.55643273, -2.75507892,
                          -2.94682783, -3.13143746, -3.30880053, -3.47894273,
                          -3.64199587, -3.79820337, -3.94785746])
    # Create interpolation function
    logdebfun = interp1d(debyex, debyelogf, kind='cubic', bounds_error=False,
                         fill_value=np.nan)
    logfval = logdebfun(x)
    # Check for extrapolated values indicated by NaN
    #   - replace with linear extrapolation
    logfval = np.where(x > debyex[-1], debyelogf[-1] + (x - debyex[-1]) *
                       (debyelogf[-1]-debyelogf[-2])/(debyex[-1]-debyex[-2]),
                       logfval)
    # Exponentiate to get integral value
    return np.exp(logfval)

def press_vinet(v,param_d):
  V0 = param_d['V0']
  K0 = param_d['K0']
  Kp = param_d['K0p']
  x = (v/V0)**(1./3)
  Vinet = 3*K0*(1.-x)*x**(-2)*np.exp(3./2.*(Kp - 1.)*(1. - x))
  return Vinet

def Ptot_powerlaw(volume,temperature,param_d,RefT):
  gamma= param_d['gamma0']*(volume/param_d['V0'])**param_d['q']
  theta = param_d['theta0']*np.exp((-1)*(gamma-param_d['gamma0'])/param_d['q'])
  Debye_Int = debye_fun(theta/RefT)
  # print gamma,theta,Debye_Int
  # print "C_DP in therm",param_d['C_DP']
  P_th_ref = param_d['C_DP']*gamma*RefT*Debye_Int/volume
  # print "P_th_ref",P_th_ref
  # print volume
  P_th = (param_d['C_DP']*gamma*temperature*debye_fun(theta/temperature))/volume
  # print "P_th",P_th
  # print "thermal pressure", P_th - P_th_ref
  # print "Natom",param_d['Natom']
  return P_th - P_th_ref

def Ptot_shiftedpowlaw(volume,temperature,param_d):
  x = (volume/param_d['V0'])**(1./3)
  gammaV = param_d['gamma0']*x**(3*param_d['q'])+1./2
  thetaV = param_d['theta0']*x**(-3./2)*np.exp(param_d['gamma0']/param_d['q']*((1-x**(3.*param_d['q']))))

  debye_Int = debye_fun(thetaV/temperature)
  P_th = (param_d['C_DP']*temperature*gammaV/volume*debye_Int)*param_d['P_conv_Fac']
  return P_th

def MGD_PowerLaw(volume, temperature,param_d,Natom = 4):

  # if np.isscalar(temperature):
  #   temperature = temperature*np.ones(len(volume))

  #assert len(volume)==len(temperature), 'temperature should be a scalar or its length should be equal to volume'
#set parameters into dictionary
  set_const(param_d)
  if Natom != 4:
    set_const(param_d)
    set_material_const(Natom,param_d)
  Vinet = press_vinet(volume,param_d)
  #print "Vinet", Vinet
  RefT = 300.0 #unit is K
  Ptot = Ptot_powerlaw(volume,temperature,param_d,RefT)
  print "P_tot",Ptot
  MGD = Vinet + Ptot
  return MGD

def MGD_PowerLawShifted(volume, temperature, param_d,Natom=4):
  set_constshifted(param_d)
  if Natom!=4:
    set_material_constshifted(Natom,param_d)
  #param_d['kB'] = 8.6173324e-5 #eV per K
  #param_d['P_conv_Fac'] = 160.217657#GPa in 1 eV/Ang^3
  #param_d['C_DP'] = 3*param_d['kB'] *param_d['Natom']#Dulong-Petit limit for Cv
  print param_d
  Vinet = press_vinet(volume,param_d)
  x = (volume/param_d['V0'])**(1./3)
  #gammaV = param_d['gamma0']*x**(3*param_d['q'])+1./2
  #thetaV = param_d['theta0']*x**(-3./2)*np.exp(param_d['gamma0']/param_d['q']*((1-x**(3.*param_d['q']))))

  #debye_Int = debye_fun(thetaV/temperature)
  #P_th = (param_d['C_DP']*temperature*gammaV/volume*debye_Int)*param_d['P_conv_Fac']
  P_th = Ptot_shiftedpowlaw(volume,temperature,param_d)
  #compute P(V,T)
  MGD = Vinet + P_th
  return MGD

"""
def MGD_PowerLaw(volume, temperature,p_eos, Natom): # MGD_PowerLaw
  
  if np.isscalar(temperature):
    temperature = temperature*np.ones(len(volume))

  assert len(p_eos)==6, 'EOS parameter array must have correct length of 6'
  assert len(volume)==len(temperature), 'temperature should be a scalar or its length should be equal to volume'

  kB = 13.806488*10**(-24)  #J per K
  P_conv_Fac= 160.217657*6.24*10**(18) #GPa in 1 J/Ang^3
  C_DP = 3*kB*Natom*P_conv_Fac