#!/usr/bin/env python3
"""iHARP batch annotation + ER/AhR semi-quantitative calculation entry point (self-contained).

Full pipeline:
    1. read input xlsx (Name + Calc. MW + Area(Max.) / Group Area)
    2. each compound name -> local ToxCast match -> CASRN + endpoint + AC50 + HitCall
    3. CASRN -> cluster table -> CLUSTER -> curve table -> slope/intercept
    4. concentration = (Area - intercept) / slope; Concentration(ng/mL)_adjust = concentration x ear_coefficient
    5. filter endpoints by pathway, assemble endpoint details
    6. call the self-contained calculation modules to run RQmodified
    7. output 4 files, named as "input-file stem + suffix":
       {stem}_summary.csv / {stem}_endpoint_details.csv /
       {stem}_RQmodified_results.csv / {stem}_RQmodified_sum.csv / {stem}_warnings.txt

Function separation principle (cost consideration):
    - this script = pure local computation (ToxCast + concentration + RQmodified), zero cost, calls no MCP.
    - pubchemlite physicochemical fields (XLogP/BioPathway/ToxicityInfo/KnownUse/DisorderDisease)
      plus CASRN and fastgpt classification fields (MainCategory/SecondaryCategory1/2):
      looked up by the agent via MCP and merged in through --enrich-map. The script itself calls no MCP.
    - fastgpt classification has cost, fully independent and on-demand; pubchemlite goes through local
      MySQL at zero cost, but is still an MCP query merged in via --enrich-map.

Usage:
    python batch_annotate.py --input AhR_CS_test.xlsx --pathway ahr \
        --ear-coefficient 0.00033 --output-dir <output_dir>
    python batch_annotate.py --input ER_CS_test.xlsx --pathway er \
        --ear-coefficient 1.0 --output-dir <output_dir> \
        --enrich-map <supplementary_fields_looked_up_via_mcp.csv>

Parameters:
    --input           input xlsx (columns: Name, Calc. MW, Area(Max.) or Group Area)
    --pathway         ahr or er
    --ear-coefficient concentration adjustment coefficient, default 1.0 (tests commonly use 0.00033)
    --output-dir      output directory, default current directory
    --enrich-map      optional, CSV of supplementary fields looked up by the agent via MCP.
                      columns may include: Name, CASRN, MainCategory, SecondaryCategory1,
                      SecondaryCategory2, XLogP, BioPathway, ToxicityInfo,
                      KnownUse, DisorderDisease (any subset).
                      CASRN overrides the ToxCast match (authoritative source: invitrodb MCP).

Data dependencies (bundled with the skill under assets/):
    - assets/ToxCast_AC50/               ToxCast AC50 data
    - assets/AhR_cluster.xlsx + AhR_curve.xlsx     AhR calibration tables
    - assets/ER_cluster.xlsx + ER_curve.xlsx       ER calibration tables

Note: standard library + in-skill modules only, no third-party dependencies.
"""

from __future__ import annotations

import argparse
import csv
import json
import time
from pathlib import Path
from typing import Any

import ahr_rqmodified
import er_rqmodified
import query_toxcast
import xlsx_reader

RULE_VERSION = "iharp_skill_v1"

_SKILL_ROOT = Path(__file__).resolve().parent.parent
_ASSETS = _SKILL_ROOT / "assets"

# ===== output fields (aligned with the original project; removed legacy EAR/EAR_Rank/Risk_Level fields) =====
SUMMARY_BASE = [
    "Input", "CompoundName", "CASRN",
    "MainCategory", "SecondaryCategory1", "SecondaryCategory2",
    "MatchedEndpoint",
    "XLogP", "BioPathway", "ToxicityInfo", "KnownUse", "DisorderDisease",
    "Endpoint", "AC50",
    "EndpointCount", "ActiveEndpointCount",
    "Concentration(ng/mL)", "Concentration(ng/mL)_adjust",
]

DETAIL_BASE = [
    "Input", "CompoundName", "CASRN",
    "MainCategory", "SecondaryCategory1", "SecondaryCategory2",
    "Endpoint", "AC50", "HitCall",
    "Concentration(ng/mL)", "Concentration(ng/mL)_adjust",
]

