"""Human-animal comparison from Ensembl release 116.

Reproduces ProteinIQ's analysis for the guide "How much DNA do humans share
with other animals?". It queries Ensembl's public MySQL server (anonymous,
read-only) and writes:

  species-summary.csv        one row per species: ortholog presence, one-to-one
                             orthologs, median protein identity, genome alignment
  human-gene-orthologs.csv.gz  one row per human protein-coding gene with the
                             relationship type and protein identity per species
  paired-comparisons.csv     for genes with one-to-one orthologs in two species,
                             how often each protein is closer to the human one
  pan-compara-summary.csv    ortholog presence across animals, plants, fungi and
                             protists from Ensembl's pan-taxonomic comparison
  summary.json               the summaries plus provenance

Requirements: Python 3.10+ and pymysql (pip install pymysql).
Usage: python analyze.py [--cache raw] [--out .]
"""

from __future__ import annotations

import argparse
import csv
import gzip
import json
import statistics
from datetime import date
from pathlib import Path

import pymysql

RELEASE = 116
HOST = "ensembldb.ensembl.org"
CORE_DB = "homo_sapiens_core_116_38"
COMPARA_DB = "ensembl_compara_116"
CHROMOSOMES = [str(i) for i in range(1, 23)] + ["X", "Y", "MT"]

# (key, display name, scientific name, Ensembl ortholog MLSS id, LASTZ net MLSS id or None)
SPECIES = [
    ("chimpanzee", "Chimpanzee", "Pan troglodytes", 103890, 1098),
    ("bonobo", "Bonobo", "Pan paniscus", 103886, 1030),
    ("gorilla", "Gorilla", "Gorilla gorilla", 103885, 1037),
    ("orangutan", "Orangutan", "Pongo abelii", 147000, 1978),
    ("macaque", "Rhesus macaque", "Macaca mulatta", 120910, 1562),
    ("pig", "Pig", "Sus scrofa", 102510, 1039),
    ("dog", "Dog", "Canis lupus familiaris", 147003, None),
    ("horse", "Horse", "Equus caballus", 110530, 1347),
    ("cat", "Cat", "Felis catus", 153563, None),
    ("cow", "Cow", "Bos taurus", 153357, None),
    ("elephant", "Elephant", "Loxodonta africana", 50936, 691),
    ("mouse", "Mouse", "Mus musculus", 144756, 1960),
    ("rat", "Rat", "Rattus norvegicus", 153506, None),
    ("dolphin", "Dolphin", "Tursiops truncatus", 50966, 700),
    ("platypus", "Platypus", "Ornithorhynchus anatinus", 135319, 1767),
    ("chicken", "Chicken", "Gallus gallus", 148730, 2043),
    ("frog", "Western clawed frog", "Xenopus tropicalis", 148731, None),
    ("zebrafish", "Zebrafish", "Danio rerio", 104747, 1113),
    ("fruit_fly", "Fruit fly", "Drosophila melanogaster", 153076, None),
    ("nematode", "Nematode worm", "Caenorhabditis elegans", 50964, None),
    ("yeast", "Baker's yeast", "Saccharomyces cerevisiae", 50962, None),
]

TYPE_CODES = {"ortholog_one2one": "1:1", "ortholog_one2many": "1:m", "ortholog_many2many": "m:m"}

WGA_TAGS = [
    "ref_genome_length",
    "ref_genome_coverage",
    "ref_coding_exon_length",
    "ref_matches",
    "ref_mis_matches",
]


# Head-to-head protein identity comparisons: (species, reference species)
PAIRS = [
    ("dog", "mouse"),
    ("cat", "mouse"),
    ("pig", "mouse"),
    ("horse", "mouse"),
    ("cow", "mouse"),
    ("elephant", "mouse"),
    ("dolphin", "mouse"),
    ("macaque", "mouse"),
    ("rat", "mouse"),
    ("dog", "cat"),
]

