"""Human-pig similarity figures for ProteinIQ's pig guide.

Reads the gene-level Ensembl 116 ortholog table published with the
human-animal guide and, optionally, the LASTZ alignment statistics that
Ensembl stores for human-referenced pairwise genome alignments.

Usage:
    python analyze_pig.py --out .                 # gene-level figures only
    python analyze_pig.py --out . --lastz         # also query Ensembl (needs pymysql)
    python analyze_pig.py --out . --lastz-json lastz-tags.json   # reuse a saved query
"""

from __future__ import annotations

import argparse
import bisect
import csv
import gzip
import json
import statistics
from pathlib import Path

HERE = Path(__file__).resolve().parent
GENE_TABLE = HERE.parent / "human-animal-dna-similarity" / "human-gene-orthologs.csv.gz"

# Columns in the gene table and their Ensembl 116 LASTZ method-link-species-set ids.
SPECIES = {
    "chimpanzee": 1098,
    "macaque": 1562,
    "pig": 1039,
    "horse": 1347,
    "mouse": 1960,
}

# Identity bins (percent of the human protein length), lower bound inclusive.
BINS = [
    ("Under 50%", 0, 50),
    ("50–70%", 50, 70),
    ("70–80%", 70, 80),
    ("80–90%", 80, 90),
    ("90–95%", 90, 95),
    ("95–99%", 95, 99),
    ("99% or more", 99, 100.001),
]

# Human transgenes in the 69-edit donor pig of Anand et al. (2023), plus
# examples discussed in the article.
GENES_OF_INTEREST = [
    "CD46", "CD55", "THBD", "PROCR", "CD47", "TNFAIP3", "HMOX1",
    "INS", "CFTR", "MYH7", "TP53", "ALB",
]


def as_float(value: str) -> float | None:
    try:
        return float(value)
    except ValueError:
        return None


def one_to_one(rows: list[dict], species: str) -> dict[str, float]:
    out = {}
    for row in rows:
        identity = as_float(row[f"{species}_identity"])
        if row[f"{species}_type"] == "1:1" and identity is not None:
            out[row["gene_id"]] = identity
    return out


def paired(a: dict[str, float], b: dict[str, float]) -> dict:
    shared = sorted(set(a) & set(b))
    a_closer = sum(1 for g in shared if a[g] > b[g])
    b_closer = sum(1 for g in shared if a[g] < b[g])
    ties = len(shared) - a_closer - b_closer
    n = len(shared)
    return {
        "genes": n,
        "first_closer": a_closer,
        "second_closer": b_closer,
        "ties": ties,
        "pct_first_closer": round(100 * a_closer / n, 1),
        "pct_second_closer": round(100 * b_closer / n, 1),
        "pct_ties": round(100 * ties / n, 1),
    }


def query_lastz() -> dict:
    import pymysql  # optional dependency

    conn = pymysql.connect(
        host="ensembldb.ensembl.org", user="anonymous", port=3306, db="ensembl_compara_116"
    )
    tags = {}
    with conn.cursor() as cur:
        for species, mlss in SPECIES.items():
            cur.execute(
                "SELECT tag, value FROM method_link_species_set_tag "
                "WHERE method_link_species_set_id = %s",
                (mlss,),
            )
            tags[species] = {"mlss_id": mlss, **{t: v for t, v in cur.fetchall()}}
    conn.close()
    return tags


