from __future__ import annotations

from typing import Mapping

import numpy as np

from .io import REQUIRED_BANDS


FEATURE_NAMES = (
    "rgb_r_mean",
    "rgb_g_mean",
    "rgb_b_mean",
    "rgb_exg_mean",
    "rgb_vari_mean",
    "rgb_brightness_std",
    "rgb_edge_mean",
    "ms_green_mean",
    "ms_red_mean",
    "ms_red_edge_mean",
    "ms_nir_mean",
    "ndvi_mean",
    "ndre_mean",
    "gndvi_mean",
    "ndvi_std",
    "ndre_std",
    "shadow_fraction",
    "saturation_fraction",
)


def normalised_difference(high: np.ndarray, low: np.ndarray) -> np.ndarray:
    high_value = np.asarray(high, dtype=np.float32)
    low_value = np.asarray(low, dtype=np.float32)
    denominator = high_value + low_value
    return np.divide(
        high_value - low_value,
        denominator,
        out=np.full_like(high_value, np.nan),
        where=np.abs(denominator) > 1e-8,
    )


def compute_indices(bands: Mapping[str, np.ndarray]) -> dict[str, np.ndarray]:
    if set(bands) != set(REQUIRED_BANDS):
        raise ValueError(f"Multispectral bands must be exactly {REQUIRED_BANDS}")
    arrays = {name: np.asarray(value, dtype=np.float32) for name, value in bands.items()}
    shapes = {value.shape for value in arrays.values()}
    if len(shapes) != 1 or len(next(iter(shapes))) != 2:
        raise ValueError("Multispectral bands must be aligned two-dimensional arrays")
    return {
        "ndvi": normalised_difference(arrays["nir"], arrays["red"]),
        "ndre": normalised_difference(arrays["nir"], arrays["red_edge"]),
        "gndvi": normalised_difference(arrays["nir"], arrays["green"]),
    }


def _safe_mean(value: np.ndarray, valid: np.ndarray) -> float:
    selected = np.asarray(value, dtype=np.float64)[valid]
    return float(selected.mean()) if selected.size else float("nan")


def _safe_std(value: np.ndarray, valid: np.ndarray) -> float:
    selected = np.asarray(value, dtype=np.float64)[valid]
    return float(selected.std()) if selected.size else float("nan")


def _edge_magnitude(brightness: np.ndarray) -> np.ndarray:
    horizontal = np.abs(np.diff(brightness, axis=1, prepend=brightness[:, :1]))
    vertical = np.abs(np.diff(brightness, axis=0, prepend=brightness[:1, :]))
    return np.sqrt(horizontal * horizontal + vertical * vertical)


def extract_patch_features(rgb: np.ndarray, bands: Mapping[str, np.ndarray]) -> dict[str, float]:
    rgb_value = np.asarray(rgb, dtype=np.float32)
    if rgb_value.ndim != 3 or rgb_value.shape[2] != 3:
        raise ValueError("RGB input must have shape (rows, columns, 3)")
    arrays = {name: np.asarray(value, dtype=np.float32) for name, value in bands.items()}
    indices = compute_indices(arrays)
    if any(value.shape != rgb_value.shape[:2] for value in arrays.values()):
        raise ValueError("RGB and multispectral inputs must share one registered grid")

    red_rgb, green_rgb, blue_rgb = (rgb_value[:, :, index] for index in range(3))
    brightness = rgb_value.mean(axis=2)
    exg = 2.0 * green_rgb - red_rgb - blue_rgb
    vari_denominator = green_rgb + red_rgb - blue_rgb
    vari = np.divide(
        green_rgb - red_rgb,
        vari_denominator,
        out=np.full_like(green_rgb, np.nan),
        where=np.abs(vari_denominator) > 1e-8,
    )
    edge = _edge_magnitude(brightness)
    valid = np.isfinite(rgb_value).all(axis=2)
    valid &= np.logical_and.reduce([np.isfinite(value) for value in arrays.values()])
    valid &= np.logical_and.reduce([np.isfinite(value) for value in indices.values()])
    if int(valid.sum()) < 16:
        raise ValueError("At least 16 valid co-registered pixels are required per patch")
    shadow = brightness < 0.08
    saturated = (rgb_value.max(axis=2) >= 0.995) | np.logical_or.reduce([value >= 0.98 for value in arrays.values()])

    values = {
        "rgb_r_mean": _safe_mean(red_rgb, valid),
        "rgb_g_mean": _safe_mean(green_rgb, valid),
        "rgb_b_mean": _safe_mean(blue_rgb, valid),
        "rgb_exg_mean": _safe_mean(exg, valid),
        "rgb_vari_mean": _safe_mean(vari, valid),
        "rgb_brightness_std": _safe_std(brightness, valid),
        "rgb_edge_mean": _safe_mean(edge, valid),
        "ms_green_mean": _safe_mean(arrays["green"], valid),
        "ms_red_mean": _safe_mean(arrays["red"], valid),
        "ms_red_edge_mean": _safe_mean(arrays["red_edge"], valid),
        "ms_nir_mean": _safe_mean(arrays["nir"], valid),
        "ndvi_mean": _safe_mean(indices["ndvi"], valid),
        "ndre_mean": _safe_mean(indices["ndre"], valid),
        "gndvi_mean": _safe_mean(indices["gndvi"], valid),
        "ndvi_std": _safe_std(indices["ndvi"], valid),
        "ndre_std": _safe_std(indices["ndre"], valid),
        "shadow_fraction": float(shadow[valid].mean()),
        "saturation_fraction": float(saturated[valid].mean()),
    }
    if tuple(values) != FEATURE_NAMES or not np.isfinite(list(values.values())).all():
        raise RuntimeError("Feature extraction produced an invalid or reordered feature contract")
    return {name: round(float(value), 8) for name, value in values.items()}


def feature_vector(features: Mapping[str, float]) -> np.ndarray:
    if set(features) != set(FEATURE_NAMES):
        raise ValueError("Feature row differs from the published feature contract")
    value = np.asarray([features[name] for name in FEATURE_NAMES], dtype=np.float64)
    if value.shape != (len(FEATURE_NAMES),) or not np.isfinite(value).all():
        raise ValueError("Feature vector contains non-finite values")
    return value
