#!/usr/bin/env python3

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

def Main():

    x = [2**n for n in range(3,11)]

    for to in [1.1, 1.01, 1.001, 1]:
        print(f't_o = {to}')

        te = 1

        y = []

        for L in x:

            filling=L//2

            u = filling

            E, V = ModelSolution(L, to, te)

            Lambda = CorrMatrix(E, V, u, filling)
            S = Entropy(Lambda)

            y.append(S)
            print(f'{L: <4} {S:.12f}')

        plt.plot(x, y, label=f'${to:.3f}$', lw=0.5, marker='.')

    plt.xscale('log')
    plt.xlabel('$L$')
    plt.ylabel('$S$')
    plt.legend(title='$t_o$')
    plt.show()

    return

def ModelSolution(L, to, te):
    M = np.zeros((L, L))

    for l in range(0, L-1, 2):
        M[l, l+1] = to
        M[l+1, l] = to

    for l in range(1, L-1, 2):
        M[l, l+1] = te
        M[l+1, l] = te

    E, V = la.eigh(-M)
    # E is the vector of eigenvalues, from smallest to largest
    # V is the matrix such that V * diag(E) * V.T = M

    return E, V

def CorrMatrix(E, V, u, fill):
    C = np.zeros((u, u))

    for i in range(u):
        for j in range(i, u):
            C[i, j] = np.sum(V[i, :fill].conj() * V[j, :fill])
            C[j, i] = C[i, j]

    Lambda, _ = la.eigh(C)

    return Lambda

def Entropy(Lambda):

    Lambda = np.array([l for l in Lambda if 1e-12 < l < 1-1e-12])
    S = -np.sum( Lambda*np.log(Lambda) + (1-Lambda)*np.log(1-Lambda) )

    return S

Main()