ENRICH_COLS = [
    "CASRN", "MainCategory", "SecondaryCategory1", "SecondaryCategory2",
    "XLogP", "BioPathway", "ToxicityInfo", "KnownUse", "DisorderDisease",
]


def _find_col(headers: dict[str, Any], *candidates: str) -> str | None:
    for cand in candidates:
        for col, val in headers.items():
            if val is not None and str(val).strip().lower() == cand.strip().lower():
                return col
    return None


def read_input_rows(xlsx_path: str) -> list[dict[str, Any]]:
    rows = xlsx_reader.read_sheet(xlsx_path)
    if not rows:
        return []
    header = rows[0]
    name_col = _find_col(header, "Name")
    mw_col = _find_col(header, "Calc. MW")
    area_col = _find_col(header, "Area (Max.)", "Area (Max)", "Group Area", "Area")
    if name_col is None:
        raise ValueError(f"input xlsx is missing the Name column, actual header: {header}")
    records: list[dict[str, Any]] = []
    for r in rows[1:]:
        name = str(r.get(name_col) or "").strip() if name_col else ""
        if not name:
            continue
        records.append({
            "name": name,
            "mw": r.get(mw_col) if mw_col else "",
            "area": r.get(area_col) if area_col else "",
        })
    return records


def load_enrich_map(path: str | None) -> dict[str, dict[str, str]]:
    """Read the Name -> supplementary-fields mapping."""
    if not path:
        return {}
    result: dict[str, dict[str, str]] = {}
    with open(path, newline="", encoding="utf-8-sig") as f:
        for row in csv.DictReader(f):
            name = str(row.get("Name") or "").strip()
            if not name:
                continue
            result[name] = {c: str(row.get(c) or "").strip() for c in ENRICH_COLS}
    return result


def _header_map(rows: list[dict[str, Any]]) -> dict[str, str]:
    """Header dict {col_letter: header_name} -> {header_name_upper: col_letter}."""
    if not rows:
        return {}
    hdr = rows[0]
    return {str(v).strip().upper(): k for k, v in hdr.items() if v}


def load_calibration(pathway: str) -> tuple[dict[str, str], dict[str, tuple[float, float]]]:
    if pathway == "ahr":
        cluster_f = _ASSETS / "AhR_cluster.xlsx"
        curve_f = _ASSETS / "AhR_curve.xlsx"
    else:
        cluster_f = _ASSETS / "ER_cluster.xlsx"
        curve_f = _ASSETS / "ER_curve.xlsx"

    # the cluster table column layout differs by pathway (AhR: B=CASRN/D=CLUSTER; ER: C=CASRN/G=CLUSTER),
    # so columns must be located by header name, not hard-coded column letters.
    cluster_rows = xlsx_reader.read_sheet(str(cluster_f))
    chmap = _header_map(cluster_rows)
    cas_col = chmap.get("CASRN")
    clu_col = chmap.get("CLUSTER")
    if not cas_col or not clu_col:
        raise ValueError(f"{cluster_f.name} is missing CASRN/CLUSTER columns, header: {chmap}")

    cas2cls: dict[str, str] = {}
    for r in cluster_rows[1:]:
        cas = str(r.get(cas_col) or "").strip()
        clu = str(r.get(clu_col) or "").strip()
        if cas and clu:
            cas2cls[cas] = clu

    # the curve table layout is identical for both (CLUSTER/slope/intercept), also located by header.
    curve_rows = xlsx_reader.read_sheet(str(curve_f))
    cumap = _header_map(curve_rows)
    cclu_col = cumap.get("CLUSTER")
    slope_col = cumap.get("SLOPE")
    inter_col = cumap.get("INTERCEPT")

    cls2para: dict[str, tuple[float, float]] = {}
    for r in curve_rows[1:]:
        clu = str(r.get(cclu_col) or "").strip() if cclu_col else ""
        try:
            slope = float(r.get(slope_col)) if slope_col else None
        except (TypeError, ValueError):
            slope = None
        try:
            inter = float(r.get(inter_col)) if inter_col else None
        except (TypeError, ValueError):
            inter = None
        if clu and slope is not None and inter is not None:
            cls2para[clu] = (slope, inter)
    return cas2cls, cls2para


