"""Layer 5 change point detection."""

from __future__ import annotations

import numpy as np
import ruptures as rpt
from scipy.stats import wasserstein_distance
from sklearn.decomposition import PCA

from cognitive_signal.models.types import (
    ChangePointOutput,
    Phase,
    SemanticChangePoint,
    SignalVector,
)

SCALAR_SIGNALS = [
    "valence",
    "arousal",
    "agency",
    "self_focus",
    "social_orientation",
    "action_ratio",
]


class ChangePointDetector:
    def detect_signal_changepoints(
        self,
        values: list[float],
        method: str = "pelt",
        penalty: float = 3.0,
    ) -> list[int]:
        if not values:
            return []

        series = np.asarray(values, dtype=np.float32).reshape(-1, 1)
        n_samples = series.shape[0]
        if n_samples < 3:
            return []

        if method == "pelt":
            model = rpt.Pelt(model="rbf").fit(series)
            raw_changepoints = model.predict(pen=penalty)
        elif method == "binseg":
            model = rpt.Binseg(model="rbf").fit(series)
            raw_changepoints = model.predict(pen=penalty)
        else:
            raise ValueError("method must be either 'pelt' or 'binseg'")

        # Strip terminal index n_samples and keep only internal cps.
        changepoints = sorted(
            {
                int(index)
                for index in raw_changepoints
                if 0 < int(index) < n_samples
            }
        )
        return changepoints

    def detect_semantic_changepoints(
        self,
        embeddings: list[np.ndarray],
        window: int = 10,
        quantile: float = 0.85,
    ) -> list[SemanticChangePoint]:
        if not 0.0 <= quantile <= 1.0:
            raise ValueError("quantile must be in [0.0, 1.0]")
        if window < 1:
            raise ValueError("window must be >= 1")

        matrix = self._as_matrix(embeddings)
        n_entries, n_features = matrix.shape
        if n_entries < (2 * window + 1):
            return []
        n_components = min(5, n_entries - 1, n_features)
        if n_components < 1:
            return []

        reduced = PCA(n_components=n_components, random_state=42).fit_transform(matrix)
        candidates: list[tuple[int, float]] = []

        for timestep in range(window, n_entries - window + 1):
            before = reduced[timestep - window : timestep]
            after = reduced[timestep : timestep + window]
            per_dimension = [
                wasserstein_distance(before[:, dimension], after[:, dimension])
                for dimension in range(n_components)
            ]
            average_distance = float(np.mean(per_dimension))
            candidates.append((timestep, average_distance))

        if not candidates:
            return []

        distances = np.asarray([distance for _, distance in candidates], dtype=np.float32)
        threshold = float(np.quantile(distances, quantile))
        return [
            SemanticChangePoint(timestep=timestep, distance=distance, method="wasserstein-pca")
            for timestep, distance in candidates
            if distance > threshold
        ]

    def detect_all(
        self,
        signal_vectors: list[SignalVector],
        embeddings: list[np.ndarray],
        semantic_quantile: float = 0.85,
    ) -> ChangePointOutput:
        if signal_vectors and embeddings and len(signal_vectors) != len(embeddings):
            raise ValueError("signal_vectors and embeddings must have the same length.")

        signal_changepoints: dict[str, list[int]] = {}
        for signal_name in SCALAR_SIGNALS:
            values = [float(getattr(vector, signal_name)) for vector in signal_vectors]
            signal_changepoints[signal_name] = self.detect_signal_changepoints(values)

        semantic_changepoints = self.detect_semantic_changepoints(
            embeddings,
            quantile=semantic_quantile,
        )

        all_points = []
        for indices in signal_changepoints.values():
            all_points.extend(indices)
        all_points.extend(item.timestep for item in semantic_changepoints)
        merged_changepoints = self._merge_nearby_points(all_points, radius=3)

        return ChangePointOutput(
            signal_changepoints=signal_changepoints,
            semantic_changepoints=semantic_changepoints,
            merged_changepoints=merged_changepoints,
        )

    def annotate_phases(self, changepoints: list[int], n_entries: int) -> list[Phase]:
        if n_entries <= 0:
            return []

        clean_points = sorted({point for point in changepoints if 0 < point < n_entries})
        if not clean_points:
            return [Phase(start=0, end=n_entries - 1, phase_id=0)]

        phases: list[Phase] = []
        start = 0
        phase_id = 0
        for point in clean_points:
            end = point - 1
            if end >= start:
                phases.append(Phase(start=start, end=end, phase_id=phase_id))
                phase_id += 1
            start = point
        if start <= n_entries - 1:
            phases.append(Phase(start=start, end=n_entries - 1, phase_id=phase_id))
        return phases

    @staticmethod
    def _as_matrix(embeddings: list[np.ndarray]) -> np.ndarray:
        if not embeddings:
            return np.empty((0, 0), dtype=np.float32)
        matrix = np.vstack(
            [np.asarray(embedding, dtype=np.float32).reshape(-1) for embedding in embeddings]
        )
        if matrix.ndim != 2:
            raise ValueError("embeddings must form a 2D matrix.")
        return matrix

    @staticmethod
    def _merge_nearby_points(points: list[int], radius: int = 3) -> list[int]:
        if not points:
            return []
        sorted_points = sorted(set(points))
        merged = [sorted_points[0]]
        for point in sorted_points[1:]:
            if point - merged[-1] <= radius:
                continue
            merged.append(point)
        return merged


__all__ = ["ChangePointDetector", "SCALAR_SIGNALS"]
