from __future__ import annotations

import uuid
from dataclasses import dataclass, field
from numbers import Real
from typing import Any, Iterable

from models import EvidenceRecord, MetricRecord, PhaseName, PhaseRecord, RunContext
from storage import LocalSessionStore, utc_now

_TRAIN_LOSS_KEYS = ("train/loss", "training/loss")
_VAL_LOSS_KEYS = ("val/loss", "validation/loss")
_TRAIN_ACCURACY_KEYS = ("train/accuracy", "training/accuracy", "train/acc")
_VAL_ACCURACY_KEYS = ("val/accuracy", "validation/accuracy", "val/acc")
_EPSILON = 1e-12
_PHASE_DISPLAY_NAMES: dict[PhaseName, str] = {
    "initialization_adaptation": "Initialization Adaptation",
    "rapid_learning": "Rapid Learning",
    "stabilization": "Stabilization",
    "stagnation": "Stagnation",
    "overfitting_onset": "Overfitting Onset",
    "divergence": "Divergence",
}
_PHASE_SUMMARIES: dict[PhaseName, str] = {
    "initialization_adaptation": (
        "Early optimization steps are still settling before a durable learning regime forms."
    ),
    "rapid_learning": "The primary training signal is improving quickly from one step to the next.",
    "stabilization": "Training is still improving, but the step-to-step gains have moderated.",
    "stagnation": "Recent training progress has flattened into a near-plateau.",
    "overfitting_onset": (
        "Training continues to improve while validation quality regresses and the gap widens."
    ),
    "divergence": "The optimization trajectory has reversed sharply away from its best recent state.",
}


def _normalize_metric_keys(values: Iterable[str], *, label: str) -> tuple[str, ...]:
    normalized_values: list[str] = []
    for raw_value in values:
        if not isinstance(raw_value, str):
            raise TypeError(f"{label} values must be strings.")
        candidate = raw_value.strip()
        if not candidate:
            raise ValueError(f"{label} values must not be empty.")
        if candidate not in normalized_values:
            normalized_values.append(candidate)
    if not normalized_values:
        raise ValueError(f"{label} must include at least one metric key.")
    return tuple(normalized_values)


def _validate_non_negative(name: str, value: Real) -> float:
    if isinstance(value, bool) or not isinstance(value, Real):
        raise TypeError(f"{name} must be a real number.")
    normalized_value = float(value)
    if normalized_value < 0.0:
        raise ValueError(f"{name} must not be negative.")
    return normalized_value


def _validate_positive_int(name: str, value: Any) -> int:
    if isinstance(value, bool) or not isinstance(value, int):
        raise TypeError(f"{name} must be an integer.")
    if value <= 0:
        raise ValueError(f"{name} must be a positive integer.")
    return int(value)


@dataclass(frozen=True, slots=True)
class PhaseAnalyzerThresholds:
    initialization_points: int = 2
    rapid_learning_min_relative_change: float = 0.18
    stagnation_max_relative_change: float = 0.01
    overfitting_min_gap: float = 0.12
    overfitting_min_train_relative_change: float = 0.015
    overfitting_min_validation_regression: float = 0.015
    divergence_min_relative_regression: float = 0.15
    divergence_from_best_ratio: float = 0.25

    def __post_init__(self) -> None:
        object.__setattr__(
            self,
            "initialization_points",
            _validate_positive_int("initialization_points", self.initialization_points),
        )
        for field_name in (
            "rapid_learning_min_relative_change",
            "stagnation_max_relative_change",
            "overfitting_min_gap",
            "overfitting_min_train_relative_change",
            "overfitting_min_validation_regression",
            "divergence_min_relative_regression",
            "divergence_from_best_ratio",
        ):
            object.__setattr__(
                self,
                field_name,
                _validate_non_negative(field_name, getattr(self, field_name)),
            )


