from __future__ import annotations

import dataclasses
import json
import math
from dataclasses import dataclass
from enum import IntEnum
from pathlib import Path
from typing import Literal, cast

from . import _aci_rssa

SourceCode = Literal["Mcnp", "D1sUned"]

#: A coordinate axis, as named by the track columns.
Axis = Literal["x", "y", "z"]

#: Version of the JSON written by `FileParameters.save_to_json()`. Must match
#: `PARAMETERS_JSON_FORMAT_VERSION` on the Rust side. Version 1 (aci_rssa <= 0.0.6)
#: parsed the surface records with the wrong layout, so the values it stored are
#: silently wrong rather than merely differently shaped: they are rejected on load.
PARAMETERS_JSON_FORMAT_VERSION = 2

#: For a surface aligned with `axis`, the two perpendicular coordinates in cyclic
#: (x -> y -> z -> x) order. These are the coordinates a cylinder's centre is given in,
#: and the two axes an axis-aligned plane is plotted on.
PERPENDICULAR_AXES: dict[Axis, tuple[Axis, Axis]] = {
    "x": ("y", "z"),
    "y": ("z", "x"),
    "z": ("x", "y"),
}


class SurfaceType(IntEnum):
    """MCNP surface types, as recorded in an RSSA header.

    The value is the index of the surface's mnemonic in MCNP's standard surface-card
    order. Codes 1 (`p`), 4 (`pz`) and 15 (`cz`) are confirmed against the files in
    `resources/`; the rest follow the same table. Unknown codes are accepted (they
    become an `UNKNOWN_<n>` member) rather than raising, since a file may well use a
    surface this package has no geometry handling for.
    """

    P = 1
    PX = 2
    PY = 3
    PZ = 4
    SO = 5
    S = 6
    SX = 7
    SY = 8
    SZ = 9
    C_X = 10
    C_Y = 11
    C_Z = 12
    CX = 13
    CY = 14
    CZ = 15
    K_X = 16
    K_Y = 17
    K_Z = 18
    KX = 19
    KY = 20
    KZ = 21
    SQ = 22
    GQ = 23
    TX = 24
    TY = 25
    TZ = 26

    @classmethod
    def _missing_(cls, value: object) -> SurfaceType | None:
        if not isinstance(value, int):
            return None
        pseudo_member = int.__new__(cls, value)
        pseudo_member._name_ = f"UNKNOWN_{value}"
        pseudo_member._value_ = value
        # setdefault so repeated lookups of the same unknown code return the same
        # object, keeping `is` comparisons and dict lookups well behaved.
        return cast(
            "SurfaceType", cls._value2member_map_.setdefault(value, pseudo_member)
        )

    @property
    def mnemonic(self) -> str:
        """The MCNP surface mnemonic, e.g. `"cz"` or `"c/z"`."""
        if self.name.startswith("UNKNOWN_"):
            return f"type {self.value}"
        return self.name.lower().replace("_", "/")


#: Surface types this package can derive a 2D plotting projection for.
_PLANES = frozenset({SurfaceType.P, SurfaceType.PX, SurfaceType.PY, SurfaceType.PZ})
_CYLINDERS = frozenset(
    {
        SurfaceType.C_X,
        SurfaceType.C_Y,
        SurfaceType.C_Z,
        SurfaceType.CX,
        SurfaceType.CY,
        SurfaceType.CZ,
    }
)
#: The coordinate axis each axis-aligned surface type is aligned with: the axis of
#: revolution for cylinders, the normal direction for planes.
_ALIGNED_AXES: dict[SurfaceType, Axis] = {
    SurfaceType.PX: "x",
    SurfaceType.PY: "y",
    SurfaceType.PZ: "z",
    SurfaceType.C_X: "x",
    SurfaceType.C_Y: "y",
    SurfaceType.C_Z: "z",
    SurfaceType.CX: "x",
    SurfaceType.CY: "y",
    SurfaceType.CZ: "z",
}


