#!/usr/bin/env python
""" Hirsh-Fye algorithm to solve the quantum impurity model
    
"""
import os, time
from pylab import *
from scipy import *
from scipy import weave
import scipy
scipy.pkgload('lib')  # To use blas functions directly

def g0_2D(tau, G0):
    "Initialization of ising g0(tau,tau') from G0(tau-tau'). Must be antiperiodic in time."
    L = len(tau)-1
    print 'L=', L
    g0 = zeros((L,L), dtype=float)
    for i in range(L):
        for j in range(L):
            if i>=j:
                g0[i,j] = -G0[i-j]  # antiperiodic boundary conditions for fermions
            else:
                g0[i,j] = G0[L+i-j] # antiperiodic boundary conditions for fermions    
    return g0

def CleanUpdate(g, g0, vn):
    " Goes over all ising spins and recalculates g from g0 according to the current ising configuration."
    L  = len(vn)
    sigma = [1,-1]
    A = zeros((L,L),dtype=float)
    for i in range(2):
        # a = e^V-1
        a = [exp(vn[l]*sigma[i])-1 for l in range(L)]
        # Matrix A = 1 + (1-g)(e^V-1) == (1+a)*I - g*a
        for l1 in range(L):
            for l2 in range(L): A[l1,l2] = -g0[l1,l2]*a[l2]
            A[l1,l1] += 1 + a[l1]
        # Solves system of linear equations A*g0[i]=g[i]. This is faster than computing inverse.
        g[i] = solve(A, g0)

def DetRatio(it, g, a):
    # One of the most important subroutines in QMC.
    # Computes the probability for the MC step - the ration of determinants between ising configurations.
    # it -- time slice index
    a[0] = exp(-2*vn[it]) - 1
    a[1] = exp( 2*vn[it]) - 1
    Det_up = 1+(1-g[0,it,it])*a[0]
    Det_dn = 1+(1-g[1,it,it])*a[1]
    return Det_up*Det_dn

def AcceptMove(it, g, a, vn, accepted):
    # Most important subroutine in QMC.
    # Updates the quantities (green's function) since the new ising configuration is accepted.
    # This function needs most of the time and should be heavily optimized! Most of the time is currently spend in blas
    # routine daxpy when doing rank1 update.
    vn[it] *= -1   # perform spin-flip
    L = len(vn)
    #if accepted%ndirty==0: CleanUpdate(g,g0,vn)
    #else # dirty update
    for s in range(2):     # both members of pair need to be updated
        # prefactor b = a/(1+a(1-g_pp))
        b = a[s]/(1+a[s]*(1-g[s,it,it]))
        # this computes (g-1)_{l,it}
        x0 = copy(g[s,:,it])  # first we set x0[l] = g_{l,it}
        x0[it] -= 1     # now be sutract delta_{l,il} to get x0[l] = g_{l,it}-delta_{l,it}
        x1 = g[s,it,:]
        # Here we calculate b*(g-1) x g , where x is tensor product
        #g[s] += lib.blas.fblas.dger(b,x0,x1)
        g[s]=lib.blas.fblas.dger(b,x0,x1,a=g[s],overwrite_a=1)
    return (accepted+1)

def AcceptMoveCPP(it, g, a, x0, x1, vn, accepted):
    # Optimized using weave
    # Most important subroutine in QMC. 
    # Updates the quantities (green's function) since the new ising configuration is accepted.
    # This function needs most of the time and should be heavily optimized! Most of the time is currently spend in blas
    # routine daxpy when doing rank1 update.
    vn[it] *= -1   # perform spin-flip
    L = len(vn)
    b = 0.0
    for s in range(2):     # both members of pair need to be updated
        code="""
          using namespace std;
          b = a(s)/(1+a(s)*(1-g(s,it,it))); // prefactor b = a/(1+a(1-g_pp))
          for (int l=0; l<L; l++){
              x0(l) = g(s,l,it);
              x1(l) = g(s,it,l);
          }
          x0(it) -= 1.;
          return_val = b;
        """
        b = weave.inline(code, ['it','s','a','g','L','x0','x1','b'],
                         type_converters=weave.converters.blitz, compiler = 'gcc')

        # Here we calculate b*(g-1) x g , where x is tensor product
        g[s]=lib.blas.fblas.dger(b,x0,x1,a=g[s],overwrite_a=1)
    return accepted+1
 
