from scipy import *
from pylab import *
from scipy import weave
import time

def ferm(x):
    if x>100.:
        return 0
    if x<-100:
        return 1
    return 1/(exp(x)+1.)

def Momentum2D(Nk, t, T):
    b1 = 2*pi*array([1,0])
    b2 = 2*pi*array([0,1])
    epsk = zeros((Nk*Nk))   # dispersion
    klist = zeros((Nk*Nk,2))
    index = zeros((Nk,Nk),dtype=int)
    tt = -2*t
    ii=0
    for i1 in range(Nk):
        for i2 in range(Nk):
            k = (i1 * b1 + i2 * b2)/(Nk + 0.0)
            cx = cos(k[0])
            cy = cos(k[1])
            epsk[ii] = tt*(cx+cy)
            klist[ii,:] = k
            index[i1,i2] = ii
            ii+=1
    fm = array([ferm(epsk[ik]/T) for ik in range(len(epsk))])

    ik_m_iq = zeros( (len(klist), len(klist)), dtype=int ) # index for k-q
    #two_pi = 2*pi
    #for ik in range(len(klist)):
    #    for iq in range(len(klist)):
    #        kq = klist[ik]-klist[iq]
    #        i1 = int(round(kq[0]/two_pi*Nk))
    #        i2 = int(round(kq[1]/two_pi*Nk))
    #        # i1 and i2 are between 0...Nk
    #        while (i1 >= Nk) : i1 -= Nk
    #        while (i1 <  0)  : i1 += Nk
    #        while (i2 >= Nk) : i2 -= Nk
    #        while (i2  <  0) : i2 += Nk
    #        ik_m_iq[ik,iq] = index[i1,i2]
    #        #print "%3d  %3d  %3d" % (ik, iq, ik_m_iq[ik,iq])
    code="""
    using namespace std;
    int Nlist = klist.shape()(0);
    double two_pi = 2*M_PI;
    for (int ik=0; ik<Nlist; ik++){
       for (int iq=0; iq<Nlist; iq++){
          int i1 = round((klist(ik,0)-klist(iq,0))*Nk/two_pi);
          int i2 = round((klist(ik,1)-klist(iq,1))*Nk/two_pi);
          while (i1 >= Nk)  i1 -= Nk;
          while (i1 <  0)   i1 += Nk;
          while (i2 >= Nk)  i2 -= Nk;
          while (i2  <  0)  i2 += Nk;
          int ikq = index(i1,i2);
          ik_m_iq(ik,iq) = ikq;
          //cout<<ik<<" "<<iq<<" "<<ikq<<endl;
       }
    }
    """
    weave.inline(code, ['klist', 'ik_m_iq', 'Nk', 'index'],
                 type_converters=weave.converters.blitz, compiler='gcc')
    return (epsk, klist, ik_m_iq, fm)

