"""Case-level diagnosis rules.

``published_rule`` is an exact replica of the rule in the original desktop
application (toolv9), which produced the published results. LeukoAI 1.1 keeps
this rule unchanged and adds ``review_flags``: warnings for situations where
the rule is known to be unreliable. Changes to the rule itself must wait
until they can be validated on new cases.
"""
from __future__ import annotations

from dataclasses import dataclass, field

from .models import DX1, DX2, DX3

# Images per patient in the published external validation: 2,973 / 64.
VALIDATION_MEAN_IMAGES = 46
MIN_RECOMMENDED_IMAGES = 10


@dataclass
class ImageResult:
    name: str
    p1: list  # Model 1 probabilities, order DX1
    p2: list  # Model 2 probabilities, order DX2
    p3: list  # Model 3 probabilities, order DX3
    quality: dict = field(default_factory=dict)

    @property
    def k1(self):
        return max(range(len(self.p1)), key=self.p1.__getitem__)

    @property
    def k2(self):
        return max(range(len(self.p2)), key=self.p2.__getitem__)

    @property
    def k3(self):
        return max(range(len(self.p3)), key=self.p3.__getitem__)


@dataclass
class CaseResult:
    three_class: str          # ACUTE LEUKEMIA / CHRONIC LEUKEMIA / NO LEUKEMIA
    five_class: str           # ALL / AML / CLL / CML / NO LEUKEMIA (as the app reports it)
    images_of_interest: list  # up to 3 file names, as the app shows them
    n_images: int
    n_leukemic: int           # Model 1 leukemia images
    n_acute: int              # Model 2 acute images
    n_chronic: int            # Model 2 chronic images
    votes: dict               # Model 3 image votes per class
    threshold: int
    flags: list = field(default_factory=list)


def published_rule(images: list[ImageResult], threshold: int = 3) -> CaseResult:
    """Exact replica of toolv9 ``clicked()`` decision logic."""
    leukemia = 0
    bcls = []                  # (rounded Model 1 prob, file) for leukemic images
    scls = [[], [], [], []]    # (rounded Model 2 prob, file) per Model 2 class
    fcls = [0, 0, 0, 0, 0]     # Model 3 votes
    for im in images:
        k1, k2, k3 = im.k1, im.k2, im.k3
        if k1 == 1:
            leukemia += 1
            bcls.append((round(im.p1[k1], 2), im.name))
        scls[k2].append((round(im.p2[k2], 2), im.name))
        fcls[k3] += 1

    interesting = []
    if leukemia > threshold:
        if len(scls[1]) > 1:
            sub = fcls[2:4]
            five = DX3[sub.index(max(sub)) + 2]
            three = "CHRONIC LEUKEMIA"
            interesting = [x for _, x in sorted(scls[1], reverse=True)][:3]
        else:
            sub = fcls[:2]
            five = DX3[sub.index(max(sub))]
            three = "ACUTE LEUKEMIA"
            interesting = [x for _, x in sorted(bcls, reverse=True)][:3]
    else:
        three = five = "NO LEUKEMIA"

    return CaseResult(
        three_class=three, five_class=five, images_of_interest=interesting,
        n_images=len(images), n_leukemic=leukemia,
        n_acute=len(scls[0]), n_chronic=len(scls[1]),
        votes=dict(zip(DX3, fcls)), threshold=threshold,
    )


@dataclass
class Flag:
    code: str
    severity: str   # "info" | "caution" | "warning"
    message: str


