"""Turn the raw downloads (see fetch_data.py) into small tidy CSVs in data/ used by part2_brains.py.

If a raw file is missing the corresponding derived CSV already present in data/ is kept, so the
analysis can be re-run from the derived tables alone.
"""
import csv
import os
import re
import html
import numpy as np

import phylo
from fetch_data import RAW

HERE = os.path.dirname(os.path.abspath(__file__))
DATA = os.path.join(HERE, "..", "data")

# Species names in Burger et al. (2019) that differ from the Upham et al. (2019) tree.
SYNONYMS = {
    "Physeter_catodon": "Physeter_macrocephalus", "Cebus_apella": "Sapajus_apella",
    "Cebus_libidinosus": "Sapajus_libidinosus", "Cebus_xanthosternos": "Sapajus_xanthosternos",
    "Cebus_nigritus": "Sapajus_nigritus", "Uncia_uncia": "Panthera_uncia",
    "Equus_burchellii": "Equus_quagga", "Equus_asinus": "Equus_africanus",
    "Taurotragus_oryx": "Tragelaphus_oryx", "Taurotragus_derbianus": "Tragelaphus_derbianus",
    "Hexaprotodon_liberiensis": "Choeropsis_liberiensis", "Otaria_flavescens": "Otaria_byronia",
    "Monachus_schauinslandi": "Neomonachus_schauinslandi", "Felis_manul": "Otocolobus_manul",
    "Puma_yagouaroundi": "Herpailurus_yagouaroundi", "Profelis_aurata": "Caracal_aurata",
    "Zaglossus_bruijni": "Zaglossus_bruijnii", "Bunopithecus_hoolock": "Hoolock_hoolock",
    "Callithrix_pygmaea": "Cebuella_pygmaea", "Galago_demidoff": "Galagoides_demidoff",
    "Manis_tricuspis": "Phataginus_tricuspis", "Alcelaphus_lichtensteinii": "Alcelaphus_buselaphus",
}


def _w(path, rows, cols=None):
    cols = cols or list(rows[0].keys())
    with open(path, "w", newline="", encoding="utf8") as f:
        w = csv.DictWriter(f, fieldnames=cols, extrasaction="ignore")
        w.writeheader()
        w.writerows(rows)


# Records in Burger et al. (2019) that are demonstrably wrong (kept in the CSV, flagged, and
# excluded from the analyses; a sensitivity run keeps them).
FLAGS = {
    "Phocoena_phocoena": "brain 1,735 g / body 142 kg is the Crile & Quiring (1940) 'porpoise' record (see Boddy "
                         "et al. 2012 database); other harbour-porpoise records are ~500-525 g (Haug 1987; Marino 2003)",
    "Bos_taurus": "body mass 46 kg is not an adult domestic cow",
    "Rattus_norvegicus": "brain 8 g is ~4x the usual ~2 g for this species",
    "Mus_musculoides": "brain 1.0 g at 8 g body mass is 4x the genus median EQ",
    # added in revision (referee point M1): the single record (no sample size) pairs an adult-sized brain
    # (1,126 g) with a body mass of 67.5 kg; adults of this species weigh about 180-300 kg and calves about
    # 40 kg (Duke OBIS-SEAMAP species account; IWDG species profile: average adult 200 kg). With an adult
    # body mass the same brain gives EQ 2.1-2.9, not 5.6.
    "Lagenorhynchus_albirostris": "body mass 67.5 kg is a quarter to a third of adult mass (180-300 kg); "
                                  "single record with no sample size",
}


def well_sampled(n_brain):
    """True if the Burger et al. record states that at least three brains were measured.  The column holds
    integers, '', '>1', '>1,<10' and '>10'; only an integer >= 3 or '>10' counts."""
    s = str(n_brain).strip()
    if s == ">10":
        return True
    try:
        return float(s) >= 3
    except ValueError:
        return False


def _match_tips(names, tipnames, lower=False):
    """Map Burger binomials to tree tips. Tree tips may carry suffixes (family/order or subspecies);
    the exact binomial is preferred, otherwise the first tip whose first two tokens match."""
    exact, first = {}, {}
    for t in tipnames:
        parts = t.strip("_").split("_")
        key = "_".join(parts[:2])
        if len(parts) == 2:
            exact[key] = t
        first.setdefault(key, t)
    out, used = {}, set()
    for nm in names:
        cands = [nm, SYNONYMS.get(nm, "")]
        if nm == "Physeter_catodon":
            cands.append("Physeter_catodon")
        tip = None
        for c in cands:
            if not c:
                continue
            k = c.lower() if lower else c
            tip = exact.get(k) or first.get(k)
            if tip:
                break
        if tip and tip not in used:
            used.add(tip)
            out[nm] = tip
    return out


