from triqs.gf import *
from h5 import *
from numpy import *
import scipy as sp

def get_mu(norbs, U, beta):
    with HDFArchive("../../../final/final_data_basic_calcs/"+str(norbs)+"orbCTHYB_U"+str(U)+"_b"+str(beta)+".h5",'r') as A:
            mu = A['mu']
    return mu

def get_S_w(norbs, U, beta, mode):
    with HDFArchive("../../../final/final_data_basic_calcs/"+str(norbs)+"orbCTHYB_U"+str(U)+"_b"+str(beta)+".h5",'r') as A:
        S_w = A["Sigma_w_"+mode]
    return S_w 


def get_df(omega, beta):
    return -beta/(2+exp(omega*beta)+exp(-omega*beta))
def phi(eps):
    phi_return = 0
    if abs(eps)<=1:
        phi_return = 2/(3*pi)*sqrt(1-(eps)**2)**3
    return phi_return 

def get_G(omega, eps, mu, S_w):
        return 1/(omega - eps + mu - S_w(omega)[0,0])

    
    
def func(omega, eps, beta,  mu, S_w):
    return -phi(eps)*get_df(omega, beta)/pi*(imag(get_G(omega, eps, mu, S_w)))**2/phi(0)


def _infunc(x,func,gfun,hfun,more_args):
    a = gfun
    b = hfun
    myargs = (x,) + more_args
    return sp.integrate.quad(func,a,b,args=myargs, limit = 50)[0]
def custom_dblquad(func, a, b, gfun, hfun, args=(), epsabs=1.49e-8, 
                   epsrel=1.49e-8, maxp1=50, limit=50):
    return sp.integrate.quad(_infunc, a, b, (func, gfun, hfun, args), 
                          epsabs=epsabs, epsrel=epsrel, maxp1=maxp1, limit=limit)







#fix parameters:
norbs = 3
U = 5.0
beta = [100, 90, 80, 70, 60, 50, 40, 35, 30, 25, 22, 20, 17, 15, 13, 12, 11, 10]


name_list_3 = ['up_0', 'down_0', 'up_1', 'down_1', 'up_2', 'down_2']
name_list_2 = ['up_0', 'down_0', 'up_1', 'down_1']



cm = zeros([len(beta), 2*norbs])
for run in range(9,10):
    mode = "ME_"+str(run)
    outfile = open("../../conductivity_U5.0_3orbs.txt", "w")
    outfile.write("norbs/U/beta/run/up_0/down_0/up_1/down_1/up_2/down_2/together/\n")
    for j in range(len(beta)):
            b = beta[j]
            S_w = get_S_w(norbs, U, beta[j], mode)
            mu = get_mu(norbs, U, beta[j])

            for k in range(2*norbs):
                name = eval("name_list_"+str(norbs))[k]
                a = custom_dblquad(func, -1, 1, -20/beta[j], 20/beta[j], args = (beta[j],  mu, S_w[name]))
                integral = a[0]
                cm[j][k] = integral

            outfile.write("%i/%.1f/%i/%s/%.8e/%.8e/%.8e/%.8e/%.8e/%.8e/%.8e/\n"%(norbs, U, b, mode, cm[j][0], cm[j][1], cm[j][2], cm[j][3], cm[j][4], cm[j][5],  sum(cm[j][:])))
            print(beta[j])
    outfile.close()
print("done")