"""
Tree integrity check -- run BEFORE any Cohort B inferential analysis.

Verifies, in this order:
  1. how many species resolved / how many excluded, with the exclusion list
  2. correspondence between tree tip labels and the CORAL dataset
  3. that the CORAL-name -> binomial mapping is one-to-one where it matters
  4. that the tree is usable: ultrametric-ish, positive branch lengths, no
     duplicate tips

Nothing here makes an analytical choice. It reports; it does not repair.
Unresolved names are LISTED and EXCLUDED, never silently substituted by
synonyms (protocol v5, section 5.3.1).

Usage:
    python3 tree_check.py tree1.nwk [tree2.nwk ...]
"""

import os as _os
HERE = _os.path.dirname(_os.path.abspath(__file__))
ROOT_DIR = _os.path.dirname(HERE)


import json
import sys
from collections import Counter

import dendropy
import pandas as pd

MAPPING = ROOT_DIR + "/data/coral_name_mapping.json"
COUNTS = ROOT_DIR + "/data/coral_counts.pkl"


def load_trees(paths):
    tips = {}
    for p in paths:
        t = dendropy.Tree.get(path=p, schema="newick", preserve_underscores=True)
        for leaf in t.leaf_node_iter():
            tips[leaf.taxon.label.strip()] = p
    return tips


def main(paths):
    mapping = json.load(open(MAPPING))               # CORAL name -> binomial
    counts = pd.read_pickle(COUNTS)
    coral_species = list(counts.index)

    binom = {s: mapping.get(s) for s in coral_species}
    unmapped = sorted([s for s, b in binom.items() if not b])
    wanted = sorted({b for b in binom.values() if b})

    tips = load_trees(paths)
    tip_names = set(tips)

    resolved = sorted(set(wanted) & tip_names)
    missing = sorted(set(wanted) - tip_names)
    extra = sorted(tip_names - set(wanted))

    print("=" * 62)
    print("TREE INTEGRITY CHECK -- Cohort B")
    print("=" * 62)
    print(f"CORAL focal species (post-aggregation) : {len(coral_species)}")
    print(f"  with an extractable binomial          : {len(wanted)}")
    print(f"  without one (-> plan B / excluded)    : {len(unmapped)}")
    print(f"Tree tips supplied                      : {len(tip_names)}")
    print()
    print(f"RESOLVED (usable in the phylogenetic analysis) : {len(resolved)}")
    print(f"EXCLUDED, unresolved by TimeTree               : {len(missing)}")
    print(f"Tips in tree but not in dataset                : {len(extra)}")
    pct = 100.0 * len(resolved) / max(len(coral_species), 1)
    print(f"Coverage of the focal-species set              : {pct:.1f}%")

    # one-to-one check on the names that actually matter
    back = Counter(b for s, b in binom.items() if b in set(resolved))
    collisions = {b: n for b, n in back.items() if n > 1}
    print()
    print(f"Binomials receiving >1 CORAL name (collisions)  : {len(collisions)}")
    for b, n in list(collisions.items())[:10]:
        originals = [s for s, x in binom.items() if x == b]
        print(f"   {b} <- {originals}")
    print("Mapping is one-to-one over resolved species." if not collisions
          else "NOT one-to-one -- collisions above must be resolved or excluded.")

    # tree usability
    print()
    dup = [k for k, v in Counter(tips.values()).items() if False]  # placeholder
    for p in paths:
        t = dendropy.Tree.get(path=p, schema="newick", preserve_underscores=True)
        lens = [e.length for e in t.edges() if e.length is not None]
        leaves = [l.taxon.label for l in t.leaf_node_iter()]
        d = [k for k, n in Counter(leaves).items() if n > 1]
        neg = sum(1 for x in lens if x is not None and x < 0)
        depths = [l.distance_from_root() for l in t.leaf_node_iter()]
        print(f"{p}: {len(leaves)} tips | duplicate labels: {len(d)} | "
              f"negative branch lengths: {neg} | "
              f"root-to-tip depth {min(depths):.1f}-{max(depths):.1f}")

    with open(ROOT_DIR + "/data/tree_excluded_species.txt", "w") as f:
        f.write("# unresolved by TimeTree -- excluded, not substituted\n")
        f.write("\n".join(missing) + "\n")
        f.write("# no extractable binomial\n")
        f.write("\n".join(unmapped) + "\n")
    print()
    print("Exclusion list written to tree_excluded_species.txt")
    print()
    print("VERDICT:", "PASS -- proceed to Cohort B analysis" if not collisions
          and len(resolved) > 0 else "REVIEW REQUIRED before proceeding")


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(__doc__)
    else:
        main(sys.argv[1:])
