from __future__ import annotations

import json
from collections.abc import Iterable
from dataclasses import dataclass
from pathlib import Path

import numpy as np

WINDOW_SIZE = 50
FEATURE_NAMES = (
    "force_slope_kn_per_cycle",
    "force_shift_kn",
    "force_std_kn",
    "max_deviation_mm",
    "last_deviation_mm",
    "force_range_kn",
)

DEFAULT_CENTER = np.asarray([0.05, 2.0, 0.8, 0.12, 0.12, 3.0], dtype=np.float32)
DEFAULT_SCALE = np.asarray([0.05, 2.0, 0.8, 0.08, 0.08, 3.0], dtype=np.float32)
DEFAULT_WEIGHT = np.asarray([2.8, 1.6, 0.5, 0.8, 0.6, 0.4], dtype=np.float32)
DEFAULT_BIAS = np.float32(-0.4)


def as_telemetry_array(values: Iterable) -> np.ndarray:
    array = np.asarray(values, dtype=np.float32)
    if array.ndim == 2:
        array = array[None, ...]
    if array.ndim != 3 or array.shape[1:] != (WINDOW_SIZE, 2):
        raise ValueError(f"telemetry must have shape [batch,{WINDOW_SIZE},2]")
    if not np.isfinite(array).all():
        raise ValueError("telemetry contains non-finite values")
    return array


def extract_features_numpy(values: Iterable) -> np.ndarray:
    telemetry = as_telemetry_array(values)
    force = telemetry[:, :, 0]
    deviation = telemetry[:, :, 1]
    x = np.arange(WINDOW_SIZE, dtype=np.float32)
    x_centered = x - x.mean()
    slope = (force * x_centered).sum(axis=1) / np.square(x_centered).sum()
    shift = force[:, -10:].mean(axis=1) - force[:, :10].mean(axis=1)
    std = force.std(axis=1)
    max_deviation = deviation.max(axis=1)
    last_deviation = deviation[:, -1]
    force_range = force.max(axis=1) - force.min(axis=1)
    return np.stack(
        [slope, shift, std, max_deviation, last_deviation, force_range],
        axis=1,
    ).astype(np.float32)


@dataclass(frozen=True)
class DriftPrediction:
    probability: float
    drift_detected: bool
    features: dict[str, float]
    runtime: str
    model_version: str

    def as_dict(self) -> dict:
        return {
            "probability": self.probability,
            "drift_detected": self.drift_detected,
            "features": self.features,
            "runtime": self.runtime,
            "model_version": self.model_version,
        }


class NumpyDriftRuntime:
    """Dependency-light reference runtime used by the CPU-only Space."""

    def __init__(self, weights_path: str | Path | None = None) -> None:
        if weights_path is None:
            weights_path = (
                Path(__file__).resolve().parents[3] / "models" / "artifacts" / "weights.json"
            )
        path = Path(weights_path)
        if path.exists():
            payload = json.loads(path.read_text())
            self.center = np.asarray(payload["feature_center"], dtype=np.float32)
            self.scale = np.asarray(payload["feature_scale"], dtype=np.float32)
            self.weight = np.asarray(payload["linear_weight"], dtype=np.float32)
            self.bias = np.float32(payload["linear_bias"])
            self.model_version = payload["model_version"]
        else:
            self.center = DEFAULT_CENTER
            self.scale = DEFAULT_SCALE
            self.weight = DEFAULT_WEIGHT
            self.bias = DEFAULT_BIAS
            self.model_version = "bootstrap-untrained"

    def predict_batch(self, values: Iterable) -> tuple[np.ndarray, np.ndarray]:
        features = extract_features_numpy(values)
        normalized = (features - self.center) / self.scale
        logits = normalized @ self.weight + self.bias
        probability = 1.0 / (1.0 + np.exp(-logits))
        return probability.astype(np.float32), features

    def predict(self, values: Iterable, threshold: float = 0.5) -> DriftPrediction:
        probabilities, features = self.predict_batch(values)
        if len(probabilities) != 1:
            raise ValueError("predict expects exactly one telemetry window")
        probability = float(probabilities[0])
        return DriftPrediction(
            probability=round(probability, 6),
            drift_detected=probability >= threshold,
            features={
                name: round(float(value), 6)
                for name, value in zip(FEATURE_NAMES, features[0], strict=True)
            },
            runtime="numpy-reference",
            model_version=self.model_version,
        )