"""
Build running-matched change/prechange trial-index pickles.

No source generator for this Figure 4/5 asset was found in the repository, so
this helper is a reconstruction rather than exact-preserved notebook logic. It
recreates the expected input shape for the downstream notebook code: for each
active/passive condition, select nonshared change and prechange trials matched
on the corresponding baseline running speed.
"""

import argparse
import pickle
from pathlib import Path
from typing import Dict, Optional, Sequence, Tuple

import numpy as np
import pandas as pd


CONDITIONS = ("active", "passive")
FLASHES = ("change", "prechange")
DEFAULT_GROUP_COLUMNS = ("session_id", "image_name")


def read_table(path: Path) -> pd.DataFrame:
    table = pd.read_csv(path)
    unnamed_columns = [col for col in table.columns if col.startswith("Unnamed:")]
    if unnamed_columns:
        table = table.drop(columns=unnamed_columns)
    return table


def write_pickle(data: object, output_file: Path) -> None:
    output_file.parent.mkdir(parents=True, exist_ok=True)
    with output_file.open("wb") as file:
        pickle.dump(data, file)


def require_columns(table: pd.DataFrame, columns: Sequence[str]) -> None:
    missing = [column for column in columns if column not in table.columns]
    if missing:
        raise ValueError(f"Missing required columns: {missing}")


def as_bool(series: pd.Series) -> pd.Series:
    if series.dtype == bool:
        return series.fillna(False)
    return series.replace({"True": True, "False": False}).fillna(False).astype(bool)


def get_candidate_masks(stim_table: pd.DataFrame) -> Tuple[pd.Series, pd.Series]:
    require_columns(
        stim_table,
        ["is_change", "is_prechange", "hit", "is_shared", "engaged", "image_name", "session_id"],
    )
    hit = as_bool(stim_table["hit"])
    engaged = as_bool(stim_table["engaged"])
    nonshared = ~as_bool(stim_table["is_shared"])

    change = as_bool(stim_table["is_change"]) & hit & engaged & nonshared
    prechange = as_bool(stim_table["is_prechange"]) & hit & engaged & nonshared
    return change, prechange


def nearest_neighbor_match(
    change_trials: pd.DataFrame,
    prechange_trials: pd.DataFrame,
    speed_column: str,
    max_speed_difference: Optional[float] = None,
) -> Tuple[Sequence[int], Sequence[int]]:
    """
    Pair change and prechange trials by baseline running speed.

    The condition with fewer valid trials determines the number of pairs. Each
    trial from that smaller set is kept and paired with the unused trial from
    the other condition that has the closest value in ``speed_column``.
    """

    change_trials = change_trials.dropna(subset=[speed_column])
    prechange_trials = prechange_trials.dropna(subset=[speed_column])
    if len(change_trials) == 0 or len(prechange_trials) == 0:
        return [], []

    if len(change_trials) <= len(prechange_trials):
        anchor = change_trials.sort_values([speed_column], kind="mergesort")
        pool = prechange_trials.copy()
        anchor_is_change = True
    else:
        anchor = prechange_trials.sort_values([speed_column], kind="mergesort")
        pool = change_trials.copy()
        anchor_is_change = False

    selected_change = []
    selected_prechange = []
    for anchor_index, anchor_row in anchor.iterrows():
        if len(pool) == 0:
            break

        differences = (pool[speed_column] - anchor_row[speed_column]).abs()
        matched_index = differences.sort_values(kind="mergesort").index[0]
        matched_difference = differences.loc[matched_index]
        if max_speed_difference is not None and matched_difference > max_speed_difference:
            continue

        if anchor_is_change:
            selected_change.append(int(anchor_index))
            selected_prechange.append(int(matched_index))
        else:
            selected_change.append(int(matched_index))
            selected_prechange.append(int(anchor_index))

        pool = pool.drop(index=matched_index)

    return selected_change, selected_prechange


def build_condition_matches(
    stim_table: pd.DataFrame,
    condition: str,
    group_columns: Sequence[str] = DEFAULT_GROUP_COLUMNS,
    max_speed_difference: Optional[float] = None,
) -> Dict[str, Sequence[int]]:
    speed_column = f"{condition}_baseline_running_speed"
    require_columns(stim_table, [speed_column, *group_columns])

    change_mask, prechange_mask = get_candidate_masks(stim_table)
    candidate_rows = stim_table.loc[change_mask | prechange_mask, list(group_columns)]

    matched_change_indices = []
    matched_prechange_indices = []
    for _, group in candidate_rows.drop_duplicates().iterrows():
        group_mask = pd.Series(True, index=stim_table.index)
        for column in group_columns:
            group_mask &= stim_table[column] == group[column]

        change_trials = stim_table.loc[group_mask & change_mask]
        prechange_trials = stim_table.loc[group_mask & prechange_mask]
        change_indices, prechange_indices = nearest_neighbor_match(
            change_trials,
            prechange_trials,
            speed_column=speed_column,
            max_speed_difference=max_speed_difference,
        )
        matched_change_indices.extend(change_indices)
        matched_prechange_indices.extend(prechange_indices)

    return {
        "change": sorted(matched_change_indices),
        "prechange": sorted(matched_prechange_indices),
    }


def build_running_matched_trial_indices(
    stim_table_file: Path,
    output_file: Path,
    group_columns: Sequence[str] = DEFAULT_GROUP_COLUMNS,
    max_speed_difference: Optional[float] = None,
) -> Dict[str, Dict[str, Sequence[int]]]:
    stim_table = read_table(stim_table_file)

    matched_indices = {
        condition: build_condition_matches(
            stim_table,
            condition=condition,
            group_columns=group_columns,
            max_speed_difference=max_speed_difference,
        )
        for condition in CONDITIONS
    }

    write_pickle(matched_indices, output_file)
    return matched_indices


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--stim-table-file", type=Path, required=True)
    parser.add_argument("--output-file", type=Path, required=True)
    parser.add_argument(
        "--group-columns",
        nargs="+",
        default=list(DEFAULT_GROUP_COLUMNS),
        help="Columns within which to match trials. Defaults to session_id image_name.",
    )
    parser.add_argument(
        "--max-speed-difference",
        type=float,
        default=None,
        help="Optional maximum allowed baseline-running-speed difference for a matched pair.",
    )
    return parser


def main(argv: Optional[Sequence[str]] = None) -> None:
    args = build_parser().parse_args(argv)
    matched_indices = build_running_matched_trial_indices(
        stim_table_file=args.stim_table_file,
        output_file=args.output_file,
        group_columns=args.group_columns,
        max_speed_difference=args.max_speed_difference,
    )

    for condition in CONDITIONS:
        change_count = len(matched_indices[condition]["change"])
        prechange_count = len(matched_indices[condition]["prechange"])
        print(f"{condition}: {change_count} change, {prechange_count} prechange")


if __name__ == "__main__":
    main()