from __future__ import annotations

from collections import Counter

import numpy as np
import pytest

from cognitive_signal.data.generate_test_data import generate_synthetic_entries
from cognitive_signal.hmm.segmenter import OBSERVATION_FEATURE_NAMES, StateSegmenter
from cognitive_signal.models.types import SignalVector
from cognitive_signal.signals.extractors import SignalExtractor


def _build_phase_coded_signal_vectors() -> list[SignalVector]:
    rng = np.random.default_rng(42)
    centers = [
        np.concatenate([np.full(120, 1.6), np.zeros(264)]).astype(np.float32),
        np.concatenate([np.zeros(120), np.full(120, 1.6), np.zeros(144)]).astype(np.float32),
        np.concatenate([np.zeros(240), np.full(120, 1.6), np.zeros(24)]).astype(np.float32),
    ]
    scalar_profiles = [
        (0.55, 0.42, 0.75, 0.32, 0.28, 0.82),
        (0.22, 0.58, 0.45, 0.54, 0.68, 0.42),
        (0.08, 0.30, 0.25, 0.80, 0.35, 0.20),
    ]

    vectors: list[SignalVector] = []
    for index in range(60):
        phase = 0 if index < 20 else 1 if index < 40 else 2
        embedding_noise = rng.normal(0.0, 0.08, size=384).astype(np.float32)
        embedding = centers[phase] + embedding_noise
        base = scalar_profiles[phase]
        scalar_noise = rng.normal(0.0, 0.03, size=6)

        vectors.append(
            SignalVector(
                embedding=embedding,
                valence=float(base[0] + scalar_noise[0]),
                arousal=float(base[1] + scalar_noise[1]),
                agency=float(base[2] + scalar_noise[2]),
                self_focus=float(base[3] + scalar_noise[3]),
                social_orientation=float(base[4] + scalar_noise[4]),
                action_ratio=float(base[5] + scalar_noise[5]),
            )
        )
    return vectors


@pytest.fixture(scope="module")
def phase_vectors() -> list[SignalVector]:
    return _build_phase_coded_signal_vectors()


def _dominant_state(sequence: list[int]) -> int:
    return Counter(sequence).most_common(1)[0][0]


def test_prepare_observations_shape_and_standardization(
    phase_vectors: list[SignalVector],
) -> None:
    segmenter = StateSegmenter()
    observations = segmenter.prepare_observations(phase_vectors)

    assert observations.shape == (60, 11)
    assert np.isfinite(observations).all()
    means = np.mean(observations, axis=0)
    stds = np.std(observations, axis=0)
    assert np.all(np.abs(means) < 1e-5)
    non_constant = stds > 1e-8
    assert np.all(np.abs(stds[non_constant] - 1.0) < 1e-4)


def test_fit_predict_outputs_valid_shapes(phase_vectors: list[SignalVector]) -> None:
    segmenter = StateSegmenter()
    observations = segmenter.prepare_observations(phase_vectors)
    result = segmenter.fit_predict(observations, n_states=3, n_iterations=200)

    assert len(result.state_sequence) == 60
    assert result.transition_matrix.shape == (3, 3)
    assert result.state_means.shape == (3, 11)
    assert result.state_covariances.shape[0] == 3
    assert np.isfinite(result.log_likelihood)


def test_transition_matrix_rows_sum_to_one(phase_vectors: list[SignalVector]) -> None:
    segmenter = StateSegmenter()
    observations = segmenter.prepare_observations(phase_vectors)
    result = segmenter.fit_predict(observations, n_states=3, n_iterations=200)
    row_sums = np.sum(result.transition_matrix, axis=1)
    assert row_sums.tolist() == pytest.approx([1.0, 1.0, 1.0], abs=1e-6)


def test_select_n_states_bic_reasonable_range(phase_vectors: list[SignalVector]) -> None:
    segmenter = StateSegmenter()
    observations = segmenter.prepare_observations(phase_vectors)
    n_states = segmenter.select_n_states(observations, max_states=6)
    assert 2 <= n_states <= 4


def test_parameter_count_respects_trainable_flags() -> None:
    n_states = 3
    n_features = 11
    with_stmc = StateSegmenter._count_free_parameters(n_states, n_features, "stmc")
    without_st = StateSegmenter._count_free_parameters(n_states, n_features, "mc")
    assert without_st < with_stmc
    expected_reduction = (n_states - 1) + (n_states * (n_states - 1))
    assert with_stmc - without_st == expected_reduction


