"""SW1L equations solving exactly the linear terms
==================================================

(:mod:`fluidsim.solvers.sw1l.exactlin.solver`)



"""

import numpy as np


from fluidsim.base.setofvariables import SetOfVariables


from fluidsim.solvers.sw1l.solver import InfoSolverSW1L
from fluidsim.solvers.sw1l.solver import Simul as SimulSW1L


from fluiddyn.util import mpi


class InfoSolverSW1LExactLin(InfoSolverSW1L):
    """Information about the solver SW1L."""

    def _init_root(self):
        super()._init_root()

        sw1l = "fluidsim.solvers.sw1l"

        self.module_name = sw1l + ".exactlin.solver"
        self.class_name = "Simul"
        self.short_name = "SW1Lexlin"

        classes = self.classes

        classes.State.module_name = sw1l + ".exactlin.state"
        classes.State.class_name = "StateSW1LExactLin"


class Simul(SimulSW1L):
    """A solver of the shallow-water 1 layer equations (SW1L)"""

    InfoSolver = InfoSolverSW1LExactLin

    def __init__(self, params):
        if params.beta != 0:
            raise NotImplementedError("Do not use this solver for beta-plane!")

        super().__init__(params)

    def tendencies_nonlin(self, state_spect=None, old=None):
        oper = self.oper
        fft2 = oper.fft2

        if state_spect is None:
            state_phys = self.state.state_phys
            state_spect = self.state.state_spect
        else:
            state_phys = self.state.return_statephys_from_statespect(state_spect)

        ux = state_phys.get_var("ux")
        uy = state_phys.get_var("uy")
        eta = state_phys.get_var("eta")
        rot = state_phys.get_var("rot")

        # compute the nonlinear terms for ux, uy and eta
        N1x = +rot * uy
        N1y = -rot * ux
        gradu2_x_fft, gradu2_y_fft = oper.gradfft_from_fft(
            fft2(ux ** 2 + uy ** 2) / 2
        )

        Nx_fft = fft2(N1x) - gradu2_x_fft
        Ny_fft = fft2(N1y) - gradu2_y_fft

        jx_fft = fft2(eta * ux)
        jy_fft = fft2(eta * uy)
        Neta_fft = -oper.divfft_from_vecfft(jx_fft, jy_fft)

        # self.verify_tendencies(state_spect, state_phys,
        #                        Nx_fft, Ny_fft, Neta_fft)

        # compute the nonlinear terms for q, ap and am
        (Nq_fft, Np_fft, Nm_fft) = self.oper.qapamfft_from_uxuyetafft(
            Nx_fft, Ny_fft, Neta_fft
        )

        # Nq_fft = self.oper.create_arrayK(value=0)
        # Np_fft = self.oper.create_arrayK(value=0)
        # Nm_fft = self.oper.create_arrayK(value=0)

        oper.dealiasing(Nq_fft, Np_fft, Nm_fft)

        if old is None:
            tendencies_fft = SetOfVariables(
                like=self.state.state_spect, info="tendencies_nonlin"
            )
        else:
            tendencies_fft = old
        tendencies_fft.set_var("q_fft", Nq_fft)
        tendencies_fft.set_var("ap_fft", Np_fft)
        tendencies_fft.set_var("am_fft", Nm_fft)

        if self.params.forcing.enable:
            tendencies_fft += self.forcing.get_forcing()

        return tendencies_fft

    def compute_freq_complex(self, key):
        K2 = self.oper.K2
        if key == "q_fft":
            omega = self.oper.create_arrayK(value=0)
        elif key == "ap_fft":
            omega = 1.0j * np.sqrt(self.params.f ** 2 + self.params.c2 * K2)
        elif key == "am_fft":
            omega = -1.0j * np.sqrt(self.params.f ** 2 + self.params.c2 * K2)
        return omega

    # def verify_tendencies(
    #     self, state_spect, state_phys, Nx_fft, Ny_fft, Neta_fft
    # ):
    #     # for verification conservation energy
    #     # compute the linear terms
    #     oper = self.oper
    #     ux = state_phys.get_var("ux")
    #     uy = state_phys.get_var("uy")
    #     eta = state_phys.get_var("eta")

    #     q_fft = state_spect.get_var("q_fft")
    #     ap_fft = state_spect.get_var("ap_fft")
    #     am_fft = state_spect.get_var("am_fft")
    #     a_fft = ap_fft + am_fft
    #     div_fft = oper.divfft_from_apamfft(ap_fft, am_fft)

    #     eta_fft = oper.etafft_from_qfft(q_fft) + oper.etafft_from_afft(a_fft)

    #     dx_c2eta_fft, dy_c2eta_fft = oper.gradfft_from_fft(
    #         self.params.c2 * eta_fft
    #     )
    #     LCx = self.params.f * uy
    #     LCy = -self.params.f * ux
    #     Lx_fft = oper.fft2(LCx) - dx_c2eta_fft
    #     Ly_fft = oper.fft2(LCy) - dy_c2eta_fft
    #     Leta_fft = -div_fft

    #     # compute the full tendencies
    #     Fx_fft = Lx_fft + Nx_fft
    #     Fy_fft = Ly_fft + Ny_fft
    #     Feta_fft = Leta_fft + Neta_fft
    #     oper.dealiasing(Fx_fft, Fy_fft, Feta_fft)

    #     # test : ux, uy, eta ---> q, ap, am
    #     (Fq_fft, Fp_fft, Fm_fft) = self.oper.qapamfft_from_uxuyetafft(
    #         Fx_fft, Fy_fft, Feta_fft
    #     )
    #     # test : q, ap, am ---> ux, uy, eta
    #     (Fx2_fft, Fy2_fft, Feta2_fft) = self.oper.uxuyetafft_from_qapamfft(
    #         Fq_fft, Fp_fft, Fm_fft
    #     )
    #     print(np.max(abs(Fx2_fft - Fx_fft)))
    #     print(np.max(abs(Fy2_fft - Fy_fft)))
    #     print(np.max(abs(Feta2_fft - Feta_fft)))
    #     Fx_fft = Fx2_fft
    #     Fy_fft = Fy2_fft
    #     Feta_fft = Feta2_fft

    #     (Fq2_fft, Fp2_fft, Fm2_fft) = self.oper.qapamfft_from_uxuyetafft(
    #         Fx2_fft, Fy2_fft, Feta2_fft
    #     )
    #     print(np.max(abs(Fq2_fft - Fq_fft)))
    #     print(np.max(abs(Fp2_fft - Fp_fft)))
    #     print(np.max(abs(Fm2_fft - Fm_fft)))

    #     Fx = oper.ifft2(Fx_fft)
    #     Fy = oper.ifft2(Fy_fft)
    #     Feta = oper.ifft2(Feta_fft)
    #     A = (
    #         Feta * (ux ** 2 + uy ** 2) / 2
    #         + (1 + eta) * (ux * Fx + uy * Fy)
    #         + self.params.c2 * eta * Feta
    #     )
    #     A_fft = oper.fft2(A)
    #     if mpi.rank == 0:
    #         print("should be zero =", A_fft[0, 0])


