#!/usr/bin/env python3

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

### Gamma function
from scipy.special import gamma

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

    dists = ['uniform', 'normal', '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)
                ### Calculate the average of the error sigma
                err[i, j, k] = ###

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

    ### Plotting the data
    for j in range(len(alpha)):
        a = alpha[j]
        for i in range(len(dists)):
            plt.plot(1/x, err[i][j], marker='.', linestyle='--', linewidth=0.5, label=dists[i])
            # plt.semilogx(...)
            # plt.loglog(...)
        plt.xlabel('1/m')
        plt.ylabel('$\\bar{\\sigma}$')
        plt.legend()
        plt.title(f'$\\alpha = {a}$')
        plt.show()
        plt.close()

    return

def Prediction(m, alpha):
    if alpha == 1:
        return 
    else:
        return 

def Entanglement_Entropy(s, alpha):
    if alpha == 1:
        return 
    else:
        return 

def Entanglement_Spectrum(m, dist):

    if dist == 'uniform':
        ###
    elif dist == 'normal':
        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':
        ###
    else:
        ### try another distribution here         ###
        ### for example, try to break convergence ###

    ### Normalize the state ###

    _, s, _ = la.svd(A)

    return s

Main()