def SaveMeasurements(g, stored, nn, Gtave):
    L = shape(g)[1]
    
    # Saving Green's function
    Gt=zeros(L+1,dtype=float)
    for s in range(2):
        for l0 in range(L):
            for l1 in range(L):
                if (l0>=l1):
                    Gt[l0-l1] += -g[s,l0,l1]  # antiperiodic boundary conditions and 
                else:
                    Gt[L+l0-l1] += g[s,l0,l1] # -isgn convention in QMC community
    Gt *= (1./L/2.); # normalization because there were L^2 pairs for L time slices.
    Gt[L] = -Gt[0]-1.; # Sets the G(beta) due to analytical knowledge

    
    # Double occupancy
    nnt=0; # double occupancy can be calculated just like average of n*n over ising configurations
    # the reason is that the action is quadratic in the representation of ising spins (non-interacting particles)
    for l in range(L): nnt += (1-g[0,l,l])*(1-g[1,l,l])
    nnt *= (1./L)
    
    return (stored+1, nn+nnt, Gtave+Gt)


def noninteracting_G0(beta, N):
    # This is exact for half-filling non-interacting case on bethe lattice
    # 
    # On the real axis, the Green's function is:
    #  G(w) = 2*(w - sqrt(w**2-1))
    #  on imaginary axis we have
    #  G(omn) = 2*i*omn * (1-sqrt(1+1/omn**2))
    #  
    om = arange(1,2*N,2)*pi/beta  # Matsubara frequency points
    G0 = zeros(len(om), dtype=complex)
    for iw,w in enumerate(om):
        #Delta = 0.5j*w*(1-sqrt(1+1/w**2)) # semicircular
        #G0[iw] = 1/(w*1j-Delta)
        G0[iw] = 2j*w*(1-sqrt(1+1/w**2))
    return (om, G0)

def InverseFourier(Giom, om, tau):
    "Inverse Fourier computes G(tau) from G(iom)"
    beta = pi/om[0]
    Gtau = zeros(len(tau), dtype=float)
    for it,t in enumerate(tau):
        dsum=0
        for iw,w in enumerate(om):
            dsum += cos(w*t)*Giom[iw].real + sin(w*t)*(Giom[iw].imag+1./w)
        Gtau[it] = 2*dsum/beta-0.5
    return Gtau

def Fourier(Gt, tau, om, mm=50):
    "Fourier computes G(iom) from G(tau)"
    # First get the spline for smooth G(iom)
    gtk = interpolate.splrep(tau,Gt)
    # We need much larger mesh to push Nayquid frequency to large enough energy
    tt = linspace(0,beta,int(mm*L+1))
    gtt = interpolate.splev(tt,gtk)
    
    Giom = zeros(len(om), dtype=complex)
    for iw,w in enumerate(om):
        rp = gtt*cos(tt*w)
        ip = gtt*sin(tt*w)
        Giom[iw] = integrate.trapz(rp,tt)+ integrate.trapz(ip,tt)*1j
    return Giom


