from __future__ import print_function
import math
import scipy
import scipy.optimize as opt
import sys
import numpy as np
import scipy.linalg as sl
import subprocess
    
def correlation_2L(U,t,occ,imp_model=False):
   """
   Second parametrization for the Hubbard dimer,
   see Senjean et al. Theoretical Chemistry Accounts (2018) 137:169 (or Carrascal paper). Careful : both have a corrigendum...
   imp_model = True will set u = U/4t instead of U/2t.
   """
   occ_tmp = 0
   if U == 0:
     U =0.000001
   if occ == 1:
     occ = 0.99999
     occ_tmp = 1
   sign = lambda x: x and (1, -1)[x<0]
   rho = abs(occ - 1)
   # For the impurity case, u = U/4t and not U/2t as in the fully-interacting dimer.
   if imp_model:
     u = U/(4*t)
   else:
     u = U/(2*t)
   # Impurity correlation energy
   a_12 = 0.5*(1 - rho)
   a_21 = 0.5*np.sqrt(rho*(1- rho)/2.0)
   a_11 = a_21*(1 + 1.0/rho)
   a_22 = a_12*0.5
   a_1 = a_11 + u*a_12
   a_2 = a_21 + u*a_22
   da12_drho = - 0.5
   da21_drho = (1 - 2*rho)/(8*np.sqrt((1 - rho)*rho/2))
   da11_drho = da21_drho*(1 + 1.0/rho) - a_21/rho**2
   da22_drho = - 0.25
   da1_drho = da11_drho + u*da12_drho
   da2_drho = da21_drho + u*da22_drho
   N = (1 - rho)*(1 + rho*(1 + (1 + rho)**3*u*a_1))
   D = 1 + (1 + rho)**3*u*a_2
   dN_drho = - 1 + (1 - 2*rho)*(1 + (1 + rho)**3*u*a_1) + rho*u*(1 - rho)*(1 + rho)**2*(3*a_1 + (1 + rho)*da1_drho)
   dD_drho = u*(1 + rho)**2*(3*a_2 + (1 + rho)*da2_drho)
   g0 = np.sqrt((1 - rho)*(1 + rho*(1 + (1 + rho)**3*u*a_1))/(1 + (1 + rho)**3*u*a_2))
   dg0_drho = (dN_drho - g0**2*dD_drho)/(2*g0*D)
   Y0 = np.sqrt(1 - g0**2 - rho**2)
   dh_dg0 = g0*(g0**4 + 3*g0**2*rho**2 + 2*rho**2*(rho**2 - 1 - Y0))/(2*(g0**2 + rho**2)**2*Y0)
   j = (1 - rho)*(1 + rho)**3*u**2
   k = (3*rho/2 - 1 + rho*(1 + rho)**3*u*a_2)*a_12 - rho*(1 + (1 + rho)**3*u*a_1)*a_22
   l = 2*g0*(1 + (1 + rho)**3*u*a_2)**2
   q = j*k/l
   g1 = g0 + (u*dh_dg0 - 1)*q
   Y1 = np.sqrt(1 - g1**2 - rho**2)
   h = (g1**2*(1 - Y1) + 2*rho**2)/(2*(g1**2 + rho**2))
   f = -2*t*(g1 - u*h)
   Ts = - 2*t*np.sqrt(occ*(2 - occ))
   EHx = (u*2*t)*(1.0+occ*((0.5*occ)-1.0))
   Ec = f - Ts - EHx

   # Derivative with respect to U
   da1_du = a_12
   da2_du = a_22
   v = 2*Y0*(g0**2 + rho**2)**2
   w = g0*(g0**4 + 3*g0**2*rho**2 + 2*rho**2*(rho**2 - 1 - Y0))
   dj_du = 2*(1 - rho)*(1 + rho)**3*u
   G = N/D
   dN_du = (1 - rho)*rho*(1 + rho)**3*(a_1 + u*da1_du)
   dD_du = (1 + rho)**3*(a_2 + u*da2_du)
   dG_du = (dN_du*D - N*dD_du)/D**2
   dg0_du = dG_du/(2*np.sqrt(G))
   dk_du = a_12*(rho*(1 + rho)**3*(a_2 + u*da2_du)) - a_22*(rho*(1 + rho)**3*(a_1 + u*da1_du))
   dl_du = 4*g0*(1 + (1 + rho)**3*u*a_2)*(1 + rho)**3*(a_2 + u*da2_du) + 2*dg0_du*(1 + (1 + rho)**3*u*a_2)**2
   dw_du = dg0_du*(g0**4 + 3*g0**2*rho**2 + 2*rho**2*(rho**2 - 1 - Y0) + g0*(4*g0**3 + 6*g0*rho**2 + 2*rho**2*g0/Y0))
   dv_du = g0*(g0**2 + rho**2)*dg0_du*(-2*(g0**2 + rho**2)/Y0 + 8*Y0)
   dh_dg1 = g1*(g1**4 + 3*g1**2*rho**2 + 2*rho**2*(rho**2 - 1 - Y1))/(2*(g1**2 + rho**2)**2*Y1)
   dq_du = ( (dj_du*k + j*dk_du)*l - j*k*dl_du )/l**2
   ddh_ddug0 = (dw_du*v - w*dv_du)/v**2
   dg1_du = dg0_du + (dh_dg0 + u*ddh_ddug0)*q + (u*dh_dg0 - 1)*dq_du
   dh_du = dg1_du*dh_dg1
   df_du = -2*t*(dg1_du - h - u*dh_du)
   dEc_dU =(u/U)*df_du - (1 + occ*(0.5*occ - 1))*(u*2*t/U)

   # Derivative with respect to t
   dEc_dt = Ec/t - U*dEc_dU/t

   # Derivative with respect to n
   if occ_tmp == 1:
     dEc_dn = 0
   else:
     drho_dn = sign(occ - 1)
     dTs_dn = - 2*t*(1 - occ)/np.sqrt(occ*(2 - occ))
         # Note that P = j*k and Q = l in Eq. 155 of Ref. Senjean et al. Theoretical Chemistry Accounts (2018) 137:169
     dP_drho = (3*(1 - rho)*(1 + rho)**2 - (1 + rho)**3)*u**2*((3*rho/2 - 1 + rho*(1 + rho)**3*u*a_2)*a_12 \
               - rho*(1 + (1 + rho)**3*u*a_1)*a_22) \
               + (1 - rho)*(1 + rho)**3*u**2*((3/2.0 + 3*u*(1 + rho)**2*rho*a_2 \
               + u*(1 + rho)**3*(a_2 + rho*da2_drho))*a_12 \
               + (3*rho/2 - 1 + rho*(1 + rho)**3*u*a_2)*da12_drho - (rho*da22_drho + a_22)*(1 + (1 + rho)**3*u*a_1) \
               - rho*a_22*(3*(1 + rho)**2*u*a_1 + (1 + rho)**3*u*da1_drho))
     dQ_drho = 2*dg0_drho*(1 + (1 + rho)**3*u*a_2)**2 + 4*g0*(1 + (1 + rho)**3*u*a_2)*u*(3.0*(1 + rho)**2*a_2 + (1 + rho)**3*da2_drho)
     dq_drho = (dP_drho*l - j*k*dQ_drho)/l**2
     ddh_ddrhog0 = (- dg0_drho*(g0**2 + rho**2) + 4*g0*(g0*dg0_drho + rho))*(2*rho**2 + g0**2*(1 - Y0))/(g0**2 + rho**2)**3 \
                   - g0*(4*rho + 2*g0*dg0_drho*(1 - Y0) + g0**2*(g0*dg0_drho + rho)/Y0)/(g0**2 + rho**2)**2 \
                   - (g0*dg0_drho + rho)*(2*g0*(1 - Y0) + g0**3/Y0)/(g0**2 + rho**2)**2 \
                   + (2*dg0_drho*(1 - Y0) + 2*g0*rho/Y0 + 5*g0**2*dg0_drho/Y0 + g0**3*(g0*dg0_drho + rho)/Y0**3)/(2*(g0**2 + rho**2))
     dg1_drho = dg0_drho + (u*dh_dg0 - 1)*dq_drho + u*ddh_ddrhog0*q
     dh1_drho = (1/(2*(g1**2 + rho**2)))*(4*rho + 2*g1*dg1_drho*(1 - Y1) + g1**2*(g1*dg1_drho + rho)/Y1) \
               - (g1*dg1_drho + rho)*(2*rho**2 + g1**2*(1 - Y1))/(g1**2 + rho**2)**2
     df_drho = -2*t*(dg1_drho - u*dh1_drho)
     dEHxc_dn = drho_dn*df_drho - dTs_dn
     dEc_dn = drho_dn*df_drho - dTs_dn - (u*2*t)*(occ - 1)
   return Ec, dEc_dU, dEc_dt, dEc_dn, f