@dataclass(frozen=True, slots=True)
class SurfaceParameters:
    """A surface tracks were recorded on, with its MCNP coefficients.

    `parameters` holds the coefficients exactly as MCNP stores them internally, which
    is not always how they are written on a surface card: notably, cylinders and
    spheres store the **squared** radius (a 30 cm `cz` cylinder has `parameters ==
    [900.0]`). The properties below do that conversion.
    """

    id: int
    #: Extra macrobody-facet info; only meaningful when `kjaq == 1`, `-1` otherwise.
    info: int
    surface_type: SurfaceType
    num_parameters: int
    parameters: list[float]

    @property
    def mnemonic(self) -> str:
        """The MCNP surface mnemonic, e.g. `"cz"`."""
        return self.surface_type.mnemonic

    @property
    def is_plane(self) -> bool:
        return self.surface_type in _PLANES

    @property
    def is_cylinder(self) -> bool:
        return self.surface_type in _CYLINDERS

    @property
    def axis(self) -> Axis | None:
        """The coordinate axis this surface is aligned with - the axis of revolution
        for a cylinder, the normal direction for an axis-aligned plane - or `None` for
        surfaces with no such axis (e.g. a general `p` plane)."""
        return _ALIGNED_AXES.get(self.surface_type)

    @property
    def center(self) -> tuple[float, float]:
        """Where a cylinder's axis sits, as the two coordinates perpendicular to
        `axis` in cyclic order (so a `c/z` cylinder gives `(x0, y0)`). The `cx`/`cy`/`cz`
        forms are centred on the origin by definition."""
        self._require(self.is_cylinder, "a centre")
        if self.surface_type in (SurfaceType.CX, SurfaceType.CY, SurfaceType.CZ):
            return (0.0, 0.0)
        # c/x, c/y and c/z are written as `x0 y0 r`, and stored here as the two centre
        # coordinates followed by the squared radius. No fixture with an off-axis
        # cylinder is available to confirm the order, which is why `plot_surface()`
        # checks the geometry against the tracks themselves.
        self._require(len(self.parameters) >= 3, "a centre")
        return (self.parameters[0], self.parameters[1])

    @property
    def radius(self) -> float:
        """The cylinder's radius, in cm. MCNP stores its square."""
        self._require(self.is_cylinder and bool(self.parameters), "a radius")
        return math.sqrt(self.parameters[-1])

    @property
    def normal(self) -> tuple[float, float, float]:
        """The plane's normal vector `(a, b, c)`, normalized to unit length."""
        self._require(self.is_plane and bool(self.parameters), "a normal")
        if self.surface_type is SurfaceType.P:
            a, b, c = self.parameters[0], self.parameters[1], self.parameters[2]
            length = math.sqrt(a * a + b * b + c * c)
            if length == 0.0:
                raise ValueError(f"surface {self.id} has a degenerate (zero) normal")
            return (a / length, b / length, c / length)
        axis = self.axis
        return (float(axis == "x"), float(axis == "y"), float(axis == "z"))

    @property
    def offset(self) -> float:
        """The plane's `d` coefficient: the plane is `normal . (x, y, z) == offset`."""
        self._require(self.is_plane and bool(self.parameters), "an offset")
        if self.surface_type is not SurfaceType.P:
            return self.parameters[0]
        a, b, c, d = self.parameters[:4]
        length = math.sqrt(a * a + b * b + c * c)
        return d / length

    def describe(self) -> str:
        """A one-line, human-readable summary, e.g.
        `"surface 63: cz (cylinder about the z axis, radius 164.30 cm)"`."""
        summary = f"surface {self.id}: {self.mnemonic}"
        if self.is_cylinder:
            center = self.center
            about = f"cylinder about the {self.axis} axis, radius {self.radius:.2f} cm"
            if center != (0.0, 0.0):
                about += f", centred at {center}"
            return f"{summary} ({about})"
        if self.is_plane:
            if self.surface_type is SurfaceType.P:
                a, b, c = self.normal
                return f"{summary} (plane {a:.4g}x + {b:.4g}y + {c:.4g}z = {self.offset:.4g})"
            return f"{summary} (plane {self.axis} = {self.offset:.4g})"
        if self.parameters:
            coefficients = ", ".join(f"{p:.4g}" for p in self.parameters)
            return f"{summary} ({coefficients})"
        return summary

    def _require(self, condition: bool, what: str) -> None:
        if not condition:
            raise ValueError(
                f"surface {self.id} is of type {self.mnemonic!r}, which does not have"
                f" {what}"
            )


