from pathlib import Path
from typing import Iterator

import numpy as np
import pandas as pd
import pyarrow.parquet as pq

# A manifest is a plain text file listing one parquet path per line, used to
# name a curated subset of files (e.g. a train/holdout pool) without copying
# or symlinking the underlying parquet files. Lines are resolved relative to
# the manifest's own directory, so the manifest stays valid if the whole
# dataset tree is moved or copied elsewhere intact.
MANIFEST_SUFFIX = ".manifest"


def _read_manifest(path: Path) -> list[Path]:
    files = []
    for line in path.read_text().splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        resolved = (path.parent / line).resolve()
        if not resolved.is_file():
            raise FileNotFoundError(f"{path} lists missing file: {resolved}")
        files.append(resolved)
    if not files:
        raise FileNotFoundError(f"manifest {path} lists no files")
    return files


def find_parquet_files(path: str | Path) -> list[Path]:
    p = Path(path)
    if p.suffix == MANIFEST_SUFFIX:
        return _read_manifest(p)
    if p.is_dir():
        files = sorted(p.glob("*.parquet"))
        if not files:
            raise FileNotFoundError(f"no .parquet files found in {p}")
        return files
    return [p]


def _df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]:
    return {
        "event_id": df["event_id"].to_numpy(),
        "pdg": df["pdg"].to_numpy(dtype=np.int32),
        "pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32),
        "pre_E": df["pre_E"].to_numpy(dtype=np.float32),
        "pre_dir": df[["pre_dx", "pre_dy", "pre_dz"]].to_numpy(dtype=np.float32),
        "material": df["material"].to_numpy(dtype=object),
        "layer_id": df["layer_id"].to_numpy(dtype=np.int32),
        "n_sec": df["child_track_ids"].apply(len).to_numpy(dtype=np.int32),
        "e_sec": df["e_sec"].to_numpy(dtype=np.float32),
        "step_length": df["step_length"].to_numpy(dtype=np.float32),
        "post_E": df["post_E"].to_numpy(dtype=np.float32),
        "delta_e": (df["pre_E"] - df["post_E"]).to_numpy(dtype=np.float32),
        "edep": df["edep"].to_numpy(dtype=np.float32),
        "post_dir": df[["post_dx", "post_dy", "post_dz"]].to_numpy(dtype=np.float32),
        "post_pos": df[["post_x", "post_y", "post_z"]].to_numpy(dtype=np.float32),
    }


def load_steps(path: str | Path) -> dict[str, np.ndarray]:
    return _df_to_dict(pd.read_parquet(path))


def load_event_ids(path: str | Path) -> np.ndarray:
    """Read only the event_id column — cheap scan for split assignment."""
    return pd.read_parquet(path, columns=["event_id"])["event_id"].to_numpy()


def iter_file_chunks(path: str | Path) -> Iterator[dict[str, np.ndarray]]:
    """Yield one parquet row-group at a time so a large file never fully loads."""
    pf = pq.ParquetFile(path)
    for i in range(pf.num_row_groups):
        yield _df_to_dict(pf.read_row_group(i).to_pandas())


_COND_COLS = [
    "event_id",
    "pdg",
    "pre_x",
    "pre_y",
    "pre_z",
    "pre_E",
    "pre_dx",
    "pre_dy",
    "pre_dz",
    "material",
    "layer_id",
    "child_track_ids",
    "e_sec",
]


def _cond_df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]:
    return {
        "event_id": df["event_id"].to_numpy(),
        "pdg": df["pdg"].to_numpy(dtype=np.int32),
        "pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32),
        "pre_E": df["pre_E"].to_numpy(dtype=np.float32),
        "pre_dir": df[["pre_dx", "pre_dy", "pre_dz"]].to_numpy(dtype=np.float32),
        "material": df["material"].to_numpy(dtype=object),
        "layer_id": df["layer_id"].to_numpy(dtype=np.int32),
        "n_sec": df["child_track_ids"].apply(len).to_numpy(dtype=np.int32),
        "e_sec": df["e_sec"].to_numpy(dtype=np.float32),
    }


def iter_cond_chunks(path: str | Path) -> Iterator[dict[str, np.ndarray]]:
    """Yield conditioning-only row-groups (no post-step columns read from disk)."""
    pf = pq.ParquetFile(path)
    for i in range(pf.num_row_groups):
        yield _cond_df_to_dict(pf.read_row_group(i, columns=_COND_COLS).to_pandas())


def build_index_maps(
    data: dict[str, np.ndarray],
) -> tuple[dict[int, int], dict[str, int]]:
    pdg_vals = sorted(int(v) for v in np.unique(data["pdg"]))
    mat_vals = sorted(str(v) for v in np.unique(data["material"]))
    return (
        {v: i for i, v in enumerate(pdg_vals)},
        {v: i for i, v in enumerate(mat_vals)},
    )


def build_index_maps_from_files(
    files: list[Path],
) -> tuple[dict[int, int], dict[str, int]]:
    """Scan only pdg and material columns across all files (2-column read)."""
    pdg_vals: set[int] = set()
    mat_vals: set[str] = set()
    for path in files:
        df = pd.read_parquet(path, columns=["pdg", "material"])
        pdg_vals.update(int(v) for v in df["pdg"].unique())
        mat_vals.update(str(v) for v in df["material"].unique())
    return (
        {v: i for i, v in enumerate(sorted(pdg_vals))},
        {v: i for i, v in enumerate(sorted(mat_vals))},
    )
