from __future__ import annotations

from typing import Any, Mapping

import numpy as np

from .io import REQUIRED_BANDS


REQUIRED_METADATA = (
    "schemaVersion",
    "dataProvenance",
    "captureId",
    "fieldId",
    "flightId",
    "captureDate",
    "sensorModel",
    "scaleFactor",
    "processingLevel",
    "measurementScale",
    "calibrationMethod",
    "crs",
    "gridTransform",
    "rgbToMultispectralRegistered",
    "coRegistrationRmsePx",
    "bandMapValidated",
)


def quality_check_capture(
    rgb: np.ndarray,
    bands: Mapping[str, np.ndarray],
    metadata: Mapping[str, Any],
    *,
    maximum_coregistration_rmse_px: float = 1.0,
) -> dict[str, Any]:
    failures: list[str] = []
    rgb_value = np.asarray(rgb, dtype=np.float32)
    arrays = {name: np.asarray(value, dtype=np.float32) for name, value in bands.items()}
    for key in REQUIRED_METADATA:
        if key not in metadata:
            failures.append(f"metadata:{key}")
    if set(arrays) != set(REQUIRED_BANDS):
        failures.append("required_band_set")
    shapes = {value.shape for value in arrays.values()}
    if rgb_value.ndim != 3 or rgb_value.shape[-1:] != (3,):
        failures.append("rgb_shape")
    if len(shapes) != 1 or not shapes or len(next(iter(shapes))) != 2:
        failures.append("multispectral_alignment")
    elif rgb_value.ndim == 3 and rgb_value.shape[:2] != next(iter(shapes)):
        failures.append("rgb_multispectral_alignment")

    all_arrays = [rgb_value, *arrays.values()]
    finite_fraction = min((float(np.isfinite(value).mean()) for value in all_arrays), default=0.0)
    out_of_range_fraction = max((float(((value < 0) | (value > 1)).mean()) for value in all_arrays), default=1.0)
    saturation_fraction = max((float((value >= 0.98).mean()) for value in all_arrays), default=1.0)
    brightness = rgb_value.mean(axis=2) if rgb_value.ndim == 3 and rgb_value.shape[-1:] == (3,) else np.zeros((1, 1))
    shadow_fraction = float((brightness < 0.08).mean())
    horizontal = np.abs(np.diff(brightness, axis=1)).mean() if brightness.shape[1] > 1 else 0.0
    vertical = np.abs(np.diff(brightness, axis=0)).mean() if brightness.shape[0] > 1 else 0.0
    edge_energy = float(horizontal + vertical)
    if finite_fraction < 0.995:
        failures.append("finite_fraction")
    if out_of_range_fraction > 0.001:
        failures.append("reflectance_or_rgb_range")
    if saturation_fraction > 0.05:
        failures.append("saturation_fraction")
    if shadow_fraction > 0.60:
        failures.append("shadow_fraction")
    if edge_energy < 0.001:
        failures.append("rgb_blur_or_flat_scene")
    if metadata.get("rgbToMultispectralRegistered") is not True:
        failures.append("rgb_multispectral_registration")
    if metadata.get("bandMapValidated") is not True:
        failures.append("band_map_validation")
    if metadata.get("processingLevel") != "orthomosaic_reflectance_rgb_coregistered":
        failures.append("processing_level")
    if metadata.get("measurementScale") not in {"relative_reflectance", "surface_reflectance"}:
        failures.append("measurement_scale")
    if metadata.get("dataProvenance") not in {"synthetic", "field"}:
        failures.append("data_provenance")
    try:
        rmse = float(metadata.get("coRegistrationRmsePx"))
        if not np.isfinite(rmse) or rmse < 0 or rmse > maximum_coregistration_rmse_px:
            failures.append("co_registration_rmse")
    except (TypeError, ValueError):
        failures.append("co_registration_rmse")
        rmse = float("nan")
    transform = metadata.get("gridTransform")
    if not isinstance(transform, list) or len(transform) != 6 or not all(isinstance(value, (int, float)) for value in transform):
        failures.append("grid_transform")

    return {
        "schemaVersion": 1,
        "passed": not failures,
        "failures": sorted(set(failures)),
        "metrics": {
            "finiteFraction": round(finite_fraction, 6),
            "outOfRangeFraction": round(out_of_range_fraction, 6),
            "saturationFraction": round(saturation_fraction, 6),
            "shadowFraction": round(shadow_fraction, 6),
            "rgbEdgeEnergy": round(edge_energy, 6),
            "coRegistrationRmsePx": round(rmse, 6) if np.isfinite(rmse) else None,
            "shape": list(rgb_value.shape[:2]) if rgb_value.ndim == 3 else None,
        },
        "interpretation": "QC pass permits research-model inference only; it is not plant-disease confirmation.",
    }
