# -*- coding: utf-8 -*-

import numpy as np
import scipy.sparse as sp
import scipy.sparse.linalg as spla

def binconf(c,L):
    "binary representation of a conf integer"
    return np.binary_repr(c,L)
    
def ket(conf,L):
    "quantum ket formatting with first site on the left"
    sconf = binconf(conf,L)
    return "|"+"".join([s for s in sconf[::-1]])+">"

def readsite(conf,i):
    "reads quantum number on site i"
    return (conf&(1<<i))>>i

def conf2int(conf):
    "[1,0,0,1,0] --> 01001 beware of the order"
    return int("".join(str(c) for c in conf[::-1]),base=2)
        
def int2conf(intconf,L):
    "01001 --> [1,0,0,1,0] beware of the order"
    return [ readsite(intconf,i) for i in range(L) ]
    
class IsingHilbert(object):
    "Hilbert space for Ising in a transverse field model"
    
    def __init__(self, L=1):
        "description is a list containing the local hilbert space at site i"
        self.L = L
        self.hilbertsize = 2**self.L
        # no need for that here, because each conf is its own index for Ising
        # but we introduce it for latter purpose
        self.Confs = { i:i for i in range(self.hilbertsize) }            

    def __repr__(self):
        "prints the first 30 confs of hilbert space"
        from itertools import islice
        num = min(30,self.hilbertsize)
        s = "Printing first "+str(num)+" configurations among "+str(self.hilbertsize)+"\n"
        s += ", ".join([ket(c,self.L) for c in list(islice(self.Confs.keys(),num))])
        return s

    def index(self,conf):
        "returns the index of a configuration"
        return self.Confs[conf]
        
def cpp_convert(x,p):
    if type(x) is complex or type(x) is np.complex64 or type(x) is np.complex128:
        return "("+str(round(x.real,p))+","+str(round(x.imag,p))+")"
    else:
        return str(round(x,p))

class Wavefunction(np.ndarray):

    def __new__(cls,hilbert,dtype=float):
        obj = np.ndarray.__new__(cls,(hilbert.hilbertsize,),dtype)
        obj.hilbert = hilbert
        return obj

    def __array_finalize__(self, obj):
        if obj is None: return
        self.hilbert = getattr(obj,'hilbert',None)

    def setvalue(self, conf, value):
        "Sets the coefficient in front of conf"
        self[self.hilbert.index(conf)] = value

    def getvalue(self, conf):
        "gets the coefficient in front of conf"
        return self[self.hilbert.index(conf)]

    def readable(self,precision=4,eps=1e-4):
        from itertools import islice
        num = min(30,self.hilbert.hilbertsize)
        L = self.hilbert.L
        s = ""
        for conf,i in islice(self.hilbert.Confs.items(),num):
            coef = self[i]
            if abs(coef)>eps:
                symb = " + "
                if type(coef) is np.float64 and np.sign(coef) == -1.0: 
                    symb = " - "
                    coef = abs(coef)
                s += symb + cpp_convert(coef,precision) + ket(conf,L)
        if s[:3] == " + ": s = s[3:]
        return "|psi> = "+s
        
    def norm(self):
        return np.linalg.norm(self)
    
    def normalize(self):
        norm = self.norm()
        try: self /= norm
        except: print("can't normalize a vector of norm", norm)

    def randomize(self,a=-1.0,b=1.0):
        self[:] = (b-a)*np.random.random_sample(self.size).reshape(self.shape)+a
        if self.dtype == np.complex128:
            self.imag = (b-a)*np.random.random_sample(self.size).reshape(self.shape)+a
        self.normalize()
    
    def measure(self,op,hermitic=True):
        "computes mean-value of an hermitic operator"
        res = scalar(self,op*self)
        if abs(res.imag) > 1e-15:
            print("Warning, issue with hermiticity")
        return res.real

def scalar(psi1,psi2):
    "computes overlaps of psi1 and psi2 with complex conjugation on psi1"
    return np.vdot(psi1,psi2)

def zeroWf(hilbert,dtype=float):
    psi = Wavefunction(hilbert,dtype=dtype)
    psi[:] = 0.0
    return psi

