"""
Utilities for testing
"""

from typing import Any, Dict, List, Union

import torch
import numpy as np
from astropy.cosmology import FlatLambdaCDM as FlatLambdaCDM_AP
from lenstronomy.Data.pixel_grid import PixelGrid
from lenstronomy.LensModel.lens_model import LensModel

from caustics.lenses import ThinLens, EPL, NFW, ThickLens, PixelatedConvergence
from caustics.light import Sersic, Pixelated
from caustics.utils import meshgrid
from caustics.sims import Simulator
from caustics.cosmology import FlatLambdaCDM
from .models import mock_from_file

__all__ = ("mock_from_file",)


def setup_simulator(
    cosmo_static=False, use_nfw=True, simulator_static=False, batched_params=False, device=None
):
    n_pix = 20

    class Sim(Simulator):
        def __init__(self, name="simulator"):
            super().__init__(name)
            if simulator_static:
                self.add_param("z_s", 1.0)
            else:
                self.add_param("z_s", None)
            z_l = 0.5
            self.cosmo = FlatLambdaCDM(h0=0.7 if cosmo_static else None, name="cosmo")
            if use_nfw:
                self.lens = NFW(
                    self.cosmo, z_l=z_l, name="lens"
                )  # NFW  wactually depend on cosmology, so a better test for Parametrized
            else:
                self.lens = EPL(self.cosmo, z_l=z_l, name="lens")
            self.sersic = Sersic(name="source")
            self.thx, self.thy = meshgrid(0.04, n_pix, device=device)
            self.n_pix = n_pix
            self.to(device=device)

        def forward(self, params):
            (z_s,) = self.unpack(params)
            alphax, alphay = self.lens.reduced_deflection_angle(
                x=self.thx, y=self.thy, z_s=z_s, params=params
            )
            bx = self.thx - alphax
            by = self.thy - alphay
            return self.sersic.brightness(bx, by, params)

    # default simulator params
    z_s = torch.tensor([1.0, 1.5])
    sim_params = [z_s]
    # default cosmo params
    h0 = torch.tensor([0.68, 0.75])
    cosmo_params = [h0]
    # default lens params
    if use_nfw:
        x0 = torch.tensor([0.0, 0.1])
        y0 = torch.tensor([0.0, 0.1])
        m = torch.tensor([1e12, 1e13])
        c = torch.tensor([10, 5])
        lens_params = [x0, y0, m, c]
    else:
        x0 = torch.tensor([0, 0.1])
        y0 = torch.tensor([0, 0.1])
        q = torch.tensor([0.9, 0.8])
        phi = torch.tensor([-0.56, 0.8])
        b = torch.tensor([1.5, 1.2])
        t = torch.tensor([1.2, 1.0])
        lens_params = [x0, y0, q, phi, b, t]
    # default source params
    x0s = torch.tensor([0, 0.1])
    y0s = torch.tensor([0, 0.1])
    qs = torch.tensor([0.9, 0.8])
    phis = torch.tensor([-0.56, 0.8])
    n = torch.tensor([1.0, 4.0])
    Re = torch.tensor([0.2, 0.5])
    Ie = torch.tensor([1.2, 10.0])
    source_params = [x0s, y0s, qs, phis, n, Re, Ie]

    if not batched_params:
        sim_params = [_x[0] for _x in sim_params]
        cosmo_params = [_x[0] for _x in cosmo_params]
        lens_params = [_x[0] for _x in lens_params]
        source_params = [_x[0] for _x in source_params]

    sim = Sim()
    # Set device when not None
    if device is not None:
        sim = sim.to(device=device)
        sim_params = [_p.to(device=device) for _p in sim_params]
        cosmo_params = [_p.to(device=device) for _p in cosmo_params]
        lens_params = [_p.to(device=device) for _p in lens_params]
        source_params = [_p.to(device=device) for _p in source_params]

    return sim, (sim_params, cosmo_params, lens_params, source_params)


