"""
This module defines abstract base classes for commensurability analysis
of a galactic potential's phase space. The "dimensionality" associated
with the class corresponds with the dimensionality of the relevant orbits.

This module also defines user-facing analysis classes for 2D and 3D
commensurability analysis using the tessellation subpackage.
"""

from __future__ import annotations

import inspect
import warnings
from abc import abstractmethod
from collections.abc import Mapping as MappingABC
from math import prod
from multiprocessing import Pool
from pathlib import Path
from typing import Any, Callable, Generator, Mapping, Optional, Sequence, Union

import astropy.coordinates as c
import astropy.units as u
import h5py
import numpy as np
import pidgey
from more_itertools import chunked
from pidgey import AgamaBackend, GalaBackend, GalpyBackend, get_backend_from
from pidgey.base import Backend
from tqdm import tqdm

from .evaluation import Evaluation
from .interactive import InteractivePlot2D, InteractivePlot3D, InteractivePlotBase
from .utils import collapse_coords, make_quantity


class AnalysisBase(MappingABC):
    """
    Base class for analyzing commensurate orbits within galactic potentials.

    This class provides methods for evaluating orbits, constructing images of
    phase space slices, and saving/loading analysis data.

    Attributes:
        ic_function (Callable[..., c.SkyCoord]): Function to generate initial conditions for orbits.
        ic_values (Mapping[str, Sequence[float]]): Values to pass into ic_function.
        axis_names (List[str]): Names of the parameters for the initial condition function.
        shape (Tuple[int]): Shape of the measures array.
        size (int): Size of the measures array.
        potential (Any): Potential generated by potential_function.
        dt (u.Quantity): Time step for orbit integration.
        steps (int): Number of integration steps.
        pattern_speed (u.Quantity): Pattern speed for orbit integration.
        measures (np.ndarray): Array to store orbit measures.
    """

    @staticmethod
    @abstractmethod
    def evaluate(orbit: c.SkyCoord) -> Evaluation:
        """
        Evaluate an orbit and return an evaluation object.

        Args:
            orbit (c.SkyCoord): Integrated orbit coordinates.

        Returns:
            Evaluation: Evaluation object containing analysis results.
        """
        pass

    def __init__(
        self,
        ic_function: Callable[..., c.SkyCoord],
        values: Mapping[str, Sequence[float]],
        /,
        potential_function: Callable[[], Any],
        dt: Union[float, np.ndarray, u.Quantity],
        steps: Union[int, np.ndarray],
        *,
        pattern_speed: Union[float, u.Quantity] = 0.0,
        backend: Optional[Union[str, Backend]] = None,
        progressbar: bool = True,
        pidgey_chunksize: int = 1,
        _blank_measures: bool = False,
    ) -> None:
        """
        Initialize AnalysisBase instance.

        Args:
            ic_function (Callable[..., c.SkyCoord]): Function to generate initial conditions for orbits.
            values (Mapping[str, Sequence[float]]): Values to pass into ic_function.
            potential_function (Callable[[], Any]): Function to generate the potential for orbit integration.
            dt (Union[float, np.ndarray, u.Quantity]): Time step for orbit integration.
            steps (Union[int, np.ndarray]): Number of integration steps.
            pattern_speed (Union[float, u.Quantity], optional): Pattern speed for orbit integration (default 0.0).
            backend (Optional[Union[str, Backend]], optional): Backend for orbit computation.
            progressbar (bool, optional): Whether to show progress bar during image construction (default True).
            pidgey_chunksize (int, optional): Chunk size for orbit integration (default 1).
        """
        self.ic_function = ic_function
        self.ic_values = values
        argspec = inspect.getfullargspec(ic_function)
        self.axis_names = argspec.args
        if len(self.ic_values) != len(self.axis_names) or not all(
            name in self.axis_names for name in self.ic_values
        ):
            raise ValueError("values and ic_function signature do not correspond")
        self.shape = tuple(len(self.ic_values[name]) for name in self.axis_names)
        self.size = prod(self.shape)

        self.potential_function = potential_function
        self.potential = potential_function()
        self.dt = make_quantity(dt, u.Gyr)
        if self.dt.value <= 0:
            raise ValueError("dt must be greater than 0")
        self.steps = steps
        if self.steps <= 0:
            raise ValueError("steps must be greater than 0")

        self.pattern_speed = make_quantity(pattern_speed, u.km / u.s / u.kpc)
        if isinstance(backend, str):
            if backend == "agama":
                backend = AgamaBackend()
            elif backend == "gala":
                backend = GalaBackend()
            elif backend == "galpy":
                backend = GalpyBackend()
            else:
                raise ValueError(f"Unrecognized backend: {backend}")
        self.backend = backend or get_backend_from(self.potential)
        if not self.backend:
            raise TypeError(f"Unrecognized potential: {self.potential}")

        if pidgey_chunksize <= 0:
            raise ValueError("chunksize must be greater than 0")
        if pidgey_chunksize >= self.size:
            pidgey_chunksize = self.size
            # raise ValueError("chunksize must be less than total number of starting coordinates")

        self.measures = np.zeros(self.shape)
        if not _blank_measures:
            self._construct_image(pidgey_chunksize, progressbar)

    def __len__(self) -> int:
        """
        Returns the size of the measures array.

        Returns:
            int: The size of the measures array.
        """
        return self.size

    def __iter__(self) -> Generator[tuple[int, ...], None, None]:
        """
        Iterate over the measures array's indices.

        Yields:
            Tuple[float]: Tuple of the indices for each grid point.
        """
        for pixel in np.ndindex(self.shape):
            yield pixel

    def __getitem__(self, key: tuple[int, ...]) -> tuple[c.SkyCoord, float]:
        """
        Get the measure at a specific grid point.

        Args:
            key: Tuple of indices for the grid point.

        Returns:
            Tuple[SkyCoord, float]: Tuple of the initial condition and measure at the grid point.
        """
        if not isinstance(key, tuple):
            raise KeyError("key must be a tuple")
        args = [self.ic_values[ax][i] for i, ax in zip(key, self.axis_names)]
        return self.ic_function(*args), self.measures[key]

    def _construct_image(self, pidgey_chunksize: int = 1, progressbar: bool = True):
        """
        Construct an image of a slice of phase space by integrating orbits and evaluating them.

        Args:
            pidgey_chunksize (int, optional): Chunk size for batching orbit integration (default 1).
            progressbar (bool, optional): Whether to show progress bar during construction (default True).
        """
        for pixels in tqdm(
            chunked(np.ndindex(self.shape), pidgey_chunksize),
            desc=f"with {pidgey_chunksize=}",
            total=self.size // pidgey_chunksize,
            disable=not progressbar,
        ):
            coords = []
            for pixel in pixels:
                params = [self.ic_values[ax][i] for i, ax in zip(pixel, self.axis_names)]
                coord = self.ic_function(*params)
                coords.append(coord)
            coords = collapse_coords(coords)

            orbits = self.backend.compute_orbit(
                coords,
                self.potential,
                self.dt,
                self.steps,
                pattern_speed=self.pattern_speed,
            )
            for pixel, orbit in tqdm(
                zip(pixels, orbits),
                desc="commensurability evaluation",
                total=pidgey_chunksize,
                disable=not progressbar,
                leave=False,
            ):
                self.measures[pixel] = self.evaluate(orbit).measure

    def save(self, path: Any):
        """
        Save the analysis data to an HDF5 file.

        Args:
            path: Path to the HDF5 file.
        """
        path = Path(path)
        if not path.parent.exists():
            print("Parent directory does not exist; creating directory.")
            path.parent.mkdir(parents=True)

        # store image mapping function source
        icsource = inspect.getsource(self.ic_function)
        icsource = icsource.replace(self.ic_function.__name__, "ic_function", 1)

        # store potential function source
        potsource = inspect.getsource(self.potential_function)
        potsource = potsource.replace(self.potential_function.__name__, "potential_function", 1)

        attrs = dict(
            icfunc=np.void(icsource.encode("utf8")),
            potfunc=np.void(potsource.encode("utf8")),
            dt=self.dt,
            steps=self.steps,
            pattern_speed=self.pattern_speed,
            backend=np.void(self.backend.__class__.__name__.encode("utf8")),
        )
        with h5py.File(path, "w") as f:
            dset = f.create_dataset(self.__class__.__name__, data=self.measures)
            for attr, value in attrs.items():
                dset.attrs[attr] = value
            for attr, value in self.ic_values.items():
                dset.attrs[attr] = value

    @classmethod
    def read_from_hdf5(cls, path: Any) -> AnalysisBase:
        """
        Read analysis data from an HDF5 file.

        Args:
            path: Path to the HDF5 file.

        Returns:
            AnalysisBase: Instance of AnalysisBase class with loaded data.
        """
        with h5py.File(path, "r") as f:
            dset = f[cls.__name__]

            if "icfunc" in dset.attrs:
                icsource = dset.attrs["icfunc"].tobytes().decode("utf8")
                namespace: dict[str, Any] = {}
                exec(icsource, {"u": u, "c": c}, namespace)
                ic_function = namespace["ic_function"]
            else:
                warnings.warn("No potential function defined.")

                def ic_function() -> None:
                    pass

            axis_names = inspect.getfullargspec(ic_function).args
            values = {name: dset.attrs[name] for name in axis_names}

            if "potfunc" in dset.attrs:
                potsource = dset.attrs["potfunc"].tobytes().decode("utf8")
                namespace = {}
                exec(potsource, {"u": u, "c": c}, namespace)
                potential_function = namespace["potential_function"]
            else:
                warnings.warn("No potential function defined.")

                def potential_function() -> None:
                    pass

            backend_cls = getattr(pidgey, dset.attrs["backend"].tobytes().decode("utf8"))
            analysis = cls(
                ic_function,
                values,
                potential_function=potential_function,
                dt=dset.attrs["dt"],
                steps=dset.attrs["steps"],
                pattern_speed=dset.attrs["pattern_speed"],
                backend=backend_cls(),
                _blank_measures=True,
            )
            analysis.measures = dset[()]
        return analysis


