#!/usr/bin/env python3

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

def Main():

    L = 10

    hx = [0.7, 0.1]
    hz = [0, 0]

    reps = 50

    bins = 100
    bin_lims = np.linspace(0, 1, bins+1)
    bin_spacing = bin_lims[1]-bin_lims[0]
    x = (bin_lims[1:] +  bin_lims[:-1])/2

    y = GetSpacings(reps, SpinChain, L, hx, hz)
    y = np.histogram(y/max(y), bins=bin_lims)[0].astype('float64')
    y /= np.sum(y)

    x_av = np.sum(y*x)

    yRM = GetSpacings(reps, RandomMatrix, L)
    yRM *= x_av/np.average(yRM)
    yRM = np.histogram(yRM, bins=bin_lims)[0].astype('float64')
    yRM /= np.sum(yRM)

    yW = (x/x_av)*np.exp(-(x/x_av)**2)
    yW /= np.sum(yW)

    plt.plot(x, y/bin_spacing,   marker='.', lw=0.5, label='Ising')
    plt.plot(x, yRM/bin_spacing, marker='.', lw=0.5, label='RM')
    plt.plot(x, yW/bin_spacing,  marker='.', lw=0.5, label='Wigner')

    plt.xlabel('$\\delta \\varepsilon/max(\\delta \\varepsilon)$')
    plt.ylabel('$p(\\delta \\varepsilon)$')
    plt.xlim(0, 0.4)
    plt.legend()
    plt.savefig('spacings.pdf')

    return

def GetSpacings(reps, fun, *args):
    y = np.array([])
    for j in range(reps):

        E = ED(fun(*args))
        s = np.diff(E)

        y = np.append(y, s)
    return y

def SpinChain(L, hx, hz):

    H = np.zeros((2**L, 2**L))

    # Get the k-th spin of integer configuration n
    def get_spin(n, k):
        b = (n & 1 << k) != 0
        return 2*b-1

    # "flip" the k-th spin of integer configuration n
    def flip_spin(n, k):
        return n ^ (1 << k)

    hx = hx[0]+hx[1]*np.random.randn(L)
    hz = hz[0]+hz[1]*np.random.randn(L)

    # Iterate over all configurations
    for i in range(2**L):

        s2 = get_spin(i, 0)

        for l in range(L-1):

            s1 = s2
            s2 = get_spin(i, l+1)

            H[i, i] -= s1*hz[l]

            j = flip_spin(i, l)
            H[i, j] -= hx[l]

            if s1 == s2:
                H[i, i] -= 1
            else:
                H[i, i] += 1

        H[i, i] -= s2*hz[-1]
        j = flip_spin(i, L-1)
        H[i, j] -= hx[-1]

    return H

def RandomMatrix(L):
    R = np.zeros((2**L, 2**L))
    for i in range(2**L):
        R[i, i] = np.random.normal() # scale=1
        for j in range(i+1, 2**L):
            R[i, j] = np.random.normal()
            R[j, i] = R[i, j]
    return R

def ED(H):
    s = H.shape[0]
    return la.eigh(H, eigvals_only=True, subset_by_index=[s//12, s-s//12])

Main()