"""
tebdIsing4.py
Use TEBD to compute ground state of transverse field Ising model.
2014-09-17
"""

import numpy as np
from cmath import *
from mpstest17 import *
import matplotlib.pyplot as plt
import time

def main():
    test0()

def test0():
    (J,hx) = (0,1)
    (nsteps,chi) = (100,5)
    for N in [3]:
        test2(J,hx,N,1.0/nsteps,nsteps,chi)

def test1():
    print(getU(1,1,1))
    a = np.random.rand(3,3)
    show(a,"a")
    show(np.linalg.eig(a)[0],"eignenvalues of a")
    ea = expm(a)
    show(ea,"e^a")
    show(np.linalg.eig(ea)[0],"eigenvalues of e^a")
    show(np.log(np.linalg.eig(ea)[0]),"log of eigenvalues of e^a")
    print(s(0),"I")
    print(s(1),"X")
    print(s(2),"Y")
    print(s(3),"Z")

def test2(J,hx,N,t,nsteps,chi):
    print("\nStarting TEBD for N = %d Ising model with parameters:"%N)
    print("    (J,hx,N) = (%.4f,%.4f,%d)"%(J,hx,N))
    print("    (t,nsteps,chi) = (%f,%d,%d)"%(t,nsteps,chi))
    startTime = time.clock()
    mps = tebdIsing(J,hx,N,t,nsteps,chi)
    runTime = time.clock()-startTime
    print("Simulation completed in %f seconds"%runTime)
    gs = getStateOBC(mps)
    gs = gs/np.linalg.norm(gs) # Normalize state
##    show(gs,"Ground State")
    print("Ground state =")
    displayState(gs)
    startTime = time.clock()
    energyBF = getEnergyBruteForce(J,hx,N,gs).real
    print("Energy = %f, Energy per spin = %f"%(energyBF,energyBF/N))
    runTime = time.clock() - startTime
    print("Brute force energy computed in %f seconds"%runTime)
    startTime = time.clock()
    energy = getEnergy(J,hx,N,mps).real
    print("Energy = %f, Energy per spin = %f"%(energy,energy/N))
    runTime = time.clock() - startTime
    print("MPO energy computed in %f seconds"%runTime)

def test3():
    print(pairHamiltonianMPO())

def isingH(J,hx,N):
    """
    Full matrix representation of ising model hamiltonian.
    """
    X = 3
    Z = 1
    pairs = np.zeros((2**N,2**N),dtype=complex)
    for i in xrange(N-1):
        pairs += pauli([X,X],[i,i+1],N)
    fields = np.zeros((2**N,2**N),dtype=complex)
    for i in xrange(N):
        fields += pauli([Z],[i],N)
    return -J*pairs-hx*fields

def tebdIsing(J,hx,N,t,nsteps,chi):
    """
    Run TEBD algorithm on 1D N-spin transverse field Ising model.
    Uses open boundary conditions and imaginary time evolution.

    Parameters
    ----------
    J : float
        Pair-wise interaction energy.
    hx : float
        Magnetic energy in transverse B-field.
    N : int        
        Number of spins.
    t : float
        Timestep of each iteration.
    nsteps : int
        Number of time evolution iterations simulated.
    
    Returns
    -------
    groundState : list
        MPS representation of ground state.
    energies : (nsteps) ndarray
        Energies at each timestep.
    """
    
    # Initiate system with random MPS state.
    d = 2
    state = randomMPSOBC(N,chi,d,real=True)
##    state = allEqual(N)
    print("Intial state")
    displayState(getStateOBC(state))

    # Initilise list of energies at each iteration
    
    # Compute time evolution operators.
    U = getU(J,hx,t) # Acts on pairs of spins in middle
    Ub = getUb(J,hx,t) # Acts on only one boundary spin

    energies = []
    # Run iteration nstep times
    for step in xrange(nsteps):
        # First half evolution
        # Evolve first two spins
        state[0],state[1] = leftPair(state[0],state[1],U,chi,d)
        # Evolve middle spins
        for i in xrange(2,N-2,2):
            state[i],state[i+1] = middlePair(state[i],state[i+1],U,chi,d)
        # Evolve last spin pair (or single spin if odd)
##        show(state[-1].shape,"state[-1]")
        if N % 2 and N > 2: # N is odd
            state[-1] = rightSingle(state[-1],Ub)
##            print("odd")
        elif N > 2: # N is even
            state[-2],state[-1] = rightPair(state[-2],state[-1],U,chi,d)
