"""
Build the master_sessions_table.csv used by figure notebooks.

The base rows come from the AllenSDK VBN 0.5.0 ecephys session table with
abnormal sessions retained. When a master stimulus table is provided, this
helper also adds the notebook-required engaged behavior summary columns:

    engaged_dprime
    engaged_hitcount
    strategy

The strategy label comes from the external behavior-model summary used by the
original notebooks when --behavior-summary-file is provided.
"""

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

import numpy as np
import pandas as pd

from notebook_utils import get_session_engaged_dprime, get_session_engaged_hit_count


DEFAULT_MANIFEST = "visual-behavior-neuropixels_project_manifest_v0.5.0.json"
BEHAVIOR_COLUMNS = ("engaged_dprime", "engaged_hitcount")


def load_cache(cache_dir: Path, manifest: str = DEFAULT_MANIFEST, cache_source: str = "local"):
    """Load a VisualBehaviorNeuropixelsProjectCache without import-time side effects."""
    from allensdk.brain_observatory.behavior.behavior_project_cache.\
        behavior_neuropixels_project_cache import VisualBehaviorNeuropixelsProjectCache

    if cache_source == "s3":
        cache = VisualBehaviorNeuropixelsProjectCache.from_s3_cache(cache_dir=cache_dir)
    else:
        cache = VisualBehaviorNeuropixelsProjectCache.from_local_cache(cache_dir=cache_dir)

    if manifest:
        cache.load_manifest(manifest)
    return cache


def reset_index_column(table: pd.DataFrame, column_name: str) -> pd.DataFrame:
    """Return a copy with `column_name` as a normal column."""
    table = table.copy()
    if column_name in table.columns:
        return table

    table = table.reset_index()
    if column_name in table.columns:
        return table

    first = table.columns[0]
    return table.rename(columns={first: column_name})


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


def read_behavior_summary(path: Path) -> pd.DataFrame:
    suffix = path.suffix.lower()
    if suffix in (".pkl", ".pickle"):
        return pd.read_pickle(path)

    table = pd.read_csv(path)
    unnamed_columns = [col for col in table.columns if str(col).startswith("Unnamed")]
    if "ecephys_session_id" not in table.columns and unnamed_columns:
        table = table.rename(columns={unnamed_columns[0]: "ecephys_session_id"})
        unnamed_columns = unnamed_columns[1:]
    if unnamed_columns:
        table = table.drop(columns=unnamed_columns)
    return table


def load_session_table(
    cache=None,
    session_table_file: Optional[Path] = None,
) -> pd.DataFrame:
    if session_table_file is not None:
        return read_table(session_table_file)

    if cache is None:
        raise ValueError("Either cache or session_table_file is required.")

    return reset_index_column(
        cache.get_ecephys_session_table(filter_abnormalities=False),
        "ecephys_session_id",
    )


def add_engaged_behavior_metrics(
    sessions: pd.DataFrame,
    stim_table: pd.DataFrame,
) -> pd.DataFrame:
    sessions = sessions.copy()
    dprimes = []
    hit_counts = []
    for _, session in sessions.iterrows():
        session_id = session["ecephys_session_id"]
        try:
            dprimes.append(get_session_engaged_dprime(stim_table, session_id))
            hit_counts.append(get_session_engaged_hit_count(stim_table, session_id))
        except Exception:
            dprimes.append(np.nan)
            hit_counts.append(np.nan)

    sessions["engaged_dprime"] = dprimes
    sessions["engaged_hitcount"] = hit_counts
    return sessions


def add_strategy_labels(
    sessions: pd.DataFrame,
    behavior_summary: Optional[pd.DataFrame] = None,
) -> pd.DataFrame:
    sessions = sessions.copy()
    if behavior_summary is None:
        if "strategy" not in sessions.columns:
            sessions["strategy"] = np.nan
        return sessions

    summary = behavior_summary.copy()
    if "ecephys_session_id" in summary.columns:
        summary = summary.set_index("ecephys_session_id", drop=False)

    strategy_column = "strategy_labels"
    if strategy_column not in summary.columns:
        if "strategy" in summary.columns:
            strategy_column = "strategy"
        else:
            raise ValueError(
                "Behavior summary must include a strategy_labels or strategy column."
            )

    sessions["strategy"] = "visual"
    timing_sessions = summary.index[summary[strategy_column] == "timing"]
    sessions.loc[
        sessions["ecephys_session_id"].isin(timing_sessions),
        "strategy",
    ] = "timing"
    return sessions


def build_master_sessions_table(
    output_file: Path,
    cache_dir: Optional[Path] = None,
    session_table_file: Optional[Path] = None,
    stim_table_file: Optional[Path] = None,
    behavior_summary_file: Optional[Path] = None,
    session_ids: Optional[Sequence[int]] = None,
    manifest: str = DEFAULT_MANIFEST,
    cache_source: str = "local",
) -> pd.DataFrame:
    cache = None
    if session_table_file is None:
        if cache_dir is None:
            raise ValueError("Either --cache-dir or --session-table-file is required.")
        cache = load_cache(cache_dir, manifest=manifest, cache_source=cache_source)

    sessions = load_session_table(cache=cache, session_table_file=session_table_file)
    sessions = reset_index_column(sessions, "ecephys_session_id")

    if session_ids is not None:
        session_ids = [int(session_id) for session_id in session_ids]
        sessions = sessions[sessions["ecephys_session_id"].isin(session_ids)]

    if stim_table_file is not None:
        stim_table = read_table(stim_table_file)
        sessions = add_engaged_behavior_metrics(sessions, stim_table)
    else:
        for column in BEHAVIOR_COLUMNS:
            if column not in sessions.columns:
                sessions[column] = np.nan

    behavior_summary = None
    if behavior_summary_file is not None:
        behavior_summary = read_behavior_summary(behavior_summary_file)
    sessions = add_strategy_labels(sessions, behavior_summary)

    output_file.parent.mkdir(parents=True, exist_ok=True)
    sessions.to_csv(output_file, index=False)
    return sessions


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--cache-dir", type=Path)
    parser.add_argument("--session-table-file", type=Path)
    parser.add_argument("--stim-table-file", type=Path)
    parser.add_argument("--behavior-summary-file", type=Path)
    parser.add_argument("--output-file", type=Path, required=True)
    parser.add_argument("--session-ids", nargs="+", type=int)
    parser.add_argument("--manifest", default=DEFAULT_MANIFEST)
    parser.add_argument("--cache-source", choices=["local", "s3"], default="local")
    return parser


def main(argv: Optional[Sequence[str]] = None) -> None:
    args = build_parser().parse_args(argv)
    build_master_sessions_table(
        output_file=args.output_file,
        cache_dir=args.cache_dir,
        session_table_file=args.session_table_file,
        stim_table_file=args.stim_table_file,
        behavior_summary_file=args.behavior_summary_file,
        session_ids=args.session_ids,
        manifest=args.manifest,
        cache_source=args.cache_source,
    )


if __name__ == "__main__":
    main()