if __name__ == '__main__':
    beta = 16.       # inverse temperature
    N = 1500         # number of matsubara points
    L = 64           # number of time slices
    U = 2.           # interaction U
    nwarm0 = 100     # how many sweeps to warm up
    nmeas0 = 10      # how often to meassure (in sweeps)
    nsteps = 100000  # number of all MC steps
    ncout  = 10000   # how often to print info
    
    nwarmup = int(nwarm0*L)
    measure = int(nmeas0*L)

    print 'U=', U, 'beta=', beta, 'L=', L
    print 'nwarmup=', nwarmup
    print 'nsteps=', nsteps
    print 'measure=', measure
    
    (om, G0start) = noninteracting_G0(beta, N)  # G0(iom)
    
    
    S_ising = sign(rand(L)-0.5)     # random ising spins [+1 or -1]
    # previous ising configuration, if existsing
    if os.path.exists('ising.dat') and os.path.getsize('ising.dat')>0:
        S_ising_new = map(float,open('ising.dat', 'r').readlines()[0].split())
        print 'S1, s2=', len(S_ising_new), len(S_ising)
        if len(S_ising_new)==len(S_ising):
            S_ising = S_ising_new
            print 'Started from old configuration of spins'
    

    G0iom = G0start
    tau = linspace(0,beta,L+1) # We also need the enpoints
    G0tau = InverseFourier(G0iom, om, tau) # G0(tau)
    
    xlam = arccosh(exp(0.5*(tau[1]-tau[0])*U)) # lambda
    vn = array(S_ising)*xlam

    # g0 in 2D representation
    g0 = g0_2D(tau, G0tau)
    # g in 2D representation
    g = zeros((2,L,L), dtype=float)
    # g computed from g0 by inverse A
    CleanUpdate(g, g0, vn)

    # A few arrays for faster computation
    a=zeros(2,dtype=float)
    x0 = zeros(L, dtype=float)
    x1 = zeros(L, dtype=float)

    
    accepted=0 # accepted steps
    stored=0   # stored measurements
    st1=0      # time used
    st2=0      # time used
    st3=0      # time used
    
    nn=0                           # double occupancy
    Gtave = zeros(L+1,dtype=float) # <G(tau)>
    
    print "%-2s %-6s %-6s %8s %8s %8s" % ('#', 'accpt', 'strd', 't-try', 't-accpt', 't-measr')
    
    for istep in range(nsteps):
        
        t1=time.time()
        
        it = int(rand()*L)
        Det = DetRatio(it, g, a)
        P = Det/(1+Det)
        
        t2=time.time(); st1 += t2-t1
    
        if P>rand():
            #accepted = AcceptMove(it, g, a, vn, accepted)
            accepted = AcceptMoveCPP(it, g, a, x0, x1, vn, accepted)
    
        t3 = time.time()
        st2 += t3-t2
        
        if istep > nwarmup and (istep-nwarmup)%measure==0:
            (stored, nn, Gtave) = SaveMeasurements(g, stored, nn, Gtave)
        
        if istep%ncout==1:
            print "%-2d %-6d %-6d %8.4f %8.4f %8.4f" % (istep/10000, accepted, stored, st1/istep*1e5, st2/istep*1e5, st3/istep*1e5)
            
    
        t4 = time.time()
        st3 += t4-t3
        

    # Save ising configuration
    S_ising = (vn/xlam)
    fis = open('ising.dat', 'w')
    print >> fis, "%2.0f "*len(S_ising) % tuple(S_ising)
    fis.close()
    
    
    Gtave *= 1./stored
    nn *= 1./stored
    
    print 'double-occupancy=', nn
    
    subplot(2,2,1)
    plot(tau, Gtave, 'o-', label='G(tau)')
    plot(tau, G0tau, label='G0(tau)')
    legend(loc='best')

    Giom = Fourier(Gtave, tau, om, 400)

    Sigma = 1/G0iom-1/Giom
    
    subplot(2,2,2)
    plot(om, Giom.imag, label='Im(G)')
    plot(om, G0iom.imag, label='Im(G0)')
    axis([0,20,-1.8,0])
    legend(loc='best')
    
    subplot(2,2,3)
    plot(om, Sigma.imag, label='Im(Sig)')
    Ly = min(Sigma[:10].imag)
    axis([0,20,Ly,0])
    legend(loc='best')
    show()