import os

import autofit as af
import autolens as al
import autolens.plot as aplt

from test_autogalaxy.simulators.imaging import instrument_util

test_path = "{}/../../".format(os.path.dirname(os.path.realpath(__file__)))


def pixel_scale_from_instrument(instrument):
    """
    Returns the pixel scale from an instrument type based on real observations.

    These options are representative of VRO, Euclid, HST, over-sampled HST and Adaptive Optics image.

    Parameters
    ----------
    instrument : str
        A string giving the resolution of the desired instrument (VRO | Euclid | HST | HST_Up | AO).
    """
    if instrument in "vro":
        return (0.2, 0.2)
    elif instrument in "euclid":
        return (0.1, 0.1)
    elif instrument in "hst":
        return (0.05, 0.05)
    elif instrument in "hst_up":
        return (0.03, 0.03)
    elif instrument in "ao":
        return (0.01, 0.01)
    else:
        raise ValueError("An invalid instrument was entered - ", instrument)


def grid_from_instrument(instrument):
    """
    Returns the `Grid` from an instrument type based on real observations.

    These options are representative of VRO, Euclid, HST, over-sampled HST and Adaptive Optics image.

    Parameters
    ----------
    instrument : str
        A string giving the resolution of the desired instrument (VRO | Euclid | HST | HST_Up | AO).
    """
    if instrument in "vro":
        return al.GridIterate.uniform(shape_2d=(80, 80), pixel_scales=0.2)
    elif instrument in "euclid":
        return al.GridIterate.uniform(shape_2d=(120, 120), pixel_scales=0.1)
    elif instrument in "hst":
        return al.GridIterate.uniform(shape_2d=(200, 200), pixel_scales=0.05)
    elif instrument in "hst_up":
        return al.GridIterate.uniform(shape_2d=(300, 300), pixel_scales=0.03)
    elif instrument in "ao":
        return al.GridIterate.uniform(shape_2d=(800, 800), pixel_scales=0.01)
    else:
        raise ValueError("An invalid instrument was entered - ", instrument)


def psf_from_instrument(instrument):
    """
    Returns the *PSF* from an instrument type based on real observations.

    These options are representative of VRO, Euclid, HST, over-sampled HST and Adaptive Optics image.

    Parameters
    ----------
    instrument : str
        A string giving the resolution of the desired instrument (VRO | Euclid | HST | HST_Up | AO).
    """
    if instrument in "vro":
        return al.Kernel.from_gaussian(
            shape_2d=(31, 31), sigma=0.5, pixel_scales=0.2, renormalize=True
        )

    elif instrument in "euclid":
        return al.Kernel.from_gaussian(
            shape_2d=(31, 31), sigma=0.1, pixel_scales=0.1, renormalize=True
        )
    elif instrument in "hst":
        return al.Kernel.from_gaussian(
            shape_2d=(31, 31), sigma=0.05, pixel_scales=0.05, renormalize=True
        )
    elif instrument in "hst_up":
        return al.Kernel.from_gaussian(
            shape_2d=(31, 31), sigma=0.05, pixel_scales=0.03, renormalize=True
        )
    elif instrument in "ao":
        return al.Kernel.from_gaussian(
            shape_2d=(31, 31), sigma=0.025, pixel_scales=0.01, renormalize=True
        )

    else:
        raise ValueError("An invalid instrument was entered - ", instrument)


