PlantAI-rag-new/bioflow_kernel/scripts/reference_classification_from_tree.py

#!/usr/bin/env python3
"""Assign target-species subfamilies from a reference-labelled joint tree."""

from __future__ import annotations

import argparse
import csv
import json
from collections import Counter
from datetime import datetime
from pathlib import Path

from Bio import Phylo

TOPOLOGY_SUPPORT_BOOTSTRAP = 70.0
STRONG_TOPOLOGY_BOOTSTRAP = 90.0


def read_reference_subfamily(path: Path, reference_prefix: str) -> dict[str, str]:
    labels: dict[str, str] = {}
    with path.open(encoding="utf-8") as handle:
        for raw in handle:
            line = raw.strip()
            if not line or line.startswith("#"):
                continue
            parts = line.split()
            if len(parts) < 2:
                continue
            gene, subfamily = parts[0], parts[1]
            labels[gene] = subfamily
            labels[gene.lower()] = subfamily
            labels[f"{reference_prefix}|{gene}"] = subfamily
            labels[f"{reference_prefix}|{gene}".lower()] = subfamily
    return labels


def lookup_subfamily(labels: dict[str, str], tip_name: str, default: str = "Reference_unclassified") -> str:
    if tip_name in labels:
        return labels[tip_name]
    lowered = tip_name.lower()
    if lowered in labels:
        return labels[lowered]
    if "|" in tip_name:
        gene = tip_name.split("|", 1)[1]
        if gene in labels:
            return labels[gene]
        if gene.lower() in labels:
            return labels[gene.lower()]
    return default


def is_reference_tip(tip_name: str, labels: dict[str, str], reference_prefix: str) -> bool:
    if tip_name.startswith(f"{reference_prefix}|"):
        return True
    return lookup_subfamily(labels, tip_name, default="") != ""


