"""Turn the real fly PN -> KC connections into the wiring matrix of our mushroom body model.

    # A. already in the repo: data/connectome/kc_glomerulus_R.csv (extracted from male-cns v1.0)
    python experiments/connectome_wiring.py

    # B. rebuild the CSV yourself through neuPrint (free token, a few MB):
    export NEUPRINT_APPLICATION_CREDENTIALS=<your token>
    python experiments/connectome_wiring.py --source neuprint

    # C. rebuild it from the public Feather files (no token, ~1.1 GB download):
    python experiments/connectome_wiring.py --source files --files-folder downloads/

Our e-nose has 8 sensor channels -> 16 PNs (ON and OFF). The fly has ~50 glomeruli.
Which glomerulus stands in for which sensor PN is a design choice, written to
channel_map.csv so it can be read and edited:
  * co2+  -> V glomerulus: the fly's CO2 receptor neurons (Gr21a/Gr63a) go only to V
             (Suh et al. 2004). This one has a biological reason.
  * all other PNs -> the glomeruli that send the most synapses to KCs, in that order.
             This part is an engineering choice, not biology.
Outputs in data/connectome/: kc_glomerulus_R.csv, channel_map.csv, kc_wiring.npy, connectome_summary.png
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

import numpy as np
import pandas as pd

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from enose import connectome  # noqa: E402
from enose.common import output_dir, plt, save_figure  # noqa: E402
from enose.simulate import ODOR_CHANNELS  # noqa: E402

FOLDER = Path(__file__).resolve().parents[1] / "data" / "connectome"
BIOLOGICAL: dict[str, str] = {"co2+": "V"}


def pn_names() -> list[str]:
    """Same order as MushroomBodyClassifier.on_off: all ON PNs, then all OFF PNs."""
    return [f"{name}+" for name in ODOR_CHANNELS] + [f"{name}-" for name in ODOR_CHANNELS]


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--source", choices=("csv", "neuprint", "files"), default="csv")
    parser.add_argument("--files-folder", type=Path, default=Path("downloads"))
    parser.add_argument("--side", default="R", help="R or L brain hemisphere")
    arguments = parser.parse_args()
    folder = output_dir(FOLDER)
    table_path = folder / f"kc_glomerulus_{arguments.side}.csv"

    if arguments.source == "neuprint":
        table = connectome.from_neuprint(side=arguments.side)
    elif arguments.source == "files":
        table = connectome.from_flat_files(arguments.files_folder, side=arguments.side)
    else:
        table = pd.read_csv(table_path)
    if arguments.source != "csv":
        table.to_csv(table_path, index=False)
        print(f"已寫入 {table_path}")

    per_glomerulus = table.groupby("glomerulus")["synapses"].sum().sort_values(ascending=False)
    kc_count = table["kc_body_id"].nunique()
    print(f"{kc_count} 個 Kenyon 細胞、{len(per_glomerulus)} 個嗅小球、{int(table['synapses'].sum())} 個 PN→KC 突觸")

    names = pn_names()
    mapping: dict[str, tuple[str, str]] = {}
    for pn, glomerulus in BIOLOGICAL.items():
        if glomerulus in per_glomerulus.index:
            mapping[pn] = (glomerulus, "生物對應")
        else:
            print(f"注意：{glomerulus} 嗅小球在這個資料裡沒有連到 KC，{pn} 改用一般規則")
    remaining = iter(g for g in per_glomerulus.index if g not in {g for g, _ in mapping.values()})
    for pn in names:
        if pn not in mapping:
            mapping[pn] = (next(remaining), "依 KC 輸入量排序（工程選擇）")
    channel_map = pd.DataFrame([{"pn": pn, "glomerulus": mapping[pn][0], "reason": mapping[pn][1],
                                 "synapses_to_kc": int(per_glomerulus[mapping[pn][0]])} for pn in names])
    channel_map.to_csv(folder / "channel_map.csv", index=False)
    print(channel_map.to_string(index=False))

    wiring, kc_ids = connectome.wiring_matrix(table, list(channel_map["glomerulus"]))
    np.save(folder / "kc_wiring.npy", wiring)
    inputs_per_kc = (wiring > 0).sum(axis=1)
    print(f"\n接線矩陣 {wiring.shape[0]} 個 KC × {wiring.shape[1]} 個 PN；"
          f"每個 KC 平均接 {inputs_per_kc.mean():.1f} 個所選嗅小球（中位數 {np.median(inputs_per_kc):.0f}）")

    all_inputs = table.groupby("kc_body_id")["glomerulus"].nunique()
    figure, axes = plt.subplots(1, 2, figsize=(12, 4))
    axes[0].hist(all_inputs, bins=range(0, all_inputs.max() + 2), color="#5eead4", edgecolor="black")
    axes[0].set_xlabel("每個 KC 接收幾個嗅小球的輸入")
    axes[0].set_ylabel("KC 數量")
    axes[0].set_title(f"真實果蠅（{arguments.side} 側，全部嗅小球）")
    top = per_glomerulus.head(25)
    axes[1].bar(top.index, top.to_numpy(), color=["#fbbf24" if g in set(channel_map["glomerulus"]) else "#a78bfa"
                                                 for g in top.index])
    axes[1].tick_params(axis="x", rotation=70)
    axes[1].set_title("各嗅小球送到 KC 的突觸數（黃＝被電子鼻 PN 使用）")
    save_figure(figure, folder / "connectome_summary.png")


if __name__ == "__main__":
    main()