"""Count virus species with a recorded human host under current ICTV taxonomy.

Reproduces the ProteinIQ analysis in content/guides/how-many-viruses-infect-humans.mdx.

Inputs (download into one folder; raw files are not stored in the repository):
  virushostdb.tsv        https://www.genome.jp/ftp/db/virushostdb/virushostdb.tsv
  virus_genome_type.tsv  https://www.genome.jp/ftp/db/virushostdb/virus_genome_type.tsv
  VMR_MSL41.v1.20260729.xlsx
                         https://ictv.global/sites/default/files/VMR/VMR_MSL41.v1.20260729.xlsx
The published run used the Virus-Host DB files dated 2026-10-04 (dbrel.txt:
RefSeq release 237, GenBank release 273.0) and ICTV VMR MSL41.v1 (VMR release
2026-07-29, DOI 10.5281/zenodo.21694279), all downloaded on 2026-10-23.

Run (Python 3.11+, openpyxl for the VMR workbook):
  python analyze.py --input-dir RAW --output-dir .

Method in brief. A Virus-Host DB record counts when its host taxid is 9606
(Homo sapiens). Each record's virus taxid is mapped to a species name through
virus_genome_type.tsv, then to an ICTV MSL41 species by exact name, falling back
to any GenBank accession shared with the VMR. Records that map to no ICTV
species (NCBI "unclassified" groups, provisional names) are counted separately,
as are ICTV species whose VMR host source is bacteria (phages of gut bacteria
sampled from people). Each human-host ICTV species takes its strongest evidence
tier across its human records:
  disease      a human record names a disease (Virus-Host DB DISEASE or KEGG DISEASE)
  literature   otherwise, a human record cites a PubMed ID
  annotation   otherwise, the host comes only from database annotation
                 (RefSeq, GenBank via NCBI Virus, UniProt)
ICTV's own host-source column then screens the list: a species is counted only when
its VMR host source includes "vertebrates" as a whole token ("invertebrates" does not
match); species with human records but only bacterial, archaeal or invertebrate VMR
hosts (gut phages and CRESS viruses sampled from people, an amoeba virus) are listed
as excluded. One NCBI name is mapped by hand because ICTV split the species: NCBI
"Orthoflavivirus encephalitidis" (RefSeq NC_001672, tick-borne encephalitis virus
Neudoerfl) is ICTV "Orthoflavivirus neudoerflense" (VMR U27495).
A species is "also recorded in another vertebrate" when any Virus-Host DB record
of the same ICTV species has a non-human vertebrate host. A recorded host is a
reported association, not proof of replication, disease or natural infection.
"""

import argparse
import csv
import hashlib
import json
import re
from collections import Counter, defaultdict
from pathlib import Path

import openpyxl

HUMAN = "9606"
VMR_FILE = "VMR_MSL41.v1.20260729.xlsx"
TIERS = ("disease", "literature", "annotation")
ACCESSION = re.compile(r"\b([A-Z]{1,2}_?\d{5,8})(?:\.\d+)?\b")
GENOME_GROUPS = {
    "dsDNA": "dsDNA",
    "dsDNA-RT": "Reverse-transcribing",
    "ssRNA-RT": "Reverse-transcribing",
    "dsRNA": "dsRNA",
    "ssRNA(+)": "ssRNA(+)",
    "ssRNA(-)": "ssRNA(-) and ambisense",
    "ssRNA(+/-)": "ssRNA(-) and ambisense",
    "ssRNA(-); ssRNA(+/-)": "ssRNA(-) and ambisense",
}


MANUAL_MAP = {"orthoflavivirus encephalitidis": "Orthoflavivirus neudoerflense"}


def vmr_has_vertebrate_host(hosts):
    tokens = {token.strip().replace(" (S)", "") for host in hosts for token in host.split(",")}
    return "vertebrates" in tokens


def genome_group(genome):
    groups = {
        GENOME_GROUPS.get(token) or ("ssDNA" if token.startswith("ssDNA") else "Other")
        for token in (part.strip() for part in genome.split(";"))
    }
    return groups.pop() if len(groups) == 1 else "Other"


def read_tsv(path):
    with open(path, encoding="utf-8", newline="") as stream:
        return list(csv.DictReader(stream, delimiter="\t"))


def read_vmr(path):
    sheet = openpyxl.load_workbook(path, read_only=True)["VMR MSL41"]
    rows = sheet.iter_rows(values_only=True)
    header = [str(cell).strip() if cell else "" for cell in next(rows)]
    column = {name: header.index(name) for name in ("Realm", "Family", "Species", "Genome", "Host source", "Virus GENBANK accession")}
    species, by_accession = {}, {}
    for row in rows:
        name = row[column["Species"]]
        if not name:
            continue
        name = name.strip()
        entry = species.setdefault(name, {"realm": row[column["Realm"]] or "", "family": row[column["Family"]] or "", "genome": set(), "host": set()})
        entry["genome"].add(row[column["Genome"]] or "")
        entry["host"].add(row[column["Host source"]] or "")
        for accession in ACCESSION.findall(row[column["Virus GENBANK accession"]] or ""):
            by_accession[accession] = name
    return species, by_accession


