from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

import polars as pl

from . import _aci_rssa
from .expressions import neutrons, on_surface, photons
from .header import FileParameters
from .plotting import RSSAPlot, RSSASpectraPlot
from .projection import Projection


@dataclass
class RSSA:
    """In-memory representation of an RSSA file.

    Attributes
    ----------
    parameters: FileParameters
        Parameters extracted from the RSSA file header.
    tracks: pl.DataFrame
        DataFrame containing the tracks recorded in the RSSA file, with columns
        `history`, `particle_type`, `weight`, `energy`, `time`, `x`, `y`, `z`, `u`,
        `v`, `w`, `surface_id`. `surface_id` is null for files that don't record it
        (see [`aci_rssa.header.FileParameters.records_surface_id`]). See
        [`aci_rssa.header.FileParameters`] and [`aci_rssa.scan_tracks`] for details on
        the header and for reading very large files without loading them fully into
        memory.

    Examples
    --------
    >>> from aci_rssa import RSSA
    >>> my_rssa = RSSA.read_from_file("small_cyl.w")
    >>> print(my_rssa)
    """

    parameters: FileParameters
    tracks: pl.DataFrame

    @staticmethod
    def read_from_file(path: Path | str, columns: list[str] | None = None) -> RSSA:
        """Loads an RSSA file as generated by MCNP.

        Parameters
        ----------
        path: Path | str
            The RSSA file to read.
        columns: list[str] | None
            If given, only decode these track columns (see `RSSA.tracks` for the
            available names). Decoding fewer columns is cheaper for very large files.
        """
        path = str(path)
        parameters = FileParameters.read_from_file(path)
        tracks = _aci_rssa.read_tracks(path, columns)
        return RSSA(parameters, tracks)

    def save_to_files(self, output_dir: Path | str) -> None:
        """Saves the RSSA as two files: a JSON file with the file parameters and a
        parquet file with the tracks DataFrame."""
        directory_path = Path(output_dir)
        directory_path.mkdir(parents=True, exist_ok=True)
        self.parameters.save_to_json(directory_path / "parameters.json")
        self.tracks.write_parquet(directory_path / "tracks.parquet")

    @staticmethod
    def load_from_saved_files(directory: Path | str) -> RSSA:
        """Loads the RSSA from the JSON and parquet files generated by
        `save_to_files()`.

        Parameters
        ----------
        directory: Path | str
            Directory where the `parameters.json` and `tracks.parquet` files are
            located.
        """
        directory_path = Path(directory)
        parameters = FileParameters.load_from_json(directory_path / "parameters.json")
        tracks = pl.read_parquet(directory_path / "tracks.parquet")
        return RSSA(parameters, tracks)

    def __repr__(self) -> str:
        return self.get_summary()

    def __str__(self) -> str:
        return self.get_summary()

    def get_summary(self) -> str:
        """Returns a summary of the RSSA file."""
        summary = "RSSA file was recorded using the following surfaces:\n"
        summary += self.parameters.describe_surfaces() + "\n"

        summary += f"The total number of tracks recorded is {self.parameters.nrss}.\n"
        summary += f"Neutrons: {self.neutron_tracks.shape[0]}"
        summary += f" photons: {self.photon_tracks.shape[0]}.\n"

        summary += "The simulation that produced this RSSA run "
        summary += f"{abs(self.parameters.np1)} histories.\n"
        summary += "The amount of independent histories that reached the RSSA surfaces "
        summary += f"was {self.parameters.niss}.\n"
        return summary

    @property
    def neutron_tracks(self) -> pl.DataFrame:
        """Returns the neutron tracks from the RSSA file."""
        return self.tracks.filter(neutrons())

    @property
    def photon_tracks(self) -> pl.DataFrame:
        """Returns the non-neutron (typically photon) tracks from the RSSA file."""
        return self.tracks.filter(photons())

    def plot_surface(
        self,
        surface_id: int | None = None,
        *,
        projection: Projection | None = None,
        verify: bool = True,
    ) -> RSSAPlot:
        """Returns an `RSSAPlot` of a single recorded surface, with the projection
        derived from that surface's geometry as stored in the file header.

        This is the only entry point to a current map, and deliberately so: the
        projection is what makes the map's per-area normalization correct, so it comes
        from the surface rather than from a guess. It needs no knowledge of how the
        surface is oriented, it uses a cylinder's real axis and centre, and on a file
        with several surfaces it keeps only the tracks that crossed the one asked for -
        rastering several surfaces together is meaningless.

        Parameters
        ----------
        surface_id: int | None
            Which recorded surface to plot, matched irrespective of sign. May be
            omitted only when the file recorded a single surface; otherwise the
            available ids are listed in the error.
        projection: Projection | None
            Overrides the projection derived from the header - for a surface type
            `Projection.from_surface()` has no geometry for, or when the header's
            coefficients need reinterpreting. The tracks are still restricted to
            `surface_id`, and the projection still carries the full geometry, so the
            bin areas remain areas on the surface.
        verify: bool
            Check a sample of the tracks against the surface's equation, so that a
            mismatch between the header and this package's reading of it fails loudly.

        Examples
        --------
        >>> rssa = RSSA.read_from_file("multiple_surfs.w")
        >>> plot = rssa.plot_surface(63).with_bin_width(10)
        >>> plot.particle_current(1e17).show()

        See `Projection.from_surface()` for the projection used per surface type.
        """
        return RSSAPlot.for_surface(
            self.tracks,
            self.parameters,
            surface_id,
            projection=projection,
            verify=verify,
        )

    def plot_spectra(self, surface_id: int | None = None) -> RSSASpectraPlot:
        """Returns an instance of RSSASpectraPlot to plot energy spectra.

        Pass `surface_id` to restrict the spectrum to one recorded surface. Unlike
        `plot_surface()`, leaving it out on a multi-surface file is not an error: a
        spectrum combining every surface is still a meaningful quantity, since it
        needs no geometry.
        """
        plot = RSSASpectraPlot(self.tracks, self.parameters)
        if surface_id is not None:
            surface = self.parameters.get_surface(surface_id)
            plot = plot.filter(on_surface(self.parameters, surface.id))
        return plot