def extract_burger_and_tree():
    src = os.path.join(RAW, "zenodo/burger2019_BrainAllometry_Supplement_Data.csv")
    tre_a = os.path.join(RAW, "trees/alvarez2022_4705sp_mean.nwk")
    tre_u = os.path.join(RAW, "trees/mamphy_mcc.tre")
    if not (os.path.exists(src) and os.path.exists(tre_a) and os.path.exists(tre_u)):
        print("  (raw Burger/tree files missing; keeping existing derived tables)")
        return
    rows = list(csv.DictReader(open(src, encoding="utf8", errors="replace")))
    names = [r["Binomial"] for r in rows]
    # primary tree: Alvarez-Carretero et al. 2022 (branch lengths in units of 100 Myr -> Myr)
    ra = phylo.parse_newick(open(tre_a).read())
    for nd in phylo.iter_nodes(ra):
        nd.length *= 100.0
    ma = _match_tips(names, [t.name for t in phylo.tips(ra)], lower=True)
    # secondary tree: Upham et al. 2019 MCC (robustness only; see report for its problems in Cetacea)
    ru = phylo.read_nexus_tree(tre_u)
    mu = _match_tips(names, [t.name for t in phylo.tips(ru)])
    out = []
    for r in rows:
        nm = r["Binomial"]
        out.append(dict(species=nm, order=r["order"], family=r["family"], genus=r["genus"],
                        brain_g=float(r["Mean_brain_mass_g"]), body_g=float(r["Mean_body_mass_g"]),
                        n_brain=r["Sample_size.brain"], well_sampled=int(well_sampled(r["Sample_size.brain"])),
                        source=r["BrainReference1"],
                        tip_alvarez2022=ma.get(nm, ""), tip_upham2019=mu.get(nm, ""),
                        flag=FLAGS.get(nm, "")))
    _w(os.path.join(DATA, "mammals_extant_burger2019.csv"), out)
    for root, m, fn in ((ra, ma, "mammal_tree_pruned_alvarez2022.nwk"), (ru, mu, "mammal_tree_pruned_upham2019.nwk")):
        pr = phylo.prune(root, set(m.values()))
        rename = {v: k for k, v in m.items()}
        for t in phylo.tips(pr):
            t.name = rename[t.name]
        open(os.path.join(DATA, fn), "w").write(phylo.to_newick(pr))
    print(f"  Burger: {len(out)} species; {len(ma)} matched to the Alvarez-Carretero 2022 tree, "
          f"{len(mu)} to the Upham 2019 tree")


# Ages (Ma) for the proboscidean endocasts of Benoit et al. (2019) Table 1. 'V' = stated in the text of that
# paper; 'L' = general stratigraphic range from the literature, not re-verified in this study.
BENOIT_AGES = {
    "Elephas maximus": (0.0, "V", "extant"),
    "Loxodonta africana": (0.0, "V", "extant"),
    "Palaeoloxodon falconeri": (0.4, "L", "Middle Pleistocene, Spinagallo Cave, Sicily (age range ~0.2-0.6 Ma)"),
    "Palaeoloxodon antiquus": (0.2, "L", "Middle-Late Pleistocene"),
    "Mammuthus primigenius": (0.03, "L", "Late Pleistocene"),
    "Mammuthus meridionalis": (2.0, "L", "Early Pleistocene"),
    "Mammuthus columbi": (0.05, "L", "Late Pleistocene"),
    "Mammut americanum": (0.05, "L", "Late Pleistocene"),
    "Stegodon insignis": (8.0, "V", "'unknown Miocene locality in the Himalayas' (Benoit et al.); age very uncertain (23-2 Ma)"),
    "Zygolophodon borsoni": (3.5, "V", "Pliocene, Mammal Neogene zone 15"),
    "Palaeomastodon beadnelli": (31.0, "V", "early Oligocene, Fayum"),
    "Moeritherium lyonsi": (36.0, "V", "late Eocene, Qasr el Sagha Fm, Fayum"),
    "Prorastomus sirenoides": (50.0, "V", "early Eocene (~50 Ma); sirenian"),
    "Seggeurius amourensis": (50.0, "V", "early Eocene (~50 Ma); hyracoid"),
}