def lastz_summary(tags: dict) -> list[dict]:
    rows = []
    for species, t in tags.items():
        if "ref_matches" not in t:
            continue
        matches = int(t["ref_matches"])
        mismatches = int(t["ref_mis_matches"])
        exon_len = int(t["ref_coding_exon_length"])
        rows.append({
            "species": species,
            "mlss_id": t["mlss_id"],
            "pct_human_genome_in_alignment": round(
                100 * int(t["ref_genome_coverage"]) / int(t["ref_genome_length"]), 1
            ),
            "pct_human_coding_exon_bases_aligned": round(100 * (matches + mismatches) / exon_len, 1),
            "pct_identical_among_aligned_coding_exon_bases": round(
                100 * matches / (matches + mismatches), 1
            ),
            "other_genome_length_bp": int(t["non_ref_genome_length"]),
        })
    return rows


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--out", type=Path, default=HERE)
    parser.add_argument("--lastz", action="store_true", help="query Ensembl's public MySQL server")
    parser.add_argument("--lastz-json", type=Path, help="reuse saved LASTZ tags")
    args = parser.parse_args()
    args.out.mkdir(parents=True, exist_ok=True)

    with gzip.open(GENE_TABLE, "rt", encoding="utf8") as handle:
        rows = list(csv.DictReader(handle))

    identity = {s: one_to_one(rows, s) for s in SPECIES}
    by_name = {row["gene_name"]: row for row in rows}

    summary: dict = {
        "source": "Ensembl release 116 gene-level ortholog table, queried 2026-09-25",
        "human_protein_coding_genes": len(rows),
        "one_to_one": {},
    }
    for species, values in identity.items():
        v = sorted(values.values())
        summary["one_to_one"][species] = {
            "genes": len(v),
            "median_identity": round(statistics.median(v), 1),
            "pct_identical_100": round(100 * sum(1 for x in v if x >= 100) / len(v), 1),
            "pct_at_least_99": round(100 * sum(1 for x in v if x >= 99) / len(v), 1),
            "pct_at_least_95": round(100 * sum(1 for x in v if x >= 95) / len(v), 1),
            "pct_below_70": round(100 * bisect.bisect_left(v, 70) / len(v), 1),
        }

    pig_types = {}
    for row in rows:
        pig_types[row["pig_type"] or "none"] = pig_types.get(row["pig_type"] or "none", 0) + 1
    summary["pig_relationship_types"] = pig_types

    summary["paired"] = {
        "pig_vs_chimpanzee": paired(identity["pig"], identity["chimpanzee"]),
        "pig_vs_macaque": paired(identity["pig"], identity["macaque"]),
        "pig_vs_mouse": paired(identity["pig"], identity["mouse"]),
    }

    with open(args.out / "identity-bins.csv", "w", newline="", encoding="utf8") as handle:
        writer = csv.writer(handle)
        writer.writerow(["bin"] + [f"{s}_pct" for s in ("chimpanzee", "pig", "mouse")])
        for label, low, high in BINS:
            line = [label]
            for s in ("chimpanzee", "pig", "mouse"):
                v = identity[s].values()
                line.append(round(100 * sum(1 for x in v if low <= x < high) / len(v), 1))
            writer.writerow(line)

    pig_sorted = sorted(identity["pig"].values())
    with open(args.out / "genes-of-interest.csv", "w", newline="", encoding="utf8") as handle:
        writer = csv.writer(handle)
        writer.writerow([
            "gene", "pig_type", "pig_identity", "pig_percentile_among_one_to_one",
            "chimpanzee_identity", "mouse_type", "mouse_identity",
        ])
        for gene in GENES_OF_INTEREST:
            row = by_name.get(gene)
            if row is None:
                continue
            pig_id = as_float(row["pig_identity"])
            pct = (
                round(100 * bisect.bisect_left(pig_sorted, pig_id) / len(pig_sorted), 1)
                if pig_id is not None and row["pig_type"] == "1:1" else ""
            )
            writer.writerow([
                gene, row["pig_type"], row["pig_identity"], pct,
                row["chimpanzee_identity"], row["mouse_type"], row["mouse_identity"],
            ])

    tags = None
    if args.lastz_json:
        tags = json.loads(args.lastz_json.read_text(encoding="utf8"))
    elif args.lastz:
        tags = query_lastz()
        (args.out / "lastz-tags.json").write_text(json.dumps(tags, indent=2), encoding="utf8")
    if tags:
        summary["lastz"] = lastz_summary(tags)

    (args.out / "summary.json").write_text(json.dumps(summary, indent=2), encoding="utf8")
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()