def basisvectorWf(hilbert,index,dtype=float):
    psi = zeroWf(hilbert,dtype=dtype)
    psi[index] = 1.0
    return psi

def randomWf(hilbert,dtype=float):
    psi = Wavefunction(hilbert,dtype=dtype)
    psi.randomize()
    return psi

def copy(psi):
    res = Wavefunction(psi.hilbert,dtype=psi.dtype)
    res[:] = psi[:]
    return res
    
class Operator(object):
    
    def __mul__(self,psi):
        res =  zeroWf(psi.hilbert,dtype=psi.dtype)
        for conf, i in psi.hilbert.Confs.items():
            coef, newconf = self.Apply(conf)
            j = psi.hilbert.Confs[newconf]
            res[j] += coef * psi[i]
        return res

    def __str__(self):
        return self.name
        
class Sz(Operator):
    def __init__(self,site):
        self.site = site
        self.name = "Sz_"+str(site)
    def Apply(self,conf):
        return readsite(conf,self.site)-0.5, conf

class SzSz(Operator):
    def __init__(self,s1,s2):
        self.s1 = s1
        self.s2 = s2
        self.name = "Sz_"+str(s1)+"Sz_"+str(s2)
    def Apply(self,conf):
        si = readsite(conf,self.s1)-0.5
        sj = readsite(conf,self.s2)-0.5
        return si*sj, conf

class Sx(Operator):
    def __init__(self,site):
        self.site = site
        self.name = "Sx_"+str(site)
    def Apply(self,conf):
        return 0.5, conf^(1<<self.site)

class SxSx(Operator):
    def __init__(self,s1,s2):
        self.s1 = s1
        self.s2 = s2
        self.name = "Sx_"+str(s1)+"Sx_"+str(s2)
    def Apply(self,conf):
# to be written by students
        pass
        
class Spinflip(Operator):
    def __init__(self,i,j):
        self.i = i
        self.j = j
        self.name = "[Sx_"+str(i)+"Sx_"+str(j)+" + Sy_"+str(i)+"Sy_"+str(j)+"]"
    def Apply(self,conf):
        if readsite(conf,self.i) != readsite(conf,self.j):
            return 0.5, conf^((1<<self.i)^(1<<self.j))
        else: return 0.0, conf