def simulator_from_instrument(instrument):
    """
    Returns the *Simulator* from an instrument type based on real observations.

    These options are representative of VRO, Euclid, HST, over-sampled HST and Adaptive Optics image.

    Parameters
    ----------
    instrument : str
        A string giving the resolution of the desired instrument (VRO | Euclid | HST | HST_Up | AO).
    """

    grid = grid_from_instrument(instrument=instrument)
    psf = psf_from_instrument(instrument=instrument)

    if instrument in "vro":
        return al.SimulatorImaging(
            exposure_time_map=al.Array.full(fill_value=100.0, shape_2d=grid.shape_2d),
            psf=psf,
            background_sky_map=al.Array.full(fill_value=1.0, shape_2d=grid.shape_2d),
            add_poisson_noise=True,
        )
    elif instrument in "euclid":
        return al.SimulatorImaging(
            exposure_time_map=al.Array.full(fill_value=2260.0, shape_2d=grid.shape_2d),
            psf=psf,
            background_sky_map=al.Array.full(fill_value=1.0, shape_2d=grid.shape_2d),
            add_poisson_noise=True,
        )
    elif instrument in "hst":
        return al.SimulatorImaging(
            exposure_time_map=al.Array.full(fill_value=2000.0, shape_2d=grid.shape_2d),
            psf=psf,
            background_sky_map=al.Array.full(fill_value=1.0, shape_2d=grid.shape_2d),
            add_poisson_noise=True,
        )
    elif instrument in "hst_up":
        return al.SimulatorImaging(
            exposure_time_map=al.Array.full(fill_value=2000.0, shape_2d=grid.shape_2d),
            psf=psf,
            background_sky_map=al.Array.full(fill_value=1.0, shape_2d=grid.shape_2d),
            add_poisson_noise=True,
        )
    elif instrument in "ao":
        return al.SimulatorImaging(
            exposure_time_map=al.Array.full(fill_value=1000.0, shape_2d=grid.shape_2d),
            psf=psf,
            background_sky_map=al.Array.full(fill_value=1.0, shape_2d=grid.shape_2d),
            add_poisson_noise=True,
        )
    else:
        raise ValueError("An invalid instrument was entered - ", instrument)


def simulate_imaging_from_instrument(instrument, dataset_name, galaxies):

    # Simulate the imaging data, remembering that we use a special image which ensures edge-effects don't
    # degrade our modeling of the telescope optics (e.al. the PSF convolution).

    grid = instrument_util.grid_from_instrument(instrument=instrument)

    simulator = simulator_from_instrument(instrument=instrument)

    # Use the input galaxies to setup a tracer, which will generate the image for the simulated imaging data.
    tracer = al.Tracer.from_galaxies(galaxies=galaxies)

    imaging = simulator.from_tracer_and_grid(tracer=tracer, grid=grid)

    # Now, lets output this simulated imaging-data to the test_autoarray/simulator folder.
    test_path = "{}/../../".format(os.path.dirname(os.path.realpath(__file__)))

    dataset_path = f"dataset/imaging/{dataset_name}/{instrument}"

    imaging.output_to_fits(
        image_path=f"{dataset_path}/image.fits",
        psf_path=f"{dataset_path}/psf.fits",
        noise_map_path=f"{dataset_path}/noise_map.fits",
        overwrite=True,
    )

    plotter = aplt.Plotter(output=aplt.Output(path=dataset_path, format="png"))
    sub_plotter = aplt.SubPlotter(output=aplt.Output(path=dataset_path, format="png"))

    aplt.Imaging.subplot_imaging(imaging=imaging, sub_plotter=sub_plotter)

    aplt.Imaging.individual(
        imaging=imaging,
        plot_image=True,
        plot_noise_map=True,
        plot_psf=True,
        plot_signal_to_noise_map=True,
        plotter=plotter,
    )

    aplt.Tracer.subplot_tracer(tracer=tracer, grid=grid, sub_plotter=sub_plotter)

    aplt.Tracer.individual(
        tracer=tracer,
        grid=grid,
        plot_image=True,
        plot_source_plane=True,
        plot_convergence=True,
        plot_potential=True,
        plot_deflections=True,
        plotter=plotter,
    )


def load_test_imaging(dataset_name, instrument, name=None):

    pixel_scales = instrument_util.pixel_scale_from_instrument(instrument=instrument)

    dataset_path = f"dataset/imaging/{dataset_name}/{instrument}"

    return al.Imaging.from_fits(
        image_path=f"{dataset_path}/image.fits",
        psf_path=f"{dataset_path}/psf.fits",
        noise_map_path=f"{dataset_path}/noise_map.fits",
        pixel_scales=pixel_scales,
        name=name,
    )