class MPAnalysisBase(AnalysisBase):
    """
    Base class for analyzing commensurate orbits within galactic potentials.
    This makes use of the multiprocessing module for parallelization of orbit evaluation.

    This class provides methods for evaluating orbits, constructing images of
    phase space slices, and saving/loading analysis data.

    Attributes:
        ic_function (Callable[..., c.SkyCoord]): Function to generate initial conditions for orbits.
        ic_values (Mapping[str, Sequence[float]]): Values to pass into ic_function.
        axis_names (List[str]): Names of the parameters for the initial condition function.
        shape (Tuple[int]): Shape of the measures array.
        size (int): Size of the measures array.
        potential (Any): Potential generated by potential_function.
        dt (u.Quantity): Time step for orbit integration.
        steps (int): Number of integration steps.
        pattern_speed (u.Quantity): Pattern speed for orbit integration.
        measures (np.ndarray): Array to store orbit measures.
    """

    @staticmethod
    @abstractmethod
    def __eval__(orbit: c.SkyCoord) -> float:
        return 0.0

    def __init__(
        self,
        ic_function: Callable[..., c.SkyCoord],
        values: Mapping[str, Sequence[float]],
        /,
        potential_function: Callable[[], Any],
        dt: Union[float, np.ndarray, u.Quantity],
        steps: Union[int, np.ndarray],
        *,
        pattern_speed: Union[float, u.Quantity] = 0.0,
        backend: Optional[Union[str, Backend]] = None,
        progressbar: bool = True,
        pidgey_chunksize: int = 1,
        mp_chunksize: int = 1,
        _blank_measures: bool = False,
    ) -> None:
        super().__init__(
            ic_function,
            values,
            potential_function,
            dt,
            steps,
            pattern_speed=pattern_speed,
            backend=backend,
            progressbar=progressbar,
            pidgey_chunksize=pidgey_chunksize,
            _blank_measures=True,
        )
        if not _blank_measures:
            self._construct_image_with_mp(pidgey_chunksize, mp_chunksize, progressbar)

    def _construct_image_with_mp(
        self, pidgey_chunksize: int = 1, mp_chunksize: int = 1, progressbar: bool = True
    ):
        for pixels in tqdm(
            chunked(np.ndindex(self.shape), pidgey_chunksize),
            desc=f"with {pidgey_chunksize=}",
            total=self.size // pidgey_chunksize,
            disable=not progressbar,
        ):
            coords = []
            for pixel in pixels:
                params = [self.ic_values[ax][i] for i, ax in zip(pixel, self.axis_names)]
                coord = self.ic_function(*params)
                coords.append(coord)
            coords = collapse_coords(coords)

            orbits = self.backend.compute_orbit(
                coords,
                self.potential,
                self.dt,
                self.steps,
                pattern_speed=self.pattern_speed,
            )
            with Pool() as p:
                values = list(
                    tqdm(
                        p.imap(self.__eval__, orbits, chunksize=mp_chunksize),
                        total=pidgey_chunksize,
                        leave=False,
                    )
                )
            for pixel, value in zip(pixels, values):
                self.measures[pixel] = value