class GeneralOp(list):
    """
        A list of 2-element in the format [ [coef, Operator] ]
        corresponding to \sum_<i,j> coef_ij * Op_ij
        warning: assumes hermitic operators
    """
    def __init__(self,iterable=[]):
        list.__init__(self,iterable)
        self.having_matrix = False
        self.storage = 0

    def __str__(self):
        if not len(self): return ""
        s = "("+str(self[0][0])+")*"+str(self[0][1])
        for coef, op in self[1:]:
            s += "+("+str(coef)+")*"+str(op)
        return s
        
    def __iadd__(self,other):
        return GeneralOp(self[:]+other[:])

    def __add__(self,other):
        return GeneralOp(self[:]+other[:])

    def __mul__(self,psi):
        res = copy(psi)
        if self.having_matrix:
            if self.sparse:
                res[:] = self.matrix.dot(psi)
            else:
                res[:] = np.dot(self.matrix,psi)            
        else:
            res[:] = 0.0
            for coef, op in self:
                res += coef * (op*psi)
        return res

    def createMatrix(self,hilbert,dtype=float,sparse=False):
        if self.having_matrix: return
        self.sparse = sparse
        if hilbert.hilbertsize < 2: return
        if not sparse:
            self.createFullMatrix(hilbert,dtype)
        else:
            self.createSparseMatrix(hilbert,dtype)
        self.having_matrix = True

    def createFullMatrix(self,hilbert,dtype=float):
        size = hilbert.hilbertsize
        self.matrix = np.zeros(shape=(size,size), dtype = dtype)
        for index in np.arange(size):
            self.matrix[:,index] = self*basisvectorWf(hilbert,index)
        self.storage = size**2

    def createSparseMatrix(self,hilbert,dtype=float):
        size = hilbert.hilbertsize
        self.matrix = sp.csr_matrix((size,size))
        for conf, i in hilbert.Confs.items():
            row, col, data = [], [], []
            for coef, op in self:
                rescoef, newconf = op.Apply(conf)
                if not rescoef == 0.0:
                    row.append(i)
                    col.append(hilbert.Confs[newconf]) 
                    data.append(coef*rescoef)
            self.matrix = self.matrix + \
                           sp.csr_matrix((data,(row,col)),\
                                        shape=(size,size),dtype=dtype)
        self.storage = self.matrix.nnz
        
    def lanczos(self,hilbert,with_eigvec=False,num=1,verbose=False):
        "Lanczos algorithm, computes the ground-state in option."
        def lanczosSpectrum(a,b):
            tridiag = np.diag(a)
            tridiag+= np.diag(b,k=-1)
            return np.linalg.eigvalsh(tridiag)

        def lanczosGS(a,b):
            tridiag = np.diag(a)
            tridiag+= np.diag(b,k=-1)
            ev, vec = np.linalg.eigh(tridiag)
            return vec[:,0]

        numiteration = 100
        a, b = [], []
        u2 = randomWf(hilbert)
        if with_eigvec: gs = copy(u2) # save the initial random state

        for i in range(numiteration):
            a.append(u2.measure(self))
            if verbose:
                print(i+1,lanczosSpectrum(a,b)[:num])
            u3 = self*u2 - a[i]*u2
            if i>0: u3 -= b[i-1]*u1
            b.append(u3.norm())
            u3 /= b[i]
            u1 = u2
            u2 = u3
            
        if not with_eigvec: 
            return lanczosSpectrum(a,b[:-1])[:num]
        # start again from the SAME initial random WF
        u2 = copy(gs)
        gs[:] = 0.0
        gs_krylov = lanczosGS(a,b[:-1])
        a,b = [],[]

        for i in range(numiteration):
            gs += gs_krylov[i]*u2  # recomputing u2 on the fly
            a.append(u2.measure(self))
            u3 = self*u2 - a[i]*u2
            if i>0: u3 -= b[i-1]*u1
            b.append(u3.norm())
            u3 /= b[i]
            u1 = u2
            u2 = u3
            
        return lanczosSpectrum(a,b[:-1])[:num], gs
        
    def diagonalize(self,num=1,with_eigvec=True):
        "Assumes an hermitic operator by default and that matrix has been created"
        if self.having_matrix:
            if self.sparse:
                if with_eigvec:
                    self.En, self.Phin = spla.eigsh(self.matrix,k=num,which='SA',return_eigenvectors=True)
                else:
                    self.En = np.sort(spla.eigsh(self.matrix,k=num,which='SA',return_eigenvectors=False))
            else:
                if with_eigvec:
                    self.En, self.Phin = np.linalg.eigh(self.matrix)
                else:
                    self.En = np.linalg.eigvalsh(self.matrix)

    def Energies(self,hilbert,dtype=float,num=2,algorithm="Arnoldi"):
        """
         computes num energies and returns it.
         Algorithms are : Fulldiag, Lanczos or Arnoldi
        """
        if algorithm == "Lanczos":
            return self.lanczos(hilbert,with_eigvec=False,num=num)
        if algorithm == "Arnoldi" or algorithm == "Fulldiag":
            sparse = False if algorithm == "Fulldiag" else True
            if not self.having_matrix: self.createMatrix(hilbert,dtype,sparse)
            self.diagonalize(num=num,with_eigvec=False)
            return self.En[:num]

    def Groundstate(self,hilbert,dtype=float,algorithm="Arnoldi"):
        """
         computes the groundstate and returns it.
         Algorithms are : Fulldiag, Lanczos or Arnoldi
        """
        if algorithm == "Lanczos":
            self.En, gs = self.lanczos(hilbert,with_eigvec=True,num=1)
        if algorithm == "Arnoldi" or algorithm == "Fulldiag":
            sparse = False if algorithm == "Fulldiag" else True
            if not self.having_matrix: self.createMatrix(hilbert,dtype,sparse)
            self.diagonalize()
            gsIndex = np.argmin(self.En)
            gs = Wavefunction(hilbert,dtype=dtype)
            gs[:] = self.Phin[:, gsIndex]
        return gs