#!/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,
        'triangle': '1010010000000000',
        'cols':     '01'*(N//2),
    }

    # Ising term
    H0 = SpinHamiltonian(Lx, Ly, 0)

    # integers corresponding to bit configurations
    StateInt = {s: int(States[s], 2) for s in States}
    # Ising energies of configurations
    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)

    # Operators of number of spins up and down
    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:

        # Full Hamiltonian
        H = SpinHamiltonian(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)
            # Plot Sz magnetization
            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, vec, n, dt):
    # E: eigenvalues of T
    # U: eigenvectors of T
    # V: Lanczos vectors
    E, U, V = Lanczos(H, vec, n)

    E = np.diag(np.exp(-1j * dt * E))
    ### Complete here
    M = ...
    return M

def Lanczos(H, v, n):

    energies = []
    norms = []
    V = [v]

    # @ is matrix or matrix-vector multiplication of numpy/scipy arrays
    def energy(vec):
        Hv = H @ vec
        ...
        return Hv

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

    ...

    V = np.array(V)

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

# Returns spin Hamiltonian matrix
def SpinHamiltonian(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')

        # off-diagonal elements implemented
        # diagonal terms missing
        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]
                ...

    return csr_matrix(H, dtype=np.float64)

Main()