"""
Build per-unit receptive-field array files used by Figure 1.

The original script read per-session RF metrics from a hard-coded directory and
saved one ``<unit_id>.npy`` array per significant on-screen RF. This version
keeps that behavior but makes cache, RF metric, and output paths explicit.
"""

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

import numpy as np
import pandas as pd


DEFAULT_MANIFEST = "visual-behavior-neuropixels_project_manifest_v0.5.0.json"
DEFAULT_P_VALUE_THRESHOLD = 0.001


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 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 load_rf_metrics(
    session_id: int,
    rf_metrics_dir: Optional[Path] = None,
    rf_metrics_file: Optional[Path] = None,
) -> pd.DataFrame:
    if rf_metrics_file is not None:
        metrics = read_table(rf_metrics_file)
        if "ecephys_session_id" in metrics.columns:
            metrics = metrics[metrics["ecephys_session_id"] == session_id]
        return metrics

    if rf_metrics_dir is None:
        raise ValueError("Either --rf-metrics-dir or --rf-metrics-file is required.")

    return read_table(rf_metrics_dir / f"{session_id}.csv")


def get_good_rf_metrics(
    rf_metrics: pd.DataFrame,
    p_value_threshold: float = DEFAULT_P_VALUE_THRESHOLD,
    require_on_screen: bool = True,
) -> pd.DataFrame:
    good_rfs = rf_metrics[rf_metrics["p_value_rf"] < p_value_threshold]
    if require_on_screen:
        good_rfs = good_rfs[good_rfs["on_screen_rf"]]
    return good_rfs.set_index("unit_id")


def build_rf_arrays_for_session(
    session,
    rf_metrics: pd.DataFrame,
    output_dir: Path,
    p_value_threshold: float = DEFAULT_P_VALUE_THRESHOLD,
    require_on_screen: bool = True,
) -> Sequence[int]:
    from brain_observatory_utilities.datasets.electrophysiology.\
        receptive_field_mapping import ReceptiveFieldMapping_VBN

    rf = ReceptiveFieldMapping_VBN(session)
    good_rfs = get_good_rf_metrics(
        rf_metrics,
        p_value_threshold=p_value_threshold,
        require_on_screen=require_on_screen,
    )

    output_dir.mkdir(parents=True, exist_ok=True)
    saved_unit_ids = []
    for unit_id in good_rfs.index.values:
        unit_rf = rf.get_receptive_field(unit_id)
        np.save(output_dir / f"{unit_id}.npy", unit_rf)
        saved_unit_ids.append(int(unit_id))
    return saved_unit_ids


def build_rf_arrays(
    session_id: int,
    cache_dir: Path,
    output_dir: Path,
    rf_metrics_dir: Optional[Path] = None,
    rf_metrics_file: Optional[Path] = None,
    p_value_threshold: float = DEFAULT_P_VALUE_THRESHOLD,
    require_on_screen: bool = True,
    manifest: str = DEFAULT_MANIFEST,
    cache_source: str = "local",
) -> Sequence[int]:
    cache = load_cache(cache_dir, manifest=manifest, cache_source=cache_source)
    session = cache.get_ecephys_session(ecephys_session_id=session_id)
    rf_metrics = load_rf_metrics(
        session_id,
        rf_metrics_dir=rf_metrics_dir,
        rf_metrics_file=rf_metrics_file,
    )
    return build_rf_arrays_for_session(
        session,
        rf_metrics,
        output_dir,
        p_value_threshold=p_value_threshold,
        require_on_screen=require_on_screen,
    )


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--session-id", type=int, required=True)
    parser.add_argument("--cache-dir", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--rf-metrics-dir", type=Path)
    parser.add_argument("--rf-metrics-file", type=Path)
    parser.add_argument("--p-value-threshold", type=float, default=DEFAULT_P_VALUE_THRESHOLD)
    parser.add_argument("--include-off-screen", action="store_true")
    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_rf_arrays(
        session_id=args.session_id,
        cache_dir=args.cache_dir,
        output_dir=args.output_dir,
        rf_metrics_dir=args.rf_metrics_dir,
        rf_metrics_file=args.rf_metrics_file,
        p_value_threshold=args.p_value_threshold,
        require_on_screen=not args.include_off_screen,
        manifest=args.manifest,
        cache_source=args.cache_source,
    )


if __name__ == "__main__":
    main()