def u(twot,bigu):
    return bigu/twot
def nu(twot,deltav):
    return deltav/twot
def omega(nu,u):
    return np.sqrt(3.0*(1.0+(nu**2))+(u**2))
def theta(u,nu,w):
    return (1.0/3.0)*np.arccos((u/(w**3))*(9.0*(nu**2-0.5)-u**2))
def r(twot,bigu,deltav):
    return np.sqrt(3.0*(twot**2 + deltav**2) + bigu**2)
def thetatilde(twot,bigu,deltav):
    return (1.0/3.0)*np.arccos((9.0*bigu*(deltav**2-(twot**2)/2.0) - bigu**3)/(r(twot,bigu,deltav)**3))
def E(twot,bigu,deltav,i): # i = 0 --> GS (should be same as above, not tested), i = 1 --> ES.
    return 2.0*bigu/3.0 + (2.0*r(twot,bigu,deltav)/3.0)*np.cos(thetatilde(twot,bigu,deltav) + 2.0*np.pi*(i+1.0)/3.0)
def dE_dv(twot,bigu,deltav,i):
    return 2.0*deltav*(2.0*bigu/3.0 + (2.0*r(twot,bigu,deltav)/3.0)*np.cos(thetatilde(twot,bigu,deltav) + 2.0*np.pi*(i+1.0)/3.0))/(3.0*(2.0*bigu/3.0 + (2.0*r(twot,bigu,deltav)/3.0)*np.cos(thetatilde(twot,bigu,deltav) + 2.0*np.pi*(i+1.0)/3.0))**2 - 4.0*bigu*(2.0*bigu/3.0 + (2.0*r(twot,bigu,deltav)/3.0)*np.cos(thetatilde(twot,bigu,deltav) + 2.0*np.pi*(i+1.0)/3.0)) + bigu**2 - twot**2 - deltav**2)
