#!/usr/bin/env python3

import matplotlib.pyplot as plt
import scipy.linalg as la
import numpy as np

def Main():

    h = 1

    x = 2**np.arange(4, 11)

    energy_gap = []
    S = []

    print('L\tΔ')
    for L in x:

        E, X, Y = ModelSolution(L, h)

        energy_gap.append( 2*E[0] )
        print(f'{L}\t{energy_gap[-1]}')

    p, cov = np.polyfit(np.log(1/x), np.log(energy_gap), 1, cov=True)
    err = np.sqrt(np.diag(cov))
    if err[0]/p[0] < 0.05:
        plt.loglog(1/x, np.exp(p[1])*x**-p[0], lw=0.5, color='tab:gray')

    plt.loglog(1/x, energy_gap, marker='.', linestyle='')
    plt.xlabel('$1/L$')
    plt.ylabel('$\\delta$')
    plt.title(f'$\\alpha = {p[0]:.2f}$')
    plt.savefig('E_vs_L.pdf')
    plt.close()

    print('\nl\tS')
    l = x[:-1]
    for u in l:
        S.append( Entropy(CorrMatrix(X, Y, u)) )
        print(f'{u}\t{S[-1]}')

    X = np.log(2*L/np.pi*np.sin(np.pi*l/L))/6

    p, cov = np.polyfit(X, S, 1, cov=True)
    err = np.sqrt(np.diag(cov))
    if err[0]/p[0] < 0.05:
        plt.plot(X, np.polyval(p, X), lw=0.5, color='tab:gray')

    plt.plot(X, S, marker='.', linestyle='')
    plt.xlabel('$\\log(2L/\\pi \\sin(\\pi l/L))/6$')
    plt.ylabel('$S_l$')
    plt.title(f'$c = {p[0]:.2f}$')
    plt.savefig('S_vs_l.pdf')
    plt.close()

    return

def ModelSolution(L, h):
    A = np.zeros((L, L))
    B = np.zeros((L, L))

    for l in range(L-1):
        B[l, l+1] = -0.5
        B[l+1, l] = 0.5

        A[l, l+1] = -0.5
        A[l+1, l] = -0.5
        A[l, l] = h
    A[-1, -1] = h

    M = np.zeros((2*L, 2*L))
    M[:L,:L] = -A
    M[L:,L:] = A
    M[:L,L:] = B
    M[L:,:L] = -B

    E, U = la.eigh(M, subset_by_index=[L, 2*L-1])
    # E is the vector of eigenvalues, from smallest to largest
    # U is the matrix such that U * diag(E) * U^T = M

    Y = U[:L, :]
    X = U[L:, :]

    return E, X, Y

def CorrMatrix(X, Y, u):
    A = X[:u, :]
    B = Y[:u, :]

    M = np.zeros((2*A.shape[0], 2*A.shape[0]))
    M[:u,:u] = np.dot(A, A.T)
    M[u:,u:] = np.dot(B, B.T)
    M[:u,u:] = np.dot(A, B.T)
    M[u:,:u] = np.dot(B, A.T)

    return M

def Entropy(C):
    Lambda = la.eigh(C, eigvals_only=True, subset_by_value=[1e-4, 1-1e-4])
    return -np.sum( Lambda*np.log(Lambda) )

Main()