#!/usr/bin/env python3

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

def Main():

    L = 11

    sets = {
        'I': {
            'hx': [-0.9, 0],
            'hz': [0, 0],
        },
        'NI': {
            'hx': [-1.1, 0],
            'hz': [0.73, 0],
        },
        'C': {
            'hx': [0, 0],
            'hz': [-0.47, 0.1],
        },
    }

    Sx = OSx(L)
    Sz = OSz(L)

    figx, axx = plt.subplots(1, 1)
    figz, axz = plt.subplots(1, 1)

    for l, s in sets.items():

        E, V = ED(SpinChain(L, s['hx'], s['hz']))

        sx = V.T @ Sx @ V
        sz = V.T @ Sz @ V

        axx.plot(E/L, np.diag(sx)/L, '.', markersize=0.3, label=l)
        axz.plot(E/L, np.diag(sz)/L, '.', markersize=0.3, label=l)

    for fig, ax, lbl in [(figx, axx, 'x'), (figz, axz, 'z')]:
        ax.legend(bbox_to_anchor=(0.99, 0.5), loc='center left')
        ax.set_xlabel('$\\varepsilon$')
        ax.set_ylabel(f'$\\sigma_{lbl}$')
        fig.savefig(f'S{lbl}.pdf')

    return

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 ED(H):
    s = H.shape[0]
    return la.eigh(H, subset_by_index=[s//12, s-s//12])

def OSz(L):
    sz = [2*n.bit_count()-L for n in range(2**L)]
    return np.diag(sz)

def OSx(L):
    sx = np.array( [[0, 1], [1, 0]] )
    Sx = np.zeros((2**L, 2**L))
    for l in range(L):
        Sx += np.kron(np.eye(2**l),
                np.kron(sx, np.eye(2**(L-l-1)))
            )
    return Sx

Main()