def Ebar(twot,bigu,deltavbar):
    return E(twot,bigu,deltavbar*bigu) 
def E0_1(twot,deltav):
    return - np.sqrt(twot**2/4.0 + (deltav**2/4.0))
def Ts_func(n,twot):
    return -twot*np.sqrt(n*(2.0-n))
def EHx_func(n,bigu):
    return bigu*(1.0+n*((0.5*n)-1.0))
def EH_func(n,bigu):
    return bigu*(1.0+(n-1.0)**2)
def dE0_1_dv(twot,deltav):
    return - deltav/(2*np.sqrt(twot**2 + deltav**2))
# ensemble energy (ground and first excited state)
def Eens_1(twot,bigu,deltav,w):
    return (1-w)*E(twot,bigu,deltav,0) + w*E(twot,bigu,deltav,1)
# ensemble energy (ground and second excited state)
def Eens_2(twot,bigu,deltav,w):
    return (1-w)*E(twot,bigu,deltav,0) + w*E(twot,bigu,deltav,2)

if __name__ == "__main__":
 t = 1
 twot = 2*t
 wlist = [0.0,0.1,0.2,0.3,0.4,0.5]
 Ulist = [1.0,2.0,5.0,10.0,100.0,1000.0]
 for w in wlist:
  for bigu in Ulist:
   LF_gs_gace_edft = open("Ew1_Ew2_fctdeltav_Usurt" + str(bigu) + "_w" + str(w), "w")
   v = range(int(-bigu-100),int(bigu+100))
   number_of_values     = 3
   format_string_str    = number_of_values*"{:>8} "
   format_string_real   = number_of_values*"{:8.4f} "
   print(format_string_str.format("deltav","Ew_1","Ew_2"),file = LF_gs_gace_edft)
   for deltav in v:
    print(format_string_real.format(deltav,Eens_1(twot,bigu,deltav,w),Eens_2(twot,bigu,deltav,w)),file = LF_gs_gace_edft)