def Momentum3D(Nk, t, T):
    b1 = 2*pi*array([1,1,-1.])
    b2 = 2*pi*array([1,-1,1.])
    b3 = 2*pi*array([-1,1,1.])
    epsk = zeros((Nk*Nk*Nk))
    klist = zeros((Nk*Nk*Nk,3))
    index = zeros((Nk,Nk,Nk),dtype=int)
    tt = -4*t
    ii=0
    for i1 in range(Nk):
        for i2 in range(Nk):
            for i3 in range(Nk):
                k = (i1*b1 + i2*b2 + i3*b3)/(Nk+0.0)
                cx = cos(k[0]/2.)
                cy = cos(k[1]/2.)
                cz = cos(k[2]/2.)
                epsk[ii] = tt*(cx*cy+cx*cz+cy*cz)
                klist[ii,:] = k
                index[i1,i2,i3] = ii
                ii += 1
    fm = array([ferm(epsk[ik]/T) for ik in range(len(epsk))])
    ik_m_iq = zeros((len(klist),len(klist)), dtype=int)
    R1 = 0.5*array([1,1,0])
    R2 = 0.5*array([1,0,1])
    R3 = 0.5*array([0,1,1])
    #two_pi = 2*pi
    #for ik,k in enumerate(klist):
    #    for iq,q in enumerate(klist):
    #        kq = k-q
    #        i1 = round(dot(kq,R1)/two_pi * Nk)
    #        i2 = round(dot(kq,R2)/two_pi * Nk)
    #        i3 = round(dot(kq,R3)/two_pi * Nk)
    #        i1, i2, i3 = int(i1), int(i2), int(i3)
    #        # back to first BZ
    #        while (i1 >= Nk) : i1 -= Nk
    #        while (i1 <  0)  : i1 += Nk
    #        while (i2 >= Nk) : i2 -= Nk
    #        while (i2  <  0) : i2 += Nk
    #        while (i3 >= Nk) : i3 -= Nk
    #        while (i3  <  0) : i3 += Nk
    #        ikq = index[i1,i2,i3]
    #        ik_m_iq[ik,iq] = ikq
    #        #print "%3d %3d  %3d" % (ik,iq,ikq)

    code="""
    #include <blitz/array.h>
    using namespace std;
    int Nlist = klist.shape()(0);
    double two_pi = 2*M_PI;
    blitz::TinyVector<double,3> _R1_, _R2_, _R3_;
    for (int i=0; i<3; i++){
      _R1_(i) = R1(i);
      _R2_(i) = R2(i);
      _R3_(i) = R3(i);
    }
    for (int ik=0; ik<Nlist; ik++){
      for (int iq=0; iq<Nlist; iq++){
         blitz::TinyVector<double,3> kq;
         for (int i=0; i<3; i++) kq(i) = klist(ik,i)-klist(iq,i);
         int i1 = round(blitz::dot(kq,_R1_)/two_pi * Nk);
         int i2 = round(blitz::dot(kq,_R2_)/two_pi * Nk);
         int i3 = round(blitz::dot(kq,_R3_)/two_pi * Nk);
          while (i1 >= Nk)  i1 -= Nk;
          while (i1 <  0)   i1 += Nk;
          while (i2 >= Nk)  i2 -= Nk;
          while (i2  <  0)  i2 += Nk;
          while (i3 >= Nk)  i3 -= Nk;
          while (i3  <  0)  i3 += Nk;
          int ikq = index(i1,i2,i3);
          ik_m_iq(ik,iq) = ikq;
          //cout<<ik<<" "<<iq<<" "<<ikq<<endl;
      }
    }
    """
    weave.inline(code, ['klist', 'ik_m_iq', 'Nk', 'index', 'R1', 'R2', 'R3'],
                 type_converters=weave.converters.blitz, compiler='gcc')
    return (epsk, klist, ik_m_iq, fm)

def GetQPolarization(klist, fm, ik_m_iq, epsk, om, igamma):
    P1 = zeros( (len(klist),len(om)), dtype=complex)
    P2 = zeros( (len(klist),len(om)), dtype=complex)
    Nlist = len(klist)
    for iq in range(Nlist):
        for ik in range(Nlist):
            ikq = ik_m_iq[ik,iq]
            c1 = (1-fm[ik])*fm[ikq]/Nlist
            c2 = fm[ik]*(1-fm[ikq])/Nlist
            w = 1./(om-epsk[ik]+epsk[ikq]+igamma)
            P1[iq,:] += c1*w
            P2[iq,:] += c2*w
    return (P1, P2)