def setup_image_simulator(cosmo_static=False, batched_params=False, device=None):
    n_pix = 20

    class Sim(Simulator):
        def __init__(self, name="test"):
            super().__init__(name)
            pixel_scale = 0.04
            z_l = 0.5
            self.z_s = torch.tensor(1.0)
            self.cosmo = FlatLambdaCDM(h0=0.7 if cosmo_static else None, name="cosmo")
            self.epl = EPL(self.cosmo, z_l=z_l, name="lens")
            self.kappa = PixelatedConvergence(
                pixel_scale,
                self.cosmo,
                z_l=z_l,
                shape=(n_pix, n_pix),
                name="kappa",
            )
            self.source = Pixelated(
                x0=0.0,
                y0=0.0,
                pixelscale=pixel_scale / 2,
                shape=(n_pix, n_pix),
                name="source",
            )
            self.thx, self.thy = meshgrid(pixel_scale, n_pix, device=device)
            self.n_pix = n_pix
            self.to(device=device)

        def forward(self, params):
            alphax, alphay = self.epl.reduced_deflection_angle(
                x=self.thx, y=self.thy, z_s=self.z_s, params=params
            )
            alphax_h, alphay_h = self.kappa.reduced_deflection_angle(
                x=self.thx, y=self.thy, z_s=self.z_s, params=params
            )
            bx = self.thx - alphax - alphax_h
            by = self.thy - alphay - alphay_h
            return self.source.brightness(bx, by, params)

    # default cosmo params
    h0 = torch.tensor([0.68, 0.75])
    # default lens params
    x0 = torch.tensor([0, 0.1])
    y0 = torch.tensor([0, 0.1])
    q = torch.tensor([0.9, 0.8])
    phi = torch.tensor([-0.56, 0.8])
    b = torch.tensor([1.5, 1.2])
    t = torch.tensor([1.2, 1.0])
    # default kappa params
    kappa = torch.randn([2, n_pix, n_pix])
    source = torch.randn([2, n_pix, n_pix])

    cosmo_params = [h0]
    lens_params = [x0, y0, q, phi, b, t]
    if not batched_params:
        cosmo_params = [_x[0] for _x in cosmo_params]
        lens_params = [_x[0] for _x in lens_params]
        kappa = kappa[0]
        source = source[0]

    sim = Sim()
    # Set device when not None
    if device is not None:
        sim = sim.to(device=device)
        cosmo_params = [_p.to(device=device) for _p in cosmo_params]
        lens_params = [_p.to(device=device) for _p in lens_params]
        kappa = kappa.to(device=device)
        source = source.to(device=device)

    return sim, (cosmo_params, lens_params, [kappa], [source])


def get_default_cosmologies(device=None):
    cosmology = FlatLambdaCDM("cosmo")
    cosmology_ap = FlatLambdaCDM_AP(100 * cosmology.h0.value, cosmology.Om0.value, Tcmb0=0)

    if device is not None:
        cosmology = cosmology.to(device=device)
    return cosmology, cosmology_ap


def setup_grids(res=0.05, n_pix=100, device=None):
    # Caustics setup
    thx, thy = meshgrid(res, n_pix, device=device)
    if device is not None:
        thx = thx.to(device=device)
        thy = thy.to(device=device)

    # Lenstronomy setup
    fov = res * n_pix
    ra_at_xy_0, dec_at_xy_0 = ((-fov + res) / 2, (-fov + res) / 2)
    transform_pix2angle = np.array([[1, 0], [0, 1]]) * res
    kwargs_pixel = {
        "nx": n_pix,
        "ny": n_pix,  # number of pixels per axis
        "ra_at_xy_0": ra_at_xy_0,
        "dec_at_xy_0": dec_at_xy_0,
        "transform_pix2angle": transform_pix2angle,
    }
    pixel_grid = PixelGrid(**kwargs_pixel)
    thx_ls, thy_ls = pixel_grid.coordinate_grid(n_pix, n_pix)
    return thx, thy, thx_ls, thy_ls


def alpha_test_helper(lens, lens_ls, z_s, x, kwargs_ls, atol, rtol, device=None):
    thx, thy, thx_ls, thy_ls = setup_grids(device=device)
    alpha_x, alpha_y = lens.reduced_deflection_angle(thx, thy, z_s, x)
    alpha_x_ls, alpha_y_ls = lens_ls.alpha(thx_ls, thy_ls, kwargs_ls)
    print(np.sum(np.abs(1 - alpha_x.cpu().numpy() / alpha_x_ls) > 1e-3))
    assert np.allclose(alpha_x.cpu().numpy(), alpha_x_ls, rtol, atol)
    assert np.allclose(alpha_y.cpu().numpy(), alpha_y_ls, rtol, atol)


def Psi_test_helper(lens, lens_ls, z_s, x, kwargs_ls, atol, rtol, device=None):
    thx, thy, thx_ls, thy_ls = setup_grids(device=device)
    Psi = lens.potential(thx, thy, z_s, x)
    Psi_ls = lens_ls.potential(thx_ls, thy_ls, kwargs_ls)
    # Potential is only defined up to a constant
    Psi -= Psi.min()
    Psi_ls -= Psi_ls.min()
    assert np.allclose(Psi.cpu().numpy(), Psi_ls, rtol, atol)


def kappa_test_helper(lens, lens_ls, z_s, x, kwargs_ls, atol, rtol, device=None):
    thx, thy, thx_ls, thy_ls = setup_grids(device=device)
    kappa = lens.convergence(thx, thy, z_s, x)
    kappa_ls = lens_ls.kappa(thx_ls, thy_ls, kwargs_ls)
    assert np.allclose(kappa.cpu().numpy(), kappa_ls, rtol, atol)


def lens_test_helper(
    lens: Union[ThinLens, ThickLens],
    lens_ls: LensModel,
    z_s,
    x,
    kwargs_ls: List[Dict[str, Any]],
    rtol,
    atol,
    test_alpha=True,
    test_Psi=True,
    test_kappa=True,
    device=None,
):
    if device is not None:
        lens = lens.to(device=device)
        z_s = z_s.to(device=device)
        x = x.to(device=device)

    if test_alpha:
        alpha_test_helper(lens, lens_ls, z_s, x, kwargs_ls, atol, rtol, device=device)

    if test_Psi:
        Psi_test_helper(lens, lens_ls, z_s, x, kwargs_ls, atol, rtol, device=device)

    if test_kappa:
        kappa_test_helper(lens, lens_ls, z_s, x, kwargs_ls, atol, rtol, device=device)
