#################
#   AA Solvers  #
#################
from itertools import product

import numpy as np
from scipy.linalg import solve


def check_each_direction(n, angs, ifprint=True):
    """returns a list of the index of elements of n which do not have adequate
    toy angle coverage. The criterion is that we must have at least one sample
    in each Nyquist box when we project the toy angles along the vector n"""
    checks = np.array([])
    P = np.array([])
    if ifprint:
        pass
    for _k, i in enumerate(n):
        np.linalg.norm(i)
        X = np.dot(angs, i)
        if (np.abs(np.max(X) - np.min(X)) < 2.0 * np.pi) or (
            np.abs(np.max(X) - np.min(X)) / len(X) > np.pi
        ):
            if ifprint:
                pass
            checks = np.append(checks, i)
            P = np.append(P, (2.0 * np.pi - np.abs(np.max(X) - np.min(X))))
    if ifprint:
        pass
    return checks, P


def solver(AA, N_max, symNx=2, throw_out_modes=False):
    """Constructs the matrix A and the vector b from a timeseries of toy
    action-angles AA to solve for the vector x = (J_0, J_1, J_2, S...) where
    x contains all Fourier components of the generating function with |n|<N_max"""

    # Find all integer component n_vectors which lie within sphere of radius N_max
    # Here we have assumed that the potential is symmetric x->-x, y->-y, z->-z
    # This can be relaxed by changing symN to 1
    # Additionally due to time reversal symmetry S_n = -S_-n so we only consider
    # "half" of the n-vector-space

    angs = unroll_angles(AA.T[3:].T, np.ones(3))

    symNz = 2
    NNx = range(-N_max, N_max + 1, symNx)
    NNy = range(-N_max, N_max + 1, symNz)
    NNz = range(-N_max, N_max + 1, symNz)
    n_vectors = np.array(
        [
            [i, j, k]
            for (i, j, k) in product(NNx, NNy, NNz)
            if (
                not (i == 0 and j == 0 and k == 0)  # exclude zero vector
                and (
                    k > 0  # northern hemisphere
                    or (k == 0 and j > 0)  # half of x-y plane
                    or (k == 0 and j == 0 and i > 0)
                )  # half of x axis
                and np.sqrt(i * i + j * j + k * k) <= N_max
            )
        ]
    )  # inside sphere

    check_each_direction(n_vectors, angs)

    if throw_out_modes:
        n_vectors = np.delete(n_vectors, check_each_direction(n_vectors, angs), axis=0)

    n = len(n_vectors) + 3
    b = np.zeros(shape=(n,))
    a = np.zeros(shape=(n, n))

    a[:3, :3] = len(AA) * np.identity(3)

    for i in AA:
        a[:3, 3:] += 2.0 * n_vectors.T[:3] * np.cos(np.dot(n_vectors, i[3:]))
        a[3:, 3:] += (
            4.0
            * np.dot(n_vectors, n_vectors.T)
            * np.outer(
                np.cos(np.dot(n_vectors, i[3:])), np.cos(np.dot(n_vectors, i[3:]))
            )
        )
        b[:3] += i[:3]
        b[3:] += 2.0 * np.dot(n_vectors, i[:3]) * np.cos(np.dot(n_vectors, i[3:]))

    a[3:, :3] = a[:3, 3:].T

    return np.array(solve(a, b)), n_vectors


def unroll_angles(A, sign):
    """Unrolls the angles, A, so they increase continuously"""
    n = np.array([0, 0, 0])
    P = np.zeros(np.shape(A))
    P[0] = A[0]
    for i in range(1, len(A)):
        n = (
            n
            + ((A[i] - A[i - 1] + 0.5 * sign * np.pi) * sign < 0)
            * np.ones(3)
            * 2.0
            * np.pi
        )
        P[i] = A[i] + sign * n
    return P


def angle_solver(AA, timeseries, N_max, sign, symNx=2, throw_out_modes=False):
    """Constructs the matrix A and the vector b from a timeseries of toy
    action-angles AA to solve for the vector x = (theta_0, theta_1, theta_2, omega_1,
    omega_2, omega_3, dSdx..., dSdy..., dSdz...) where x contains all derivatives
    of the Fourier components of the generating function with |n| < N_max"""

    # First unroll angles
    angs = unroll_angles(AA.T[3:].T, sign)

    # Same considerations as above
    symNz = 2
    NNx = range(-N_max, N_max + 1, symNx)
    NNy = range(-N_max, N_max + 1, symNz)
    NNz = range(-N_max, N_max + 1, symNz)
    n_vectors = np.array(
        [
            [i, j, k]
            for (i, j, k) in product(NNx, NNy, NNz)
            if (
                not (i == 0 and j == 0 and k == 0)  # exclude zero vector
                and (
                    k > 0  # northern hemisphere
                    or (k == 0 and j > 0)  # half of x-y plane
                    or (k == 0 and j == 0 and i > 0)
                )  # half of x axis
                and np.sqrt(i * i + j * j + k * k) <= N_max  # inside sphere
            )
        ]
    )

    if throw_out_modes:
        n_vectors = np.delete(n_vectors, check_each_direction(n_vectors, angs), axis=0)

    nv = len(n_vectors)
    n = 3 * nv + 6

    b = np.zeros(shape=(n,))
    a = np.zeros(shape=(n, n))

    a[:3, :3] = len(AA) * np.identity(3)
    a[:3, 3:6] = np.sum(timeseries) * np.identity(3)
    a[3:6, :3] = a[:3, 3:6]
    a[3:6, 3:6] = np.sum(timeseries * timeseries) * np.identity(3)

    for i, j in zip(angs, timeseries):
        a[6 : 6 + nv, 0] += -2.0 * np.sin(np.dot(n_vectors, i))
        a[6 : 6 + nv, 3] += -2.0 * j * np.sin(np.dot(n_vectors, i))
        a[6 : 6 + nv, 6 : 6 + nv] += 4.0 * np.outer(
            np.sin(np.dot(n_vectors, i)), np.sin(np.dot(n_vectors, i))
        )

        b[:3] += i
        b[3:6] += j * i

        b[6 : 6 + nv] += -2.0 * i[0] * np.sin(np.dot(n_vectors, i))
        b[6 + nv : 6 + 2 * nv] += -2.0 * i[1] * np.sin(np.dot(n_vectors, i))
        b[6 + 2 * nv : 6 + 3 * nv] += -2.0 * i[2] * np.sin(np.dot(n_vectors, i))

    a[6 + nv : 6 + 2 * nv, 1] = a[6 : 6 + nv, 0]
    a[6 + 2 * nv : 6 + 3 * nv, 2] = a[6 : 6 + nv, 0]
    a[6 + nv : 6 + 2 * nv, 4] = a[6 : 6 + nv, 3]
    a[6 + 2 * nv : 6 + 3 * nv, 5] = a[6 : 6 + nv, 3]
    a[6 + nv : 6 + 2 * nv, 6 + nv : 6 + 2 * nv] = a[6 : 6 + nv, 6 : 6 + nv]
    a[6 + 2 * nv : 6 + 3 * nv, 6 + 2 * nv : 6 + 3 * nv] = a[6 : 6 + nv, 6 : 6 + nv]

    a[:6, :] = a[:, :6].T

    return np.array(solve(a, b))
