#!/usr/bin/env python3

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

def Main():

    to = 1
    te = 1

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

    for L in x:
        u = L//2

        E, V = ModelSolution(L, to, te)
        Lambda = CorrMatrix(E, V, u)
        S = Entropy(Lambda)

        y.append(S)

    plt.plot(x, y)
    plt.show()

    return

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

    ### Fill in M here ###

    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):
    C = np.zeros((u, u))

    ### Fill in C here ###

    Lambda, _ = la.eigh(C)
    return Lambda

def Entropy(Lambda):
    # removing the problematic limits Lambda -> 0,1
    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) )

    ### optionally, calculate the entanglement spectrum s_i ###
    # return s_i, S
    return S

Main()