@dataclass(frozen=True, slots=True)
class PhaseAnalyzerConfig:
    thresholds: PhaseAnalyzerThresholds = field(default_factory=PhaseAnalyzerThresholds)
    replace_existing_outputs: bool = True
    train_loss_metric_keys: tuple[str, ...] = field(default_factory=lambda: _TRAIN_LOSS_KEYS)
    val_loss_metric_keys: tuple[str, ...] = field(default_factory=lambda: _VAL_LOSS_KEYS)
    train_accuracy_metric_keys: tuple[str, ...] = field(
        default_factory=lambda: _TRAIN_ACCURACY_KEYS
    )
    val_accuracy_metric_keys: tuple[str, ...] = field(default_factory=lambda: _VAL_ACCURACY_KEYS)

    def __post_init__(self) -> None:
        if not isinstance(self.thresholds, PhaseAnalyzerThresholds):
            raise TypeError("thresholds must be a PhaseAnalyzerThresholds instance.")
        if not isinstance(self.replace_existing_outputs, bool):
            raise TypeError("replace_existing_outputs must be a boolean.")
        object.__setattr__(
            self,
            "train_loss_metric_keys",
            _normalize_metric_keys(self.train_loss_metric_keys, label="train_loss_metric_keys"),
        )
        object.__setattr__(
            self,
            "val_loss_metric_keys",
            _normalize_metric_keys(self.val_loss_metric_keys, label="val_loss_metric_keys"),
        )
        object.__setattr__(
            self,
            "train_accuracy_metric_keys",
            _normalize_metric_keys(
                self.train_accuracy_metric_keys,
                label="train_accuracy_metric_keys",
            ),
        )
        object.__setattr__(
            self,
            "val_accuracy_metric_keys",
            _normalize_metric_keys(
                self.val_accuracy_metric_keys,
                label="val_accuracy_metric_keys",
            ),
        )


@dataclass(frozen=True, slots=True)
class PhaseAnalysisResult:
    phases: tuple[PhaseRecord, ...]


@dataclass(frozen=True, slots=True)
class _MetricSeries:
    metric_key: str
    records: tuple[MetricRecord, ...]
    direction_multiplier: int

    def normalized(self, record: MetricRecord) -> float:
        return float(record.value) * float(self.direction_multiplier)


@dataclass(frozen=True, slots=True)
class _PhaseSignalPoint:
    train_record: MetricRecord
    normalized_train_value: float
    train_relative_change: float
    regression_from_best: float
    validation_record: MetricRecord | None
    normalized_validation_value: float | None
    validation_relative_change: float | None
    generalization_gap: float | None


@dataclass(frozen=True, slots=True)
class _PhaseSegment:
    phase_name: PhaseName
    start_index: int
    end_index: int