def review_flags(res: CaseResult, images: list[ImageResult], skipped: list[str] | None = None) -> list[Flag]:
    """Warnings for situations where the published rule is known to be unreliable."""
    f: list[Flag] = []
    t = res.threshold
    v = res.votes

    if res.n_images == 0:
        f.append(Flag("NO_IMAGES", "warning", "No readable images were found, so no analysis was possible."))
        return f

    if res.n_images <= t:
        f.append(Flag("TOO_FEW_IMAGES", "warning",
                      f"Only {res.n_images} image(s) were analysed. With the threshold set to {t}, "
                      f"the tool cannot report leukemia for {t} or fewer images, whatever they show."))
    elif res.n_images < MIN_RECOMMENDED_IMAGES:
        f.append(Flag("FEW_IMAGES", "caution",
                      f"Only {res.n_images} images were analysed. The published validation used about "
                      f"{VALIDATION_MEAN_IMAGES} images per patient; results with fewer fields are less reliable."))

    if abs(res.n_leukemic - (t + 1)) <= 1:
        f.append(Flag("BORDERLINE_COUNT", "caution",
                      f"{res.n_leukemic} image(s) were classed as leukemic, close to the decision point "
                      f"(more than {t}). One image more or less would change the result."))

    if res.three_class == "NO LEUKEMIA":
        m3_leuk = {k: n for k, n in v.items() if k != "NO LEUKEMIA"}
        top = max(m3_leuk, key=m3_leuk.get)
        if m3_leuk[top] > v["NO LEUKEMIA"] and m3_leuk[top] > t:
            f.append(Flag("MODEL_DISAGREEMENT", "warning",
                          f"The subtype model classed {m3_leuk[top]} of {res.n_images} images as {top}, "
                          f"although the screening model did not reach the leukemia threshold."))

    if res.three_class == "CHRONIC LEUKEMIA" and res.n_acute > res.n_chronic:
        f.append(Flag("ACUTE_CHRONIC_CONFLICT", "warning",
                      f"Reported as chronic because {res.n_chronic} images were classed chronic, although "
                      f"{res.n_acute} images were classed acute. The rule calls a case chronic once 2 images are chronic."))

    if res.three_class == "ACUTE LEUKEMIA":
        a, b = v["ALL"], v["AML"]
        if a == b:
            f.append(Flag("SUBTYPE_TIE", "warning",
                          f"ALL and AML received equal votes ({a} each). The subtype is undetermined; "
                          f"the original tool reports ALL in a tie."))
        else:
            f.append(Flag("ALL_AML_LIMITATION", "caution",
                          "ALL vs AML separation was weak in validation (F1 0.36 for ALL, 0.57 for AML). "
                          "Treat the subtype as unconfirmed."))
        if res.five_class == "AML" or a == b:
            f.append(Flag("APML_NOT_DETECTED", "info",
                          "APML was merged into AML during validation and is not reported separately. "
                          "If APML is suspected clinically, act on standard urgent pathways."))

    if res.three_class == "CHRONIC LEUKEMIA":
        c, m = v["CLL"], v["CML"]
        if c == m:
            f.append(Flag("SUBTYPE_TIE", "warning",
                          f"CLL and CML received equal votes ({c} each). The subtype is undetermined; "
                          f"the original tool reports CLL in a tie."))

    poor = [im.name for im in images if im.quality and not im.quality.get("ok", True)]
    if poor:
        f.append(Flag("IMAGE_QUALITY", "warning",
                      f"{len(poor)} image(s) do not look like 1000x MGG-stained smear fields and were still "
                      f"included, as in the original tool: {', '.join(poor[:5])}{' ...' if len(poor) > 5 else ''}"))

    if skipped:
        f.append(Flag("SKIPPED_FILES", "caution",
                      f"{len(skipped)} file(s) were not analysed: {', '.join(skipped[:5])}{' ...' if len(skipped) > 5 else ''}"))
    return f


def display_subtype(res: CaseResult, flags: list[Flag]) -> str:
    """Subtype wording for LeukoAI 1.1: ties are shown as undetermined instead of ALL/CLL."""
    if any(fl.code == "SUBTYPE_TIE" for fl in flags):
        return "ALL or AML (undetermined)" if res.three_class == "ACUTE LEUKEMIA" else "CLL or CML (undetermined)"
    return res.five_class