def extract_benoit():
    src = os.path.join(RAW, "epmc/PMC6597534.xml")
    if not os.path.exists(src):
        return
    s = open(src, encoding="utf8").read()
    m = re.search(r"<table-wrap[ >].*?</table-wrap>", s, flags=re.S)
    out = []
    for r in re.findall(r"<tr[^>]*>(.*?)</tr>", m.group(0), flags=re.S):
        cells = re.findall(r"<t[dh][^>]*?(?:/>|>(.*?)</t[dh]>)", r, flags=re.S)
        # drop footnote superscripts before stripping tags
        cells = [html.unescape(re.sub(r"<[^>]+>", "", re.sub(r"<sup>.*?</sup>", "", c))).strip() for c in cells]
        if len(cells) < 5 or cells[0] not in BENOIT_AGES:
            continue
        age, flag, note = BENOIT_AGES[cells[0]]
        out.append(dict(species=cells[0].replace(" ", "_"), age_Ma=age, age_flag=flag, age_note=note,
                        ecv_cm3=float(cells[1]), brain_g=float(cells[2]), body_g=float(cells[3]),
                        clade="Proboscidea" if cells[0].split()[0] not in ("Prorastomus", "Seggeurius") else "other Afrotheria",
                        source="Benoit et al. 2019 Sci Rep 9:9323 Table 1"))
    _w(os.path.join(DATA, "fossil_proboscideans_benoit2019.csv"), out)
    print(f"  Benoit et al.: {len(out)} proboscidean (and outgroup) endocasts")


def extract_bertrand():
    src = os.path.join(RAW, "zenodo/bertrand2022_Placental_data_final.csv")
    if not os.path.exists(src):
        return
    rows = list(csv.DictReader(open(src, encoding="utf8", errors="replace")))
    out = []
    for r in rows:
        try:
            ecv = float(r["Brain_volume_cm3"]); bm = float(r["Body_mass_g"])
        except ValueError:
            continue
        fad = float(r["FAD_Ma"]) if r["FAD_Ma"] not in ("", "NA") else np.nan
        lad = float(r["LAD_Ma"]) if r["LAD_Ma"] not in ("", "NA") else np.nan
        out.append(dict(species=r["Species"], period=r["Period"], group=r["Group"], order=r["Order"],
                        fad_Ma=fad, lad_Ma=lad, age_Ma=np.nanmean([fad, lad]) if not (np.isnan(fad) and np.isnan(lad)) else np.nan,
                        body_g=bm, ecv_cm3=ecv, status=r["Status"]))
    _w(os.path.join(DATA, "fossil_placentals_bertrand2022.csv"), out)
    print(f"  Bertrand: {len(out)} taxa with endocranial volume and body mass")


def _cell_text(cell):
    """Text of a docx table cell with superscript runs (footnote marks) removed."""
    parts = []
    for p in cell.paragraphs:
        for run in p.runs:
            if run.font.superscript:
                continue
            parts.append(run.text)
        parts.append(" ")
    return re.sub(r"\s+", " ", "".join(parts)).strip()


def _num(s):
    s = s.replace(",", "").replace("–", "").replace("-", "").strip()
    try:
        return float(s)
    except ValueError:
        return np.nan


def extract_waugh():
    import docx
    src = os.path.join(RAW, "epmc/waugh2021/pone.0257803.s003.docx")
    if not os.path.exists(src):
        return
    d = docx.Document(src)
    t = d.tables[2]  # S3 Table: fossil cetacean brain and body mass and EQ estimates
    out = []
    for r in t.rows[1:]:
        c = [_cell_text(x) for x in r.cells]
        fam, sp, spec = c[0], c[1], c[2]
        age, ecv, brain, body_wt, body_b = _num(c[3]), _num(c[6]), _num(c[7]), _num(c[8]), _num(c[9])
        out.append(dict(family=fam, species=sp, specimen=spec, age_Ma=age, ecv_cm3=ecv,
                        brain_g=brain, body_kg_waugh=body_wt, body_kg_boessenecker=body_b))
    _w(os.path.join(DATA, "fossil_cetaceans_waugh2021.csv"), out)
    print(f"  Waugh & Thewissen S3: {len(out)} fossil cetacean specimens")