def sha256(path):
    return hashlib.sha256(path.read_bytes()).hexdigest()


def main():
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--input-dir", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    args = parser.parse_args()

    records = read_tsv(args.input_dir / "virushostdb.tsv")
    ncbi_species = {row["virus_tax_id"]: row["species"] for row in read_tsv(args.input_dir / "virus_genome_type.tsv")}
    vmr, vmr_by_accession = read_vmr(args.input_dir / VMR_FILE)
    vmr_lower = {name.lower(): name for name in vmr}

    def ictv_species(record):
        ncbi_name = ncbi_species.get(record["virus tax id"], "").lower()
        name = vmr_lower.get(ncbi_name) or MANUAL_MAP.get(ncbi_name)
        if name:
            return name, "name"
        for accession in ACCESSION.findall(record["refseq id"]):
            if accession in vmr_by_accession:
                return vmr_by_accession[accession], "accession"
        return None, None

    human_records = [row for row in records if row["host tax id"] == HUMAN]
    tier = {}
    match_method = Counter()
    human_virus_taxa = Counter()
    unmapped = Counter()
    for record in human_records:
        name, method = ictv_species(record)
        if not name:
            unmapped[ncbi_species.get(record["virus tax id"]) or record["virus name"]] += 1
            continue
        match_method[method] += 1
        human_virus_taxa[name] += 1
        if record["DISEASE"] or record["KEGG DISEASE"]:
            level = "disease"
        elif record["pmid"]:
            level = "literature"
        else:
            level = "annotation"
        if name not in tier or TIERS.index(level) < TIERS.index(tier[name]):
            tier[name] = level

    other_vertebrate = defaultdict(set)
    for record in records:
        if record["host tax id"] in ("", "1", HUMAN) or "Vertebrata" not in record["host lineage"]:
            continue
        name, _ = ictv_species(record)
        if name in tier:
            other_vertebrate[name].add(record["host name"])

    bacterial = {name for name in tier if not vmr_has_vertebrate_host(vmr[name]["host"])}
    counted = sorted(set(tier) - bacterial)

    rows = []
    for name in counted:
        entry = vmr[name]
        genome = "; ".join(sorted(entry["genome"]))
        rows.append({
            "ictv_species": name,
            "realm": entry["realm"],
            "family": entry["family"],
            "genome": genome,
            "genome_group": genome_group(genome),
            "evidence_tier": tier[name],
            "human_virus_taxa": human_virus_taxa[name],
            "other_vertebrate_hosts": len(other_vertebrate[name]),
            "vmr_host_source": "; ".join(sorted(entry["host"])),
        })

    args.output_dir.mkdir(parents=True, exist_ok=True)
    with open(args.output_dir / "human-host-species.csv", "w", encoding="utf-8", newline="") as stream:
        writer = csv.DictWriter(stream, fieldnames=list(rows[0]), lineterminator="\n")
        writer.writeheader()
        writer.writerows(rows)

    with open(args.output_dir / "unmapped-ncbi-names.csv", "w", encoding="utf-8", newline="") as stream:
        writer = csv.writer(stream, lineterminator="\n")
        writer.writerow(["ncbi_species_name", "human_host_records"])
        writer.writerows(sorted(unmapped.items()))

    summary = []
    def add(group, value, subset):
        subset = list(subset)
        summary.append({
            "group": group,
            "value": value,
            "species": len(subset),
            **{tier_name: sum(row["evidence_tier"] == tier_name for row in subset) for tier_name in TIERS},
            "also_other_vertebrate": sum(row["other_vertebrate_hosts"] > 0 for row in subset),
        })
    add("all", "ICTV species with a recorded human host", rows)
    for group in ("genome_group", "realm"):
        for value, _ in Counter(row[group] for row in rows).most_common():
            add(group, value, (row for row in rows if row[group] == value))
    with open(args.output_dir / "summary.csv", "w", encoding="utf-8", newline="") as stream:
        writer = csv.DictWriter(stream, fieldnames=list(summary[0]), lineterminator="\n")
        writer.writeheader()
        writer.writerows(summary)

    report = {
        "inputs": {name: sha256(args.input_dir / name) for name in ("virushostdb.tsv", "virus_genome_type.tsv", VMR_FILE)},
        "human_host_records": len(human_records),
        "human_host_ncbi_species": len({ncbi_species.get(r["virus tax id"]) for r in human_records}),
        "records_mapped": sum(match_method.values()),
        "ictv_species_with_human_host": len(tier),
        "excluded_no_vertebrate_host_species": sorted(bacterial),
        "counted_species": len(counted),
        "unmapped_ncbi_names": len(unmapped),
        "unmapped_human_host_records": sum(unmapped.values()),
        "vmr_species_total": len(vmr),
    }
    print(json.dumps(report, indent=2))
    for line in summary:
        print(line)


if __name__ == "__main__":
    main()
