"""Reproduce ProteinIQ's descriptive mammal chromosome-count analysis.

Python 3.10+, standard library only. Run without arguments for local reproduction;
--check compares outputs without writing; --refresh-inputs retrieves pinned sources.
No rounded ACC haploid field, inferred missing zero, or taxonomic synonym is used.
"""
from __future__ import annotations

import argparse
import csv
import hashlib
import io
import json
import re
import statistics
import sys
from collections import Counter
from pathlib import Path
from urllib.request import urlopen
from xml.etree import ElementTree as ET
from zipfile import ZipFile

HERE = Path(__file__).resolve().parent
ACC_COMMIT = "15afa843279faddf0e5161ab3476b2348afdbfdb"
ACC_URL = f"https://raw.githubusercontent.com/cromanpa94/ACC/{ACC_COMMIT}/category/ACCDB_Mar24.csv"
SUPPLEMENT_URL = "https://www.ebi.ac.uk/europepmc/webservices/rest/PMC6590138/supplementaryFiles"
SUPPLEMENT_MEMBER = "EVO-73-511-s001.zip"
TABLE_MEMBER = "supp.table.1.docx"
TARGETS = {"Cat": 38, "Human": 46, "Horse": 64, "Dog / gray wolf": 78}


def sha256(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


def csv_text(rows: list[dict], fields: list[str] | None = None) -> str:
    stream = io.StringIO(newline="")
    writer = csv.DictWriter(stream, fieldnames=fields or list(rows[0]), lineterminator="\n")
    writer.writeheader()
    writer.writerows(rows)
    return stream.getvalue()


def parse_counts(value: str) -> tuple[int, ...]:
    """Deduplicate repeated explicit observations; reject ranges and annotations."""
    if not value.strip():
        return ()
    parts = value.split(",")
    if any(not re.fullmatch(r"[1-9][0-9]*", p.strip()) for p in parts):
        raise ValueError(f"Not an explicit positive integer list: {value!r}")
    return tuple(sorted({int(p) for p in parts}))


def extract_acc(data: bytes) -> list[dict]:
    rows = csv.DictReader(io.StringIO(data.decode("utf-8-sig")))
    # source_record is one-based CSV data-record position, excluding the header.
    return [{"source_record": i, **r} for i, r in enumerate(rows, 1) if r["class"] == "Mammalia"]


def extract_blackmon(docx: bytes) -> list[dict]:
    ns = {"w": "http://schemas.openxmlformats.org/wordprocessingml/2006/main"}
    with ZipFile(io.BytesIO(docx)) as archive:
        root = ET.fromstring(archive.read("word/document.xml"))
    result = []
    for table_id, table in enumerate(root.findall(".//w:tbl", ns), 1):
        for row_id, row in enumerate(table.findall("w:tr", ns), 1):
            cells = ["".join(n.text or "" for n in c.findall(".//w:t", ns))
                     for c in row.findall("w:tc", ns)]
            if cells[0] == "species":
                continue
            if len(cells) != 5 or len(parse_counts(cells[1])) != 1:
                raise ValueError(f"Unexpected Table S1 row: {cells}")
            result.append(dict(table=table_id, table_row=row_id, species=cells[0],
                               female_haploid=cells[1]))
    return result


def indexed(rows: list[dict], field: str) -> dict[str, dict]:
    result = {}
    for row in rows:
        # The only name transformation: underscores to spaces and trim.
        name = row[field].replace("_", " ").strip()
        if name in result:
            raise ValueError(f"Duplicate species key: {name}")
        result[name] = row
    return result


def classify(acc: dict | None, blackmon: dict | None) -> str:
    if acc is None:
        return "blackmon_only"
    if blackmon is None:
        return "acc_only"
    diploid = parse_counts(acc["ploidy_2N"])
    haploid = parse_counts(acc["ploidy_N"])
    reference_n = int(blackmon["female_haploid"])
    if not diploid:
        return "missing_acc_diploid"
    if reference_n * 2 not in diploid:
        return "count_disagreement"
    if len(diploid) > 1:
        return "multiple_acc_counts_including_reference"
    if haploid and haploid != (reference_n,):
        return "acc_haploid_conflict"
    return "included"


def describe(values: list[int]) -> dict:
    if not values:
        raise ValueError("Empty sample")
    frequencies = Counter(values)
    mode_frequency = max(frequencies.values())
    return dict(n=len(values), median=statistics.median(values),
                modes=sorted(k for k, v in frequencies.items() if v == mode_frequency),
                mode_frequency=mode_frequency, minimum=min(values), maximum=max(values))


def rank_counts(values: list[int], threshold: int) -> dict:
    if not values:
        raise ValueError("Empty sample")
    below = sum(v < threshold for v in values)
    equal = values.count(threshold)
    return dict(count=threshold, n=len(values), below=below, equal=equal,
                above=len(values) - below - equal,
                below_percent=round(100 * below / len(values), 6),
                equal_percent=round(100 * equal / len(values), 6))


def analyze(acc_rows: list[dict], blackmon_rows: list[dict]) -> tuple[dict, list[dict]]:
    acc = indexed(acc_rows, "Species_corrected")
    blackmon = indexed(blackmon_rows, "species")
    rows = []
    for species in sorted(acc.keys() | blackmon.keys()):
        a, b = acc.get(species), blackmon.get(species)
        state = classify(a, b)
        n = int(b["female_haploid"]) if b else None
        rows.append(dict(species=species, decision=state,
                         acc_source_record=a["source_record"] if a else "",
                         acc_order=a["order"] if a else "",
                         acc_family=a["family"] if a else "",
                         acc_genus=a["genus"] if a else "",
                         acc_diploid_raw=a["ploidy_2N"] if a else "",
                         acc_haploid_raw=a["ploidy_N"] if a else "",
                         acc_diploid_distinct=";".join(map(str, parse_counts(a["ploidy_2N"]))) if a else "",
                         blackmon_table=b["table"] if b else "",
                         blackmon_table_row=b["table_row"] if b else "",
                         blackmon_female_haploid=n if b else "",
                         blackmon_female_diploid=2*n if b else "",
                         analysis_diploid=2*n if state == "included" else "",
                         acc_reference=a["Reference"] if a else ""))
    strict = [r["analysis_diploid"] for r in rows if r["decision"] == "included"]
    # Alternative treatments are explicit, not asserted corrections.
    no_haploid_screen = [r["blackmon_female_diploid"] for r in rows
                         if r["decision"] in {"included", "acc_haploid_conflict"}]
    female_supported = [r["blackmon_female_diploid"] for r in rows
                        if r["decision"] in {"included", "acc_haploid_conflict",
                                             "multiple_acc_counts_including_reference"}]
    all_blackmon = [2 * int(r["female_haploid"]) for r in blackmon_rows]
    scenarios = {}
    for label, values in [("strict_agreement", strict),
                          ("ignore_acc_haploid_conflicts", no_haploid_screen),
                          ("allow_multiple_acc_counts", female_supported),
                          ("blackmon_table_without_acc_screen", all_blackmon)]:
        scenarios[label] = {**describe(values), "human_comparison": rank_counts(values, 46)}
    frequencies = Counter(strict)
    summary = dict(analysis_date="2026-09-19", counting_unit="One published species label per record; female diploid complement",
                   acc_mammal_records=len(acc), blackmon_records=len(blackmon),
                   exact_name_matches=len(acc.keys() & blackmon.keys()),
                   decisions=dict(sorted(Counter(r["decision"] for r in rows).items())),
                   primary=describe(strict),
                   targets={label: rank_counts(strict, count) for label, count in TARGETS.items()},
                   frequencies=[dict(chromosomes=k, species=v) for k,v in sorted(frequencies.items())],
                   bins=[dict(lower=lo, upper=lo+9, species=sum(lo <= v <= lo+9 for v in strict))
                         for lo in range(10, 110, 10)],
                   order_composition=dict(sorted(Counter(r["acc_order"] for r in rows if r["decision"] == "included").items())),
                   sensitivity=scenarios)
    if sum(r["species"] for r in summary["bins"]) != len(strict):
        raise ValueError("Histogram bins do not cover the sample")
    return summary, rows


def refresh_inputs() -> None:
    provenance = json.loads((HERE / "provenance.json").read_text(encoding="utf-8"))
    with urlopen(ACC_URL, timeout=60) as response:
        acc = response.read()
    if sha256(acc) != provenance["sources"]["acc"]["original_sha256"]:
        raise ValueError("ACC source checksum mismatch")
    with urlopen(SUPPLEMENT_URL, timeout=60) as response:
        archive_bytes = response.read()
    with ZipFile(io.BytesIO(archive_bytes)) as archive:
        supplement = archive.read(SUPPLEMENT_MEMBER)
    if sha256(supplement) != provenance["sources"]["blackmon"]["supplement_sha256"]:
        raise ValueError("Blackmon supplement checksum mismatch")
    with ZipFile(io.BytesIO(supplement)) as archive:
        docx = archive.read(TABLE_MEMBER)
    if sha256(docx) != provenance["sources"]["blackmon"]["table_docx_sha256"]:
        raise ValueError("Blackmon Table S1 checksum mismatch")
    snapshots = {"acc-mammals.csv": csv_text(extract_acc(acc)),
                 "blackmon-table-s1.csv": csv_text(extract_blackmon(docx))}
    for name, text in snapshots.items():
        if sha256(text.encode()) != provenance["input_sha256"][name]:
            raise ValueError(f"Extracted snapshot mismatch: {name}")
        (HERE / name).write_bytes(text.encode())


def results() -> dict[str, str]:
    provenance = json.loads((HERE / "provenance.json").read_text(encoding="utf-8"))
    for name, checksum in provenance["input_sha256"].items():
        if sha256((HERE / name).read_bytes()) != checksum:
            raise ValueError(f"Input checksum mismatch: {name}")
    with (HERE / "acc-mammals.csv").open(encoding="utf-8", newline="") as stream:
        acc = list(csv.DictReader(stream))
    with (HERE / "blackmon-table-s1.csv").open(encoding="utf-8", newline="") as stream:
        blackmon = list(csv.DictReader(stream))
    summary, rows = analyze(acc, blackmon)
    return {"summary.json": json.dumps(summary, indent=2, ensure_ascii=False) + "\n",
            "species-audit.csv": csv_text(rows),
            "distribution.csv": csv_text(summary["frequencies"]),
            "comparisons.csv": csv_text([dict(species=k, **v) for k,v in summary["targets"].items()])}


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--check", action="store_true")
    parser.add_argument("--refresh-inputs", action="store_true")
    args = parser.parse_args()
    if args.refresh_inputs:
        refresh_inputs()
    for name, text in results().items():
        output = HERE / name
        if args.check:
            if not output.exists() or output.read_bytes() != text.encode():
                sys.exit(f"Output mismatch: {name}")
        else:
            output.write_bytes(text.encode())
    print("All four analysis outputs verified." if args.check else "Wrote all four analysis outputs.")


if __name__ == "__main__":
    main()