import numpy as np
import pandas as pd

from giant.constants import COND_DIM, X_DIM
from giant.data import setup_cache
from giant.data.dataset import StreamingStepsDataset, make_event_split
from giant.data.transforms import Normalizer


def test_make_event_split_sizes():
    rng = np.random.default_rng(42)
    event_ids = rng.integers(0, 50, size=1000)
    train_set, val_set = make_event_split(event_ids, val_fraction=0.2)
    unique = np.unique(event_ids)
    assert len(train_set) + len(val_set) == len(unique)


def test_make_event_split_no_overlap():
    rng = np.random.default_rng(7)
    event_ids = rng.integers(0, 50, size=1000)
    train_set, val_set = make_event_split(event_ids, val_fraction=0.2)
    assert train_set.isdisjoint(val_set)


def test_make_event_split_no_empty_sets():
    rng = np.random.default_rng(0)
    event_ids = rng.integers(0, 20, size=500)
    train_set, val_set = make_event_split(event_ids, val_fraction=0.2)
    assert len(train_set) > 0
    assert len(val_set) > 0


def test_make_event_split_reproducible():
    event_ids = np.arange(100)
    a_tr, a_val = make_event_split(event_ids, val_fraction=0.1, seed=42)
    b_tr, b_val = make_event_split(event_ids, val_fraction=0.1, seed=42)
    assert a_tr == b_tr
    assert a_val == b_val


# ── StreamingStepsDataset: cross-file event_id offsetting ──────────────────


def _steps_df(event_ids, n_per_event=3, pre_E=100.0):
    """A schema-complete but minimal steps DataFrame — no secondaries, so
    `require_secondaries=True` never needs the per-secondary list columns."""
    rows = []
    for eid in event_ids:
        for s in range(n_per_event):
            rows.append(
                {
                    "event_id": eid,
                    "pdg": 11,
                    "pre_x": 0.0,
                    "pre_y": 0.0,
                    "pre_z": 0.0,
                    "pre_E": pre_E,
                    "pre_dx": 0.0,
                    "pre_dy": 0.0,
                    "pre_dz": 1.0,
                    "material": "G4_AIR",
                    "layer_id": s,
                    "child_track_ids": [],
                    "e_sec": 0.0,
                    "step_length": 1.0,
                    "post_E": pre_E * 0.9,
                    "edep": pre_E * 0.1,
                    "post_dx": 0.0,
                    "post_dy": 0.0,
                    "post_dz": 1.0,
                    "post_x": 0.0,
                    "post_y": 0.0,
                    "post_z": 1.0,
                }
            )
    return pd.DataFrame(rows)


def _dummy_normalizer(width):
    norm = Normalizer()
    norm.mean = np.zeros(width, dtype=np.float32)
    norm.std = np.ones(width, dtype=np.float32)
    return norm


def test_streaming_dataset_offsets_colliding_event_ids_across_files(tmp_path):
    """Two files that each restart event_id from 0 (one Geant4 job per file,
    see scripts/steps_to_parquet.py) must not have their same-numbered events
    collapsed together: every row from every file must show up in exactly one
    of train/val, and the number of distinct events must be the sum across
    files, not the union of raw ids."""
    n_events, n_per_event = 5, 3
    path_a = tmp_path / "a.parquet"
    path_b = tmp_path / "b.parquet"
    _steps_df(range(n_events), n_per_event=n_per_event).to_parquet(path_a)
    _steps_df(range(n_events), n_per_event=n_per_event).to_parquet(path_b)
    files = [path_a, path_b]

    unique_ids, counts = setup_cache.compute_event_index_from_files(files)
    assert len(unique_ids) == 2 * n_events

    train_events, val_events = make_event_split(unique_ids, val_fraction=0.4, seed=0)
    assert train_events.isdisjoint(val_events)

    pdg_map, mat_map = {11: 0}, {"G4_AIR": 0}
    cond_norm = _dummy_normalizer(COND_DIM)
    tgt_norm = _dummy_normalizer(X_DIM)

    def _count_rows(split_events):
        ds = StreamingStepsDataset(
            files=files,
            split_events=split_events,
            pdg_map=pdg_map,
            mat_map=mat_map,
            cond_normalizer=cond_norm,
            target_normalizer=tgt_norm,
            batch_size=4,
            shuffle=False,
            conditioning="embedding",
        )
        return sum(len(batch[0]) for batch in ds)

    n_train = _count_rows(train_events)
    n_val = _count_rows(val_events)

    total_rows = 2 * n_events * n_per_event
    assert n_train + n_val == total_rows
    assert n_train == int(counts[np.isin(unique_ids, list(train_events))].sum())
    assert n_val == int(counts[np.isin(unique_ids, list(val_events))].sum())