def GetQSigma(U, P1, P2, klist, fm, ik_m_iq, epsk, om, igamma):
    
    iSg = zeros((len(klist), len(om)), dtype=complex)
    #Nlist = len(klist)
    #for ik in range(Nlist):
    #    for iq in range(Nlist):
    #        ikq = ik_m_iq[ik,iq]
    #        for iom in range(len(om)):
    #            x = om[iom]-epsk[ikq]
    #            if (x>=om[0] and x<om[-1]):
    #                iz = int( (x-om[0])/(om[-1]-om[0])*(len(om)-1) ) # CORRECT
    #                if (iz > len(om)-2): iz = len(om)-2  # CORRECT
    #                # interpolation in [iz,iz+1]  
    #                p = (x-om[iz])/(om[iz+1]-om[iz])
    #                # linear interpolation of P(om-eps(k-q))
    #                P1t = P1[iq,iz] + (P1[iq,iz+1]-P1[iq,iz])*p
    #                P2t = P2[iq,iz] + (P2[iq,iz+1]-P2[iq,iz])*p
    #                iSg[ik,iom] += (1-fm[ikq])*P1t + fm[ikq]*P2t

    code="""
    using namespace std;
    int Nlist = klist.shape()(0);
    int nom = om.size();
    double om0 = om(0), omn = om(nom-1);
    for (int ik=0; ik<Nlist; ik++){
       for (int iq=0; iq<Nlist; iq++){
          int ikq = ik_m_iq(ik,iq);
          for (int iom=0; iom<om.size(); iom++){
             double x = om(iom)-epsk(ikq);
             if (x>=om0 && x<omn){
                int iz = (x-om0)/(omn-om0)*(nom-1); // simple linera interpolation
                if (iz >= om.size()-2) iz=nom-2;
                double p = (x-om(iz))/(om(iz+1)-om(iz));
                // linear interpolation
                complex<double> P1t = P1(iq,iz) + (P1(iq,iz+1)-P1(iq,iz))*p;
                complex<double> P2t = P2(iq,iz) + (P2(iq,iz+1)-P2(iq,iz))*p;
                iSg(ik,iom) += (1-fm(ikq))*P1t + fm(ikq)*P2t;
             }
          }
       }
    }
    """
    weave.inline(code, ['om','klist','ik_m_iq','fm','epsk','P1','P2','iSg'],
                 type_converters=weave.converters.blitz, compiler='gcc')
    
    return iSg * U**2 / len(klist)


def GetGloc(klist, epsk, om, iSg, igamma):
    Gloc = zeros(len(om), dtype=complex)
    for iom in range(len(om)):
        dsum =0.0
        for ik in range(len(klist)):
            dsum += 1/(om[iom]-epsk[ik]-iSg[ik,iom]+igamma)
        Gloc[iom] = dsum/len(klist)
    return Gloc

if __name__ == '__main__':
    T = 0.05
    Nw = 200
    L = 4
    U = 3.0
    Q2D = False
    
    # 2D case
    if (Q2D):
        Nk = 30
        t = 1./4
        gamma = 0.03
    else:
    # 3D case
        t = 1./(4*sqrt(3.))
        Nk = 12
        gamma = 0.05
    
    om = linspace(-L,L,Nw)
        
    igamma = gamma*1j

    t1 = time.time()

    if (Q2D):
        (epsk, klist, ik_m_iq, fm) = Momentum2D(Nk,t,T)
    else:
        (epsk, klist, ik_m_iq, fm) = Momentum3D(Nk,t,T)
    
    t2 = time.time()

    (P1, P2) = GetQPolarization(klist, fm, ik_m_iq, epsk, om, gamma*1j)
    t3 = time.time()
    
    print 'dt=', t2-t1, ' ', t3-t2

    iSg = GetQSigma(U, P1, P2, klist, fm, ik_m_iq, epsk, om, gamma*1j)  # CORRECT

    Gloc = GetGloc(klist, epsk, om, iSg, gamma*1j)
    
    lSg = zeros(shape(iSg[0]), dtype=complex)
    for iq in range(len(klist)):
        lSg += iSg[iq]
    lSg *= 1./len(klist)
        
    for iq in range(len(klist)):
        plot(om, -iSg[iq].imag)
    plot(om, -lSg.imag, 'r', lw=3)
    show()

    plot(om, -Gloc.imag/pi)
    show()
    
    sys.exit(0)
    g=0
    for i in range(len(epsk)):
        g += 1/(om-epsk[i]+igamma)
    g *= 1./len(epsk)
    
    plot(om, -g.imag/pi)
    show()