import time
from scipy import *
from pylab import *
from scipy import linalg
from scipy.linalg import blas
from scipy import interpolate
from scipy import integrate
from numba import jit

def noninteracting_G0(beta, N):
    # See previous lecture, where we show the Hilbert transform to be
    # w(z) = 2*(z - sign(Re(z))*sqrt(z^2-1))
    # use z = i*om
    om = arange(1,2*N,2)*pi/beta
    return (om, 2j*(om-sqrt(om**2 + 1))) # Bethe lattice, half-filling

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 g0_2D(tau, G0):
    L = len(tau)-1
    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]  # minus because G_{qmc} is -G, where G is usual green's function
            else:
                g0[i,j] = G0[L+i-j]
    return g0

def CleanUpdate(g, g0, vn):
    L = len(vn)
    A = zeros((L,L), dtype=float)
    for i,s in enumerate([1,-1]):
        a = exp(vn*s)-1  # a = e^V-1
        # Matrix A = 1 + (1-g0)(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]
        g[i] = linalg.solve(A,g0)

#@jit(nopython=True, cache=True)
def DetRatio(p, g, vn):
    a = [exp(-2*vn[p])-1, exp(2*vn[p])-1]
    Det_up = 1+a[0]*(1-g[0][p,p])
    Det_dn = 1+a[1]*(1-g[1][p,p])
    return (Det_up*Det_dn, a)


def AcceptMove(p, g, a, vn, x0, x1):
    vn[p] *= -1 # Flip spin
    L = len(vn)
    for i,s in enumerate([1,-1]):
        b = a[i]/(1+a[i]*(1-g[i][p,p]))
        x0 = copy(g[i][:,p])
        x0[p] -= 1.
        x1 = g[i][p,:]
        #gnew = g[i][:,:]
        #for l1 in range(L):
        #    for l2 in range(L):
        #        gnew[l1,l2] += b*x0[l1]*x1[l2]
        blas.dger(b,x0,x1,a=g[i],overwrite_a=1)
        #print 'gdiff=', sum(abs(gnew-g[i]))
        #print 'g[',i,']=', g[i]
        
def SaveMeasurements(g, nn, nd, Gtave):
    L = len(g[0])
    # Green's function
    Gt = zeros(L+1, dtype=float)
    for s in range(2):
        for i in range(L):
            for j in range(L):
                if (i>=j):
                    Gt[i-j] += -g[s][i,j]
                else:
                    Gt[L+i-j] += g[s][i,j]
    Gt *= 1./(2.*L)
    # G(tau=0^+) = -<c c^+> = -1+n   while
    # G(tau=0^-) = <c^+ c> = n , hence G(beta)=-G(tau=0^-) = -n
    # therefore Gt[L] + Gt[0] = -1
    Gt[L] = -Gt[0]-1.
    # density
    nnd=0
    for l in range(L): nnd += (1-g[0][l,l]) + (1-g[1][l,l])
    nnd *= (1./L)
    # double occupancy
    nnt=0
    for l in range(L): nnt += (1-g[0][l,l])*(1-g[1][l,l])
    nnt *= (1./L)
    return (nn+nnt, nd+nnd, Gtave+Gt)

def Fourier(Gtave, tau, beta, om, mm=400):
    gtk = interpolate.UnivariateSpline(tau, Gtave, s=0)
    tt = linspace(0,beta,int(mm*L+1))
    gtt = gtk(tt)
    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.simps(rp,tt) + integrate.simps(ip,tt)*1j
    return Giom

def HirshFye(om, G0iom, tau, S_ising, xlam):

    L = len(S_ising)
    vn = array(S_ising)*xlam

    G0tau = InverseFourier(G0iom, om, tau)
    g0 = g0_2D(tau, G0tau)
    #g = [zeros((L,L), dtype=float, order='F'), zeros((L,L), dtype=float, order='F')] # for spin_up and spin_down
    g = [zeros((L,L), dtype=float, order='C'), zeros((L,L), dtype=float, order='C')] # for spin_up and spin_down

    CleanUpdate(g, g0, vn)

    x0 = zeros(L, dtype=float)
    x1 = zeros(L, dtype=float)
    
    accepted = 0
    stored=0
    st1=0
    st2=0
    st3=0
    
    nn=0
    nd=0
    Gtave = zeros(L+1, dtype=float)

    print "%-2s %-6s %-6s %8s %8s %8s" % ('#', 'accpt', 'strd', 't-try', 't-accpt', 't-measr')
    
    for istep in range(nsteps):
        t1 = time.clock()
        
        p = int(rand()*L) # flipping p spin
        (rhor,a) = DetRatio(p, g, vn)
        if (rhor<0): print 'Sign problem!!!!'

        t2 = time.clock(); st1 += t2-t1
        
        if abs(rhor) > rand(): # metropolis
            AcceptMove(p, g, a, vn, x0, x1)
            accepted += 1

        t3 = time.clock(); st2 += t3-t2

        if istep>nwarmup and (istep-nwarmup)%measure==0:
            (nn, nd, Gtave) = SaveMeasurements(g, nn, nd, Gtave)
            stored += 1

        t4 = time.clock(); st3 += t4-t3
        
        if istep%ncout==1:
            print "%-2d %-6d %-6d %8.4f %8.4f %8.4f" % (istep/ncout, accepted, stored, st1/istep*1e5, st2/istep*1e5, st3/istep*1e5)
    
    Gtave *= 1./stored
    nn *= 1./stored
    nd *= 1./stored
    
    print 'density=', nd, 'double-occupancy=', nn
    Giom = Fourier(Gtave, tau, beta, om, 400)
    
    return Giom,Gtave,S_ising,G0tau


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

    nwarmup = int(nwarm0*L)
    measure = int(nmeas0*L)

    (om, G0start) = noninteracting_G0(beta, N)


    G0iom = G0start
    
    tau = linspace(0,beta,L+1)
    

    xlam = arccosh(exp(0.5*(tau[1]-tau[0])*U)) # lambda


    seed(0)
    S_ising = array(sign(rand(L)-0.5),dtype=int)

    G_store=[]
    Giom_old = G0iom[:]
    for itt in range(10):
        Giom,Gtave,S_ising,G0tau = HirshFye(om,G0iom,tau,S_ising,xlam)
        Sigma = 1/G0iom - 1/Giom

        # For a particular case of Bethe lattice, this is the most stable form of the DMFT-SCC
        G0iom_new = 1/(om*1j - 0.25*Giom)
        if itt>0:
            G0iom = mix * G0iom_new[:] + (1-mix) * G0_iom[:]
        else:
            G0_iom = G0iom_new[:]
        
        diff = sum(abs(Giom-Giom_old))
        Giom_old = Giom[:]
        print itt, 'Diff=', diff

        G_store.append( Gtave )
        
        subplot(2,1,1)
        plot(tau, Gtave, 'o-', label='G(tau)')
        plot(tau, G0tau, label='G0(tau)')
        legend(loc='best')

        subplot(2,1,2)
        plot(om, Sigma.imag, label='Sigma(iom)')
        Ly = min(Sigma[:10].imag)
        axis([0,20,Ly,0])
        legend(loc='best')

        show()

    for i in range(len(G_store)):
        plot(tau, G_store[i], 'o-', label='itt '+str(i))
    legend(loc='best')
    show()
