"""Station-space IMPROVER-style post-processing with explicit local grid patches."""

from __future__ import annotations

import numpy as np
import torch
from torch import nn
from torch.nn import functional as F


VARIABLES = ("temperature", "dewpoint", "wind_speed")


class ImproverAIFS(nn.Module):
    """Calibrate three sources at 569 station-centred target-grid neighborhoods."""

    def __init__(self, config: dict):
        super().__init__()
        self.variables = tuple(config.get("variables", VARIABLES))
        self.lead_hours = int(config.get("lead_hours", 240))
        self.station_count = int(config.get("station_count", 569))
        self.fuzzy_widths = torch.tensor(config["fuzzy_widths"], dtype=torch.float32)
        self.recursive_coefficient = float(config.get("recursive_coefficient", 0.18))
        self.recursive_iterations = int(config.get("recursive_iterations", 1))
        self.register_buffer("knot_hours", torch.linspace(0, self.lead_hours, 11))
        self.blend_logits = nn.Parameter(torch.zeros(len(self.variables), 11, 3))
        self.register_buffer("bias", torch.zeros(3, self.lead_hours + 1, len(self.variables), self.station_count))
        max_thresholds = max(map(len, config["thresholds"]))
        knots = torch.zeros(3, len(self.variables), max_thresholds, 7, 2)
        knots[..., 0] = torch.linspace(0, 1, 7)
        knots[..., 1] = torch.linspace(0, 1, 7)
        self.register_buffer("calibration_knots", knots)
        self.register_buffer("threshold_counts", torch.tensor([len(v) for v in config["thresholds"]]))

    @torch.no_grad()
    def fit_bias_chunk(self, forecasts: torch.Tensor, analyses: torch.Tensor, elevation_delta: torch.Tensor, station_slice: slice) -> None:
        """Fit each source from all 30 preceding daily 1200 UTC histories."""
        adjusted = forecasts.clone()  # history, lead, variable, source, station
        adjusted[:, :, 0] += -0.0098 * elevation_delta.view(1, 1, 1, -1)
        self.bias[..., station_slice] = (adjusted - analyses.unsqueeze(3)).mean(0).permute(2, 0, 1, 3)

    def correct_expected(self, forecasts: torch.Tensor, elevation_delta: torch.Tensor, source: int, station_slice: slice) -> torch.Tensor:
        """Correct station-centre forecasts shaped [date/history, lead, variable, station]."""
        corrected = forecasts.clone()
        corrected[:, :, 0] += -0.0098 * elevation_delta.view(1, 1, -1)
        return corrected - self.bias[source, :, :, station_slice].unsqueeze(0)

    def fuzzy_threshold(self, expected: torch.Tensor, thresholds: list[torch.Tensor]) -> list[torch.Tensor]:
        outputs = []
        widths = self.fuzzy_widths.to(expected.device)
        for variable, values in enumerate(thresholds):
            forecast = expected[:, :, variable].unsqueeze(2)
            threshold = values.to(expected.device).view(1, 1, -1, 1)
            outputs.append(((forecast - threshold + widths[variable]) / (2 * widths[variable])).clamp(0, 1))
        return outputs

    @staticmethod
    def neighborhood(patch_probabilities: torch.Tensor) -> torch.Tensor:
        """Reduce pre-extracted real-semantic 3x3 target-grid patches to stations."""
        if patch_probabilities.shape[-2:] != (3, 3):
            raise ValueError("station neighborhood must be an authoritative 3x3 target-grid patch")
        return patch_probabilities.mean(dim=(-2, -1))

    def recursive_filter(self, patch_probabilities: torch.Tensor) -> torch.Tensor:
        """Apply separable filtering only within each station's extracted 3x3 patch."""
        output = patch_probabilities.clone()
        coefficient = self.recursive_coefficient
        for _ in range(self.recursive_iterations):
            for axis in (-2, -1):
                for index in range(1, 3):
                    current, previous = [slice(None)] * output.ndim, [slice(None)] * output.ndim
                    current[axis], previous[axis] = index, index - 1
                    output[tuple(current)] = (1 - coefficient) * output[tuple(current)] + coefficient * output[tuple(previous)]
                for index in range(1, -1, -1):
                    current, following = [slice(None)] * output.ndim, [slice(None)] * output.ndim
                    current[axis], following[axis] = index, index + 1
                    output[tuple(current)] = (1 - coefficient) * output[tuple(current)] + coefficient * output[tuple(following)]
        return output

    @torch.no_grad()
    def fit_reliability(self, probabilities: list[torch.Tensor], analyses: torch.Tensor, thresholds: list[torch.Tensor], source: int) -> None:
        """Fit seven-bin mappings on configured station samples with full time axes."""
        defaults = torch.linspace(0, 1, 7, device=analyses.device)
        for variable, probability in enumerate(probabilities):
            truth = analyses[:, :, variable].unsqueeze(2) > thresholds[variable].view(1, 1, -1, 1)
            count = probability.shape[2]
            p = probability.permute(2, 0, 1, 3).reshape(count, -1)
            y = truth.permute(2, 0, 1, 3).reshape(count, -1).float()
            ids = torch.bucketize(p, torch.linspace(1 / 7, 6 / 7, 6, device=p.device))
            xs, ys = [], []
            for bin_index in range(7):
                mask = ids == bin_index
                samples = mask.sum(1)
                denominator = samples.clamp_min(1)
                xs.append(torch.where(samples > 0, (p * mask).sum(1) / denominator, defaults[bin_index]))
                ys.append(torch.where(samples > 0, (y * mask).sum(1) / denominator, defaults[bin_index]))
            xs = torch.stack(xs, dim=1).cummax(1).values
            ys = torch.stack(ys, dim=1).cummax(1).values.clamp(0, 1)
            self.calibration_knots[source, variable, :count, :, 0] = xs
            self.calibration_knots[source, variable, :count, :, 1] = ys

    def calibrate(self, probabilities: list[torch.Tensor], source: int) -> list[torch.Tensor]:
        outputs = []
        for variable, probability in enumerate(probabilities):
            count = probability.shape[2]
            knots = self.calibration_knots[source, variable, :count]
            x, y = knots[..., 0].contiguous(), knots[..., 1].contiguous()
            values = probability.permute(2, 0, 1, 3).reshape(count, -1).contiguous()
            ids = torch.searchsorted(x, values).clamp(1, 6)
            x0, x1 = x.gather(1, ids - 1), x.gather(1, ids)
            y0, y1 = y.gather(1, ids - 1), y.gather(1, ids)
            calibrated = y0 + (values - x0) * (y1 - y0) / (x1 - x0).clamp_min(1e-6)
            calibrated = calibrated.reshape(count, probability.shape[0], probability.shape[1], probability.shape[3]).permute(1, 2, 0, 3)
            outputs.append(calibrated.clamp(0, 1).cummin(dim=2).values)
        return outputs

    def blend_weights(self) -> torch.Tensor:
        """Interpolate the paper's 11 valid-time knots to 241 hourly leads."""
        return F.interpolate(self.blend_logits.softmax(-1).permute(0, 2, 1), size=241, mode="linear", align_corners=True).permute(0, 2, 1)

    def blend_expected(self, model_expected: torch.Tensor) -> torch.Tensor:
        weights = self.blend_weights().permute(1, 0, 2).view(1, 241, 3, 3, 1)
        return (model_expected * weights).sum(3)

    def blend_probabilities(self, source_probabilities: list[list[torch.Tensor]]) -> list[torch.Tensor]:
        weights = self.blend_weights()
        return [sum(source_probabilities[s][v] * weights[v, :, s].view(1, 241, 1, 1) for s in range(3)).clamp(0, 1)
                for v in range(3)]


def crps_from_thresholds(probabilities: np.ndarray, thresholds: np.ndarray, truth: np.ndarray) -> float:
    total = np.zeros_like(truth, dtype=np.float32)
    previous = (1.0 - probabilities[:, :, 0] - (thresholds[0] >= truth)) ** 2
    for index in range(1, len(thresholds)):
        current = (1.0 - probabilities[:, :, index] - (thresholds[index] >= truth)) ** 2
        total += 0.5 * (previous + current) * (thresholds[index] - thresholds[index - 1])
        previous = current
    return float(total.mean())