from triqs.gf import *
from h5 import HDFArchive
from triqs_cthyb import Solver
import triqs.operators.util as op
import triqs.utility.mpi as mpi
from triqs.operators import *
from numpy import * 
import triqs.utility.dichotomy as dichotomy

## getting started


beta = 100 #inverse Temperature at which we do our calculations
n_orbs = 3 #how many orbitals do we have? (each fits 2 electrons (bc spin))
U = 5.0 #How strong is our Hubbard U? 5.0 is strong, 2.5-3 is a typical value for many correlated materials
J = U/5 #Hunds-J, typically chosen to be at a constant ratio to U
half_bandwidth = 1 #energy scale
density = 4 #how many electrons are there on average per atom?
CF = 4.0 #crystal field (or for chemists ligand field), how much higher in energy one orbital can be

min_fit = 20 #tailfitting parameters, change!
max_fit = 50 #tailfitting parameters, change!

filename = "01_Tutorial"

spin_names = ['up','down']                   # Outer (non-hybridizing) blocks
orb_names = ['%s'%i for i in range(n_orbs)]  # Orbital indices
off_diag = False                             # Include orbital off-diagonal elements?
n_loops = 1                             # Number of DMFT self-consistency loops
prec = 0.01 #precision

# Initial chemical Potential
# Particle Hole Symmetry for Kanamori:
mu = (n_orbs-0.5)*U - 5*J/2*(n_orbs-1)
mu += 0.6 #this is an initial guess of the chemical potential. If you have a better guess, the calculation might converge quicker

# MC parameters
p = {}
p["max_time"] = -1
p["random_name"] = ""
p["random_seed"] = 123 * mpi.rank + 567
p["length_cycle"] = 2 #how many MC steps between measurements, reduced autocorrelation
p["n_warmup_cycles"] = 200#thermalization
p["n_cycles"] = 500 #number of measurements per core, the more, the better are the statistics
p["move_double"] = False #set to true for superconducting materials


#p["perform_tail_fit"] = True #Tailfitting is important for convergence, do one loop without, check it (see below), and then do rest of the loops with tailfit
#p["fit_max_moment"] = 4 #tailfit parameters
#p["fit_min_w"] = min_fit*2*pi/beta #tailfit parameters
#p["fit_max_w"] = max_fit*2*pi/beta #tailfit parameters



## Building the model

gf_struct = op.set_operator_structure(spin_names,n_orbs,off_diag=off_diag) #gives you the name structure

Umat, Uprimemat = op.U_matrix_kanamori(n_orbs,U,J) #gives you structure of Kanamori interacting Hamiltonian
H = op.h_int_kanamori(spin_names,orb_names,U=Umat,Uprime=Uprimemat,J_hund=J,off_diag=off_diag) #defines Hamiltonian as Operator (predefined class in Triqs)

#if crystal field is wanted, this shifts the "0" orbital up in energy
CFmat = {}
for nn in gf_struct: CFmat[nn[0]] = zeros([nn[1],nn[1]])
if (n_orbs>1):
        H += CF * ( n('up_0',0) + n('down_0',0) )
        CFmat['up_0'][0,0] = CF
        CFmat['down_0'][0,0] = CF


# Construct the solver
S = Solver(beta=beta, gf_struct=gf_struct)


# Set the hybridization function and G0_iw for the Bethe lattice
delta_iw = GfImFreq(indices=[0], beta=beta)
delta_iw << (half_bandwidth/2.0)**2 * SemiCircular(half_bandwidth) #Self consistency equation for Bethe lattice
for name, g0 in S.G0_iw: g0 << inverse(iOmega_n + mu - delta_iw) #define G_0 curly (effective Weiss field)

previous_runs = -1
#only if you want to continue previous calculation
"""if mpi.is_master_node():
        with HDFArchive(filename+"_%i.h5"%(previous_runs),'a') as Res:#a means you can read and write in the file
                S.G_iw << Res['G_iw']
                mu = Res['mu']

                # Compute new S.G0_iw with the self-consistency condition while imposing paramagnetism
                g_iw = GfImFreq(indices=[0], beta=beta)
                # Impose symmetries, to force paramagnetism. Helps with quicker convergence
                for ii in range(n_orbs):
                    g_iw << (S.G_iw['up_%s'%ii] + S.G_iw['down_%s'%ii]) / 2.0
                    S.G_iw['up_%s'%ii] << g_iw
                    S.G_iw['down_%s'%ii] << g_iw

                # Compute S.G0_iw with the self-consistency condition
                for name, g0 in S.G0_iw:
                    g0 << inverse(iOmega_n + mu - (half_bandwidth/2.0)**2 * S.G_iw[name] )

                print("Starting from previous run number %s, with chemical potential %s"%(previous_runs,mu))

#mpi.bcast is important to give the same information to each core!
previous_runs = mpi.bcast(previous_runs)
mu = mpi.bcast(mu)
S.G0_iw <<  mpi.bcast(S.G0_iw)
"""


## The DMFT loop
for i_loop in range(previous_runs + 1, previous_runs + 1 + n_loops):

        mpi.report("\n==================================================\nStarting iteration : %s\n"%i_loop)

        # Solve the impurity problem for the given interacting Hamiltonian and set of parameters
        S.solve(h_int=H, **p)

        # Save quantities of interest on the master node to an h5 archive
        if mpi.is_master_node():
            with HDFArchive(filename+"_%i.h5"%(i_loop),'a') as Results:
                Results['G_iw'] = S.G_iw
                Results['Sigma_iw'] = S.Sigma_iw
                Results['G_tau'] = S.G_tau
                Results['G0_iw'] = S.G0_iw
                Results['loops'] = i_loop

       
        mpi.report("density : %s"%S.G_iw.total_density().real)
        mpi.report("density matrix : %s"%S.G_iw.density())


        # Important part to recalculate and save the correct chemical potential
        def F(mu):
            Gnew = S.G_iw.copy()
            for nm, g in Gnew:
                g << inverse(inverse(S.G0_iw[nm]) + mu - S.Sigma_iw[nm])
            return Gnew.total_density().real

        delta_mu = dichotomy.dichotomy(function=F, x_init=0.0, y_value=density,
                                       precision_on_y=prec, delta_x=0.5, max_loops=100,
                                       x_name="Delta mu", y_name="Total Density",verbosity=3)[0]
        mu += delta_mu
        mpi.report("New chemical potential : %s"%mu)
        if mpi.is_master_node():
            with HDFArchive(filename+"_%i.h5"%(i_loop),'a') as Results:
                Results['mu'] = mu


        #symmetrization for paramagnetism
        g_iw = GfImFreq(indices=[0], beta=beta)
        for ii in range(n_orbs):
            g_iw << (S.G_iw['up_%s'%ii] + S.G_iw['down_%s'%ii]) / 2.0
            S.G_iw['up_%s'%ii] << g_iw
            S.G_iw['down_%s'%ii] << g_iw


        # Compute S.G0_iw with the self-consistency condition
        for name, g0 in S.G0_iw:
            g0 << inverse(iOmega_n + mu - (half_bandwidth/2.0)**2 * S.G_iw[name] )

