#!/usr/bin/env python3

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

def Main():

    x = [2**n for n in range(3, 12)]
    log_x = np.log(x)

    for v in [0, 0.05, 0.1, 0.3]:
        print(f'v = {v}')
        y = []
        for N in x:

            s, S = Solution(N, v)

            y.append(S)

            print(f'{N}\t{S:.4f}')

        p, cov = np.polyfit(log_x, y, 1, cov=True)

        err = np.sqrt(np.diag(cov))
        if err[0]/p[0] < 0.01:
            plt.plot(x, np.polyval(p, log_x), color='tab:gray', lw=0.5)

        plt.plot(x, y, '.', label=f'{v:.2f}')

    plt.xscale('log')
    plt.ylabel('S')
    plt.xlabel('N')
    plt.legend(title='$V$')
    plt.savefig('ex1.pdf')

    return

def Solution(N, V):
    H = np.zeros((N+1, N+1))

    for i in range(N):
        H[i, i+1] = -np.sqrt(i+1)*np.sqrt(N-i)/np.sqrt(N)
        H[i+1, i] = -H[i, i+1]
        H[i, i] = V*(N-2*i)
    H[-1, -1] = -V*N

    E, V = la.eigh(H)

    psi = V[:, 0]
    s = [p**2*np.log(p**2) for p in psi if abs(p) > 1e-14]
    S = -np.sum(s)

    return s, S

Main()