# Ensembl Genomes pan-taxonomic comparison, the release paired with Ensembl 116
PAN_HOST = "mysql-eg-publicsql.ebi.ac.uk"
PAN_PORT = 4157
PAN_DB = "ensembl_compara_pan_homology_63_116"
PAN_HUMAN_GENOME_DB_ID = 2039
# (key, display name, scientific name, pan ortholog MLSS id)
PAN_SPECIES = [
    ("mouse", "Mouse", "Mus musculus", 324616),
    ("chicken", "Chicken", "Gallus gallus", 325649),
    ("zebrafish", "Zebrafish", "Danio rerio", 308002),
    ("fruit_fly", "Fruit fly", "Drosophila melanogaster", 328453),
    ("nematode", "Nematode worm", "Caenorhabditis elegans", 303313),
    ("physcomitrium", "Moss", "Physcomitrium patens", 308000),
    ("arabidopsis", "Thale cress", "Arabidopsis thaliana", 21982),
    ("grape", "Grape", "Vitis vinifera", 328454),
    ("rice", "Rice", "Oryza sativa", 303508),
    ("tomato", "Tomato", "Solanum lycopersicum", 329008),
    ("dictyostelium", "Slime mold", "Dictyostelium discoideum", 303767),
    ("chlamydomonas", "Green alga", "Chlamydomonas reinhardtii", 308001),
    ("fission_yeast", "Fission yeast", "Schizosaccharomyces pombe", 23530),
    ("yeast", "Baker's yeast", "Saccharomyces cerevisiae", 23697),
    ("plasmodium", "Malaria parasite", "Plasmodium falciparum", 328022),
]


def connect(db: str):
    return pymysql.connect(host=HOST, user="anonymous", port=3306, database=db, read_timeout=1800)


def pan_homologies(key: str, mlss_id: int, cache: Path) -> list[dict[str, str]]:
    """Human-side rows of one pan-taxonomic ortholog set."""
    path = cache / f"pan-homologies-{key}.tsv"
    if not path.exists():
        sql = """
            SELECT gh.stable_id, h.description, hh.perc_id
            FROM homology h
            JOIN homology_member hh ON hh.homology_id = h.homology_id
            JOIN gene_member gh ON gh.gene_member_id = hh.gene_member_id
            WHERE h.method_link_species_set_id = %s AND gh.genome_db_id = %s
        """
        conn = pymysql.connect(host=PAN_HOST, user="anonymous", port=PAN_PORT, database=PAN_DB, read_timeout=1800)
        with conn, conn.cursor() as cur:
            cur.execute(sql, (mlss_id, PAN_HUMAN_GENOME_DB_ID))
            rows = cur.fetchall()
        with path.open("w", newline="") as f:
            w = csv.writer(f, delimiter="	")
            w.writerow(["human_gene", "type", "pid_human"])
            w.writerows(rows)
    with path.open() as f:
        return list(csv.DictReader(f, delimiter="	"))


def human_genes(cache: Path) -> dict[str, str]:
    """Protein-coding genes on the GRCh38 reference chromosomes, with gene names."""
    path = cache / "human-protein-coding-genes.tsv"
    if not path.exists():
        placeholders = ",".join(["%s"] * len(CHROMOSOMES))
        sql = f"""
            SELECT g.stable_id, COALESCE(x.display_label, '')
            FROM gene g
            JOIN seq_region s USING (seq_region_id)
            JOIN coord_system cs USING (coord_system_id)
            LEFT JOIN xref x ON x.xref_id = g.display_xref_id
            WHERE g.biotype = 'protein_coding'
              AND cs.name = 'chromosome' AND cs.attrib LIKE '%%default_version%%'
              AND s.name IN ({placeholders})
        """
        with connect(CORE_DB) as conn, conn.cursor() as cur:
            cur.execute(sql, CHROMOSOMES)
            rows = sorted(cur.fetchall())
        with path.open("w", newline="") as f:
            w = csv.writer(f, delimiter="\t")
            w.writerow(["gene_id", "gene_name"])
            w.writerows(rows)
    with path.open() as f:
        return {r["gene_id"]: r["gene_name"] for r in csv.DictReader(f, delimiter="\t")}