class AnalysisBase2D(MPAnalysisBase):
    """
    Base class for commensurability analysis on 2D orbits.

    This class extends MPAnalysisBase and provides additional methods for launching interactive plots.

    Attributes:
        ic_function (Callable[..., c.SkyCoord]): Function to generate initial conditions for orbits.
        ic_values (Mapping[str, Sequence[float]]): Values to pass into ic_function.
        axis_names (List[str]): Names of the parameters for the initial condition function.
        shape (Tuple[int]): Shape of the measures array.
        size (int): Size of the measures array.
        potential (Any): Potential generated by potential_function.
        dt (u.Quantity): Time step for orbit integration.
        steps (int): Number of integration steps.
        pattern_speed (u.Quantity): Pattern speed for orbit integration.
        measures (np.ndarray): Array to store orbit measures.
    """

    def launch_interactive_plot(self, x_axis: str, y_axis: str, var_axis: Optional[str] = None):
        """
        Launch an interactive plot for 2D orbits.

        Args:
            x_axis (str): Name of the x-axis parameter.
            y_axis (str): Name of the y-axis parameter.
            var_axis (Optional[str]): Name of the axis varied by scrolling (optional).
        """
        iplot: InteractivePlotBase = InteractivePlot2D(self, x_axis, y_axis, var_axis)
        iplot.show()


