#!/usr/bin/env python3

import matplotlib.pyplot as plt
import scipy.linalg as la
import numpy as np
from scipy.special import gamma

def Main():
    K_max = 2**16
    x = 2**np.arange(3, 11)

    dists = ['uniform', 'gaussian', 'chi1', 'chi2']
    alpha = [1, 2, 3]

    err = np.zeros((len(dists), len(alpha), len(x)))

    for i in range(len(dists)):
        dist = dists[i]
        print(dist)
        print(f'm\tα=1\tα=2\tα=3')
        for k in range(len(x)):
            m = x[k]

            K = K_max//m
            s = [Entanglement_Spectrum(m, dist) for k in range(K)]

            for j in range(len(alpha)):
                a = alpha[j]
                S = [Entanglement_Entropy(l, a) for l in s]
                P = Prediction(m, a)
                err[i][j][k] = np.average(abs(S/P-1))

            print(f'{m}\t{err[i][0][k]:.1e}\t{err[i][1][k]:.1e}\t{err[i][2][k]:.1e}')

    for j in range(len(alpha)):
        a = alpha[j]
        for i in range(len(dists)):
            plt.loglog(1/x, err[i][j], marker='.', linestyle='--', lw=0.5, label=dists[i])
        plt.xlabel('1/m')
        plt.ylabel('$\\bar{E}$')
        plt.legend()
        plt.title(f'$\\alpha = {a}$')
        plt.savefig(f'ex2_alpha_{a}.pdf')
        plt.close()

    return

def Prediction(m, alpha):
    if alpha == 1:
        return np.log(m)-1/2
    else:
        return np.log(m)-np.log(
                    4**alpha
                    *gamma(alpha+0.5)
                    /np.sqrt(np.pi)
                    /gamma(alpha+2)
                )/(alpha-1)

def Entanglement_Entropy(s, alpha):
    if alpha == 1:
        return -np.sum(2*s**2*np.log(s))
    else:
        return np.log(np.sum(s**(2*alpha)))/(1-alpha)

def Entanglement_Spectrum(m, dist):

    if dist == 'uniform':
        A = 2*np.random.rand(m, m)-1
    elif dist == 'gaussian':
        A = np.random.randn(m, m)
    elif dist == 'chi1':
        z = np.exp(2j*np.pi*np.random.rand(m, m))
        r = np.sqrt(np.random.chisquare(1, size=(m, m)))
        A = r*z
    elif dist == 'chi2':
        z = np.exp(2j*np.pi*np.random.rand(m, m))
        r = np.sqrt(np.random.chisquare(2, size=(m, m)))
        A = r*z

    A /= np.sqrt(np.sum(A*A.conj()))

    s = la.svd(A, compute_uv=False)

    return s

Main()