if __name__ == "__main__":

    import fluiddyn as fld

    params = Simul.create_default_params()

    params.short_name_type_run = "test"

    nh = 64
    Lh = 2 * np.pi
    params.oper.nx = nh
    params.oper.ny = nh
    params.oper.Lx = Lh
    params.oper.Ly = Lh

    delta_x = params.oper.Lx / params.oper.nx
    params.nu_8 = (
        2.0 * 10e-1 * params.forcing.forcing_rate ** (1.0 / 3) * delta_x ** 8
    )

    params.time_stepping.t_end = 2.0

    params.init_fields.type = "noise"

    params.forcing.enable = True
    params.forcing.type = "waves"

    params.output.periods_print.print_stdout = 0.25

    params.output.periods_save.phys_fields = 1.0
    params.output.periods_save.spectra = 0.5
    params.output.periods_save.spect_energy_budg = 0.5
    params.output.periods_save.increments = 0.5
    params.output.periods_save.pdf = 0.5
    params.output.periods_save.time_signals_fft = False

    params.output.periods_plot.phys_fields = 0.0

    params.output.phys_fields.field_to_plot = "div"

    sim = Simul(params)

    # sim.output.phys_fields.plot()
    sim.time_stepping.start()
    # sim.output.phys_fields.plot()

    fld.show()