def bootstrap_value(clade) -> float | None:
    for value in (getattr(clade, "confidence", None), getattr(clade, "name", None)):
        if value is None:
            continue
        try:
            return float(value)
        except (TypeError, ValueError):
            continue
    return None


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--tree", required=True, type=Path)
    parser.add_argument("--reference-subfamily", required=True, type=Path)
    parser.add_argument("--target-prefix", default="Ft")
    parser.add_argument("--reference-prefix", default="At")
    parser.add_argument("--assignment-out", required=True, type=Path)
    parser.add_argument("--subfamily-out", required=True, type=Path)
    parser.add_argument("--manifest-out", type=Path)
    parser.add_argument("--report-out", type=Path)
    args = parser.parse_args()

    ref_subfamily = read_reference_subfamily(args.reference_subfamily, args.reference_prefix)
    tree = Phylo.read(args.tree, "newick")
    terminals = tree.get_terminals()
    reference_tips = [
        t for t in terminals
        if is_reference_tip(t.name or "", ref_subfamily, args.reference_prefix)
    ]
    target_tips = [t for t in terminals if (t.name or "").startswith(f"{args.target_prefix}|")]
    if not target_tips:
        reference_names = {id(t) for t in reference_tips}
        target_tips = [t for t in terminals if id(t) not in reference_names]

    if not target_tips:
        raise SystemExit(
            f"No target tips found. Tree has no tips with prefix {args.target_prefix}| "
            "and no non-reference tips after reference detection."
        )
    if not reference_tips:
        raise SystemExit(
            f"No reference tips found. The tree must contain tips with prefix {args.reference_prefix}| "
            "or gene IDs present in the reference subfamily file."
        )

    parent: dict[object, object | None] = {tree.root: None}
    for clade in tree.find_clades(order="level"):
        for child in clade.clades:
            parent[child] = clade

    assignments: list[dict[str, str]] = []
    subfamily_rows: list[tuple[str, str]] = []

    for ref in sorted(reference_tips, key=lambda item: item.name or ""):
        subfamily_rows.append((ref.name, lookup_subfamily(ref_subfamily, ref.name or "")))

    for target in sorted(target_tips, key=lambda item: item.name or ""):
        distances = sorted(
            ((tree.distance(target, ref), ref.name, lookup_subfamily(ref_subfamily, ref.name or "")) for ref in reference_tips),
            key=lambda item: (item[0], item[1]),
        )
        nearest_distance, nearest_ref, nearest_subfamily = distances[0]
        nearest3 = distances[:3]
        nearest3_counts = Counter(item[2] for item in nearest3)
        nearest3_label, nearest3_count = nearest3_counts.most_common(1)[0]

        mixed_clade = parent.get(target)
        mixed_ref_subfamilies: list[str] = []
        mixed_bootstrap: float | None = None
        while mixed_clade is not None:
            mixed_refs = [tip for tip in mixed_clade.get_terminals() if (tip.name or "").startswith(f"{args.reference_prefix}|")]
            if not mixed_refs:
                mixed_refs = [
                    tip for tip in mixed_clade.get_terminals()
                    if is_reference_tip(tip.name or "", ref_subfamily, args.reference_prefix)
                ]
            if mixed_refs:
                mixed_ref_subfamilies = sorted({lookup_subfamily(ref_subfamily, tip.name or "") for tip in mixed_refs})
                mixed_bootstrap = bootstrap_value(mixed_clade)
                break
            mixed_clade = parent.get(mixed_clade)

        topology_subfamily = mixed_ref_subfamilies[0] if len(mixed_ref_subfamilies) == 1 else ""
        topology_supported = (
            bool(topology_subfamily)
            and mixed_bootstrap is not None
            and mixed_bootstrap >= TOPOLOGY_SUPPORT_BOOTSTRAP
        )

        if topology_subfamily and topology_subfamily == nearest_subfamily:
            if mixed_bootstrap is not None and mixed_bootstrap >= STRONG_TOPOLOGY_BOOTSTRAP:
                confidence = "high"
                note = "single_reference_subfamily_clade_strong_bootstrap"
            elif topology_supported:
                confidence = "medium"
                note = "single_reference_subfamily_clade_supported_bootstrap"
            else:
                confidence = "medium"
                note = "single_reference_subfamily_clade_low_or_missing_bootstrap"
            assigned = topology_subfamily
        elif topology_supported and topology_subfamily:
            confidence = "medium"
            note = "topology_conflict_nearest_reference_supported_clade"
            assigned = topology_subfamily
        elif nearest3_label == nearest_subfamily and nearest3_count >= 2:
            confidence = "medium"
            note = "nearest3_majority_support"
            assigned = nearest_subfamily
        else:
            confidence = "low"
            note = "nearest_reference_only_mixed_clade"
            assigned = nearest_subfamily

        assignments.append(
            {
                "target_gene": target.name,
                "assigned_subfamily": assigned,
                "confidence": confidence,
                "nearest_reference": nearest_ref,
                "nearest_reference_subfamily": nearest_subfamily,
                "nearest_reference_distance": f"{nearest_distance:.6g}",
                "nearest3_references": ";".join(f"{name}:{sub}:{dist:.6g}" for dist, name, sub in nearest3),
                "minimal_mixed_clade_ref_subfamilies": ";".join(mixed_ref_subfamilies) if mixed_ref_subfamilies else "NA",
                "minimal_mixed_clade_bootstrap": "NA" if mixed_bootstrap is None else f"{mixed_bootstrap:.6g}",
                "topology_reference_subfamily": topology_subfamily or "NA",
                "topology_bootstrap": "NA" if mixed_bootstrap is None else f"{mixed_bootstrap:.6g}",
                "nearest3_majority_subfamily": nearest3_label,
                "nearest3_majority_count": str(nearest3_count),
                "classification_conflict": "yes" if topology_supported and topology_subfamily != nearest_subfamily else "no",
                "classification_note": note,
            }
        )
        subfamily_rows.append((target.name, assigned))

    args.assignment_out.parent.mkdir(parents=True, exist_ok=True)
    with args.assignment_out.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(assignments[0]), delimiter="\t")
        writer.writeheader()
        writer.writerows(assignments)

    with args.subfamily_out.open("w", encoding="utf-8") as handle:
        for label, subfamily in sorted(subfamily_rows):
            handle.write(f"{label}\t{subfamily}\n")

    assignment_counts = Counter(row["assigned_subfamily"] for row in assignments)
    confidence_counts = Counter(row["confidence"] for row in assignments)
    note_counts = Counter(row["classification_note"] for row in assignments)

    manifest_out = args.manifest_out
    if manifest_out is None:
        manifest_out = args.assignment_out.with_name(
            args.assignment_out.name.replace(".assignment.tsv", ".manifest.json")
        )
    report_out = args.report_out
    if report_out is None:
        report_out = args.assignment_out.with_name(
            args.assignment_out.name.replace(".assignment.tsv", "_report.md")
        )

    def final_artifact_path(path: Path) -> str:
        if path.parent.name == "final":
            return f"final/{path.name}"
        return str(path)

    manifest = {
        "manifest_version": "1.0",
        "stage": "reference_classification",
        "status": "success",
        "prefix": args.assignment_out.name.split(".reference_classification.assignment.tsv")[0],
        "timestamp": datetime.now().isoformat(timespec="seconds"),
        "inputs": {
            "tree": str(args.tree),
            "reference_subfamily": str(args.reference_subfamily),
            "target_species": args.target_prefix,
            "reference_species": args.reference_prefix,
        },
        "artifacts": {
            "assignment_table": {"path": final_artifact_path(args.assignment_out), "type": "tsv"},
            "subfamily_file": {"path": final_artifact_path(args.subfamily_out), "type": "tsv"},
            "report": {"path": final_artifact_path(report_out), "type": "markdown"},
        },
        "metrics": {
            "target_gene_count": len(target_tips),
            "reference_gene_count": len(reference_tips),
            "total_tree_tips": len(terminals),
            "confidence_counts": dict(sorted(confidence_counts.items())),
            "assigned_subfamily_counts": dict(sorted(assignment_counts.items())),
            "classification_note_counts": dict(sorted(note_counts.items())),
        },
        "parameters": {
            "assignment_rule": (
                "a supported local topology clade containing a single reference subfamily "
                f"(bootstrap >= {TOPOLOGY_SUPPORT_BOOTSTRAP:g}) is prioritized over nearest-reference "
                "distance; nearest3 majority is used only when no supported single-subfamily topology "
                "clade is available. classification_note and classification_conflict expose conflicts."
            ),
            "script": "scripts/reference_classification_from_tree.py",
        },
    }
    manifest_out.write_text(json.dumps(manifest, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")

    report_lines = [
        "# Reference-Guided Subfamily Classification",
        "",
        "## Summary",
        "",
        f"- Target genes classified: {len(target_tips)}",
        f"- Reference genes: {len(reference_tips)}",
        f"- Total joint-tree tips: {len(terminals)}",
        "",
        "## Confidence Counts",
        "",
        "| Confidence | Ft genes |",
        "|---|---:|",
    ]
    for label, count in sorted(confidence_counts.items()):
        report_lines.append(f"| {label} | {count} |")
    report_lines.extend(["", "## Assigned Subfamily Counts", "", "| Assigned subfamily | Ft genes |", "|---|---:|"])
    for label, count in sorted(assignment_counts.items()):
        report_lines.append(f"| {label} | {count} |")
    report_lines.extend(["", "## Rule", ""])
    report_lines.append(
        "A local tree clade containing a single At reference subfamily is used as the primary "
        f"assignment evidence when its bootstrap is >= {TOPOLOGY_SUPPORT_BOOTSTRAP:g}. "
        "Nearest-reference and nearest3 evidence are retained in the assignment table and used "
        "as fallback evidence when the local topology is mixed or weakly supported."
    )
    report_out.write_text("\n".join(report_lines) + "\n", encoding="utf-8")

    print(f"target_tips={len(target_tips)}")
    print(f"reference_tips={len(reference_tips)}")
    print(f"assignment_out={args.assignment_out}")
    print(f"subfamily_out={args.subfamily_out}")
    print(f"manifest_out={manifest_out}")
    print(f"report_out={report_out}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())