def test_state_sequence_roughly_aligns_with_phases(
    phase_vectors: list[SignalVector],
) -> None:
    segmenter = StateSegmenter()
    observations = segmenter.prepare_observations(phase_vectors)
    result = segmenter.fit_predict(observations, n_states=3, n_iterations=250)
    sequence = result.state_sequence

    phase_1 = sequence[:20]
    phase_2 = sequence[20:40]
    phase_3 = sequence[40:60]
    dominant = [_dominant_state(phase_1), _dominant_state(phase_2), _dominant_state(phase_3)]
    assert len(set(dominant)) >= 2

    change_points = [idx for idx in range(1, len(sequence)) if sequence[idx] != sequence[idx - 1]]
    assert any(18 <= idx <= 24 for idx in change_points)
    assert any(38 <= idx <= 44 for idx in change_points)


def test_label_states_strictness() -> None:
    segmenter = StateSegmenter()
    feature_names = OBSERVATION_FEATURE_NAMES

    strong_contrast = np.array(
        [
            [0.0, 0.0, 0.0, 0.0, 0.0, 2.0, 1.0, 2.0, -1.0, 1.0, 1.0],
            [0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
            [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 2.0, 1.0, 1.0],
        ],
        dtype=np.float32,
    )
    labels_contrast = segmenter.label_states(strong_contrast, feature_names)
    assert labels_contrast[0].startswith("high-") and "-low-" in labels_contrast[0]

    near_neutral = np.array(
        [
            [0.0, 0.0, 0.0, 0.0, 0.0, 0.55, 0.50, 0.52, 0.48, 0.51, 0.50],
            [0.0, 0.0, 0.0, 0.0, 0.0, 0.50, 0.51, 0.49, 0.50, 0.50, 0.49],
            [0.0, 0.0, 0.0, 0.0, 0.0, 0.52, 0.49, 0.51, 0.52, 0.48, 0.51],
        ],
        dtype=np.float32,
    )
    labels_neutral = segmenter.label_states(near_neutral, feature_names)
    assert all(
        label.startswith("high-")
        or label.startswith("low-")
        or label.startswith("neutral-state-")
        for label in labels_neutral.values()
    )

    one_sided = np.array(
        [
            [0.0, 0.0, 0.0, 0.0, 0.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0],
            [0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
            [0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
        ],
        dtype=np.float32,
    )
    labels_one_sided = segmenter.label_states(one_sided, feature_names)
    assert labels_one_sided[0].startswith("high-")


def test_label_states_no_all_mixed() -> None:
    """With distinct states, labels should not collapse to a single fallback label."""
    segmenter = StateSegmenter()
    feature_names = OBSERVATION_FEATURE_NAMES
    state_means = np.array(
        [
            [0.0, 0.0, 0.0, 0.0, 0.0, 1.2, 0.4, 1.0, 0.2, 0.1, 0.9],
            [0.0, 0.0, 0.0, 0.0, 0.0, 0.4, 0.8, 0.3, 0.9, 0.8, 0.2],
            [0.0, 0.0, 0.0, 0.0, 0.0, 0.3, 0.2, 0.1, 1.1, 0.6, 0.1],
        ],
        dtype=np.float32,
    )
    labels = segmenter.label_states(state_means, feature_names)
    values = list(labels.values())
    assert len(set(values)) > 1
    assert any(value.startswith("high-") or value.startswith("low-") for value in values)


def test_prepare_observations_small_sample_padding() -> None:
    segmenter = StateSegmenter()
    small = _build_phase_coded_signal_vectors()[:4]
    observations = segmenter.prepare_observations(small)
    assert observations.shape == (4, 11)


def test_prepare_observations_single_sample_uses_raw_slice() -> None:
    segmenter = StateSegmenter()
    sample = _build_phase_coded_signal_vectors()[:1]
    observations = segmenter.prepare_observations(sample)
    assert observations.shape == (1, 11)


def test_hybrid_integration_with_real_extractors() -> None:
    entries = generate_synthetic_entries()
    try:
        extractor = SignalExtractor()
        signal_vectors = [extractor.extract(entry) for entry in entries]
    except Exception as exc:
        pytest.skip(f"SignalExtractor dependencies unavailable for integration test: {exc}")

    segmenter = StateSegmenter()
    observations = segmenter.prepare_observations(signal_vectors)
    result = segmenter.fit_predict(observations, n_states=3, n_iterations=150)
    assert observations.shape == (60, 11)
    assert len(result.state_sequence) == 60
