import pytest
import numpy.testing as npt
import numpy as np
import unittest
import os

from lenstronomy.Util.coolest_interface import (
    create_lenstronomy_from_coolest,
    update_coolest_from_lenstronomy,
    create_kwargs_mcmc_from_chain_list,
)
from lenstronomy.Util.coolest_read_util import (
    degree_coolest_to_radian_lenstronomy,
    ellibounds_coolest_to_lenstronomy,
    shearbounds_coolest_to_lenstronomy,
)
from lenstronomy.Util.coolest_update_util import (
    shapelet_amp_lenstronomy_to_coolest,
    folding_coolest,
)

from lenstronomy.LensModel.lens_model import LensModel
from lenstronomy.LightModel.light_model import LightModel
from lenstronomy.Data.imaging_data import ImageData
from lenstronomy.Data.psf import PSF
import lenstronomy.Util.image_util as image_util
from lenstronomy.ImSim.image_model import ImageModel
from lenstronomy.Plots.model_plot import ModelPlot

from lenstronomy.Workflow.fitting_sequence import FittingSequence


class TestCOOLESTinterface(object):
    def test_load(self):
        path = os.getcwd()
        if path[-11:] == "lenstronomy":
            path += "/test/test_Util"
        kwargs_out = create_lenstronomy_from_coolest(path + "/coolest_template")
        print(kwargs_out)
        kwargs_out = create_lenstronomy_from_coolest(path + "/coolest_template_pemd")
        kwargs_out = create_lenstronomy_from_coolest(
            path + "/coolest_template_pemd_random"
        )
        return

    def test_update(self):
        path = os.getcwd()
        if path[-11:] == "lenstronomy":
            path += "/test/test_Util"
        kwargs_result = {
            "kwargs_lens": [
                {"gamma1": 0.1, "gamma2": -0.05, "ra_0": 0.0, "dec_0": 0.0},
                {
                    "theta_E": 0.7,
                    "e1": -0.15,
                    "e2": 0.01,
                    "center_x": 0.03,
                    "center_y": 0.01,
                },
            ],
            "kwargs_source": [
                {
                    "amp": 15.0,
                    "R_sersic": 0.11,
                    "n_sersic": 3.6,
                    "center_x": 0.02,
                    "center_y": -0.03,
                    "e1": 0.1,
                    "e2": -0.2,
                },
                {
                    "amp": np.array(
                        [70.0, 33.0, 2.1, 3.9, 15.0, -16.0, 2.8, -1.7, -4.1, 0.2]
                    ),
                    "n_max": 3,
                    "beta": 0.1,
                    "center_x": 0.1,
                    "center_y": 0.0,
                },
            ],
            "kwargs_lens_light": [
                {
                    "amp": 11.0,
                    "R_sersic": 0.2,
                    "n_sersic": 3.0,
                    "center_x": 0.03,
                    "center_y": 0.01,
                    "e1": -0.15,
                    "e2": 0.01,
                },
                {
                    "amp": 12.0,
                    "R_sersic": 0.02,
                    "n_sersic": 6.0,
                    "center_x": 0.03,
                    "center_y": 0.01,
                    "e1": 0.0,
                    "e2": -0.15,
                },
            ],
            "kwargs_ps": [
                {
                    "point_amp": np.array([0.1]),
                    "ra_image": np.array([0.25]),
                    "dec_image": np.array([0.2]),
                }
            ],
        }
        update_coolest_from_lenstronomy(
            path + "/coolest_template", kwargs_result, ending="_update"
        )
        kwargs_out = create_lenstronomy_from_coolest(path + "/coolest_template_update")
        npt.assert_almost_equal(
            kwargs_out["kwargs_params"]["lens_model"][0][1]["e1"],
            kwargs_result["kwargs_lens"][1]["e1"],
            decimal=4,
        )
        npt.assert_almost_equal(
            kwargs_out["kwargs_params"]["lens_model"][0][1]["e2"],
            kwargs_result["kwargs_lens"][1]["e2"],
            decimal=4,
        )

        return

    def test_full(self):
        # use read json ; create an image ; create noise ; do fit (PSO for result + MCMC for chain)
        # create the kwargs mcmc ; upadte json
        path = os.getcwd()
        if path[-11:] == "lenstronomy":
            path += "/test/test_Util"

        kwargs_out = create_lenstronomy_from_coolest(path + "/coolest_template")

        # IMAGE specifics
        background_rms = 0.005  # background noise per pixel
        exp_time = 500.0  # exposure time (arbitrary units, flux per pixel is in units #photons/exp_time unit)
        # PSF : easier for test to create a gaussian PSF
        fwhm = 0.05  # full width at half maximum of PSF
        psf_type = "GAUSSIAN"  # 'GAUSSIAN', 'PIXEL', 'NONE'

        # lensing quantities to create an image
        lens_model_list = kwargs_out["kwargs_model"]["lens_model_list"]
        kwargs_sie = {
            "theta_E": 0.66,
            "center_x": 0.05,
            "center_y": 0,
            "e1": -0.1,
            "e2": 0.1,
        }  # parameters of the deflector lens model
        kwargs_shear = {
            "gamma1": 0.0,
            "gamma2": -0.05,
        }  # shear values to the source plane
        kwargs_lens = [kwargs_shear, kwargs_sie]
        lens_model_class = LensModel(lens_model_list)

        # Sersic parameters in the initial simulation for the source
        source_model_list = kwargs_out["kwargs_model"]["source_light_model_list"]
        kwargs_sersic = {
            "amp": 16,
            "R_sersic": 0.1,
            "n_sersic": 3.5,
            "e1": -0.1,
            "e2": 0.1,
            "center_x": 0.1,
            "center_y": 0,
        }
        kwargs_shapelets = {
            "amp": np.array([70.0, 33.0, 2.1, 3.9, 15.0, -16.0, 2.8, -1.7, -4.1, 0.2]),
            "n_max": 3,
            "beta": 0.1,
            "center_x": 0.1,
            "center_y": 0.0,
        }
        kwargs_source = [kwargs_sersic, kwargs_shapelets]
        source_model_class = LightModel(source_model_list)

        # Sersic parameters in the initial simulation for the lens light
        lens_light_model_list = kwargs_out["kwargs_model"]["lens_light_model_list"]
        kwargs_sersic_lens1 = {
            "amp": 16,
            "R_sersic": 0.6,
            "n_sersic": 2.5,
            "e1": -0.1,
            "e2": 0.1,
            "center_x": 0.05,
            "center_y": 0,
        }
        kwargs_sersic_lens2 = {
            "amp": 3,
            "R_sersic": 0.7,
            "n_sersic": 3,
            "e1": -0.1,
            "e2": 0.1,
            "center_x": 0.05,
            "center_y": 0,
        }
        kwargs_lens_light = [kwargs_sersic_lens1, kwargs_sersic_lens2]
        lens_light_model_class = LightModel(lens_light_model_list)

        numPix = 100
        kwargs_out["kwargs_data"]["background_rms"] = background_rms
        kwargs_out["kwargs_data"]["exposure_time"] = exp_time
        kwargs_out["kwargs_data"]["image_data"] = np.zeros((numPix, numPix))
        kwargs_out["kwargs_data"].pop("noise_map")

        data_class = ImageData(**kwargs_out["kwargs_data"])
        # PSF
        pixel_scale = (
            kwargs_out["kwargs_data"]["transform_pix2angle"][1][1]
            / kwargs_out["kwargs_psf"]["point_source_supersampling_factor"]
        )
        kwargs_psf = {
            "psf_type": "GAUSSIAN",
            "fwhm": fwhm,
            "pixel_size": pixel_scale,
            "truncation": 3,
        }
        kwargs_out["kwargs_psf"] = kwargs_psf
        psf_class = PSF(**kwargs_out["kwargs_psf"])

        kwargs_numerics = {
            "supersampling_factor": 1,
            "supersampling_convolution": False,
        }

        imageModel = ImageModel(
            data_class,
            psf_class,
            lens_model_class=lens_model_class,
            source_model_class=source_model_class,
            lens_light_model_class=lens_light_model_class,
            kwargs_numerics=kwargs_numerics,
        )

        # generate image
        image_model = imageModel.image(
            kwargs_lens,
            kwargs_source,
            kwargs_lens_light=kwargs_lens_light,
            kwargs_ps=None,
        )

        poisson = image_util.add_poisson(image_model, exp_time=exp_time)
        bkg = image_util.add_background(image_model, sigma_bkd=background_rms)
        image_real = image_model + poisson + bkg

        data_class.update_data(image_real)
        kwargs_out["kwargs_data"]["image_data"] = image_real

        # MODELING
        # Notes :
        # All the lines above were meant to create a mock image
        # The following is basically the only lines of code you will need
        # (after running the "create_lenstronomy_from_coolest" function) when you actually do the
        # modeling on a pre-existing image (with associated noise and psf proveded)
        band_list = [
            kwargs_out["kwargs_data"],
            kwargs_out["kwargs_psf"],
            kwargs_numerics,
        ]
        multi_band_list = [band_list]
        kwargs_data_joint = {
            "multi_band_list": multi_band_list,
            "multi_band_type": "single-band",
        }
        kwargs_constraints = {
            "joint_lens_with_light": [
                [0, 1, ["center_x", "center_y", "e1", "e2"]],
                [1, 1, ["center_x", "center_y", "e1", "e2"]],
            ]
        }
        kwargs_likelihood = {"check_bounds": True, "check_positive_flux": True}

        fitting_seq = FittingSequence(
            kwargs_data_joint,
            kwargs_out["kwargs_model"],
            kwargs_constraints,
            kwargs_likelihood,
            kwargs_out["kwargs_params"],
        )

        n_particules = 200
        n_iterations = 10
        wr = 5
        n_run_mcmc = 10
        n_burn_mcmc = 10
        fitting_kwargs_list = [
            [
                "PSO",
                {
                    "sigma_scale": 1.0,
                    "n_particles": n_particules,
                    "n_iterations": n_iterations,
                },
            ],
            [
                "MCMC",
                {
                    "n_burn": n_burn_mcmc,
                    "n_run": n_run_mcmc,
                    "walkerRatio": wr,
                    "sigma_scale": 0.01,
                },
            ],
        ]
        chain_list = fitting_seq.fit_sequence(fitting_kwargs_list)
        kwargs_result = fitting_seq.best_fit()

        modelPlot = ModelPlot(
            kwargs_data_joint["multi_band_list"],
            kwargs_out["kwargs_model"],
            kwargs_result,
        )
        modelPlot._imageModel.image_linear_solve(inv_bool=True, **kwargs_result)
        # the last 2 lines are meant for solving the linear parameters

        # use the function to save mcmc chains in userfriendly mode
        # kwargs_mcmc = create_kwargs_mcmc_from_chain_list(chain_list,kwargs_out['kwargs_model'],kwargs_out['kwargs_params'],
        #                                    kwargs_out['kwargs_data'],kwargs_out['kwargs_psf'],kwargs_numerics,
        #                                    kwargs_constraints,idx_chain=1)

        kwargs_mcmc = create_kwargs_mcmc_from_chain_list(
            chain_list,
            kwargs_out["kwargs_model"],
            kwargs_out["kwargs_params"],
            kwargs_out["kwargs_data"],
            kwargs_out["kwargs_psf"],
            kwargs_numerics,
            kwargs_constraints,
            idx_chain=1,
            likelihood_threshold=-100000,
        )
        # save the results (aka update the COOLEST json)
        update_coolest_from_lenstronomy(
            path + "/coolest_template", kwargs_result, kwargs_mcmc
        )

        return

    def test_pemd(self):
        path = os.getcwd()
        if path[-11:] == "lenstronomy":
            path += "/test/test_Util"
        kwargs_out = create_lenstronomy_from_coolest(path + "/coolest_template_pemd")
        print(kwargs_out)

        # kwargs_results to update the COOLEST template
        kwargs_result = {
            "kwargs_lens": [
                {
                    "theta_E": 0.7,
                    "e1": -0.15,
                    "e2": 0.01,
                    "gamma": 2.1,
                    "center_x": 0.03,
                    "center_y": 0.01,
                }
            ],
            "kwargs_source": [
                {
                    "amp": 15.0,
                    "R_sersic": 0.11,
                    "n_sersic": 3.6,
                    "center_x": 0.02,
                    "center_y": -0.03,
                    "e1": 0.1,
                    "e2": -0.2,
                }
            ],
            "kwargs_lens_light": [
                {
                    "amp": 11.0,
                    "R_sersic": 0.2,
                    "n_sersic": 3.0,
                    "center_x": 0.03,
                    "center_y": 0.01,
                    "e1": -0.15,
                    "e2": 0.01,
                }
            ],
        }
        # kwargs_mcmc to update the COOLEST template. In real cases, this list would be much bigger
        # as each element is a result from a given point at a given iteration of a MCMC chain

        kwargs_mcmc = {
            "args_lens": [
                [
                    {
                        "theta_E": 0.68,
                        "e1": -0.10,
                        "e2": -0.04,
                        "gamma": 1.9,
                        "center_x": 0.02,
                        "center_y": 0.10,
                    }
                ],
                [
                    {
                        "theta_E": 0.65,
                        "e1": -0.10,
                        "e2": -0.04,
                        "gamma": 1.9,
                        "center_x": 0.03,
                        "center_y": 0.01,
                    }
                ],
                [
                    {
                        "theta_E": 0.65,
                        "e1": -0.10,
                        "e2": -0.04,
                        "gamma": 1.9,
                        "center_x": 0.03,
                        "center_y": 0.01,
                    }
                ],
                [
                    {
                        "theta_E": 0.65,
                        "e1": -0.10,
                        "e2": -0.04,
                        "gamma": 1.9,
                        "center_x": 0.03,
                        "center_y": 0.01,
                    }
                ],
            ],
            "args_source": [
                [
                    {
                        "amp": 15.0,
                        "R_sersic": 0.11,
                        "n_sersic": 3.6,
                        "center_x": 0.02,
                        "center_y": -0.03,
                        "e1": 0.1,
                        "e2": -0.2,
                    }
                ],
                [
                    {
                        "amp": 15.0,
                        "R_sersic": 0.11,
                        "n_sersic": 3.6,
                        "center_x": 0.02,
                        "center_y": -0.03,
                        "e1": 0.1,
                        "e2": -0.2,
                    }
                ],
                [
                    {
                        "amp": 15.0,
                        "R_sersic": 0.11,
                        "n_sersic": 3.6,
                        "center_x": 0.02,
                        "center_y": -0.03,
                        "e1": 0.1,
                        "e2": -0.2,
                    }
                ],
                [
                    {
                        "amp": 15.0,
                        "R_sersic": 0.11,
                        "n_sersic": 3.6,
                        "center_x": 0.02,
                        "center_y": -0.03,
                        "e1": 0.1,
                        "e2": -0.2,
                    }
                ],
            ],
            "args_lens_light": [
                [
                    {
                        "amp": 11.0,
                        "R_sersic": 0.2,
                        "n_sersic": 3.0,
                        "center_x": 0.03,
                        "center_y": 0.01,
                        "e1": -0.15,
                        "e2": 0.01,
                    }
                ],
                [
                    {
                        "amp": 11.0,
                        "R_sersic": 0.2,
                        "n_sersic": 3.0,
                        "center_x": 0.03,
                        "center_y": 0.01,
                        "e1": -0.15,
                        "e2": 0.01,
                    }
                ],
                [
                    {
                        "amp": 11.0,
                        "R_sersic": 0.2,
                        "n_sersic": 3.0,
                        "center_x": 0.03,
                        "center_y": 0.01,
                        "e1": -0.15,
                        "e2": 0.01,
                    }
                ],
                [
                    {
                        "amp": 11.0,
                        "R_sersic": 0.2,
                        "n_sersic": 3.0,
                        "center_x": 0.03,
                        "center_y": 0.01,
                        "e1": -0.15,
                        "e2": 0.01,
                    }
                ],
            ],
        }
        update_coolest_from_lenstronomy(
            path + "/coolest_template_pemd",
            kwargs_result,
            ending="_update",
            kwargs_mcmc=kwargs_mcmc,
        )
        kwargs_out = create_lenstronomy_from_coolest(
            path + "/coolest_template_pemd_update"
        )
        print(kwargs_out)
        npt.assert_almost_equal(
            kwargs_out["kwargs_params"]["lens_model"][0][0]["e1"],
            kwargs_result["kwargs_lens"][0]["e1"],
            decimal=4,
        )
        npt.assert_almost_equal(
            kwargs_out["kwargs_params"]["lens_model"][0][0]["e2"],
            kwargs_result["kwargs_lens"][0]["e2"],
            decimal=4,
        )

        return

    def test_util_functions(self):
        radian = degree_coolest_to_radian_lenstronomy(None)
        radian = degree_coolest_to_radian_lenstronomy(-120.0)
        npt.assert_almost_equal(radian, np.pi / 6.0, decimal=4)

        radian = degree_coolest_to_radian_lenstronomy(120.0)
        npt.assert_almost_equal(radian, 5 * np.pi / 6.0, decimal=4)

        ellibounds_coolest_to_lenstronomy(0.6, 1.0, None, None)
        shearbounds_coolest_to_lenstronomy(0.0, 0.1, -90.0, None)

        shapelet_amp_lenstronomy_to_coolest(None)
        folding_coolest(np.array([-95.0, 95.0]))
        folding_coolest(-95.0)
        folding_coolest(95.0)

        return
