#!/usr/bin/env python3
"""Module for constructing different Hamiltonians.

This module implements functions to construct Hamiltonian matrices
for different models,
e.g. SYK_2, SYK_4, and Rosenzweig-Porter model.
"""

import numpy as np
from scipy.special import binom
from scipy import sparse
import os
import utils
from numba import njit


PATH = os.path.dirname(os.path.realpath(__file__))


def _syk2_rem(n, seed, as_csr, projected, split):
    """Return single realizationg of SYK_2 Hamiltonian.

    To create H_2 following the diagonal sum form
    we need to generate N i.i.d. v_i's
    with gassian distribution with width \delta^2 \gamma^2.
    But here \gamma = 1,
    and since we put \delta as a parameter in the Hamiltonian
    outside of the definition of SYK_2
    then the width here should be 1.

    Returns (H_pos, H_neg).

    Notes
    -----
    Uses $\langle J_{ijkl}^2 \rangle = \frac{3}{2 n^4}$ for $H_4$,
    and $\langle v_{\alpha}^2 \rangle = \gamma^2 for $H_2$.

    This way, variance of eigenvalues of $H_2$ is n\gamma^2,
    where \gamma^2 = \frac{3}{2n^4} \binom{2n}{4},
    and variance of eigenvalues of $H_4$ is $\gamma^2$.

    Parameters
    ----------
    n : int
    system size

    seed : int
    seed for the random number generator

    as_csr : bool
    whether to use sparse matrices

    projected : bool
    If true, project onto constant parity sectors,
    i.e. block diagonal form.

    split : bool
        whether to return a single, block diagonal, matrix,
        or to return two matrices, one for each block.

    Returns
    -------
    H : ndarray
        SYK_2 Hamiltonian.
    """

    np.random.seed(seed)
    dim = int(2**n)

    # Set variance of the model.
    J_REM = np.sqrt(3 * binom(2 * n, 4) / (2 * n ** 3)) / np.sqrt(n)

    # Generate diagonal disorder.
    H = J_REM * np.diag(np.random.normal(scale=1.0, size=dim))

    # Transform to sparse matrix form.
    H = J_REM * sparse.diags(np.random.normal(scale=1.0, size=dim), format="csr")

    if split:
        H = utils.split_blocks(H)

    return H


def _syk2(n, seed, couplings=None):
    """Return single realizationg of SYK_2 Hamiltonian.

    To create H_2 following the diagonal sum form
    we need to generate n i.i.d. v_i's
    with gassian distribution with width \delta^2 \gamma^2.
    But here \gamma = 1,
    and since we put \delta as a parameter in the Hamiltonian
    outside of the definition of SYK_2
    then the width here should be 1.

    Returns (H_pos, H_neg).

    Notes
    -----
    Uses $\langle J_{ijkl}^2 \rangle = \frac{3}{2 n^4}$ for $H_4$,
    and $\langle v_{\alpha}^2 \rangle = \gamma^2 for $H_2$.

    This way, variance of eigenvalues of $H_2$ is n\gamma^2,
    where \gamma^2 = \frac{3}{2n^4} \binom{2n}{4},
    and variance of eigenvalues of $H_4$ is $\gamma^2$.

    Parameters
    ----------
    n : int
    system size

    seed : int
    seed for the random number generator

    as_csr : bool
    whether to use sparse matrices

    projected : bool
    If true, project onto constant parity sectors,
    i.e. block diagonal form.

    split : bool
        whether to return a single, block diagonal, matrix,
        or to return two matrices, one for each block.

    Returns
    -------
    H : ndarray
        SYK_2 Hamiltonian.
    """
    # Load the products of Majorana matrices.
    product_of_2_majorana = utils.load_saved_majoranas(
        N=n,
        k=2,
        as_csr=True,
        projected=True,
        split=True,
    )

    # Get the total number of terms in the SYK hamiltonian
    # as the number of different products of majorana matrices.
    # This is equal to \binom{2n}{k}.
    size = len(product_of_2_majorana)

    # number_op = utils.number_operator(N=n)
    # size = len(number_op)

    if couplings is None:
        np.random.seed(seed)
        # Set variance of the model.
        J_2 = np.sqrt((2 * n - 2) * (2 * n - 3) / (8 * n ** 3))
        # Generate random masses.
        random_syk_J = 1j * J_2 * np.random.normal(scale=1.0, size=size)
        # random_syk_J = np.random.normal(scale=1.0, size=size)
    else:
        assert len(couplings) == binom(2*n, 2)
        # assert len(couplings) == n
        random_syk_J = 1j * couplings
        print(random_syk_J[:5])

    np.savetxt(f'N{n}-j2.txt', np.imag(random_syk_J))

    # Multiply each product of majoranas by a random coupling.
    H = [random_syk_J[i] * product_of_2_majorana[i]
            for i in range(size)]

    # identity = sparse.identity(2**n)
    # H = [random_syk_J * (2*number_op[i] - identity)
    #         for i in range(size)]

    H = utils.matrix_sum_multiproc(H)
    H = np.array([h.toarray() for h in utils.split_blocks(H)])

    return H


