# coding: utf-8

""" Test action-angle stuff """

from __future__ import division, print_function


# Standard library
import logging

# Third-party
import numpy as np
from astropy import log as logger
from astropy.tests.helper import quantity_allclose
from scipy.linalg import solve
import pytest

# Project
from ...integrate import DOPRI853Integrator
from ...potential import (IsochronePotential, HarmonicOscillatorPotential,
                          LeeSutoTriaxialNFWPotential)
from ...units import galactic
from ..actionangle import *
from ..core import *
from ..plot import *
from .helpers import *

logger.setLevel(logging.DEBUG)

def test_generate_n_vectors():
    # test against Sanders' method
    nvecs = generate_n_vectors(N_max=6, dx=2, dy=2, dz=2)
    nvecs_sanders = sanders_nvecs(N_max=6, dx=2, dy=2, dz=2)
    assert np.all(nvecs == nvecs_sanders)

    nvecs = generate_n_vectors(N_max=6, dx=1, dy=1, dz=1)
    nvecs_sanders = sanders_nvecs(N_max=6, dx=1, dy=1, dz=1)
    assert np.all(nvecs == nvecs_sanders)

def test_fit_isochrone():
    # integrate orbit in Isochrone potential, then try to recover it
    true_m = 2.81E11
    true_b = 11.
    potential = IsochronePotential(m=true_m, b=true_b, units=galactic)
    orbit = potential.integrate_orbit([15.,0,0,0,0.2,0], dt=2., n_steps=10000)

    fit_potential = fit_isochrone(orbit)
    m,b = fit_potential.parameters['m'].value, fit_potential.parameters['b'].value
    assert np.allclose(m, true_m, rtol=1E-2)
    assert np.allclose(b, true_b, rtol=1E-2)

def test_fit_harmonic_oscillator():
    # integrate orbit in harmonic oscillator potential, then try to recover it
    true_omegas = np.array([0.011, 0.032, 0.045])
    potential = HarmonicOscillatorPotential(omega=true_omegas, units=galactic)

    orbit = potential.integrate_orbit([15.,1,2,0,0,0], dt=2., n_steps=10000)

    fit_potential = fit_harmonic_oscillator(orbit)
    omegas = fit_potential.parameters['omega'].value
    assert np.allclose(omegas, true_omegas, rtol=1E-2)

# TODO: check on this after fixing HarmonicOscillatorPotential
def test_fit_toy_potential():
    # integrate orbit in both toy potentials, make sure correct one is chosen
    true_m = 2.81E11
    true_b = 11.
    true_potential = IsochronePotential(m=true_m, b=true_b, units=galactic)
    orbit = true_potential.integrate_orbit([15.,0,0,0,0.2,0], dt=2., n_steps=10000)

    potential = fit_toy_potential(orbit)
    for k,v in true_potential.parameters.items():
        assert quantity_allclose(v, potential.parameters[k], rtol=1E-2)

    # -----------------------------------------------------------------
    true_omegas = np.array([0.011, 0.032, 0.045])
    true_potential = HarmonicOscillatorPotential(omega=true_omegas, units=galactic)
    orbit = true_potential.integrate_orbit([15.,1,2,0,0,0], dt=2., n_steps=10000)

    potential = fit_toy_potential(orbit)

    assert quantity_allclose(potential.parameters['omega'],
                             true_potential.parameters['omega'], rtol=1E-2)

def test_check_angle_sampling():

    # frequencies
    omegas = np.array([0.21, 0.3421, 0.4968])

    # integer vectors
    nvecs = generate_n_vectors(N_max=6)

    # loop over times with known failures:
    #   - first one fails needing longer integration time
    #   - second one fails needing finer sampling
    for i,t in enumerate([np.linspace(0,50,500), np.linspace(0,8000,8000)]):
        # periods = 2*np.pi/omegas
        # print("Periods:", periods)
        # print("N periods:", t.max() / periods)

        angles = t[np.newaxis] * omegas[:,np.newaxis]
        checks,failures = check_angle_sampling(nvecs, angles)

        assert np.all(failures == i)

