"""Layer 3 semantic motion and identity drift analysis."""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
from sklearn.decomposition import PCA


def _clamp(value: float, lo: float, hi: float) -> float:
    return max(lo, min(hi, value))


def _as_matrix(embeddings: list[np.ndarray]) -> np.ndarray:
    if not embeddings:
        return np.empty((0, 0), dtype=np.float32)
    vectors = [np.asarray(embedding, dtype=np.float32).reshape(-1) for embedding in embeddings]
    sizes = {vector.size for vector in vectors}
    if len(sizes) != 1:
        raise ValueError("All embeddings must have the same dimensionality.")
    return np.vstack(vectors)


def _lag1_autocorr(values: list[float]) -> float:
    if len(values) < 2:
        return 0.0
    series = np.asarray(values, dtype=np.float32)
    previous = series[:-1]
    current = series[1:]
    if previous.size == 0 or current.size == 0:
        return 0.0
    if float(np.var(previous)) == 0.0 or float(np.var(current)) == 0.0:
        return 1.0
    corr = float(np.corrcoef(current, previous)[0, 1])
    if np.isnan(corr):
        return 1.0
    return _clamp(corr, -1.0, 1.0)


def _safe_unit_vector(vector: np.ndarray) -> np.ndarray:
    array = np.asarray(vector, dtype=np.float32).reshape(-1)
    norm = float(np.linalg.norm(array))
    if norm == 0.0:
        return np.zeros_like(array, dtype=np.float32)
    return array / norm


@dataclass(slots=True)
class IdentityShift:
    timestep: int
    magnitude: float
    persistence: float
    direction: np.ndarray


@dataclass(slots=True)
class SemanticMotionOutput:
    step_distances: list[float]
    reduced_embeddings: np.ndarray
    drift_vectors: list[np.ndarray]
    identity_shifts: list[IdentityShift]


class SemanticMotionAnalyzer:
    def compute_step_distances(self, embeddings: list[np.ndarray]) -> list[float]:
        matrix = _as_matrix(embeddings)
        if matrix.shape[0] == 0:
            return []

        distances = [0.0]
        for index in range(1, matrix.shape[0]):
            displacement = matrix[index] - matrix[index - 1]
            distances.append(float(np.linalg.norm(displacement)))
        return distances

    def reduce_dimensions(
        self, embeddings: list[np.ndarray], n_components: int = 10
    ) -> np.ndarray:
        matrix = _as_matrix(embeddings)
        n_entries, n_features = matrix.shape

        if n_entries == 0:
            return np.empty((0, 0), dtype=np.float32)
        if n_entries == 1:
            return matrix.copy()

        effective_components = min(int(n_components), n_entries - 1, n_features)
        if effective_components < 1:
            raise ValueError("n_components must result in at least one component.")

        pca = PCA(n_components=effective_components, random_state=42)
        reduced = pca.fit_transform(matrix)
        return np.asarray(reduced, dtype=np.float32)

    def compute_drift_vector(
        self, reduced: np.ndarray, window: int = 7
    ) -> list[np.ndarray]:
        matrix = np.asarray(reduced, dtype=np.float32)
        if matrix.size == 0:
            return []
        if matrix.ndim != 2:
            raise ValueError("Reduced embeddings must be a 2D array.")

        n_entries, dims = matrix.shape
        drifts: list[np.ndarray] = [np.zeros(dims, dtype=np.float32)]

        for timestep in range(1, n_entries):
            start = max(1, timestep - window + 1)
            displacements = [
                matrix[index] - matrix[index - 1] for index in range(start, timestep + 1)
            ]
            drift = np.mean(np.vstack(displacements), axis=0)
            drifts.append(np.asarray(drift, dtype=np.float32))
        return drifts

    def _detect_identity_shifts_from_precomputed(
        self,
        step_distances: list[float],
        drift_vectors: list[np.ndarray],
        epsilon: float,
        tau: float,
        window: int,
    ) -> list[IdentityShift]:
        shifts: list[IdentityShift] = []
        for timestep, drift_vector in enumerate(drift_vectors):
            magnitude = float(np.linalg.norm(drift_vector))
            segment_start = max(0, timestep - window + 1)
            persistence_segment = step_distances[segment_start : timestep + 1]
            if len(persistence_segment) < 3:
                persistence = 0.0
            else:
                persistence = abs(_lag1_autocorr(persistence_segment))

            if magnitude > epsilon and persistence > tau:
                shifts.append(
                    IdentityShift(
                        timestep=timestep,
                        magnitude=magnitude,
                        persistence=persistence,
                        direction=_safe_unit_vector(drift_vector),
                    )
                )
        return shifts

    def detect_identity_shifts(
        self,
        embeddings: list[np.ndarray],
        epsilon: float = 0.5,
        tau: float = 0.7,
        window: int = 7,
    ) -> list[IdentityShift]:
        if not embeddings:
            return []
        step_distances = self.compute_step_distances(embeddings)
        reduced = self.reduce_dimensions(embeddings)
        drift_vectors = self.compute_drift_vector(reduced, window=window)
        return self._detect_identity_shifts_from_precomputed(
            step_distances=step_distances,
            drift_vectors=drift_vectors,
            epsilon=epsilon,
            tau=tau,
            window=window,
        )

    def analyze(
        self,
        embeddings: list[np.ndarray],
        epsilon: float = 0.5,
        tau: float = 0.7,
        window: int = 7,
        n_components: int = 10,
    ) -> SemanticMotionOutput:
        step_distances = self.compute_step_distances(embeddings)
        reduced = self.reduce_dimensions(embeddings, n_components=n_components)
        drift_vectors = self.compute_drift_vector(reduced, window=window)
        identity_shifts = self._detect_identity_shifts_from_precomputed(
            step_distances=step_distances,
            drift_vectors=drift_vectors,
            epsilon=epsilon,
            tau=tau,
            window=window,
        )
        return SemanticMotionOutput(
            step_distances=step_distances,
            reduced_embeddings=reduced,
            drift_vectors=drift_vectors,
            identity_shifts=identity_shifts,
        )


__all__ = ["SemanticMotionAnalyzer", "IdentityShift", "SemanticMotionOutput"]