def _syk4(n, seed, couplings=None):
    """Return SYK_4 Hamiltonian.

    Returns (H_pos, H_neg).

    Notes
    -----
    Uses $\langle J_{ijkl}^2 \rangle = \frac{3}{2 n^4}$ for $H_4$,
    and $\langle v_{\alpha}^2 \rangle = \gamma^2 for $H_2$.

    This way, variance of eigenvalues of $H_2$ is n\gamma^2,
    where \gamma^2 = \frac{3}{2n^4} \binom{2n}{4},
    and variance of eigenvalues of $H_4$ is $\gamma^2$.

    Parameters
    -------
    n : int
    system size

    seed : int
    seed for the random number generator

    as_csr : bool
    whether to use sparse matrices

    projected : bool
    If true, project onto constant parity sectors,
    i.e. block diagonal form.

    split : bool
        whether to return a single, block diagonal, matrix,
        or to return two matrices, one for each block.

    Returns
    -------
    H : ndarray
        SYK_4 Hamiltonian.

    """
    dim = int(2**n)
    # Load the products of Majorana matrices.
    product_of_4_majorana = utils.load_saved_majoranas(
        N=n,
        k=4,
        as_csr=True,
        projected=True,
        split=True,
    )

    # Get the total number of terms in the SYK hamiltonian
    # as the number of different products of majorana matrices.
    # This is equal to \binom{2n}{k}.
    size = len(product_of_4_majorana)

    if couplings is None:
        np.random.seed(seed)
        # Set variance of random coupling in SYK_4
        # such that bandwidth is \sim 1 for large n.
        J_4 = np.sqrt(3 / (2 * n ** 4))
        # Create all the random coupling of the SYK at once.
        random_syk_J = J_4 * np.random.normal(scale=1.0, size=size)
    else:
        assert len(couplings) == binom(2*n, 4)
        random_syk_J = couplings

    np.savetxt(f'N{n}-j4.txt', random_syk_J)
    # Multiply each product of majoranas by a random coupling.
    H = [random_syk_J[i] * product_of_4_majorana[i]
            for i in range(size)]

    H = utils.matrix_sum_multiproc(H)
    H = np.array([h.toarray() for h in utils.split_blocks(H)])

    return H


def syk(k, n, seed):
    if k not in [2, 4]:
        raise ValueError("k argument must be 2 or 4.")
    if k == 2:
        return _syk2(n, seed)
    if k == 4:
        return _syk4(n, seed)


def extended_syk(n, delta, seed_h2=0, seed_h4=None):
    h2 = syk(k=2, n=n, seed=seed_h2)
    h4 = syk(k=4, n=n, seed=seed_h4)

    return h4 + delta * h2


@njit
def rp(n, gamma, eps=1.0, seed_diag=0, seed_offdiag=None):
    """Returns a Rosenzweig-Porter Hamiltonian.

    It is of the form \(H = H_0 + \delta H_v\),
    where \(H_0\) is GUE matrix,
    and \(H_v\) is a random diagonal matrix.

    Parameters
    ----------
    n : int
    Size of the Hamiltonian

    gamma : float
    Parameter to control the regime of the model

    seed0 : int
    seed for the disorder realization of the diagonal part

    seed1 : int
    seed for the disorder realization of the off-diagonal part

    eps : float
    Additional parameter of the Rosenzweig-Porter model
    """

    h_diag = np.zeros((n, n), dtype=np.complex64)
    h_offdiag = np.zeros((n, n), dtype=np.complex64)

    if seed_diag is not None:
        np.random.seed(seed_diag)
    for i in range(n):
        h_diag[i, i] = np.random.normal()

    if seed_diag is not None:
        np.random.seed(seed_offdiag)
    for i in range(n):
        for j in range(i+1, n):
            h_offdiag[i, j] = (np.random.normal() + 1j*np.random.normal()) * eps**2/(np.sqrt(2)*np.power(n, gamma/2))
            h_offdiag[j, i] = np.conj(h_offdiag[i, j])
    return h_diag + h_offdiag