class ActionsBase(object):

    @pytest.fixture(autouse=True)
    def set_tmpdir(self, tmpdir):
        self.tmpdir = tmpdir

    def test_classify(self):
        # my classify
        orb_type = self.orbit.circulation()

        # compare to Sanders'
        for j in range(self.N):
            sdrs = genfunc_3d.assess_angmom(self.w[...,j].T)
            logger.debug("APW: {}, Sanders: {}".format(orb_type[:,j], sdrs))
            assert np.all(orb_type[:,j] == sdrs)

    def test_actions(self):
        # t = self.t[::10]
        t = self.t

        N_max = 6
        for n in range(self.N):
            print("\n\n")
            print("======================= Orbit {} =======================".format(n))
            # w = self.w[:,::10,n]
            w = self.w[...,n]
            orb = self.orbit[:,n]
            circ = orb.circulation()

            # get values from Sanders' code
            print("Computing actions from genfunc...")
            s_actions,s_angles,s_freqs,toy_potential = sanders_act_ang_freq(t, w, circ, N_max=N_max)

            print("Computing actions with gala...")
            ret = find_actions(orb, N_max=N_max, toy_potential=toy_potential)
            actions = ret['actions']
            angles = ret['angles']
            freqs = ret['freqs']

            print("Action ratio: {}".format(actions / s_actions))
            print("Angle ratio: {}".format(angles / s_angles))
            print("Freq ratio: {}".format(freqs / s_freqs))

            assert np.allclose(actions.value, s_actions, rtol=1E-5)
            assert np.allclose(angles.value, s_angles, rtol=1E-5)
            assert np.allclose(freqs.value, s_freqs, rtol=1E-5)

            # logger.debug("Plotting orbit...")
            # fig = plot_orbits(w, marker='.', alpha=0.2, linestyle='none')
            # fig.savefig(str(self.plot_path.join("orbit_{}.png".format(n))))

            # fig = plot_angles(t,angles,freqs)
            # fig.savefig(str(self.plot_path.join("angles_{}.png".format(n))))

            # fig = plot_angles(t,s_angles,s_freqs)
            # fig.savefig(str(self.plot_path.join("angles_sanders_{}.png".format(n))))

            # plt.close('all')

            # print("Plots saved at:", self.plot_path)

class TestActions(ActionsBase):

    def setup(self):
        self.plot_path = self.tmpdir.mkdir("normal")

        self.units = galactic
        self.potential = LeeSutoTriaxialNFWPotential(v_c=0.2, r_s=20.,
                                                     a=1., b=0.77, c=0.55,
                                                     units=galactic)
        self.N = 8
        np.random.seed(42)
        w0 = isotropic_w0(N=self.N)
        n_steps = 20000

        # integrate orbits
        orbit = self.potential.integrate_orbit(w0, dt=2., n_steps=n_steps,
                                               Integrator=DOPRI853Integrator)
        self.orbit = orbit
        self.t = orbit.t.value
        self.w = orbit.w()

# TODO: need to fix this -- or assess whether needed?
# class TestHardActions(ActionsBase):

#     def setup(self):
#         self.plot_path = self.tmpdir.mkdir("hard")

#         self.units = (u.kpc, u.Msun, u.Myr)
#         params = {'disk': {'a': 6.5, 'b': 0.26, 'm': 65000000000.0},
#                   'bulge': {'c': 0.3, 'm': 20000000000.0},
#                   'halo': {'psi': 1.570796, 'theta': 1.570796, 'phi': 1.570796,
#                            'a': 1., 'b': 0.77, 'c': 0.61, 'r_s': 30.0, 'v_c': 0.22}}
#         self.potential = PW14Potential(**params)
#         self.N = 25

#         w0 = np.loadtxt(os.path.join(test_data_path, "w0.txt"))
#         n_steps = 200000

#         if not os.path.exists(os.path.join(test_data_path, "w_hard.npy")):
#             logger.debug("Integrating orbits")
#             t,w = self.potential.integrate_orbit(w0, dt=0.2, n_steps=n_steps, Integrator=DOPRI853Integrator)

#             logger.debug("Saving orbits")
#             np.save(os.path.join(test_data_path, "t_hard.npy"), t)
#             np.save(os.path.join(test_data_path, "w_hard.npy"), w)
#         else:
#             logger.debug("Loaded orbits")
#             t = np.load(os.path.join(test_data_path, "t_hard.npy"))
#             w = np.load(os.path.join(test_data_path, "w_hard.npy"))

#         self.t = t
#         self.w = w

def test_compare_action_prepare():

    from ..actionangle import _action_prepare, _angle_prepare

    logger.setLevel(logging.ERROR)
    AA = np.random.uniform(0., 100., size=(1000,6))
    t = np.linspace(0., 100., 1000)

    act_san,n_vectors = solver.solver(AA, N_max=6, symNx=2)
    A2,b2,n = _action_prepare(AA.T, N_max=6, dx=2, dy=2, dz=2)
    act_apw = np.array(solve(A2,b2))

    ang_san = solver.angle_solver(AA, t, N_max=6, symNx=2, sign=1)
    A2,b2,n = _angle_prepare(AA.T, t, N_max=6, dx=2, dy=2, dz=2)
    ang_apw = np.array(solve(A2,b2))

    assert np.allclose(act_apw, act_san)
    # assert np.allclose(ang_apw, ang_san)

    # TODO: this could be critical -- why don't our angles agree?