def homologies(key: str, mlss_id: int, cache: Path) -> list[dict[str, str]]:
    """All homologies in the human-species ortholog set, one row per pair.

    pid_human is homology_member.perc_id for the human member: identical
    aligned positions as a percentage of the human protein's length.
    """
    path = cache / f"homologies-{key}.tsv"
    if not path.exists():
        sql = """
            SELECT gh.stable_id, go.stable_id, h.description, h.is_high_confidence,
                   hh.perc_id, ho.perc_id
            FROM homology h
            JOIN homology_member hh ON hh.homology_id = h.homology_id
            JOIN gene_member gh ON gh.gene_member_id = hh.gene_member_id
            JOIN homology_member ho ON ho.homology_id = h.homology_id
                 AND ho.gene_member_id <> hh.gene_member_id
            JOIN gene_member go ON go.gene_member_id = ho.gene_member_id
            WHERE h.method_link_species_set_id = %s AND gh.stable_id LIKE 'ENSG0%%'
        """
        with connect(COMPARA_DB) as conn, conn.cursor() as cur:
            cur.execute(sql, (mlss_id,))
            rows = cur.fetchall()
        with path.open("w", newline="") as f:
            w = csv.writer(f, delimiter="\t")
            w.writerow(["human_gene", "target_gene", "type", "high_confidence", "pid_human", "pid_target"])
            w.writerows(["" if v is None else v for v in r] for r in rows)
    with path.open() as f:
        return list(csv.DictReader(f, delimiter="\t"))


def wga_stats(mlss_id: int) -> dict[str, int]:
    with connect(COMPARA_DB) as conn, conn.cursor() as cur:
        cur.execute(
            "SELECT tag, value FROM method_link_species_set_tag WHERE method_link_species_set_id = %s",
            (mlss_id,),
        )
        tags = dict(cur.fetchall())
    return {t: int(tags[t]) for t in WGA_TAGS}