def calc_concentration(
    area_val: Any,
    cas: str,
    cas2cls: dict[str, str],
    cls2para: dict[str, tuple[float, float]],
    ear_coefficient: float,
) -> tuple[float | None, float | None, str | None]:
    """Compute concentration, return (raw conc, conc_adjust or None, failure reason or None)."""
    try:
        area = float(area_val)
    except (TypeError, ValueError):
        return None, None, f"Area is empty or invalid: {area_val!r}"
    if cas not in cas2cls:
        return None, None, f"CAS {cas} not in cluster table"
    clu = cas2cls[cas]
    if clu not in cls2para:
        return None, None, f"CLUSTER {clu} missing curve parameters"
    slope, inter = cls2para[clu]
    if abs(slope) < 1e-15:
        return None, None, f"CLUSTER {clu} slope=0"
    conc = (area - inter) / slope
    if conc < 0:
        return None, None, f"negative concentration (Area<intercept): Area={area:.4g}, intercept={inter:.4g}, conc={conc:.4g}"
    return conc, conc * ear_coefficient, None


def pathway_endpoints(pathway: str) -> set[str]:
    if pathway == "ahr":
        return set(ahr_rqmodified.AGONIST_ENDPOINTS) | {ahr_rqmodified.VIABILITY_ENDPOINT}
    eps = set(er_rqmodified.AGONIST_ENDPOINTS) | set(er_rqmodified.ANTAGONIST_ENDPOINTS)
    eps |= set(er_rqmodified.NON_SCORING_ENDPOINTS)
    eps |= set(er_rqmodified.FALSE_POSITIVE_QC_PAIRS.values())
    eps.add(er_rqmodified.LUC_VM7_AGONIST_ENDPOINT)
    eps.add(er_rqmodified.LUC_VM7_0_1NM_QC_ENDPOINT)
    eps.add(er_rqmodified.LUC_VM7_0_1NM_QC_VIABILITY_ENDPOINT)
    return eps


def _fmt(v: Any) -> str:
    if v is None:
        return ""
    if isinstance(v, float):
        return f"{v:.12g}"
    return str(v)


