"""Layer 4 latent state segmentation using Gaussian HMM."""

from __future__ import annotations

import warnings

import numpy as np
from hmmlearn import hmm
from sklearn.decomposition import PCA
from sklearn.preprocessing import StandardScaler

from cognitive_signal.models.types import SegmentationOutput, SignalVector

try:
    from hmmlearn.base import ConvergenceWarning
except Exception:  # pragma: no cover - fallback for hmmlearn variants
    class ConvergenceWarning(UserWarning):
        pass


SCALAR_FEATURE_NAMES = [
    "valence",
    "arousal",
    "agency",
    "self_focus",
    "social_orientation",
    "action_ratio",
]
OBSERVATION_FEATURE_NAMES = [
    "emb_pc1",
    "emb_pc2",
    "emb_pc3",
    "emb_pc4",
    "emb_pc5",
    *SCALAR_FEATURE_NAMES,
]


def _readable_name(signal_name: str) -> str:
    mapping = {
        "self_focus": "self-focus",
        "social_orientation": "social",
        "action_ratio": "action",
    }
    return mapping.get(signal_name, signal_name)


class StateSegmenter:
    def prepare_observations(self, signal_vectors: list[SignalVector]) -> np.ndarray:
        if not signal_vectors:
            raise ValueError("signal_vectors must contain at least one item.")

        embedding_matrix = np.vstack(
            [np.asarray(vector.embedding, dtype=np.float32).reshape(-1) for vector in signal_vectors]
        )
        if embedding_matrix.shape[1] < 5:
            raise ValueError("Embeddings must have at least 5 dimensions.")

        n_entries = embedding_matrix.shape[0]
        if n_entries == 1:
            reduced = embedding_matrix[:, :5].copy()
        elif n_entries < 6:
            n_components = min(n_entries - 1, embedding_matrix.shape[1])
            reduced_partial = PCA(n_components=n_components, random_state=42).fit_transform(
                embedding_matrix
            )
            reduced = np.zeros((n_entries, 5), dtype=np.float32)
            reduced[:, :n_components] = reduced_partial
        else:
            reduced = PCA(n_components=5, random_state=42).fit_transform(embedding_matrix)

        scalar_matrix = np.asarray(
            [
                [
                    vector.valence,
                    vector.arousal,
                    vector.agency,
                    vector.self_focus,
                    vector.social_orientation,
                    vector.action_ratio,
                ]
                for vector in signal_vectors
            ],
            dtype=np.float32,
        )

        observations = np.concatenate([np.asarray(reduced, dtype=np.float32), scalar_matrix], axis=1)
        standardized = StandardScaler().fit_transform(observations)
        return np.asarray(standardized, dtype=np.float32)

    def fit_predict(
        self,
        observations: np.ndarray,
        n_states: int = 3,
        n_iterations: int = 100,
    ) -> SegmentationOutput:
        matrix = self._validate_observations(observations)
        if n_states < 1:
            raise ValueError("n_states must be >= 1")
        if matrix.shape[0] < n_states:
            raise ValueError("n_states cannot exceed number of observations.")

        model = self._build_model(n_states=n_states, n_iterations=n_iterations)
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", ConvergenceWarning)
            model.fit(matrix)

        state_sequence = model.predict(matrix).astype(int).tolist()
        log_likelihood = float(model.score(matrix))
        return SegmentationOutput(
            state_sequence=state_sequence,
            transition_matrix=np.asarray(model.transmat_, dtype=np.float32),
            state_means=np.asarray(model.means_, dtype=np.float32),
            state_covariances=np.asarray(model.covars_, dtype=np.float32),
            log_likelihood=log_likelihood,
        )

    def select_n_states(self, observations: np.ndarray, max_states: int = 6) -> int:
        matrix = self._validate_observations(observations)
        n_samples, n_features = matrix.shape
        if n_samples < 2:
            return 1

        upper = min(max_states, n_samples)
        best_k = 2
        best_bic = np.inf

        for n_states in range(2, upper + 1):
            model = self._build_model(n_states=n_states, n_iterations=100)
            try:
                with warnings.catch_warnings():
                    warnings.simplefilter("ignore", ConvergenceWarning)
                    model.fit(matrix)
                log_likelihood = float(model.score(matrix))
            except Exception:
                continue

            n_params = self._count_free_parameters(
                n_states=n_states,
                n_features=n_features,
                params=model.params,
            )
            bic = (-2.0 * log_likelihood) + (n_params * float(np.log(n_samples)))
            if bic < best_bic:
                best_bic = bic
                best_k = n_states

        return best_k

    def label_states(
        self, state_means: np.ndarray, feature_names: list[str]
    ) -> dict[int, str]:
        means = np.asarray(state_means, dtype=np.float32)
        if means.ndim != 2:
            raise ValueError("state_means must be a 2D array.")
        if len(feature_names) != means.shape[1]:
            raise ValueError("feature_names length must match state_means width.")

        scalar_names = list(SCALAR_FEATURE_NAMES)
        scalar_indices = [
            index for index, name in enumerate(feature_names) if name in scalar_names
        ]
        if len(scalar_indices) != len(scalar_names):
            raise ValueError(
                "feature_names must contain all scalar signal names for state labeling."
            )

        scalar_means = means[:, scalar_indices]
        cross_state_mean = np.mean(scalar_means, axis=0)
        cross_state_std = np.std(scalar_means, axis=0)

        labels: dict[int, str] = {}
        for state_id in range(means.shape[0]):
            state_scalars = scalar_means[state_id]
            with np.errstate(divide="ignore", invalid="ignore"):
                z_scores = np.where(
                    cross_state_std > 1e-8,
                    (state_scalars - cross_state_mean) / cross_state_std,
                    0.0,
                )

            high_idx = int(np.argmax(z_scores))
            low_idx = int(np.argmin(z_scores))
            high_name = _readable_name(scalar_names[high_idx])
            low_name = _readable_name(scalar_names[low_idx])

            has_high = bool(z_scores[high_idx] > 0.5)
            has_low = bool(z_scores[low_idx] < -0.5)
            if has_high and has_low:
                labels[state_id] = f"high-{high_name}-low-{low_name}"
            elif has_high:
                labels[state_id] = f"high-{high_name}"
            elif has_low:
                labels[state_id] = f"low-{low_name}"
            else:
                labels[state_id] = f"neutral-state-{state_id}"

        return labels

    @staticmethod
    def _validate_observations(observations: np.ndarray) -> np.ndarray:
        matrix = np.asarray(observations, dtype=np.float32)
        if matrix.ndim != 2:
            raise ValueError("observations must be a 2D matrix.")
        if matrix.shape[0] == 0:
            raise ValueError("observations must contain at least one row.")
        if matrix.shape[1] == 0:
            raise ValueError("observations must contain at least one feature.")
        return matrix

    @staticmethod
    def _build_model(n_states: int, n_iterations: int) -> hmm.GaussianHMM:
        return hmm.GaussianHMM(
            n_components=n_states,
            covariance_type="full",
            n_iter=n_iterations,
            random_state=42,
            params="stmc",
            init_params="stmc",
        )

    @staticmethod
    def _count_free_parameters(
        n_states: int, n_features: int, params: str
    ) -> int:
        count = 0
        if "s" in params:
            count += n_states - 1
        if "t" in params:
            count += n_states * (n_states - 1)
        if "m" in params:
            count += n_states * n_features
        if "c" in params:
            count += int(n_states * n_features * (n_features + 1) / 2)
        return int(count)


__all__ = [
    "StateSegmenter",
    "SCALAR_FEATURE_NAMES",
    "OBSERVATION_FEATURE_NAMES",
]
