"""Download the Moody et al. (2024) time-tree archive from Zenodo and extract the
age of the LUCA node (and a few other nodes) from every analysis variant in it.

Source: Alvarez-Carretero, S. (2024) "sabifo4/LUCA-divtimes: v1.0.0",
        Zenodo, doi:10.5281/zenodo.11260523  (227 MB zip; MCMCtree summary tables)

Output: data/moody_node_ages.csv  (one row per analysis variant x node)

Each `*_all_mean_est.tsv` in the archive holds, for every node of the tree, the mean
age and the 2.5% / 97.5% quantiles (in Ma) averaged over the converged MCMC chains,
plus the calibration string.  "CLK" runs sample from the joint calibration prior
WITHOUT sequence data, so they give the effective prior on each node.
"""
from __future__ import annotations

import csv
import io
import os
import re
import sys
import urllib.request
import zipfile

ZENODO_URL = ("https://zenodo.org/api/records/11260523/files/"
              "sabifo4/LUCA-divtimes-LUCAdivtimes-prerelease.zip/content")
RAW_DIR = "fermi_raw/timing/moody_zenodo"
ZIP_PATH = os.path.join(RAW_DIR, "LUCA-divtimes.zip")
EXPECTED_SIZE = 227269877

NODES = ["ROOT", "LUCA", "LECA", "TG-EUKARYA-ARCH", "TG-EUKARYA-MITO", "TG-OXYPHOTOBACTERIA",
         "CG-OXYPHOTOBACTERIA", "METAZOA", "EMBRYOPHYTA"]


def download(force=False):
    os.makedirs(RAW_DIR, exist_ok=True)
    if os.path.exists(ZIP_PATH) and os.path.getsize(ZIP_PATH) == EXPECTED_SIZE and not force:
        return ZIP_PATH
    print("downloading Moody et al. archive from Zenodo (227 MB) ...", flush=True)
    tmp = ZIP_PATH + ".part"
    with urllib.request.urlopen(ZENODO_URL, timeout=600) as r, open(tmp, "wb") as f:
        while True:
            b = r.read(1 << 20)
            if not b:
                break
            f.write(b)
    os.replace(tmp, ZIP_PATH)
    return ZIP_PATH


MAIN_DIRS = {
    "01_PAML/02_MCMCtree/": "concatenated, cross-braced (A)",
    "01_PAML/03_MCMCtree_nonbraced/": "concatenated, not cross-braced",
    "01_PAML/04_MCMCtree_fossbraced/": "concatenated, fossil-node cross-braced (B)",
    "01_PAML/05_MCMCtree_part/": "partitioned, cross-braced (A) [paper main]",
    "01_PAML/06_MCMCtree_part_nonbraced/": "partitioned, not cross-braced",
    "01_PAML/07_MCMCtree_part_fossbraced/": "partitioned, fossil-node cross-braced (B)",
}


def classify(path):
    """Return (analysis, clock, stage) labels from a path inside the archive."""
    p = path
    clock = "CLK" if ("_CLK" in p or "sum_files_prior" in p) else ("GBM" if "GBM" in p else "ILN")
    stage = "prior" if clock == "CLK" else "posterior"
    analysis = None
    for d, lab in MAIN_DIRS.items():
        if d in p and "03_additional_analyses" not in p:
            analysis = lab
    if "00_indgenes_analyses" in p:
        gene = re.search(r"/(ATP|EF|Leu|SRP|Tyr)(_GBM|_ILN)?_all_mean_est", p).group(1)
        analysis = f"single gene {gene}, cross-braced"
    elif "01_lou_analyses" in p:
        gene = re.search(r"LUCA-divtimes_no([a-z]+)", p).group(1)
        analysis = f"leave-one-out (no {gene}), cross-braced"
    elif "02_blunc_analysis" in p:
        analysis = "branch-length uncertainty (bsinBV), cross-braced"
    if analysis is None:
        analysis = "unclassified: " + p
    if "FILT" in p:
        analysis += " (filtered chains)"
    return analysis, clock, stage


def extract(zip_path):
    rows = []
    z = zipfile.ZipFile(zip_path)
    for info in z.infolist():
        n = info.filename
        if not n.endswith("_all_mean_est.tsv"):
            continue
        if "/sum_files_" not in n:      # the sum_analyses copies are duplicates
            continue
        analysis, clock, stage = classify(n)
        txt = io.TextIOWrapper(z.open(info), encoding="utf8").read().splitlines()
        seen = {}
        for line in txt[1:]:
            parts = line.split("\t")
            if len(parts) < 4:
                continue
            m = re.match(r"t_n(\d+)_(.+)$", parts[0])
            if not m:
                continue
            node = m.group(2)
            base = re.sub(r"-DUP\d*$", "", node)
            if base not in NODES:
                continue
            is_dup = node != base
            # cross-braced analyses: the mirrored node has identical numbers -> keep one;
            # non-cross-braced: the two copies are dated independently -> keep both
            k = (base, parts[1], parts[2], parts[3])
            if k in seen:
                continue
            seen[k] = True
            copy = sum(1 for kk in seen if kk[0] == base)
            rows.append(dict(analysis=analysis, clock=clock, stage=stage, node=base, copy=copy,
                             mean_Ma=float(parts[1]), q025_Ma=float(parts[2]), q975_Ma=float(parts[3]),
                             calibration=parts[4] if len(parts) > 4 else "",
                             file=n.split("/", 1)[1]))
    return rows


def main(out_csv):
    try:
        zp = download()
        rows = extract(zp)
    except Exception as e:  # network failure: fall back to the derived table shipped with the study
        if os.path.exists(out_csv):
            print(f"WARNING: could not download/parse the Zenodo archive ({e!r}); "
                  f"using the existing derived table {out_csv}", file=sys.stderr)
            return None
        raise
    rows.sort(key=lambda r: (r["node"] != "LUCA", r["node"], r["stage"], r["analysis"], r["clock"], r["copy"]))
    os.makedirs(os.path.dirname(out_csv), exist_ok=True)
    with open(out_csv, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
        w.writeheader()
        w.writerows(rows)
    return len(rows)


if __name__ == "__main__":
    here = os.path.dirname(os.path.abspath(__file__))
    print(main(os.path.join(here, "..", "data", "moody_node_ages.csv")))