#!/usr/bin/env python3

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

np.seterr(invalid='ignore')

def Main():

    cluster = 'circle'

    coords = []

    if cluster == 'horizontal':
        L = np.arange(50, 1001, 50)
        for l in L:
            coords += [[ [0, y] for y in range(l) ]]
    elif cluster == 'diagonal':
        L = np.arange(50, 1001, 50)
        for l in L:
            coords += [[ [x, x] for x in range(l) ]]
    elif cluster == 'ladder':
        L = np.arange(25, 501, 25)
        for l in L:
            coords += [[ [x, y] for y in range(l) for x in range(2) ]]
    elif cluster == 'square':
        L = np.arange(5, 61, 5)
        for l in L:
            coords += [[ [x, y] for y in range(l) for x in range(l) ]]
    elif cluster == 'disconnected_square':
        L = np.arange(10, 61, 5)
        for l in L:
            coords += [[ [2*x, 2*y] for y in range(l) for x in range(l) ]]
    elif cluster == 'circle':
        L = np.arange(5, 46, 5)
        for r in L:
            coords += [[ [x, y] for x in range(-r, r+1) for y in range(-r, r+1) if x**2 + y**2 <= r**2 ]]

    A = []
    S = []

    print('L\tA\tS')
    for i in range(len(coords)):
        A.append(len(coords[i]))
        S.append(Entropy(coords[i]))
        print(f'{L[i]}\t{A[-1]}\t{S[-1]:.2f}')

    tol = 0.005

    p, cov = np.polyfit(L, S, 1, cov=True)
    err = np.sqrt(np.diag(cov))
    if err[0]/p[0] < tol:
        plt.plot(L, np.polyval(p, L), lw=0.5, color='tab:gray', label=f'$\\sim {p[0]:.2f}L$')
        plt.legend()

    plt.plot(L, S, marker='.', linestyle='')
    plt.xlabel('$L$')
    plt.ylabel('$S$')
    plt.savefig(f'ex3_{cluster}_L.pdf')
    plt.close()

    p, cov = np.polyfit(L*np.log(L), S, 1, cov=True)
    err = np.sqrt(np.diag(cov))
    if err[0]/p[0] < tol:
        plt.plot(L*np.log(L), np.polyval(p, L*np.log(L)), lw=0.5, color='tab:gray', label=f'$\\sim {p[0]:.2f}L\\log L$')
        plt.legend()

    plt.plot(L*np.log(L), S, marker='.', linestyle='')
    plt.xlabel('$L\\log L$')
    plt.ylabel('$S$')
    plt.savefig(f'ex3_{cluster}_LlogL.pdf')
    plt.close()

    p, cov = np.polyfit(A, S, 1, cov=True)
    err = np.sqrt(np.diag(cov))
    if err[0]/p[0] < tol:
        plt.plot(A, np.polyval(p, A), lw=0.5, color='tab:gray', label=f'$\\sim {p[0]:.2f}A$')
        plt.legend()

    plt.plot(A, S, marker='.', linestyle='')
    plt.xlabel('$A$')
    plt.ylabel('$S$')
    plt.savefig(f'ex3_{cluster}_A.pdf')
    plt.close()

    return

def Entropy(coords):

    ### vectorized calculation of the matrix C
    X = np.tile([c[0] for c in coords], (len(coords), 1))
    Y = np.tile([c[1] for c in coords], (len(coords), 1))

    DX, DY = X - X.T, Y - Y.T

    W0 = np.pi*(DX + DY)/2
    W1 = np.pi*(DX - DY)/2

    ### nan_to_num to resolve the W -> 0 limit
    Z0 = np.nan_to_num(np.sin(W0)/W0, nan=1)
    Z1 = np.nan_to_num(np.sin(W1)/W1, nan=1)

    C = Z0*Z1/2

    ### To speed up execution, we don't need all eigenvalues of C,
    ### as eigenvalues close to 0,1 have a very small contribution to S.
    Lambda = la.eigh(C, eigvals_only=True, subset_by_value=[1e-3, 1-1e-3])

    S = -np.sum( Lambda*np.log(Lambda) + (1-Lambda)*np.log(1-Lambda) )

    return S

Main()