##        show(state[-1].shape,"state[-1]")
        
        # Second half evolution
        # Evolve first spin
        state[0] = leftSingle(state[0],Ub)
        # Evolve middle spins
        for i in xrange(1,N-2,2):
            state[i],state[i+1] = middlePair(state[i],state[i+1],U,chi,d)
##        show(state[-1].shape,"state[-1]")
        state[-1] = rightSingle(state[-1],Ub)
##        show(state[-1].shape,"state[-1]")
##        # Evolve last spin (or spin pair if odd)
##        show(state[-1].shape,"state[-1]")
        if N % 2 and N > 2: # N is odd
            state[-2],state[-1] = rightPair(state[-2],state[-1],U,chi,d)
        elif N > 2: # N is even and greater than 2
            state[-1] = rightSingle(state[-1],Ub)
##        energies.append(getEnergy(state))
##        if innerProductOBC(state,state) > 1e100:
##            state = [s/1e10 for s in state]
##        show(sum([np.sum(np.abs(s)) for s in state]),"sum state")
##        show(innerProductOBC(state,state),"<a|a>")
        energies.append(getEnergy(J,hx,N,state))
    plt.plot(energies)
    plt.show()
    return state

def middlePair(A,B,U,chi,d):
    """
    Evolve a pair of spins in middle.
    """
    lbd = A.shape[0] # Left bond dimension
    rbd = B.shape[2] # Right bond dimension
    theta = np.tensordot(A,U,axes=(1,2))
    theta = np.tensordot(theta,B,axes=((1,4),(0,1)))
    theta = np.reshape(theta,(lbd*d,rbd*d))
    (a,b) = efficientSVD(theta,chi)
    a = np.reshape(a,(lbd,d,a.shape[1]))
    b = np.reshape(b,(b.shape[0],d,rbd))
    return (a,b)

def leftPair(A,B,U,chi,d):
    """
    Evolve a pair of spins on left.
    """
    rbd = B.shape[2] # Right bond dimension
    theta = np.tensordot(A,U,axes=(0,2))
    theta = np.tensordot(theta,B,axes=((0,3),(0,1)))
    theta = np.reshape(theta,(d,d*rbd))
    (a,b) = efficientSVD(theta,chi)
    b = np.reshape(b,(b.shape[0],d,rbd))
    return (a,b)

def rightPair(A,B,U,chi,d):
    """
    Evolve a pair of spins on right.
    """
    lbd = A.shape[0] # Left bond dimension
##    show(A.shape,"A")
##    show(B.shape,"B")
##    show(U.shape,"U")
    theta = np.tensordot(A,U,axes=(1,2))
##    show(theta.shape,"A*U")
    theta = np.tensordot(theta,B,axes=((1,4),(0,1)))
##    show(theta.shape,"A*U*B")
    theta = np.reshape(theta,(lbd*d,d))
    (a,b) = efficientSVD(theta,chi)
    a = np.reshape(a,(lbd,d,a.shape[1]))
    return (a,b)

def leftSingle(A,Ub):
    """
    Evolve a single spin on left end.
    """
##    show(A.shape,"leftSingleA")
##    show(Ub.shape,"leftSingleUb")
    return np.tensordot(Ub,A,axes=(1,0))

def rightSingle(A,Ub):
    """
    Evolve a single spin on right end.
    """
    return np.tensordot(A,Ub,axes=(1,1))

def pairHamiltonianMPO():
    X = s(3)
    XX = np.kron(X,X)
    (a,b) = efficientSVD(XX,10)
    (a,b) = (np.reshape(a,(2,2,4)),np.reshape(b,(4,2,2)))
##    print(np.reshape(np.tensordot(a,b,(-1,0)),(4,4)))
    return (a,b)

def getEnergy(J,hx,N,mps):
    """
    Energy <a|H|a> of a state |a> by transfer matrices.

    Parameters
    ----------
    J : float
        Coupling constant.
    hx : float
        Product of magnetic moment and field.
    N : int
        Number of spins.
    mps : list
        MPS representation of state.
    """
    # Local energy
    I = np.reshape(s(0),(1,2,2,1))
    X = np.reshape(s(3),(1,2,2,1))
    Z = np.reshape(s(1),(1,2,2,1))
    localEnergy = 0
    for i in xrange(0,N-1):
        hamiltonian = [I for x in xrange(N)]
##        show((hamiltonian),"Hamiltonian MPO")
##        show(N,"N")
##        show(i,"i")
##        show(hamiltonian[i],"hamiltonian[i]")
##        show(hamiltonian[i+1],"hamiltonian[i+1]")
        hamiltonian[i] = X
        hamiltonian[i+1] = X
        hamiltonian[0] = np.reshape(hamiltonian[0],(2,2,1))