def extract_du():
    import openpyxl
    import warnings
    src = os.path.join(RAW, "zenodo/du2018_ECV_dataset.xlsx")
    if not os.path.exists(src):
        return
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        wb = openpyxl.load_workbook(src, read_only=True, data_only=True)
    rows = list(wb["Data"].iter_rows(values_only=True))
    hdr = rows[0]
    out = []
    for r in rows[1:]:
        d = dict(zip(hdr, r))
        ecv = [float(d[k]) for k in ("ecv1", "ecv2", "ecv3", "ecv4", "ecv5", "ecv6")
               if isinstance(d[k], (int, float))]
        if not ecv:
            continue
        if isinstance(d["mean.date"], (int, float)):
            age = float(d["mean.date"])
        else:
            age = 0.5 * (float(d["min.date"]) + float(d["max.date"]))
        out.append(dict(specimen=d["ID"], taxon=d["lump.taxon"], age_Ma=age, ecv_cm3=float(np.mean(ecv)),
                        n_estimates=len(ecv)))
    _w(os.path.join(DATA, "hominin_ecv_du2018.csv"), out)
    print(f"  Du et al.: {len(out)} hominin crania")


def extract_bird_reptile_neurons():
    import openpyxl
    import warnings
    out = []
    src = os.path.join(RAW, "kverkova/pnas.2121624119.sd01.xlsx")
    if os.path.exists(src):
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            wb = openpyxl.load_workbook(src, read_only=True, data_only=True)
        rows = list(wb.worksheets[0].iter_rows(values_only=True))
        for r in rows[2:]:
            if not r[1] or not isinstance(r[4], (int, float)):
                continue
            pall = r[40] if len(r) > 40 and isinstance(r[40], (int, float)) else None
            tel = r[11] if isinstance(r[11], (int, float)) else None
            out.append(dict(species=str(r[1]).strip().replace(" ", "_"), common=r[0], cls=r[2], group=r[3],
                            body_g=float(r[4]), brain_g=float(r[5]) if isinstance(r[5], (int, float)) else "",
                            brain_neurons=r[7] if isinstance(r[7], (int, float)) else "",
                            pallium_neurons=pall if pall else "", telencephalon_neurons=tel if tel else "",
                            method="isotropic fractionator", source="Kverkova et al. 2022 PNAS Dataset S1"))
    src = os.path.join(RAW, "zenodo/sol2022_neurons.csv")
    have = {o["species"] for o in out}
    if os.path.exists(src):
        rows = list(csv.DictReader(open(src, encoding="utf-8-sig"), delimiter=";"))
        for r in rows:
            nm = r["animal"].strip()
            if nm in have:
                continue
            out.append(dict(species=nm, common="", cls="Aves", group="", body_g=float(r["Body"]),
                            brain_g=float(r["AbsBrain"]), brain_neurons=float(r["Neurons"]),
                            pallium_neurons=float(r["Pall_N"]), telencephalon_neurons=float(r["Tel_N"]),
                            method="isotropic fractionator",
                            source="Sol et al. 2022 Nat Ecol Evol (compiles Olkowicz et al. 2016 PNAS)"))
    if out:
        _w(os.path.join(DATA, "neurons_birds_reptiles.csv"), out)
        print(f"  bird/reptile neuron counts: {len(out)} species")


def extract_massen_mammals():
    import openpyxl
    import warnings
    src = os.path.join(RAW, "massen/massen_brain_measures.xlsx")
    if not os.path.exists(src):
        return
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        wb = openpyxl.load_workbook(src, read_only=True, data_only=True)
    rows = list(wb.worksheets[0].iter_rows(values_only=True))
    out = []
    for r in rows[1:]:
        if not r[3] or not isinstance(r[5], (int, float)):
            continue
        if not isinstance(r[8], (int, float)):
            continue
        out.append(dict(species=str(r[3]).strip().replace(" ", "_"), common=r[1], order=r[4],
                        body_g=float(r[5]), brain_g=float(r[6]), brain_neurons=float(r[7]) * 1e6,
                        cortex_neurons=float(r[8]) * 1e6, method="isotropic fractionator",
                        source="Massen et al. 2021 Commun Biol source data (compiled from the primary IF literature)"))
    _w(os.path.join(DATA, "neurons_mammals_massen2021.csv"), out)
    print(f"  Massen et al. mammal cortical neuron counts: {len(out)} species")


