from __future__ import annotations

import argparse
from collections import defaultdict
import glob
import json
from pathlib import Path
import sys
from typing import Any

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import numpy as np

from patientphex.ontology import HPOOntology
from patientphex.sapbert import _unique_alias_records, encode_cls, rank_concepts


def load_jsonl(path: Path) -> list[dict[str, Any]]:
    return [json.loads(line) for line in path.open(encoding="utf-8") if line.strip()]


def main() -> None:
    parser = argparse.ArgumentParser(description="Ensemble verifier scores and normalize selected A-set spans.")
    parser.add_argument("--data", type=Path, required=True)
    parser.add_argument("--ontology", type=Path, required=True)
    parser.add_argument("--verifier-glob", required=True)
    parser.add_argument("--sapbert-model", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--threshold", type=float, default=0.66)
    parser.add_argument(
        "--target-fraction", type=float,
        help="If supplied, replace the raw threshold by the ensemble-score quantile that keeps this fraction.",
    )
    parser.add_argument("--lexicon-threshold", type=float, default=0.0)
    parser.add_argument("--phenotagger-threshold", type=float, default=0.90)
    parser.add_argument("--device", default="cuda")
    args = parser.parse_args()

    paths = sorted(glob.glob(args.verifier_glob))
    if not paths:
        raise FileNotFoundError(f"no verifier outputs matched {args.verifier_glob!r}")
    grouped: dict[tuple[str, int, int], list[dict[str, Any]]] = defaultdict(list)
    for path in paths:
        result = json.loads(Path(path).read_text(encoding="utf-8"))
        for record in result["predictions"]:
            grouped[(str(record["pmc_id"]), int(record["start"]), int(record["end"]))].append(record)
    expected_models = len(paths)
    incomplete = [key for key, records in grouped.items() if len(records) != expected_models]
    if incomplete:
        raise RuntimeError(f"{len(incomplete)} candidates are missing ensemble scores")
    ensemble_records: list[dict[str, Any]] = []
    for records in grouped.values():
        record = dict(records[0])
        record["ensemble_score"] = float(np.mean([float(item["score"]) for item in records]))
        ensemble_records.append(record)
    threshold = args.threshold
    if args.target_fraction is not None:
        if not 0.0 < args.target_fraction <= 1.0:
            parser.error("--target-fraction must be in (0, 1]")
        threshold = float(np.quantile(
            [record["ensemble_score"] for record in ensemble_records], 1.0 - args.target_fraction
        ))
    selected = [record for record in ensemble_records if record["ensemble_score"] >= threshold]

    ontology = HPOOntology.from_obo(args.ontology)

    def valid_identifier(value: Any) -> str | None:
        if value is None:
            return None
        identifier = ontology.canonical_id(str(value))
        return identifier if ontology.classify_id(identifier) == "valid_branch_hp" else None

    unresolved = sorted({
        str(record["text"]).lower() for record in selected
        if not (
            valid_identifier(record.get("lexicon_identifier")) is not None
            and float(record.get("lexicon_quality", 0.0)) >= args.lexicon_threshold
        ) and not (
            valid_identifier(record.get("phenotagger_identifier")) is not None
            and float(record.get("phenotagger_score", 0.0)) >= args.phenotagger_threshold
        )
    })
    ranking_by_text: dict[str, list[str]] = {}
    if unresolved:
        from transformers import AutoModel, AutoTokenizer
        tokenizer = AutoTokenizer.from_pretrained(args.sapbert_model, local_files_only=True)
        model = AutoModel.from_pretrained(args.sapbert_model, local_files_only=True).to(args.device)
        alias_ids, alias_texts = _unique_alias_records(ontology)
        alias_embeddings = encode_cls(alias_texts, tokenizer, model, device=args.device, batch_size=256)
        query_embeddings = encode_cls(unresolved, tokenizer, model, device=args.device, batch_size=256)
        rankings = rank_concepts(query_embeddings, alias_embeddings, alias_ids, device=args.device, max_concepts=5)
        ranking_by_text = dict(zip(unresolved, rankings, strict=True))

    entities_by_doc: dict[str, dict[tuple[int, int, str], dict[str, Any]]] = defaultdict(dict)
    source_counts: dict[str, int] = defaultdict(int)
    for record in selected:
        lexicon_identifier = valid_identifier(record.get("lexicon_identifier"))
        phenotagger_identifier = valid_identifier(record.get("phenotagger_identifier"))
        if (
            lexicon_identifier is not None
            and float(record.get("lexicon_quality", 0.0)) >= args.lexicon_threshold
        ):
            identifier = lexicon_identifier
            source = "lexicon"
        elif (
            phenotagger_identifier is not None
            and float(record.get("phenotagger_score", 0.0)) >= args.phenotagger_threshold
        ):
            identifier = phenotagger_identifier
            source = "phenotagger"
        else:
            identifier = ranking_by_text[str(record["text"]).lower()][0]
            source = "sapbert"
        source_counts[source] += 1
        start, end = int(record["start"]), int(record["end"])
        entities_by_doc[str(record["pmc_id"])][(start, end, identifier)] = {
            "identifier": identifier, "type": "Phenotype", "offset": start,
            "length": end - start, "text": record["text"], "note": None,
        }

    documents = load_jsonl(args.data)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.output.open("w", encoding="utf-8", newline="\n") as handle:
        for document in documents:
            entities = sorted(
                entities_by_doc[str(document["pmc_id"])].values(),
                key=lambda item: (int(item["offset"]), int(item["length"]), str(item["identifier"])),
            )
            output = {
                "pmc_id": document["pmc_id"], "pmid": document.get("pmid"),
                "entities": entities, "association": [],
            }
            handle.write(json.dumps(output, ensure_ascii=False, separators=(",", ":")) + "\n")
    print(json.dumps({
        "models": expected_models, "threshold": threshold, "target_fraction": args.target_fraction,
        "selected_candidates": len(selected), "entities": sum(len(x) for x in entities_by_doc.values()),
        "normalization_sources": dict(source_counts), "output": str(args.output),
    }))


if __name__ == "__main__":
    main()