"""Polars expressions for the track selections that need file-specific knowledge.

These exist because two very common filters can't be written correctly by someone who
only knows the column names: neutrons are identified by a packed leading-digit code
(`8`), and surface ids have to be matched irrespective of sign - and are unusable
altogether on files that don't record them per track.

Everything else (energy cuts, position windows, ...) is an ordinary Polars predicate,
so there is deliberately no helper for it:

```python
from aci_rssa import neutrons, on_surface, scan_tracks

scan_tracks(path).filter(neutrons() & (pl.col("energy") > 1.0))
```
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import polars as pl

from ._columns import NEUTRON_PARTICLE_TYPE

if TYPE_CHECKING:
    from .header import FileParameters


def neutrons() -> pl.Expr:
    """A predicate selecting neutron tracks.

    MCNP packs the particle type into the leading decimal digit of a field it shares
    with the sign of the `w` direction cosine; `8` is a neutron.
    """
    return pl.col("particle_type") == NEUTRON_PARTICLE_TYPE


def photons() -> pl.Expr:
    """A predicate selecting non-neutron tracks, which in the common case (a file
    written with `mipts` covering neutrons and photons) means photons."""
    return pl.col("particle_type") != NEUTRON_PARTICLE_TYPE


def on_surface(parameters: FileParameters, *surface_ids: int) -> pl.Expr:
    """A predicate selecting the tracks recorded on the given surfaces.

    Ids are matched irrespective of sign, since MCNP can write a surface in the header
    with the opposite sign to the one its tracks carry (`-610041` against `610041` in
    `resources/multiple_surfs.w`).

    Raises if an id isn't in the file, or if the file doesn't record which surface each
    track crossed (see `FileParameters.records_surface_id`) - filtering on an all-null
    column would silently drop every track instead of selecting some of them.
    """
    if not surface_ids:
        raise ValueError("on_surface() needs at least one surface id")

    valid_surface_ids = parameters.surface_ids
    unknown = [
        surface_id
        for surface_id in surface_ids
        if all(abs(surface_id) != abs(valid) for valid in valid_surface_ids)
    ]
    if unknown:
        raise ValueError(
            f"Some surface IDs are not valid: {unknown}."
            f" Valid IDs are: {valid_surface_ids}"
        )
    if not parameters.records_surface_id:
        raise ValueError(
            "This RSSA file does not record which surface each track crossed"
            f" (nrcd = {parameters.nrcd} > 0), so its tracks cannot be filtered by"
            f" surface. Note the file only has {len(valid_surface_ids)} surface(s):"
            f" {valid_surface_ids}."
        )
    wanted = [abs(surface_id) for surface_id in surface_ids]
    return pl.col("surface_id").abs().is_in(wanted)
