from __future__ import annotations

import logging
import math
from collections import defaultdict
from dataclasses import dataclass, field

from app.ml.feature_store import FeatureStore

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class FeatureModel:
    sample_count: int
    mean: float
    standard_deviation: float
    minimum: float
    maximum: float
    slope: float
    intercept: float

    def forecast(self, current_value: float | None = None) -> float:
        if current_value is not None:
            return current_value + self.slope
        return self.intercept + self.slope * self.sample_count

    @property
    def confidence(self) -> float:
        sample_score = self.sample_count / (self.sample_count + 2)
        scale = abs(self.mean) if abs(self.mean) > 1e-9 else 1.0
        stability_score = 1.0 / (1.0 + self.standard_deviation / scale)
        return min(0.99, max(0.05, sample_score * stability_score))


@dataclass(frozen=True)
class TrainedArtifact:
    artifact_id: str
    supported_sensors: tuple[str, ...]
    feature_models: dict[str, dict[str, FeatureModel]] = field(default_factory=dict)
    model_type: str = "statistical_baseline"


class TrainingPipeline:
    def __init__(self, store: FeatureStore) -> None:
        self._store = store
        self._artifacts: dict[str, TrainedArtifact] = {}

    def run(self, artifact_id: str) -> TrainedArtifact:
        vectors = self._store.all()
        if not vectors:
            raise ValueError("FeatureStore enthält keine Trainingsdaten.")

        samples: dict[str, dict[str, list[float]]] = defaultdict(lambda: defaultdict(list))
        for vector in vectors:
            for feature_name, raw_value in vector.values.items():
                value = float(raw_value)
                if math.isfinite(value):
                    samples[vector.sensor_id][feature_name].append(value)

        feature_models = {
            sensor_id: {
                feature_name: _fit_feature(values)
                for feature_name, values in sorted(features.items())
                if values
            }
            for sensor_id, features in sorted(samples.items())
        }
        feature_models = {
            sensor_id: models for sensor_id, models in feature_models.items() if models
        }
        if not feature_models:
            raise ValueError("Trainingsdaten enthalten keine endlichen numerischen Werte.")

        sensors = tuple(feature_models)
        artifact = TrainedArtifact(
            artifact_id=artifact_id,
            supported_sensors=sensors,
            feature_models=feature_models,
        )
        self._artifacts[artifact_id] = artifact
        logger.info("Training abgeschlossen für %s mit %d Sensoren", artifact_id, len(sensors))
        return artifact

    def export(self, artifact_id: str) -> TrainedArtifact:
        if artifact_id not in self._artifacts:
            raise KeyError(f"Artifact '{artifact_id}' nicht gefunden.")
        return self._artifacts[artifact_id]


def _fit_feature(values: list[float]) -> FeatureModel:
    sample_count = len(values)
    mean = sum(values) / sample_count
    variance = sum((value - mean) ** 2 for value in values) / sample_count
    standard_deviation = math.sqrt(variance)

    if sample_count == 1:
        slope = 0.0
        intercept = mean
    else:
        x_mean = (sample_count - 1) / 2
        denominator = sum((index - x_mean) ** 2 for index in range(sample_count))
        numerator = sum(
            (index - x_mean) * (value - mean) for index, value in enumerate(values)
        )
        slope = numerator / denominator
        intercept = mean - slope * x_mean

    return FeatureModel(
        sample_count=sample_count,
        mean=mean,
        standard_deviation=standard_deviation,
        minimum=min(values),
        maximum=max(values),
        slope=slope,
        intercept=intercept,
    )