def quantile(sorted_values: list[float], q: float) -> float:
    """Linear interpolation between closest ranks (same as statistics.quantiles 'inclusive')."""
    pos = q * (len(sorted_values) - 1)
    lo = int(pos)
    hi = min(lo + 1, len(sorted_values) - 1)
    return sorted_values[lo] + (sorted_values[hi] - sorted_values[lo]) * (pos - lo)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--cache", default="raw")
    parser.add_argument("--out", default=".")
    args = parser.parse_args()
    cache = Path(args.cache)
    out = Path(args.out)
    cache.mkdir(parents=True, exist_ok=True)

    genes = human_genes(cache)
    denominator = len(genes)
    per_gene: dict[str, dict[str, tuple[str, str]]] = {g: {} for g in genes}
    summary = []

    for key, name, sci, ortholog_mlss, lastz_mlss in SPECIES:
        any_ortholog: set[str] = set()
        high_conf: set[str] = set()
        types: dict[str, set[str]] = {}
        one2one: dict[str, float] = {}
        for r in homologies(key, ortholog_mlss, cache):
            g = r["human_gene"]
            if g not in genes or not r["type"].startswith("ortholog"):
                continue
            any_ortholog.add(g)
            types.setdefault(r["type"], set()).add(g)
            if r["high_confidence"] == "1":
                high_conf.add(g)
            if r["type"] == "ortholog_one2one":
                one2one[g] = float(r["pid_human"])
        for g in any_ortholog:
            code = next(TYPE_CODES[t] for t in TYPE_CODES if g in types.get(t, set()))
            per_gene[g][key] = (code, f"{one2one[g]:.2f}" if g in one2one else "")

        identities = sorted(one2one.values())
        row = {
            "species": name,
            "scientific_name": sci,
            "human_genes_with_any_ortholog": len(any_ortholog),
            "pct_any_ortholog": round(100 * len(any_ortholog) / denominator, 1),
            "human_genes_with_one_to_one_ortholog": len(one2one),
            "pct_one_to_one": round(100 * len(one2one) / denominator, 1),
            "human_genes_with_high_confidence_ortholog": len(high_conf),
            "pct_high_confidence": round(100 * len(high_conf) / denominator, 1),
            "median_protein_identity_one_to_one": round(statistics.median(identities), 1),
            "q1_protein_identity_one_to_one": round(quantile(identities, 0.25), 1),
            "q3_protein_identity_one_to_one": round(quantile(identities, 0.75), 1),
            "ensembl_ortholog_mlss_id": ortholog_mlss,
            "pct_human_genome_in_lastz_alignment": "",
            "pct_human_coding_exon_bases_aligned": "",
            "ensembl_lastz_mlss_id": lastz_mlss or "",
        }
        if lastz_mlss:
            s = wga_stats(lastz_mlss)
            row["pct_human_genome_in_lastz_alignment"] = round(
                100 * s["ref_genome_coverage"] / s["ref_genome_length"], 1
            )
            row["pct_human_coding_exon_bases_aligned"] = round(
                100 * (s["ref_matches"] + s["ref_mis_matches"]) / s["ref_coding_exon_length"], 1
            )
            row["human_reference_length_bp"] = s["ref_genome_length"]
        summary.append(row)
        print(f"{name:20} any {row['pct_any_ortholog']:5}%  1:1 {row['pct_one_to_one']:5}%  "
              f"median id {row['median_protein_identity_one_to_one']:5}%  "
              f"genome aligned {row['pct_human_genome_in_lastz_alignment']}")

    fields = [k for k in summary[0] if k != "human_reference_length_bp"]
    with (out / "species-summary.csv").open("w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=fields, extrasaction="ignore")
        w.writeheader()
        w.writerows(summary)

    keys = [s[0] for s in SPECIES]
    with gzip.open(out / "human-gene-orthologs.csv.gz", "wt", newline="") as f:
        w = csv.writer(f)
        w.writerow(["gene_id", "gene_name"] + [f"{k}_{c}" for k in keys for c in ("type", "identity")])
        for g in sorted(genes):
            w.writerow([g, genes[g]] + [v for k in keys for v in per_gene[g].get(k, ("", ""))])

    identity = {k: {g: float(v[1]) for g, v in ((g, per_gene[g].get(k)) for g in genes) if v and v[1]} for k in keys}
    paired = []
    for a, b in PAIRS:
        shared = identity[a].keys() & identity[b].keys()
        closer = sum(identity[a][g] > identity[b][g] for g in shared)
        farther = sum(identity[a][g] < identity[b][g] for g in shared)
        ties = len(shared) - closer - farther
        paired.append({
            "species": a,
            "reference": b,
            "genes_with_one_to_one_in_both": len(shared),
            "pct_species_closer": round(100 * closer / len(shared), 1),
            "pct_reference_closer": round(100 * farther / len(shared), 1),
            "pct_tie": round(100 * ties / len(shared), 1),
        })
    with (out / "paired-comparisons.csv").open("w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(paired[0]))
        w.writeheader()
        w.writerows(paired)

    pan = []
    for key, name, sci, mlss in PAN_SPECIES:
        rows = [r for r in pan_homologies(key, mlss, cache) if r["human_gene"] in genes and r["type"].startswith("ortholog")]
        any_ortholog = {r["human_gene"] for r in rows}
        types_by_gene: dict[str, set[str]] = {}
        for r in rows:
            types_by_gene.setdefault(r["human_gene"], set()).add(r["type"])
        one2one = {g for g, t in types_by_gene.items() if t == {"ortholog_one2one"}}
        pan.append({
            "species": name,
            "scientific_name": sci,
            "human_genes_with_any_ortholog": len(any_ortholog),
            "pct_any_ortholog": round(100 * len(any_ortholog) / denominator, 1),
            "human_genes_with_one_to_one_ortholog": len(one2one),
            "pct_one_to_one": round(100 * len(one2one) / denominator, 1),
            "pan_ortholog_mlss_id": mlss,
        })
        print(f"pan {name:18} any {pan[-1]['pct_any_ortholog']:5}%  1:1 {pan[-1]['pct_one_to_one']:5}%")
    with (out / "pan-compara-summary.csv").open("w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(pan[0]))
        w.writeheader()
        w.writerows(pan)

    reference_length = next(s["human_reference_length_bp"] for s in summary if "human_reference_length_bp" in s)
    json.dump(
        {
            "analysis_date": date.today().isoformat(),
            "ensembl_release": RELEASE,
            "databases": [CORE_DB, COMPARA_DB],
            "host": HOST,
            "pan_database": PAN_DB,
            "pan_host": f"{PAN_HOST}:{PAN_PORT}",
            "human_protein_coding_genes": denominator,
            "human_reference_length_bp": reference_length,
            "species": [{k: v for k, v in s.items() if k != "human_reference_length_bp"} for s in summary],
            "paired_comparisons": paired,
            "pan_taxonomic": pan,
        },
        (out / "summary.json").open("w"),
        indent=2,
    )


if __name__ == "__main__":
    main()