class PhaseAnalyzer:
    """Deterministic training-regime segmentation derived from persisted metrics."""

    def __init__(
        self,
        session: Any,
        *,
        config: PhaseAnalyzerConfig | None = None,
    ) -> None:
        self._context, self._store = self._resolve_session_dependencies(session)
        self._config = config or PhaseAnalyzerConfig()

    @property
    def config(self) -> PhaseAnalyzerConfig:
        return self._config

    @property
    def run_id(self) -> str:
        return self._context.run_id

    def analyze(self, *, replace_existing: bool | None = None) -> PhaseAnalysisResult:
        metrics = self._store.read_scalar_metrics(self.run_id)
        phases = self._build_phases(metrics)
        persisted_phases = self._store.write_phases(
            self.run_id,
            phases,
            replace_existing=(
                self._config.replace_existing_outputs
                if replace_existing is None
                else replace_existing
            ),
        )
        return PhaseAnalysisResult(phases=persisted_phases)

    def read_phases(
        self,
        *,
        phase_name: str | None = None,
        step: int | None = None,
        step_min: int | None = None,
        step_max: int | None = None,
        limit: int | None = None,
    ) -> tuple[PhaseRecord, ...]:
        return self._store.read_phases(
            self.run_id,
            phase_name=phase_name,
            step=step,
            step_min=step_min,
            step_max=step_max,
            limit=limit,
        )

    def _build_phases(self, metrics: tuple[MetricRecord, ...]) -> tuple[PhaseRecord, ...]:
        train_series, validation_series = self._select_metric_series(metrics)
        if train_series is None:
            return ()

        points = self._build_signal_points(train_series, validation_series)
        if not points:
            return ()

        phase_labels = tuple(
            self._classify_point(index, point) for index, point in enumerate(points)
        )
        segments = self._segment_points(phase_labels)
        return tuple(
            self._build_phase_record(
                sequence_id=sequence_id,
                segment=segment,
                segments=segments,
                points=points,
                train_series=train_series,
                validation_series=validation_series,
            )
            for sequence_id, segment in enumerate(segments)
        )

    def _select_metric_series(
        self,
        metrics: tuple[MetricRecord, ...],
    ) -> tuple[_MetricSeries | None, _MetricSeries | None]:
        train_loss = self._find_series(metrics, self._config.train_loss_metric_keys, direction=1)
        if train_loss is not None:
            return (
                train_loss,
                self._find_series(metrics, self._config.val_loss_metric_keys, direction=1),
            )

        train_accuracy = self._find_series(
            metrics,
            self._config.train_accuracy_metric_keys,
            direction=-1,
        )
        if train_accuracy is not None:
            return (
                train_accuracy,
                self._find_series(metrics, self._config.val_accuracy_metric_keys, direction=-1),
            )

        return None, None

    def _build_signal_points(
        self,
        train_series: _MetricSeries,
        validation_series: _MetricSeries | None,
    ) -> tuple[_PhaseSignalPoint, ...]:
        train_records = self._dedupe_records_by_step(train_series.records)
        if not train_records:
            return ()

        validation_records = (
            self._dedupe_records_by_step(validation_series.records)
            if validation_series is not None
            else ()
        )
        validation_index = -1
        best_train_value = train_series.normalized(train_records[0])
        points: list[_PhaseSignalPoint] = []

        for index, train_record in enumerate(train_records):
            train_value = train_series.normalized(train_record)
            if index == 0:
                train_relative_change = 0.0
            else:
                train_relative_change = self._relative_change(
                    train_series.normalized(train_records[index - 1]),
                    train_value,
                )

            regression_from_best = self._relative_regression(best_train_value, train_value)
            best_train_value = min(best_train_value, train_value)

            while (
                validation_index + 1 < len(validation_records)
                and validation_records[validation_index + 1].step is not None
                and validation_records[validation_index + 1].step <= train_record.step
            ):
                validation_index += 1

            validation_record = (
                validation_records[validation_index] if validation_index >= 0 else None
            )
            previous_validation = (
                validation_records[validation_index - 1] if validation_index >= 1 else None
            )
            normalized_validation = (
                validation_series.normalized(validation_record)
                if validation_series is not None and validation_record is not None
                else None
            )
            validation_relative_change = (
                self._relative_change(
                    validation_series.normalized(previous_validation),
                    normalized_validation,
                )
                if validation_series is not None
                and validation_record is not None
                and previous_validation is not None
                and normalized_validation is not None
                else None
            )
            generalization_gap = (
                normalized_validation - train_value
                if normalized_validation is not None
                else None
            )

            points.append(
                _PhaseSignalPoint(
                    train_record=train_record,
                    normalized_train_value=train_value,
                    train_relative_change=train_relative_change,
                    regression_from_best=regression_from_best,
                    validation_record=validation_record,
                    normalized_validation_value=normalized_validation,
                    validation_relative_change=validation_relative_change,
                    generalization_gap=generalization_gap,
                )
            )

        return tuple(points)

    def _classify_point(self, index: int, point: _PhaseSignalPoint) -> PhaseName:
        thresholds = self._config.thresholds
        if index < thresholds.initialization_points:
            return "initialization_adaptation"

        if (
            point.train_relative_change <= -thresholds.divergence_min_relative_regression
            and point.regression_from_best >= thresholds.divergence_from_best_ratio
        ):
            return "divergence"

        if (
            point.generalization_gap is not None
            and point.validation_relative_change is not None
            and point.generalization_gap >= thresholds.overfitting_min_gap
            and point.train_relative_change >= thresholds.overfitting_min_train_relative_change
            and point.validation_relative_change <= -thresholds.overfitting_min_validation_regression
        ):
            return "overfitting_onset"

        if point.train_relative_change >= thresholds.rapid_learning_min_relative_change:
            return "rapid_learning"

        if abs(point.train_relative_change) <= thresholds.stagnation_max_relative_change:
            return "stagnation"

        return "stabilization"

    @staticmethod
    def _segment_points(phase_labels: tuple[PhaseName, ...]) -> tuple[_PhaseSegment, ...]:
        if not phase_labels:
            return ()

        segments: list[_PhaseSegment] = []
        current_label = phase_labels[0]
        start_index = 0
        for index, phase_name in enumerate(phase_labels[1:], start=1):
            if phase_name == current_label:
                continue
            segments.append(
                _PhaseSegment(
                    phase_name=current_label,
                    start_index=start_index,
                    end_index=index - 1,
                )
            )
            current_label = phase_name
            start_index = index

        segments.append(
            _PhaseSegment(
                phase_name=current_label,
                start_index=start_index,
                end_index=len(phase_labels) - 1,
            )
        )
        return tuple(segments)

    def _build_phase_record(
        self,
        *,
        sequence_id: int,
        segment: _PhaseSegment,
        segments: tuple[_PhaseSegment, ...],
        points: tuple[_PhaseSignalPoint, ...],
        train_series: _MetricSeries,
        validation_series: _MetricSeries | None,
    ) -> PhaseRecord:
        start_point = points[segment.start_index]
        end_point = points[segment.end_index]
        previous_phase = segments[sequence_id - 1].phase_name if sequence_id > 0 else None
        evidence = self._build_phase_evidence(
            phase_name=segment.phase_name,
            previous_phase=previous_phase,
            start_point=start_point,
            end_point=end_point,
            train_series=train_series,
            validation_series=validation_series,
        )
        return PhaseRecord(
            id=uuid.uuid4().hex,
            run_id=self.run_id,
            sequence_id=sequence_id,
            phase_name=segment.phase_name,
            display_name=_PHASE_DISPLAY_NAMES[segment.phase_name],
            summary=_PHASE_SUMMARIES[segment.phase_name],
            step_start=start_point.train_record.step if start_point.train_record.step is not None else 0,
            step_end=end_point.train_record.step if end_point.train_record.step is not None else 0,
            timestamp=utc_now(),
            evidence=evidence,
            metadata={
                "entered_from_phase": previous_phase,
                "primary_metric_key": train_series.metric_key,
                "validation_metric_key": (
                    validation_series.metric_key if validation_series is not None else None
                ),
                "point_count": segment.end_index - segment.start_index + 1,
                "train_relative_change_at_boundary": round(start_point.train_relative_change, 6),
                "validation_relative_change_at_boundary": (
                    round(start_point.validation_relative_change, 6)
                    if start_point.validation_relative_change is not None
                    else None
                ),
                "generalization_gap_at_boundary": (
                    round(start_point.generalization_gap, 6)
                    if start_point.generalization_gap is not None
                    else None
                ),
            },
        )

    def _build_phase_evidence(
        self,
        *,
        phase_name: PhaseName,
        previous_phase: PhaseName | None,
        start_point: _PhaseSignalPoint,
        end_point: _PhaseSignalPoint,
        train_series: _MetricSeries,
        validation_series: _MetricSeries | None,
    ) -> tuple[EvidenceRecord, ...]:
        thresholds = self._phase_threshold_payload(phase_name)
        train_segment_improvement = (
            start_point.normalized_train_value - end_point.normalized_train_value
        )
        evidence_items = [
            EvidenceRecord(
                kind="transition",
                source="phase_boundary",
                summary=(
                    "Initial phase selected from the earliest training points."
                    if previous_phase is None
                    else (
                        f"Transition from {_PHASE_DISPLAY_NAMES[previous_phase]} to "
                        f"{_PHASE_DISPLAY_NAMES[phase_name]} based on deterministic metric thresholds."
                    )
                ),
                observed_value={
                    "train_relative_change": round(start_point.train_relative_change, 6),
                    "validation_relative_change": (
                        round(start_point.validation_relative_change, 6)
                        if start_point.validation_relative_change is not None
                        else None
                    ),
                    "generalization_gap": (
                        round(start_point.generalization_gap, 6)
                        if start_point.generalization_gap is not None
                        else None
                    ),
                    "regression_from_best": round(start_point.regression_from_best, 6),
                },
                threshold_value=thresholds,
                step=start_point.train_record.step,
                step_start=start_point.train_record.step,
                step_end=start_point.train_record.step,
            ),
            EvidenceRecord(
                kind="metric",
                source=train_series.metric_key,
                summary="Primary training metric trend across the phase segment.",
                observed_value={
                    "start": start_point.train_record.value,
                    "end": end_point.train_record.value,
                    "normalized_improvement": round(train_segment_improvement, 6),
                },
                step_start=start_point.train_record.step,
                step_end=end_point.train_record.step,
                metadata={"direction_multiplier": train_series.direction_multiplier},
            ),
        ]
        if validation_series is not None and end_point.validation_record is not None:
            evidence_items.append(
                EvidenceRecord(
                    kind="metric",
                    source=validation_series.metric_key,
                    summary="Most recent validation behavior aligned to this phase.",
                    observed_value={
                        "value": end_point.validation_record.value,
                        "relative_change": (
                            round(end_point.validation_relative_change, 6)
                            if end_point.validation_relative_change is not None
                            else None
                        ),
                        "generalization_gap": (
                            round(end_point.generalization_gap, 6)
                            if end_point.generalization_gap is not None
                            else None
                        ),
                    },
                    step=end_point.validation_record.step,
                    step_start=start_point.train_record.step,
                    step_end=end_point.train_record.step,
                    metadata={"direction_multiplier": validation_series.direction_multiplier},
                )
            )
        return tuple(evidence_items)

    def _phase_threshold_payload(self, phase_name: PhaseName) -> dict[str, float | int]:
        thresholds = self._config.thresholds
        if phase_name == "initialization_adaptation":
            return {"initialization_points": thresholds.initialization_points}
        if phase_name == "rapid_learning":
            return {
                "train_relative_change_min": thresholds.rapid_learning_min_relative_change,
            }
        if phase_name == "stagnation":
            return {
                "train_relative_change_max": thresholds.stagnation_max_relative_change,
            }
        if phase_name == "overfitting_onset":
            return {
                "generalization_gap_min": thresholds.overfitting_min_gap,
                "train_relative_change_min": thresholds.overfitting_min_train_relative_change,
                "validation_regression_min": thresholds.overfitting_min_validation_regression,
            }
        if phase_name == "divergence":
            return {
                "train_regression_min": thresholds.divergence_min_relative_regression,
                "regression_from_best_min": thresholds.divergence_from_best_ratio,
            }
        return {
            "stagnation_relative_change_max": thresholds.stagnation_max_relative_change,
            "rapid_learning_relative_change_min": thresholds.rapid_learning_min_relative_change,
        }

    @staticmethod
    def _find_series(
        metrics: tuple[MetricRecord, ...],
        keys: tuple[str, ...],
        *,
        direction: int,
    ) -> _MetricSeries | None:
        for key in keys:
            records = tuple(metric for metric in metrics if metric.metric_key == key)
            if records:
                return _MetricSeries(
                    metric_key=key,
                    records=records,
                    direction_multiplier=direction,
                )
        return None

    @staticmethod
    def _dedupe_records_by_step(records: tuple[MetricRecord, ...]) -> tuple[MetricRecord, ...]:
        records_by_step: dict[int, MetricRecord] = {}
        for record in records:
            if record.step is None:
                continue
            records_by_step[int(record.step)] = record
        return tuple(records_by_step[step] for step in sorted(records_by_step))

    @staticmethod
    def _relative_change(previous_value: float, current_value: float) -> float:
        return (previous_value - current_value) / max(abs(previous_value), _EPSILON)

    @staticmethod
    def _relative_regression(best_value: float, current_value: float) -> float:
        if current_value <= best_value:
            return 0.0
        return (current_value - best_value) / max(abs(best_value), _EPSILON)

    @staticmethod
    def _resolve_session_dependencies(session: Any) -> tuple[RunContext, LocalSessionStore]:
        if isinstance(session, RunContext):
            return session, LocalSessionStore(session.storage_dir)

        context = getattr(session, "context", None)
        store = getattr(session, "_store", None)
        if isinstance(context, RunContext) and isinstance(store, LocalSessionStore):
            return context, store

        raise TypeError(
            "PhaseAnalyzer expects a TrainingSession or RunContext created by the Quark SDK."
        )
