#!/usr/bin/env python

"""Compute eigenvalues and statistics for the quantized perturbed cat map.

Version 0.4, 01.06.2013

Copyright (C) 2002, 2013 Arnd Baecker.


Usage: call with N and kappa given::

   python pert_cat.py 51 0.3

"""

from __future__ import (division, print_function)

import sys
from math import pi
import numpy as np
from matplotlib import pyplot as plt

# pylint: disable=C0103


def quantum_cat(N, kappa):
    """Return unitary matrix of the quantized perturbed cat map.

    The matrix size is given by `N`and the strength of the perturbation
    by the parameter `kappa`.
    """

    print("Filling matrix ... N=%d  kappa=%g" % (N, kappa))
    # One could use the following lines:
    # I=1j                        # predefine sqrt(-1)
    # # First initialize complex matrix with NxN elements
    # mat = zeros((N, N), np.complex)
    # for k in xrange(0, N):
    #     for l in xrange(0, N):
    #         mat[k, l] = cmath.exp(2.0*I*pi/N*(k*k-k*l+l*l) + \
    #                          1j * kappa*N/2.0/pi*sin(2.0*pi/N*l))/sqrt(N)
    #
    # However, the following is much faster (thanks to Fernando Perez)
    alpha = 2.0*pi/N
    mat_fn = lambda k, l: alpha*(k*k-k*l+l*l)
    phi = np.fromfunction(mat_fn, (N, N)) + \
        (kappa/alpha) * np.sin(alpha*np.arange(N))
    mat = np.exp(1j*phi)/np.sqrt(N)    # note that 1j=sqrt(-1)
    print("# Fill done")
    return mat


def compute_eigenphases_pcat(N, kappa):
    """Compute eigenphases for the quantized perturbed cat map.

    For a given `N` and `kappa` this functions
       returns the eigenvalues and eigenphase
       of the unitary matrix U filled via quantum_cat(N,kappa).
    """
    # fill matrix U_N
    mat = quantum_cat(N, kappa)
    # determine the eigenvalues of the matrix U_N
    evals = np.linalg.eigvals(mat)

    # determine phase \in [0, 2\pi] from the eigenvalues
    phases = np.arctan2(evals.imag, evals.real) + pi
    return phases


def goe_approx(x):
    """Approximation to the GOE level spacing distribution."""
    return pi/2.0*x*np.exp(-pi/4*x*x)


def gue_approx(x):
    """Approximation to the GUE level spacing distribution."""
    return 32/pi/pi*x*x*np.exp(-4/pi*x*x)


def plot_levelspacing_distribution(phases):
    """Plot the level spacing distribution for the provided `eigenphases`.

    In addition the random matrix distributions for the GOE and GUE
    and the Poissonian distribution are plotted.

    Note, that the eigenphases are assumed to be in the interval [0, 2*pi[.

    Only for odd matrix size, a GOE distribution is expected
    in the limit of large matrix sizes.
    """
    N = len(phases)
    phases = np.concatenate([phases, np.array([phases[0]+2.0*pi])])
    phases.sort()

    # Unfold to have average spacing 1
    phases = phases*N/(2.0*pi)

    # determine Level spacing
    # (by computing the difference of the shifted eigenphases)
    spacings = phases[1:] - phases[0:N]

    print("Plottting....")
    plt.hist(spacings, bins=50, range=[0, 4], normed=True)
    x = np.linspace(0.0, 4.0, 70)
    plt.plot(x, goe_approx(x), lw=2, label="GOE")
    plt.plot(x, gue_approx(x), lw=2, label="GUE")
    plt.plot(x, np.exp(-x), lw=2, label="Poisson")
    plt.xlabel("s")
    plt.ylabel("P(s)")
    plt.title("Level spacing distribution, perturbed cat map,  N=%d" % N)
    plt.legend()
    plt.show()


###############################################################################
# If called as script, run a diagonalization and plot the result
###############################################################################

def main():
    """Run the diagonalization."""
    N = 601
    kappa = 0.3
    if len(sys.argv) == 3:
        # get values from commandline
        N = int(sys.argv[1])
        kappa = float(sys.argv[2])

    phases = compute_eigenphases_pcat(N, kappa)
    plot_levelspacing_distribution(phases)


if __name__ == "__main__":
    main()