from __future__ import annotations

import logging
import math
from dataclasses import dataclass
from typing import Sequence

from app.ml.explanation import FeatureExplanation, explain_feature
from app.ml.feature_store import FeatureVector
from app.ml.registry.model_registry import ModelRegistry
from app.ml.training import TrainedArtifact, TrainingPipeline

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class PredictionResult:
    artifact_id: str
    sensor_id: str
    predictions: dict[str, float]
    confidence: float
    model_type: str
    explanations: dict[str, FeatureExplanation]


class Predictor:
    def __init__(
        self,
        pipeline: TrainingPipeline | None = None,
        registry: ModelRegistry | None = None,
    ) -> None:
        if isinstance(pipeline, ModelRegistry) and registry is None:
            registry = pipeline
            pipeline = None
        if pipeline is None and registry is None:
            raise ValueError("Predictor erfordert TrainingPipeline oder ModelRegistry.")
        self._pipeline = pipeline
        self._registry = registry

    def predict(self, artifact_id: str, entity: FeatureVector) -> PredictionResult:
        artifact = self._get_artifact(artifact_id)
        if entity.sensor_id not in artifact.supported_sensors:
            raise ValueError(
                f"Sensor '{entity.sensor_id}' wird vom Modell '{artifact_id}' nicht unterstützt."
            )
        sensor_models = artifact.feature_models.get(entity.sensor_id, {})
        if not sensor_models:
            raise ValueError(f"Modell '{artifact_id}' enthält keine statistischen Parameter.")

        feature_names = sorted(set(sensor_models).intersection(entity.values))
        if not feature_names:
            raise ValueError(
                f"Keine Eingabemerkmale werden vom Modell '{artifact_id}' unterstützt."
            )

        predictions: dict[str, float] = {}
        explanations: dict[str, FeatureExplanation] = {}
        confidences: list[float] = []
        for feature_name in feature_names:
            model = sensor_models[feature_name]
            current_value = float(entity.values[feature_name])
            if not math.isfinite(current_value):
                raise ValueError("Vorhersagewerte müssen endlich sein.")
            predicted_value = model.forecast(current_value)
            predictions[feature_name] = predicted_value
            explanations[feature_name] = explain_feature(
                feature_name,
                current_value,
                predicted_value,
                model,
            )
            confidences.append(model.confidence)

        return PredictionResult(
            artifact_id=artifact_id,
            sensor_id=entity.sensor_id,
            predictions=predictions,
            confidence=sum(confidences) / len(confidences),
            model_type=artifact.model_type,
            explanations=explanations,
        )

    def predict_batch(
        self,
        artifact_id: str,
        entities: Sequence[FeatureVector],
    ) -> list[PredictionResult]:
        return [self.predict(artifact_id, entity) for entity in entities]

    @staticmethod
    def default_artifact(pipeline: TrainingPipeline) -> TrainedArtifact:
        artifacts = list(pipeline._artifacts)
        if not artifacts:
            raise ValueError("Kein trainiertes Modell gefunden.")
        return pipeline.export(artifacts[-1])

    def _get_artifact(self, artifact_id: str) -> TrainedArtifact:
        if self._registry is not None:
            return self._registry.load_artifact(artifact_id)
        if self._pipeline is not None:
            return self._pipeline.export(artifact_id)
        raise RuntimeError("Predictor nicht initialisiert.")