#!/usr/bin/env python3

from scipy.sparse import csr_matrix
import matplotlib.pyplot as plt
import scipy.linalg as la
import numpy as np
from pathlib import Path

def Main():

    Lx = 4
    Ly = 4
    N = Lx*Ly

    n = 20

    dt = 0.05
    tf = 5
    t = np.arange(0, tf+1e-5, dt)

    X = [0.1, 0.3, 0.5, 0.7, 0.9]
    States = {
        'ferro':    '0'*N,
        'row':      '1111000000000000',
        'T':        '1110010000000000',
        'triangle': '1010010000000000',
        'Y':        '1010010001000000',
        'diag':     '1000010000100001',
        'cols':     '01'*(N//2),
    }

    H0 = SpinChain(Lx, Ly, 0)

    StateInt = {s: int(States[s], 2) for s in States}
    E0 = {s: H0[StateInt[s], StateInt[s]]/N for s in States}

    for s in States:
        print(f'{s}: ε = {E0[s]:+.2f}')
        Draw(States[s], Lx, Ly)

    n_up, n_down = Observables(N)
    m_up = np.zeros(len(t))
    m_down = np.zeros(len(t))

    fig, ax = plt.subplots(1, 1)

    for x in X:

        H = SpinChain(Lx, Ly, x)

        for s in States:

            ax.clear()

            psi = np.zeros(2**N)
            psi[StateInt[s]] = 1

            file = f'm_x_{x}_E_{E0[s]:+.2f}_{s}.pdf'
            if Path(file).is_file():
                continue

            for i in range(len(t)):
                if i:
                    psi = TimeEvolve(H, psi, n, dt)
                m_up[i] = Measure(n_up, psi)/N
                m_down[i] = Measure(n_down, psi)/N
                Print(f'{s}:\t{x}\t{t[i]:.2f}\t{m_up[i]:+.4f}\t{m_down[i]:+.4f}')
            Print('\r')

            ax.plot(t, m_up,        label='$n_{\\uparrow}$',   lw=0.2, marker='.', ms=0.3)
            ax.plot(t, m_down,      label='$n_{\\downarrow}$', lw=0.2, marker='.', ms=0.3)
            ax.plot(t, m_up-m_down, label='$m_z$',             lw=0.2, marker='.', ms=0.3)

            ax.legend(loc='center left', bbox_to_anchor=(1, 0.5))
            ax.set_xlabel('$t$')
            ax.set_title(f'$\\mathrm{{{s}}}: ε = {E0[s]:+.2f}$')
            plt.savefig(file)

    return

def Print(mess):
    print(f'\r{" "*80}', end='')
    print(f'\r{mess}', end='')
    return

def Draw(st, Lx, Ly):
    for l in range(Ly):
        s = st[l*Lx:(l+1)*Lx]
        print('\t'+' '.join(s))
    print()
    return

def Observables(N):
    inds = range(2**N)
    def mat(diag):
        return csr_matrix((diag, (inds, inds)), dtype=int)

    n_up =   [n.bit_count() for n in inds]
    n_down = [N-n.bit_count() for n in inds]
    return mat(n_up), mat(n_down)

def Measure(O, v):
    return np.real(v.T @ O @ v)

def TimeEvolve(H, v, n, dt):
    E, U, V = Lanczos(H, v, n)
    return V @ U @ np.diag(np.exp(-1j * dt * E)) @ U.T[:, 0]

def Lanczos(H, v, n):

    a = []
    b = []
    V = [v]

    def energy(vec):
        Hv = H @ vec
        en = np.abs(vec.T.conj() @ Hv)
        a.append(en)
        return Hv

    def norm(vec):
        no = np.sqrt(np.abs(vec.T.conj() @ vec))
        b.append(no)
        V.append(vec/no)
        return


    Hv = energy(V[-1])
    vp = Hv - a[-1]*V[-1]
    norm(vp)

    for j in range(2, n):
        Hv = energy(V[-1])
        vp = Hv - a[-1]*V[-1] - b[-1] * V[-2]
        norm(vp)

    energy(V[-1])

    V = np.array(V)

    E, U = la.eigh_tridiagonal(a, b)
    return E, U, V.T

def SpinChain(Lx, Ly, x):

    N = Lx*Ly

    J = np.cos(np.pi*x/2)
    hx = -np.sin(np.pi*x/2)

    def flip_spin(n, k):
        return n ^ (1 << k)

    H = ([], ([], []))

    # We will store non-zero matrix elements of H: H[i, j] = v
    def add(i, j, v):
        H[0].append(v)
        H[1][0].append(i)
        H[1][1].append(j)

    # Iterate over all configurations
    for i in range(2**N):

        # string of bits of configuration i
        b = f'{i:b}'.rjust(N, '0')

        Jt = 0

        for l in range(Lx):
            for c in range(Ly):

                j = flip_spin(i, l+Lx*c)
                add(i, j, -hx)

                s1 = b[l+Lx*c]

                lp, cp = (l+1)%Lx, (c+1)%Ly
                ps = [lp+Lx*c, l+Lx*cp]
                for p in ps:
                    Jt += -1 if s1 == b[p] else 1
        add(i, i, J*Jt)

    return csr_matrix(H, dtype=np.float64)

Main()