"""
Build Figure 4/5 change/prechange response summary pickles.

The original Figure 4/5 notebook wrote these dictionaries in-place. This
helper keeps the same `vbn_utils` computation path while making the inputs and
outputs explicit.
"""

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

import numpy as np
import pandas as pd

import decoding_utils as du
import vbn_utils


CONDITIONS = ("active", "passive")
CORE_FLASHES = ("change", "prechange", "shared_nonchange", "nonshared_nonchange")
RUNNING_MATCHED_FLASHES = ("change", "prechange")


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 read_pickle(path: Path):
    with path.open("rb") as file:
        return pickle.load(file)


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 parse_session_ids(session_ids: Optional[Sequence[str]]) -> Optional[np.ndarray]:
    if session_ids is None:
        return None
    return np.array([int(session_id) for session_id in session_ids])


def get_quality_unit_ids_and_sessions(
    units: pd.DataFrame,
    session_ids: Optional[Sequence[str]] = None,
) -> Tuple[np.ndarray, Sequence[str]]:
    unit_filter = du.apply_unit_quality_filter(units)

    parsed_session_ids = parse_session_ids(session_ids)
    if parsed_session_ids is not None:
        unit_filter = unit_filter & units["ecephys_session_id"].isin(parsed_session_ids)

    unit_ids = units.loc[unit_filter, "unit_id"].values
    session_list = [str(session) for session in np.unique(units.loc[unit_filter, "ecephys_session_id"])]
    return unit_ids, session_list


def initialize_flash_data(unit_ids: Sequence[int]) -> dict:
    return {
        unit_id: {
            condition: {flash: [] for flash in CORE_FLASHES}
            for condition in CONDITIONS
        }
        for unit_id in unit_ids
    }


def apply_running_matched_flags(
    stim_table: pd.DataFrame,
    running_matched_indices: dict,
) -> pd.DataFrame:
    stim_table = stim_table.copy()
    for condition in CONDITIONS:
        for flash in RUNNING_MATCHED_FLASHES:
            column = f"{condition}_{flash}_baseline_running_matched"
            stim_table[column] = False
            stim_table.loc[running_matched_indices[condition][flash], column] = True
    return stim_table


def add_shared_nonshared_responses(
    flash_data: dict,
    tensor_file: Path,
    stim_table: pd.DataFrame,
    session_list: Sequence[str],
    unit_ids: Sequence[int],
    condition: str,
    baseline_length: int,
    response_window_length: int,
) -> None:
    shared_data, nonshared_data, unit_id_groups = vbn_utils.change_prechange_matched_psth(
        str(tensor_file),
        stim_table,
        session_list,
        unit_ids,
        baseline_length=baseline_length,
        resp_window_length=response_window_length,
        comparison="shared_nonshared",
    )
    for shared_response, nonshared_response, unit_id in zip(
        np.concatenate(shared_data),
        np.concatenate(nonshared_data),
        np.concatenate(unit_id_groups),
    ):
        flash_data[unit_id][condition]["shared_nonchange"] = shared_response
        flash_data[unit_id][condition]["nonshared_nonchange"] = nonshared_response


def add_change_prechange_responses(
    flash_data: dict,
    tensor_file: Path,
    stim_table: pd.DataFrame,
    session_list: Sequence[str],
    unit_ids: Sequence[int],
    condition: str,
    baseline_length: int,
    response_window_length: int,
) -> None:
    change_data, prechange_data, unit_id_groups = vbn_utils.change_prechange_matched_psth(
        str(tensor_file),
        stim_table,
        session_list,
        unit_ids,
        baseline_length=baseline_length,
        resp_window_length=response_window_length,
        comparison="change_prechange",
    )
    for change_response, prechange_response, unit_id in zip(
        np.concatenate(change_data),
        np.concatenate(prechange_data),
        np.concatenate(unit_id_groups),
    ):
        flash_data[unit_id][condition]["change"] = change_response
        flash_data[unit_id][condition]["prechange"] = prechange_response


def add_omission_responses(
    flash_data: dict,
    tensor_file: Path,
    stim_table: pd.DataFrame,
    session_list: Sequence[str],
    unit_ids: Sequence[int],
    condition: str,
    baseline_length: int,
    response_window_length: int,
) -> None:
    omission_data, unit_id_groups = vbn_utils.unit_averaged_psth(
        str(tensor_file),
        stim_table,
        session_list,
        unit_ids,
        "engaged",
        "omitted",
        baseline_length=baseline_length,
        resp_window_length=response_window_length,
    )
    for omission_response, unit_id in zip(
        np.concatenate(omission_data),
        np.concatenate(unit_id_groups),
    ):
        flash_data[unit_id][condition]["omission"] = omission_response


