from __future__ import annotations

from pathlib import Path
from typing import Any, Mapping

import numpy as np

from .features import extract_patch_features, feature_vector
from .io import load_capture, write_json
from .model import LogisticModel, model_sha256
from .quality import quality_check_capture


ALLOWED_OUTPUT_STATUSES = (
    "REJECTED_QC",
    "REJECTED_OOD",
    "NO_FLAG",
    "REVIEW_REQUIRED",
)


def infer_features(
    model: LogisticModel,
    features: Mapping[str, float],
    *,
    qc_receipt: Mapping[str, Any],
    sample_id: str,
) -> dict[str, Any]:
    base = {
        "sampleId": sample_id,
        "modelId": model.modelId,
        "modelSha256": model_sha256(model),
        "promotionStatus": model.promotionStatus,
        "automationAuthorized": False,
        "diseaseConfirmed": False,
        "treatmentRecommendationPresent": False,
    }
    if qc_receipt.get("passed") is not True:
        return {
            **base,
            "status": "REJECTED_QC",
            "abstained": True,
            "reason": "capture_quality_gate_failed",
            "qcFailures": list(qc_receipt.get("failures", [])),
            "probability": None,
        }
    vector = feature_vector(features)[None, :]
    ood = model.ood_metrics(vector)
    fraction = float(ood["fractionBeyondZLimit"][0])
    distance = float(ood["standardisedRmsDistance"][0])
    if bool(ood["rejected"][0]):
        return {
            **base,
            "status": "REJECTED_OOD",
            "abstained": True,
            "reason": "feature_distribution_outside_training_envelope",
            "probability": None,
            "ood": {
                "fractionBeyondZLimit": round(fraction, 6),
                "standardisedRmsDistance": round(distance, 6),
            },
        }
    probability = float(model.predict_proba(vector)[0, 1])
    if probability <= model.lowThreshold:
        status = "NO_FLAG"
        abstained = False
        reason = "probability_at_or_below_no_flag_threshold"
    elif probability >= model.highThreshold:
        status = "REVIEW_REQUIRED"
        abstained = False
        reason = "probability_at_or_above_review_threshold"
    else:
        status = "REVIEW_REQUIRED"
        abstained = True
        reason = "probability_between_thresholds_requires_human_review"
    assert status in ALLOWED_OUTPUT_STATUSES
    return {
        **base,
        "status": status,
        "abstained": abstained,
        "reason": reason,
        "probability": round(probability, 8),
        "thresholds": {"low": model.lowThreshold, "high": model.highThreshold},
        "ood": {
            "fractionBeyondZLimit": round(fraction, 6),
            "standardisedRmsDistance": round(distance, 6),
        },
    }


def infer_capture(
    model: LogisticModel,
    manifest_path: str | Path,
    *,
    patch_size: int = 32,
    output_path: str | Path | None = None,
) -> dict[str, Any]:
    if patch_size < 8:
        raise ValueError("patch_size must be at least eight pixels")
    rgb, bands, metadata = load_capture(manifest_path)
    qc = quality_check_capture(rgb, bands, metadata)
    receipt: dict[str, Any] = {
        "schemaVersion": 1,
        "captureId": metadata.get("captureId"),
        "dataProvenance": metadata.get("dataProvenance"),
        "modelId": model.modelId,
        "modelSha256": model_sha256(model),
        "promotionStatus": model.promotionStatus,
        "qc": qc,
        "patchSizePx": patch_size,
        "automationAuthorized": False,
        "diseaseConfirmed": False,
        "treatmentRecommendationPresent": False,
        "fieldPerformanceClaimed": False,
        "warning": "Optical research screening only. Human ground inspection and an appropriate reference standard remain required.",
    }
    if not qc["passed"]:
        receipt.update({"status": "REJECTED_QC", "patches": [], "reviewPatchCount": 0, "abstainedPatchCount": 0})
    else:
        rows, columns = rgb.shape[:2]
        patches: list[dict[str, Any]] = []
        for top in range(0, rows - patch_size + 1, patch_size):
            for left in range(0, columns - patch_size + 1, patch_size):
                slices = (slice(top, top + patch_size), slice(left, left + patch_size))
                features = extract_patch_features(
                    rgb[slices[0], slices[1], :],
                    {name: value[slices] for name, value in bands.items()},
                )
                patch_id = f"{metadata.get('captureId', 'capture')}-R{top:04d}-C{left:04d}"
                prediction = infer_features(model, features, qc_receipt=qc, sample_id=patch_id)
                prediction["window"] = {"top": top, "left": left, "height": patch_size, "width": patch_size}
                patches.append(prediction)
        statuses = {patch["status"] for patch in patches}
        if "REJECTED_OOD" in statuses:
            status = "REJECTED_OOD"
        elif "REVIEW_REQUIRED" in statuses:
            status = "REVIEW_REQUIRED"
        else:
            status = "NO_FLAG"
        receipt.update({
            "status": status,
            "patches": patches,
            "reviewPatchCount": sum(patch["status"] == "REVIEW_REQUIRED" for patch in patches),
            "abstainedPatchCount": sum(bool(patch["abstained"]) for patch in patches),
        })
    if receipt["status"] not in ALLOWED_OUTPUT_STATUSES:
        raise AssertionError("Inference produced an unpublished output status")
    if output_path is not None:
        write_json(output_path, receipt)
    return receipt