def extract_muller():
    """Specimen-level cetacean brain and body masses (independent cross-check of the delphinid EQ values)."""
    import openpyxl
    import warnings
    src = os.path.join(RAW, "zenodo/muller2019_CCExCet.xlsx")
    if not os.path.exists(src):
        return
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        wb = openpyxl.load_workbook(src, read_only=True, data_only=True)
    rows = list(wb.worksheets[0].iter_rows(values_only=True))
    spec, means = {}, {}
    for r in rows[4:]:
        if not r[0]:
            break  # end of the main cetacean block (excluded data and other mammals follow)
        if not (isinstance(r[2], (int, float)) and isinstance(r[3], (int, float))):
            continue
        (spec if r[1] else means).setdefault(r[0], []).append((float(r[2]), float(r[3])))
    out = []
    for sp in sorted(set(spec) | set(means)):
        v = spec.get(sp) or means[sp]
        out.append(dict(species=sp, n_specimens=len(v), body_kg=float(np.mean([x[0] for x in v])),
                        brain_g=float(np.mean([x[1] for x in v])),
                        source="Muller & Montgomery 2019 J Evol Biol, Table S1A (Dryad 10.5061/dryad.rm4368f)"))
    _w(os.path.join(DATA, "cetaceans_muller2019.csv"), out)
    print(f"  Muller & Montgomery: {len(out)} cetacean species")


def check_upham_trees():
    """Data-quality check (referee point m7): node ages and the placement of Lissodelphis borealis in the two
    distributed Upham et al. (2019) MCC 'target' trees, against the primary Alvarez-Carretero et al. tree."""
    import json
    files = {"upham2019_completed_5911sp": os.path.join(RAW, "trees/mamphy_mcc.tre"),
             "upham2019_DNAonly_4098sp": os.path.join(RAW, "trees/mamphy_mcc_dnaonly.tre"),
             "alvarez2022_4705sp": os.path.join(RAW, "trees/alvarez2022_4705sp_mean.nwk")}
    if not all(os.path.exists(f) for f in files.values()):
        return
    pairs = [("Balaenoptera_musculus", "Tursiops_truncatus", "crown Cetacea (baleen-toothed whale split)"),
             ("Physeter", "Tursiops_truncatus", "sperm whale - dolphin"),
             ("Lissodelphis_borealis", "Tursiops_truncatus", "Lissodelphis borealis - bottlenose dolphin (same family)"),
             ("Lissodelphis_borealis", "Balaenoptera_musculus", "Lissodelphis borealis - blue whale"),
             ("Homo_sapiens", "Pan_troglodytes", "human - chimpanzee"),
             ("Homo_sapiens", "Macaca_mulatta", "human - macaque"),
             ("Homo_sapiens", "Mus_musculus", "human - mouse")]
    out = {}
    for nm, path in files.items():
        if nm.startswith("alvarez"):
            root = phylo.parse_newick(open(path).read())
            scale = 100.0
        else:
            root = phylo.read_nexus_tree(path)
            scale = 1.0
        depth, mx = phylo.node_heights(root)
        tips = phylo.tips(root)
        low = {t.name.lower(): t for t in tips}

        def find(prefix):
            c = [t for k, t in low.items() if k.startswith(prefix.lower())]
            return c[0] if c else None

        def mrca(a, b):
            anc = set()
            n = a
            while n is not None:
                anc.add(id(n)); n = n.parent
            n = b
            while id(n) not in anc:
                n = n.parent
            return n
        d = {}
        for a, b, lab in pairs:
            ta, tb = find(a), find(b)
            if ta is not None and tb is not None:
                d[lab] = round((mx - depth[id(mrca(ta, tb))]) * scale, 2)
        t = find("Lissodelphis_borealis")
        if t is not None:
            d["sister_group_of_Lissodelphis_borealis"] = sorted(
                {"_".join(x.name.split("_")[:2]) for x in phylo.tips(t.parent)} - {"Lissodelphis_borealis", "lissodelphis_borealis"})[:4]
        out[nm] = d
    out["note"] = ("node ages in Ma; both distributed Upham et al. MCC 'target' files date the baleen-toothed whale "
                   "split at 16.6-17.7 Ma (primary tree: 33.6 Ma) and join the right whale dolphin Lissodelphis borealis to the "
                   "blue whale at 3-4 Ma, i.e. inside the baleen whales, while its split from another dolphin is placed at "
                   "the cetacean root; this is a property of those two files, checked here, not a claim about the "
                   "posterior sample as a whole")
    json.dump(out, open(os.path.join(DATA, "upham_tree_check.json"), "w"), indent=1)
    print("  Upham tree check written")


def main():
    os.makedirs(DATA, exist_ok=True)
    extract_burger_and_tree()
    check_upham_trees()
    extract_bertrand()
    extract_waugh()
    extract_benoit()
    extract_du()
    extract_bird_reptile_neurons()
    extract_massen_mammals()
    extract_muller()


if __name__ == "__main__":
    main()