import numpy as np
from astropy import cosmology as cosmo

from autolens import exc
from autolens.lens.util import lens_util
from autolens.lens import ray_tracing
from workspace_jam.scraps.stack import plane_stack
from autolens.model.inversion import pixelizations as pix


class AbstractTracerStack(ray_tracing.AbstractTracer):
    def __init__(self, planes, cosmology):
        """Abstract ray-tracer for lens systems with any number of planes and multiple grid-stacks.

        From the galaxies in the tracer's planes, the image-plane images of each grid-stack are computed.

        Parameters
        ----------
        planes : [plane_stack.Plane]
            The list of the tracer's planes in ascending redshift order.
        cosmology : astropy.cosmology
            The cosmology of the ray-tracing calculation.
        """
        super(AbstractTracerStack, self).__init__(planes=planes, cosmology=cosmology)

    @property
    def total_grid_stacks(self):
        return len(self.image_plane.grid_stacks)

    @property
    def image_plane_images(self):
        return list(
            map(
                lambda image_plane_image_1d, grid_stack: grid_stack.regular.scaled_array_2d_from_array_1d(
                    image_plane_image_1d
                ),
                self.image_plane_images_1d,
                self.image_plane.grid_stacks,
            )
        )

    @property
    def image_plane_images_for_simulation(self):
        return [
            sum(
                image_plane_image_of_plane_for_simulation[i]
                for image_plane_image_of_plane_for_simulation in self.image_plane_images_of_planes_for_simulation
            )
            for i in range(self.total_grid_stacks)
        ]

    @property
    def image_plane_images_of_planes_for_simulation(self):
        return [plane.image_plane_images_for_simulation for plane in self.planes]

    @property
    def image_plane_images_1d(self):
        return list(
            map(
                lambda image_plane_image_1d_of_planes: sum(
                    image_plane_image_1d_of_planes
                ),
                self.image_plane_images_1d_of_planes,
            )
        )

    @property
    def image_plane_images_1d_of_planes(self):
        image_plane_images_1d = [plane.image_plane_images_1d for plane in self.planes]
        image_plane_images_1d_of_planes = [[] for _ in range(self.total_grid_stacks)]
        for image_index in range(self.total_grid_stacks):
            for plane_index in range(self.total_planes):
                image_plane_images_1d_of_planes[image_index].append(
                    image_plane_images_1d[plane_index][image_index]
                )
        return image_plane_images_1d_of_planes

    @property
    def image_plane_blurring_images_1d(self):
        return list(
            map(
                lambda image_plane_blurring_image_1d_of_planes: sum(
                    image_plane_blurring_image_1d_of_planes
                ),
                self.image_plane_blurring_images_1d_of_planes,
            )
        )

    @property
    def image_plane_blurring_images_1d_of_planes(self):
        image_plane_blurring_images_1d = [
            plane.image_plane_blurring_images_1d for plane in self.planes
        ]
        image_plane_blurring_images_1d_of_planes = [
            [] for _ in range(self.total_grid_stacks)
        ]
        for image_index in range(self.total_grid_stacks):
            for plane_index in range(self.total_planes):
                image_plane_blurring_images_1d_of_planes[image_index].append(
                    image_plane_blurring_images_1d[plane_index][image_index]
                )
        return image_plane_blurring_images_1d_of_planes


class TracerImagePlaneStack(AbstractTracerStack):
    def __init__(
        self,
        lens_galaxies,
        image_plane_grid_stacks,
        borders=None,
        cosmology=cosmo.Planck15,
    ):
        """Ray tracer for a lens system with just an image-plane.

        As there is only 1 plane, there are no ray-tracing calculations. This class is therefore only used for fitting \
        image-plane galaxies with light profiles.

        This tracer has a list of grid-stacks (see grid_stack.GridStack) which are all used for ray-tracing.

        Parameters
        ----------
        lens_galaxies : [Galaxy]
            The list of lens galaxies in the image-plane.
        image_plane_grid_stacks : [grid_stacks.GridStack]
            The image-plane grid stacks which are traced. (each stack includes the regular-grid, sub-grid, \
            blurring-grid, etc.).
        borders : [masks.RegularGridBorder]
            The border of each grid-stacks's regular-grid, which is used to relocate demagnified traced pixels to the \
            source-plane border.
        cosmology : astropy.cosmology.Planck15
            The cosmology of the ray-tracing calculation.
        """
        if not lens_galaxies:
            raise exc.RayTracingException(
                "No lens galaxies have been input into the Tracer"
            )

        image_plane = plane_stack.PlaneStack(
            galaxies=lens_galaxies,
            grid_stacks=image_plane_grid_stacks,
            borders=borders,
            compute_deflections=True,
            cosmology=cosmology,
        )

        super(TracerImagePlaneStack, self).__init__(
            planes=[image_plane], cosmology=cosmology
        )