@dataclass(frozen=True, slots=True)
class FileParameters:
    """Parameters extracted from an RSSA file header."""

    #: The code that produced the file, identified from its identification record.
    source_code: SourceCode
    #: Number of histories of the simulation, given as a negative number.
    np1: int
    #: Number of tracks recorded.
    nrss: int
    #: Number of values recorded for each particle (always 11, signed: see
    #: `records_surface_id`).
    nrcd: int
    #: Number of surfaces tracks were recorded on.
    njsw: int
    #: Number of independent histories that reached the recorded surfaces.
    niss: int
    #: Number of extra (weight-window) records; always 0 for supported files.
    niwr: int
    #: Source particle type.
    mipts: int
    #: Flag indicating whether surfaces include macrobody facet info.
    kjaq: int
    #: Only present for D1S-UNED files: the index of the last recorded dump.
    last_dump: int | None
    surfaces: list[SurfaceParameters]

    @property
    def records_surface_id(self) -> bool:
        """Whether each track carries the id of the surface it crossed.

        MCNP 6 files (`nrcd < 0`) record it in the last value of every track record;
        files written with `nrcd > 0` (e.g. the D1S-UNED `resources/small_cyl.w`) put a
        direction cosine there instead and encode the surface in a packed field this
        package does not decode. For those, the `surface_id` column is all-null and
        filtering by surface is unavailable - which only matters when the file has more
        than one surface.
        """
        return self.nrcd < 0

    @property
    def surface_ids(self) -> list[int]:
        """The ids of the recorded surfaces, as written in the header.

        These can be negative (MCNP writes one of `multiple_surfs.w`'s surfaces as
        `-610041` while its tracks carry `610041`), so `get_surface()` matches on
        absolute values.
        """
        return [surface.id for surface in self.surfaces]

    def get_surface(self, surface_id: int | None = None) -> SurfaceParameters:
        """Returns the recorded surface with the given id, matched irrespective of
        sign. `None` selects the only surface of a single-surface file, and raises if
        the file has several - there is no sensible default in that case."""
        if not self.surfaces:
            raise ValueError("this RSSA file's header declares no surfaces")
        if surface_id is None:
            if len(self.surfaces) == 1:
                return self.surfaces[0]
            raise ValueError(
                f"this RSSA file recorded tracks on {len(self.surfaces)} surfaces, so"
                f" one has to be chosen explicitly; available ids: {self.surface_ids}"
            )
        for surface in self.surfaces:
            if abs(surface.id) == abs(surface_id):
                return surface
        raise ValueError(
            f"no surface with id {surface_id} in this RSSA file; available ids:"
            f" {self.surface_ids}"
        )

    def describe_surfaces(self) -> str:
        """A human-readable, one-line-per-surface summary of the recorded surfaces."""
        return "\n".join(f"  {surface.describe()}" for surface in self.surfaces)

    @staticmethod
    def read_from_file(path: Path | str) -> FileParameters:
        """Parses only the header of an RSSA file (fast: it never touches the,
        potentially huge, track data)."""
        return FileParameters._from_json(_aci_rssa.parse_header_json(str(path)))

    @staticmethod
    def _from_json(text: str) -> FileParameters:
        data = json.loads(text)
        version = data.pop("format_version", 1)
        if version != PARAMETERS_JSON_FORMAT_VERSION:
            raise ValueError(
                f"unsupported RSSA parameters format version {version} (this version of"
                f" aci_rssa writes and reads version {PARAMETERS_JSON_FORMAT_VERSION})."
                " Files written by aci_rssa 0.0.6 or earlier parsed the surface records"
                " with the wrong layout, so their surface types and parameters are"
                " wrong; re-read the original RSSA file instead of converting them."
            )
        data["surfaces"] = [
            SurfaceParameters(
                id=s["id"],
                info=s["info"],
                surface_type=SurfaceType(s["surface_type"]),
                num_parameters=s["num_parameters"],
                parameters=[float(p) for p in s["parameters"]],
            )
            for s in data["surfaces"]
        ]
        return FileParameters(**data)

    def save_to_json(self, path: Path | str) -> None:
        """Saves the file parameters to a JSON file."""
        data = {
            "format_version": PARAMETERS_JSON_FORMAT_VERSION,
            **dataclasses.asdict(self),
        }
        with open(path, "w") as outfile:
            json.dump(data, outfile, indent=2)

    @staticmethod
    def load_from_json(path: Path | str) -> FileParameters:
        """Loads the file parameters from a JSON file written by `save_to_json()`."""
        with open(path) as infile:
            return FileParameters._from_json(infile.read())