def main() -> int:
    parser = argparse.ArgumentParser(description="iHARP batch annotation + ER/AhR semi-quantitative")
    parser.add_argument("--input", required=True, help="input xlsx path")
    parser.add_argument("--pathway", required=True, choices=["ahr", "er"])
    parser.add_argument("--ear-coefficient", type=float, default=1.0, help="concentration adjustment coefficient, default 1.0")
    parser.add_argument("--output-dir", default=".", help="output directory")
    parser.add_argument("--enrich-map", help="optional CSV of supplementary fields looked up by the agent via MCP")
    parser.add_argument("--toxcast-dir", help="optional, override the ToxCast folder")
    args = parser.parse_args()

    t0 = time.time()

    if not Path(args.input).is_file():
        print(json.dumps({"ok": False, "error": f"input file not found: {args.input}"}, ensure_ascii=False))
        return 2

    try:
        records = read_input_rows(args.input)
    except Exception as exc:
        print(json.dumps({"ok": False, "error": f"failed to read input: {exc}"}, ensure_ascii=False))
        return 2

    toxcast_dir = Path(args.toxcast_dir) if args.toxcast_dir else query_toxcast.probe_toxcast_dir()
    if toxcast_dir is None or not toxcast_dir.is_dir():
        print(json.dumps({"ok": False, "error": "ToxCast folder not found"}, ensure_ascii=False))
        return 2
    tox_index = query_toxcast.load_index(toxcast_dir)
    cas2cls, cls2para = load_calibration(args.pathway)
    enrich = load_enrich_map(args.enrich_map)
    wanted = pathway_endpoints(args.pathway)

    # assemble
    calc_detail_rows: list[dict[str, str]] = []   # for the calculation modules (contains Concentration(ng/mL)_adjusted)
    out_detail_rows: list[dict[str, str]] = []    # endpoint_details output
    summary_rows: list[dict[str, str]] = []       # summary output (one row per compound)
    stats = {"total": len(records), "matched": 0, "no_cas": 0, "no_conc": 0, "no_endpoint": 0}
    warnings: list[str] = []

    for rec in records:
        name = rec["name"]
        enrich_row = enrich.get(name, {})
        enrich_cas = enrich_row.get("CASRN")

        # query ToxCast: prefer the enrich-map CAS (invitrodb authoritative), otherwise match by name
        details: list[dict[str, str]] = []
        if enrich_cas:
            details = query_toxcast.query_details(tox_index, enrich_cas)
        if not details:
            details = query_toxcast.query_details(tox_index, name)

        # final CAS: enrich CAS first, otherwise the ToxCast match result
        cas = enrich_cas or (details[0].get("CASRN", "") if details else "")
        if not cas:
            stats["no_cas"] += 1
            warnings.append(f"no CASRN obtained, skipped: {name}")
            continue

        conc, conc_adjust, conc_err = calc_concentration(rec["area"], cas, cas2cls, cls2para, args.ear_coefficient)
        if conc_adjust is None:
            stats["no_conc"] += 1
            warnings.append(f"concentration calculation failed, skipped: {name}; reason: {conc_err}")
            continue

        sel = [d for d in details if d["Endpoint"] in wanted]
        if not sel:
            stats["no_endpoint"] += 1
            warnings.append(f"no {args.pathway.upper()}-related endpoint, skipped: {name}")
            continue

        stats["matched"] += 1
        matched_endpoints = sorted({d["Endpoint"] for d in sel})
        active_count = sum(1 for d in sel if d["HitCall"].strip().lower() == "active")

        # endpoint details (calculation-module input + output)
        for d in sel:
            calc_detail_rows.append({
                "CompoundName": d["CompoundName"] or name,
                "Endpoint": d["Endpoint"],
                "AC50": d["AC50"],
                "HitCall": d["HitCall"],
                "Concentration(ng/mL)_adjusted": _fmt(conc_adjust),
            })
            out_detail_rows.append({
                "Input": name,
                "CompoundName": d["CompoundName"] or name,
                "CASRN": cas,
                "MainCategory": enrich_row.get("MainCategory", ""),
                "SecondaryCategory1": enrich_row.get("SecondaryCategory1", ""),
                "SecondaryCategory2": enrich_row.get("SecondaryCategory2", ""),
                "Endpoint": d["Endpoint"],
                "AC50": d["AC50"],
                "HitCall": d["HitCall"],
                "Concentration(ng/mL)": _fmt(conc),
                "Concentration(ng/mL)_adjust": _fmt(conc_adjust),
            })

        summary_rows.append({
            "Input": name,
            "CompoundName": d["CompoundName"] or name if sel else name,
            "CASRN": cas,
            "MainCategory": enrich_row.get("MainCategory", ""),
            "SecondaryCategory1": enrich_row.get("SecondaryCategory1", ""),
            "SecondaryCategory2": enrich_row.get("SecondaryCategory2", ""),
            "MatchedEndpoint": "; ".join(matched_endpoints),
            "XLogP": enrich_row.get("XLogP", ""),
            "BioPathway": enrich_row.get("BioPathway", ""),
            "ToxicityInfo": enrich_row.get("ToxicityInfo", ""),
            "KnownUse": enrich_row.get("KnownUse", ""),
            "DisorderDisease": enrich_row.get("DisorderDisease", ""),
            "Endpoint": "; ".join(matched_endpoints),
            "AC50": "",
            "EndpointCount": str(len(matched_endpoints)),
            "ActiveEndpointCount": str(active_count),
            "Concentration(ng/mL)": _fmt(conc),
            "Concentration(ng/mL)_adjust": _fmt(conc_adjust),
        })

    if not calc_detail_rows:
        print(json.dumps({"ok": False, "error": "no valid endpoint details to compute", "stats": stats, "warnings": warnings}, ensure_ascii=False))
        return 1

    # run RQmodified
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    reference = _ASSETS / ("er_reference_endpoint_details.csv" if args.pathway == "er" else "ahr_reference_endpoint_details.csv")

    if args.pathway == "ahr":
        rq_rows, calc_warnings, totals = ahr_rqmodified.calculate_ahrrqmodified(
            calc_detail_rows, ahr_rqmodified.read_csv(reference)
        )
        rq_results_cols = ahr_rqmodified.DETAIL_COLUMNS          # CompoundName,concentration,CC50,AhR_Score,EC,RQmodified_ag
        rq_sum_cols = ahr_rqmodified.SUMMARY_COLUMNS              # SUM_RQmodified
        rq_key = "RQmodified_ag"
    else:
        rq_rows, calc_warnings, totals = er_rqmodified.calculate_rqmodified(
            calc_detail_rows, er_rqmodified.read_csv(reference)
        )
        rq_results_cols = er_rqmodified.DETAIL_COLUMNS
        rq_sum_cols = er_rqmodified.SUMMARY_COLUMNS
        rq_key = "RQmodified_ag"

    warnings.extend(calc_warnings)

    # sort rq_rows by risk descending (empty/No Data at the end)
    def _rq_sort_key(r):
        keys = ["RQmodified_ag"] if args.pathway == "ahr" else ["RQmodified_ag", "RQmodified_ant"]
        vals = []
        for k in keys:
            x = r.get(k)
            if x not in (None, ""):
                try:
                    vals.append(float(x))
                except (TypeError, ValueError):
                    pass
        if not vals:
            return (1, 0.0)
        return (0, -max(vals))

    rq_rows.sort(key=_rq_sort_key)

    # aggregate RQmodified into summary by CompoundName (take max)
    rq_by_compound: dict[str, float] = {}
    for r in rq_rows:
        cname = str(r.get("CompoundName", "")).strip()
        val = r.get(rq_key)
        try:
            fv = float(val) if val not in (None, "") else None
        except (TypeError, ValueError):
            fv = None
        if cname and fv is not None:
            rq_by_compound[cname] = max(rq_by_compound.get(cname, -1e18), fv)

    for s in summary_rows:
        cname = s["CompoundName"]
        if args.pathway == "ahr":
            s["RQmodified_ag"] = _fmt(rq_by_compound.get(cname)) if cname in rq_by_compound else ""
        else:
            # ER: rq_rows keyed by (compound, concentration); take the compound's ag/ant max
            ag_vals = [float(r["RQmodified_ag"]) for r in rq_rows if str(r.get("CompoundName")).strip() == cname and r.get("RQmodified_ag") not in (None, "")]
            ant_vals = [float(r["RQmodified_ant"]) for r in rq_rows if str(r.get("CompoundName")).strip() == cname and r.get("RQmodified_ant") not in (None, "")]
            s["RQmodified_ag"] = _fmt(max(ag_vals)) if ag_vals else ""
            s["RQmodified_ant"] = _fmt(max(ant_vals)) if ant_vals else ""

    # summary columns (base + RQmodified)
    summary_cols = SUMMARY_BASE + (["RQmodified_ag"] if args.pathway == "ahr" else ["RQmodified_ag", "RQmodified_ant"])

    # write files: named as "input-file stem + suffix"
    input_stem = Path(args.input).stem
    detail_path = output_dir / f"{input_stem}_endpoint_details.csv"
    summary_path = output_dir / f"{input_stem}_summary.csv"
    rq_results_path = output_dir / f"{input_stem}_RQmodified_results.csv"
    rq_sum_path = output_dir / f"{input_stem}_RQmodified_sum.csv"
    warning_path = output_dir / f"{input_stem}_warnings.txt"

    with summary_path.open("w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=summary_cols)
        w.writeheader()
        for s in summary_rows:
            w.writerow({c: s.get(c, "") for c in summary_cols})

    with detail_path.open("w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=DETAIL_BASE)
        w.writeheader()
        for d in out_detail_rows:
            w.writerow({c: d.get(c, "") for c in DETAIL_BASE})

    with rq_results_path.open("w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=rq_results_cols)
        w.writeheader()
        for r in rq_rows:
            w.writerow({c: ("" if r.get(c) is None else r.get(c)) for c in rq_results_cols})

    with rq_sum_path.open("w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=rq_sum_cols)
        w.writeheader()
        w.writerow({c: totals.get(c) for c in rq_sum_cols})

    warning_path.write_text("\n".join(warnings) + ("\n" if warnings else ""), encoding="utf-8")

    elapsed = time.time() - t0
    result = {
        "ok": True,
        "pathway": args.pathway.upper(),
        "rule_version": RULE_VERSION,
        "ear_coefficient": args.ear_coefficient,
        "stats": stats,
        "files": {
            "summary": str(summary_path),
            "endpoint_details": str(detail_path),
            "rqmodified_results": str(rq_results_path),
            "rqmodified_sum": str(rq_sum_path),
            "warnings": str(warning_path),
        },
        "totals": totals,
        "elapsed_seconds": round(elapsed, 2),
    }
    print(json.dumps(result, ensure_ascii=False, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())