def build_change_prechange_response_summary(
    unit_table_file: Path,
    stim_table_file: Path,
    active_tensor_file: Path,
    passive_tensor_file: Path,
    output_file: Path,
    session_ids: Optional[Sequence[str]] = None,
    baseline_length: int = 50,
    response_window_length: int = 750,
) -> dict:
    """Build the full active/passive change, prechange, nonchange, and omission summary."""

    units = read_table(unit_table_file)
    stim_table = read_table(stim_table_file)
    unit_ids, session_list = get_quality_unit_ids_and_sessions(units, session_ids=session_ids)
    flash_data = initialize_flash_data(unit_ids)

    for condition, tensor_file in zip(CONDITIONS, [active_tensor_file, passive_tensor_file]):
        add_shared_nonshared_responses(
            flash_data,
            tensor_file,
            stim_table,
            session_list,
            unit_ids,
            condition,
            baseline_length,
            response_window_length,
        )
        add_change_prechange_responses(
            flash_data,
            tensor_file,
            stim_table,
            session_list,
            unit_ids,
            condition,
            baseline_length,
            response_window_length,
        )
        add_omission_responses(
            flash_data,
            tensor_file,
            stim_table,
            session_list,
            unit_ids,
            condition,
            baseline_length,
            response_window_length,
        )

    write_pickle(flash_data, output_file)
    return flash_data


def build_running_matched_response_summary(
    unit_table_file: Path,
    stim_table_file: Path,
    active_tensor_file: Path,
    passive_tensor_file: Path,
    running_matched_indices_file: Path,
    output_file: Path,
    session_ids: Optional[Sequence[str]] = None,
    baseline_length: int = 50,
    response_window_length: int = 750,
) -> dict:
    """Build the running-matched active/passive change and prechange summary."""

    units = read_table(unit_table_file)
    stim_table = read_table(stim_table_file)
    running_matched_indices = read_pickle(running_matched_indices_file)
    stim_table = apply_running_matched_flags(stim_table, running_matched_indices)
    unit_ids, session_list = get_quality_unit_ids_and_sessions(units, session_ids=session_ids)
    flash_data = initialize_flash_data(unit_ids)

    for condition, tensor_file in zip(CONDITIONS, [active_tensor_file, passive_tensor_file]):
        change_data, prechange_data, unit_id_groups = vbn_utils.change_prechange_matched_psth(
            str(tensor_file),
            stim_table,
            session_list,
            unit_ids,
            baseline_length=baseline_length,
            resp_window_length=response_window_length,
            comparison="change_prechange",
            match_running_speed=True,
            match_running_col_prefix=f"{condition}_",
        )

        changes = np.concatenate([change for change in change_data if change.size > 1])
        prechanges = np.concatenate(
            [prechange for prechange, change in zip(prechange_data, change_data) if change.size > 1]
        )
        concatenated_unit_ids = np.concatenate(
            [unit_group for unit_group, change in zip(unit_id_groups, change_data) if change.size > 1]
        )

        for change_response, prechange_response, unit_id in zip(
            changes,
            prechanges,
            concatenated_unit_ids,
        ):
            flash_data[unit_id][condition]["change"] = change_response
            flash_data[unit_id][condition]["prechange"] = prechange_response

    write_pickle(flash_data, output_file)
    return flash_data


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--unit-table-file", type=Path, required=True)
    parser.add_argument("--stim-table-file", type=Path, required=True)
    parser.add_argument("--active-tensor-file", type=Path, required=True)
    parser.add_argument("--passive-tensor-file", type=Path, required=True)
    parser.add_argument("--output-file", type=Path, required=True)
    parser.add_argument("--session-ids", nargs="+", help="Optional session IDs to process")
    parser.add_argument("--baseline-length", type=int, default=50)
    parser.add_argument("--response-window-length", type=int, default=750)
    parser.add_argument(
        "--running-matched",
        action="store_true",
        help="Build the running-matched change/prechange summary.",
    )
    parser.add_argument(
        "--running-matched-indices-file",
        type=Path,
        help="Required with --running-matched.",
    )
    return parser


def main(argv: Optional[Sequence[str]] = None) -> None:
    args = build_parser().parse_args(argv)

    if args.running_matched:
        if args.running_matched_indices_file is None:
            raise ValueError("--running-matched requires --running-matched-indices-file")
        build_running_matched_response_summary(
            unit_table_file=args.unit_table_file,
            stim_table_file=args.stim_table_file,
            active_tensor_file=args.active_tensor_file,
            passive_tensor_file=args.passive_tensor_file,
            running_matched_indices_file=args.running_matched_indices_file,
            output_file=args.output_file,
            session_ids=args.session_ids,
            baseline_length=args.baseline_length,
            response_window_length=args.response_window_length,
        )
    else:
        build_change_prechange_response_summary(
            unit_table_file=args.unit_table_file,
            stim_table_file=args.stim_table_file,
            active_tensor_file=args.active_tensor_file,
            passive_tensor_file=args.passive_tensor_file,
            output_file=args.output_file,
            session_ids=args.session_ids,
            baseline_length=args.baseline_length,
            response_window_length=args.response_window_length,
        )


if __name__ == "__main__":
    main()