#!/usr/bin/env python3
"""Build rank blends of independently submitted 0.921 public anchors."""

from __future__ import annotations

import hashlib
import json
from pathlib import Path

import numpy as np
import pandas as pd


UID = "StudyInstanceUID"
TARGETS = [
    "ACL",
    "MCL",
    "Medial Meniscus",
    "Lateral Meniscus",
    "Medial OA",
    "Lateral OA",
    "PF OA",
    "Effusion",
    "Synovitis",
    "Baker's",
    "Contusion",
    "Fracture",
]
SOURCES = {
    "mattia": "mattiaangeli/bend-the-knee-to-dinov3-ensembled",
    "roman": "romantamrazov/rsna-knee-dinosaur-v3",
    "legacy_mattia": "mraz1006/rsna-knee-silver-dinov3-ensemble",
}


def locate_input_file(source: str, filename: str) -> Path:
    owner, slug = source.split("/", 1)
    candidates = [
        Path("/kaggle/input") / slug / filename,
        Path("/kaggle/input/code") / slug / filename,
        Path("/kaggle/input/notebooks") / owner / slug / filename,
    ]
    for parent in (Path("/kaggle/input"), Path("/kaggle/input/code")):
        source = parent / slug
        if source.is_dir():
            candidates.extend(source.glob(f"*/{filename}"))
    unique = sorted({path.resolve() for path in candidates if path.is_file()})
    if len(unique) != 1:
        mounted = []
        for parent in (Path("/kaggle/input"), Path("/kaggle/input/code")):
            if parent.is_dir():
                mounted.extend(str(path) for path in sorted(parent.iterdir()))
        raise FileNotFoundError(
            f"expected exactly one {source}/{filename}, found {len(unique)}: {unique}; "
            f"mounted roots: {mounted}"
        )
    return unique[0]


def sha256(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def rank_columns(values: np.ndarray) -> np.ndarray:
    return pd.DataFrame(values).rank(method="average", pct=True).to_numpy(float)


def validate(frame: pd.DataFrame, expected_uids: list[str]) -> None:
    if frame.columns.tolist() != [UID, *TARGETS]:
        raise ValueError(f"unexpected submission columns: {frame.columns.tolist()}")
    if frame[UID].astype(str).tolist() != expected_uids:
        raise ValueError("submission UID order does not match test.csv")
    values = frame[TARGETS].to_numpy(float)
    if not np.isfinite(values).all() or not ((values >= 0) & (values <= 1)).all():
        raise ValueError("submission predictions must be finite and in [0, 1]")


competition = Path("/kaggle/input/competitions/rsna-knee-abnormality-detection")
if not (competition / "test.csv").is_file():
    competition = Path("/kaggle/input/rsna-knee-abnormality-detection")
test = pd.read_csv(competition / "test.csv", dtype={UID: str})
expected_uids = test[UID].astype(str).tolist()

paths = {
    name: locate_input_file(slug, "submission.csv") for name, slug in SOURCES.items()
}
frames = {
    name: pd.read_csv(path, dtype={UID: str}) for name, path in paths.items()
}
for frame in frames.values():
    validate(frame, expected_uids)

ranks = {name: rank_columns(frame[TARGETS].to_numpy(float)) for name, frame in frames.items()}
artifacts = {}


def write_candidate(filename: str, values: np.ndarray, weights: dict[str, float]) -> None:
    candidate = pd.DataFrame(rank_columns(values), columns=TARGETS)
    candidate.insert(0, UID, expected_uids)
    validate(candidate, expected_uids)
    path = Path("/kaggle/working") / filename
    candidate.to_csv(path, index=False)
    artifacts[filename] = {"weights": weights, "sha256": sha256(path)}


for roman_weight in (0.25, 0.50, 0.75):
    values = rank_columns(
        (1.0 - roman_weight) * ranks["mattia"] + roman_weight * ranks["roman"]
    )
    filename = "submission.csv" if roman_weight == 0.50 else f"submission_roman_{int(roman_weight * 100):02d}.csv"
    write_candidate(
        filename,
        values,
        {"mattia": 1.0 - roman_weight, "roman": roman_weight},
    )

write_candidate(
    "submission_three_equal.csv",
    (ranks["mattia"] + ranks["roman"] + ranks["legacy_mattia"]) / 3.0,
    {"mattia": 1 / 3, "roman": 1 / 3, "legacy_mattia": 1 / 3},
)
write_candidate(
    "submission_latest50_threeway.csv",
    0.50 * ranks["mattia"]
    + 0.25 * ranks["roman"]
    + 0.25 * ranks["legacy_mattia"],
    {"mattia": 0.50, "roman": 0.25, "legacy_mattia": 0.25},
)

receipt = {
    "status": "PUBLIC_ANCHOR_RANK_BLENDS_COMPLETE",
    "method": "per-target percentile-rank blend",
    "sources": {
        name: {"kernel": SOURCES[name], "sha256": sha256(path)}
        for name, path in paths.items()
    },
    "artifacts": artifacts,
    "rows": len(expected_uids),
}
Path("/kaggle/working/anchor_blends_receipt.json").write_text(
    json.dumps(receipt, indent=2, sort_keys=True) + "\n", encoding="utf-8"
)
print(json.dumps(receipt, indent=2, sort_keys=True))