class TracerImageSourcePlanesStack(AbstractTracerStack):
    def __init__(
        self,
        lens_galaxies,
        source_galaxies,
        image_plane_grid_stacks,
        borders=None,
        cosmology=cosmo.Planck15,
    ):
        """Ray-tracer for a lens system with two planes, an image-plane and source-plane.

        This tracer has a list of grid-stacks (see grid_stack.GridStack) which are all used for ray-tracing.

        Parameters
        ----------
        lens_galaxies : [Galaxy]
            The list of galaxies in the image-plane.
        source_galaxies : [Galaxy]
            The list of galaxies in the source-plane.
        image_plane_grid_stacks : [grid_stacks.GridStack]
            The image-plane grid stacks which are traced. (each stack includes the regular-grid, sub-grid, \
            blurring-grid, etc.).
        borders : [masks.RegularGridBorder]
            The border of each grid-stacks's regular-grid, which is used to relocate demagnified traced pixels to the \
            source-plane border.
        cosmology : astropy.cosmology.Planck15
            The cosmology of the ray-tracing calculation.
        """

        image_plane_grid_stacks = list(
            map(
                lambda data_grids: pix.setup_image_plane_pixelization_grid_from_galaxies_and_grid_stack(
                    galaxies=source_galaxies, grid_stack=data_grids
                ),
                image_plane_grid_stacks,
            )
        )

        image_plane = plane_stack.PlaneStack(
            galaxies=lens_galaxies,
            grid_stacks=image_plane_grid_stacks,
            borders=borders,
            compute_deflections=True,
        )

        source_plane_grid_stacks = image_plane.trace_grids_to_next_plane()

        source_plane = plane_stack.PlaneStack(
            galaxies=source_galaxies,
            grid_stacks=source_plane_grid_stacks,
            borders=borders,
            compute_deflections=False,
            cosmology=cosmology,
        )

        super(TracerImageSourcePlanesStack, self).__init__(
            planes=[image_plane, source_plane], cosmology=cosmology
        )


class TracerMultiPlanesStack(AbstractTracerStack):
    def __init__(
        self, galaxies, image_plane_grid_stacks, borders=None, cosmology=cosmo.Planck15
    ):
        """Ray-tracer for a lens system with any number of planes.

        To perform multi-plane ray-tracing, the cosmology that is input is used to rescale deflection-angles \
        according to the lens-geometry of the multi-plane system. All galaxies input to the tracer must therefore \
        have redshifts.

        This tracer has a list of grid-stacks (see grid_stack.GridStack) which are all used for ray-tracing.

        Parameters
        ----------
        galaxies : [Galaxy]
            The list of galaxies in the ray-tracing calculation.
        image_plane_grid_stacks : [grid_stacks.GridStack]
            The image-plane grid stacks which are traced. (each stack includes the regular-grid, sub-grid, \
            blurring-grid, etc.).
        borders : [masks.RegularGridBorder]
            The border of each grid-stacks's regular-grid, which is used to relocate demagnified traced pixels to the \
            source-plane border.
        cosmology : astropy.cosmology
            The cosmology of the ray-tracing calculation.
        """

        ordered_redshifts = lens_util.ordered_plane_redshifts_from_galaxies(
            galaxies=galaxies
        )

        galaxies_in_redshift_ordered_lists = lens_util.galaxies_in_redshift_ordered_planes_from_galaxies(
            galaxies=galaxies, plane_redshifts=ordered_redshifts
        )

        image_plane_grid_stacks = list(
            map(
                lambda grid_stack: pix.setup_image_plane_pixelization_grid_from_galaxies_and_grid_stack(
                    galaxies=galaxies, grid_stack=grid_stack
                ),
                image_plane_grid_stacks,
            )
        )

        planes = []

        for plane_index in range(0, len(ordered_redshifts)):

            if plane_index < len(ordered_redshifts) - 1:
                compute_deflections = True
            elif plane_index == len(ordered_redshifts) - 1:
                compute_deflections = False
            else:
                raise exc.RayTracingException(
                    "A galaxy was not correctly allocated its previous / next redshifts"
                )

            new_grid_stacks = image_plane_grid_stacks

            if plane_index > 0:
                for previous_plane_index in range(plane_index):

                    scaling_factor = lens_util.scaling_factor_between_redshifts_for_cosmology(
                        z1=ordered_redshifts[previous_plane_index],
                        z2=ordered_redshifts[plane_index],
                        z_final=ordered_redshifts[-1],
                        cosmology=cosmology,
                    )

                    def scale(grid):
                        return np.multiply(scaling_factor, grid)

                    if planes[previous_plane_index].deflection_stacks is not None:
                        scaled_deflections = list(
                            map(
                                lambda deflection_stack: deflection_stack.apply_function(
                                    scale
                                ),
                                planes[previous_plane_index].deflection_stacks,
                            )
                        )
                    else:
                        scaled_deflections = None

                    if scaled_deflections is not None:

                        def minus(grid, deflections):
                            return grid - deflections

                        new_grid_stacks = list(
                            map(
                                lambda grid, deflections: grid.map_function(
                                    minus, deflections
                                ),
                                new_grid_stacks,
                                scaled_deflections,
                            )
                        )

            planes.append(
                plane_stack.PlaneStack(
                    galaxies=galaxies_in_redshift_ordered_lists[plane_index],
                    grid_stacks=new_grid_stacks,
                    borders=borders,
                    compute_deflections=compute_deflections,
                    cosmology=cosmology,
                )
            )

        super(TracerMultiPlanesStack, self).__init__(planes=planes, cosmology=cosmology)