class AnalysisBase3D(MPAnalysisBase):
    """
    Base class for commensurability analysis on 3D orbits.

    This class extends MPAnalysisBase and provides additional methods for launching interactive plots.

    Attributes:
        ic_function (Callable[..., c.SkyCoord]): Function to generate initial conditions for orbits.
        ic_values (Mapping[str, Sequence[float]]): Values to pass into ic_function.
        axis_names (List[str]): Names of the parameters for the initial condition function.
        shape (Tuple[int]): Shape of the measures array.
        size (int): Size of the measures array.
        potential (Any): Potential generated by potential_function.
        dt (u.Quantity): Time step for orbit integration.
        steps (int): Number of integration steps.
        pattern_speed (u.Quantity): Pattern speed for orbit integration.
        measures (np.ndarray): Array to store orbit measures.
    """

    def launch_interactive_plot(self, x_axis: str, y_axis: str, var_axis: Optional[str] = None):
        """
        Launch an interactive plot for 3D orbits.

        Args:
            x_axis (str): Name of the x-axis parameter.
            y_axis (str): Name of the y-axis parameter.
            var_axis (Optional[str]): Name of the variable axis (optional).
        """
        iplot: InteractivePlotBase = InteractivePlot3D(self, x_axis, y_axis, var_axis)
        iplot.show()


# define user-facing analysis classes
from .tessellation import Tessellation
from .tessellation.base import TessellationBase


class TessellationAnalysis(AnalysisBase3D):
    """
    Analysis class for tessellation analysis on 3D orbits.

    This class extends AnalysisBase3D and implements the evaluate method for tessellation analysis.

    Attributes:
        ic_function (Callable[..., c.SkyCoord]): Function to generate initial conditions for orbits.
        ic_values (Mapping[str, Sequence[float]]): Values to pass into ic_function.
        axis_names (List[str]): Names of the parameters for the initial condition function.
        shape (Tuple[int]): Shape of the measures array.
        size (int): Size of the measures array.
        potential (Any): Potential generated by potential_function.
        dt (u.Quantity): Time step for orbit integration.
        steps (int): Number of integration steps.
        pattern_speed (u.Quantity): Pattern speed for orbit integration.
        measures (np.ndarray): Array to store orbit measures.
    """

    @staticmethod
    def evaluate(orbit: c.SkyCoord) -> TessellationBase:
        """
        Evaluate an orbit using the tessellation and trimming algorithm.

        Args:
            orbit (c.SkyCoord): Integrated orbit coordinates.

        Returns:
            TessellationBase: Tessellation object containing tessellation results.
        """
        return Tessellation(orbit, incremental=False)

    @staticmethod
    def __eval__(orbit: c.SkyCoord) -> float:
        return Tessellation(orbit, incremental=False).measure


class TessellationAnalysis2D(AnalysisBase2D):
    """
    Analysis class for tessellation analysis on 2D orbits.

    This class extends AnalysisBase2D and implements the evaluate method for tessellation analysis.

    Attributes:
        ic_function (Callable[..., c.SkyCoord]): Function to generate initial conditions for orbits.
        ic_values (Mapping[str, Sequence[float]]): Values to pass into ic_function.
        axis_names (List[str]): Names of the parameters for the initial condition function.
        shape (Tuple[int]): Shape of the measures array.
        size (int): Size of the measures array.
        potential (Any): Potential generated by potential_function.
        dt (u.Quantity): Time step for orbit integration.
        steps (int): Number of integration steps.
        pattern_speed (u.Quantity): Pattern speed for orbit integration.
        measures (np.ndarray): Array to store orbit measures.
    """

    @staticmethod
    def evaluate(orbit) -> TessellationBase:
        """
        Evaluate an orbit using the tessellation and trimming algorithm.

        Args:
            orbit (c.SkyCoord): Integrated orbit coordinates.

        Returns:
            TessellationBase: Tessellation object containing tessellation results.
        """
        return Tessellation(orbit.xyz[:2].T, incremental=False)

    @staticmethod
    def __eval__(orbit: c.SkyCoord) -> float:
        return Tessellation(orbit.xyz[:2].T, incremental=False).measure
