# -*- coding: utf-8 -*-
"""
test_dynamics.py — tests for the 3.1.0a1 cognitive dynamics model.

Covers:
  - state ↔ vector encoding (round-trip, simplex projection)
  - Observation construction (raw vector + from_thoughts)
  - fit() math correctness:
      * machine-epsilon recovery on full-rank gaussian inputs
      * noisy recovery is bounded
      * dirichlet (simplex) inputs: predictions exact, parameters
        non-unique (rank-deficient regime)
  - predict() consistency with fit()
  - simulate() multi-step rollout
  - suggest() — the controller — drives state to target in raw space
  - forecast_horizon() bounded for stable A
  - residual() on held-out tuples
  - save / load round-trip with .cogdyn files
  - public API exposure on styxx module
  - error paths (unfitted predict, bad shapes, etc.)

These tests run with no network access and no API keys. They use
synthetic data with known dynamics so every assertion is verifiable
against ground truth.
"""

from __future__ import annotations

import json
import math
import os
import sys
import tempfile
from pathlib import Path

import numpy as np
import pytest

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import styxx
from styxx import Thought
from styxx.dynamics import (
    CognitiveDynamics,
    Observation,
    FitResult,
    synthetic_observations,
    thought_to_state,
    state_to_thought,
    COGDYN_FORMAT,
    COGDYN_VERSION,
)
from styxx.thought import N_CATEGORIES, CATEGORIES


# ══════════════════════════════════════════════════════════════════
# Helpers
# ══════════════════════════════════════════════════════════════════

def _stable_random_A(seed: int = 42, target_radius: float = 0.7) -> np.ndarray:
    """Generate a random 6x6 matrix with spectral radius == target_radius."""
    rng = np.random.default_rng(seed)
    A = rng.normal(0, 0.15, size=(N_CATEGORIES, N_CATEGORIES))
    radius = float(np.max(np.abs(np.linalg.eigvals(A))))
    if radius > 1e-9:
        A = A * (target_radius / radius)
    return A


def _random_B(seed: int = 43, scale: float = 0.2) -> np.ndarray:
    rng = np.random.default_rng(seed)
    return rng.normal(0, scale, size=(N_CATEGORIES, N_CATEGORIES))


# ══════════════════════════════════════════════════════════════════
# State <-> vector encoding
# ══════════════════════════════════════════════════════════════════

class TestStateVectorEncoding:

    def test_thought_to_state_returns_6d_vector(self):
        t = Thought.target("reasoning", confidence=0.7)
        v = thought_to_state(t)
        assert v.shape == (N_CATEGORIES,)
        assert isinstance(v, np.ndarray)

    def test_thought_to_state_is_simplex_valid_for_target(self):
        t = Thought.target("reasoning", confidence=0.7)
        v = thought_to_state(t)
        assert abs(v.sum() - 1.0) < 1e-9
        assert all(x >= 0 for x in v)

    def test_state_to_thought_roundtrip_simplex(self):
        v = np.array([0.1, 0.4, 0.1, 0.2, 0.1, 0.1])
        t = state_to_thought(v)
        v2 = thought_to_state(t)
        assert np.allclose(v, v2, atol=1e-9)

    def test_state_to_thought_clips_negatives(self):
        v = np.array([-0.1, 0.4, 0.2, 0.2, 0.2, 0.1])
        t = state_to_thought(v)
        v2 = thought_to_state(t)
        assert all(x >= 0 for x in v2)
        assert abs(v2.sum() - 1.0) < 1e-9

    def test_state_to_thought_renormalizes(self):
        v = np.array([2.0, 3.0, 1.0, 0.0, 0.0, 0.0])  # sum = 6
        t = state_to_thought(v)
        v2 = thought_to_state(t)
        assert abs(v2.sum() - 1.0) < 1e-9

    def test_state_to_thought_zero_vector_returns_uniform(self):
        v = np.zeros(N_CATEGORIES)
        t = state_to_thought(v)
        v2 = thought_to_state(t)
        for x in v2:
            assert abs(x - 1.0 / N_CATEGORIES) < 1e-9

    def test_state_to_thought_wrong_shape_raises(self):
        with pytest.raises(ValueError):
            state_to_thought(np.array([0.5, 0.5]))

    def test_thought_to_state_rejects_non_thought(self):
        with pytest.raises(TypeError):
            thought_to_state("not a thought")  # type: ignore


# ══════════════════════════════════════════════════════════════════
# Observation
# ══════════════════════════════════════════════════════════════════

class TestObservation:

    def test_observation_holds_raw_vectors(self):
        ob = Observation(
            state_vec=np.array([0.1, 0.4, 0.1, 0.2, 0.1, 0.1]),
            action_vec=np.array([0.0, 0.7, 0.0, 0.1, 0.1, 0.1]),
            next_state_vec=np.array([0.05, 0.55, 0.1, 0.15, 0.1, 0.05]),
        )
        assert ob.state_vec.shape == (N_CATEGORIES,)
        assert ob.action_vec.shape == (N_CATEGORIES,)
        assert ob.next_state_vec.shape == (N_CATEGORIES,)

    def test_observation_from_thoughts(self):
        t1 = Thought.target("reasoning", confidence=0.7)
        t2 = Thought.target("creative", confidence=0.7)
        t3 = Thought.target("retrieval", confidence=0.7)
        ob = Observation.from_thoughts(t1, t2, t3)
        assert ob.state_vec.shape == (N_CATEGORIES,)
        assert ob.action_vec.shape == (N_CATEGORIES,)
        assert ob.next_state_vec.shape == (N_CATEGORIES,)

    def test_observation_lazy_thought_accessors(self):
        t1 = Thought.target("reasoning", confidence=0.7)
        ob = Observation.from_thoughts(t1, t1, t1)
        # The .state property returns a Thought (simplex-projected)
        assert isinstance(ob.state, Thought)
        assert isinstance(ob.action, Thought)
        assert isinstance(ob.next_state, Thought)
        assert ob.state.primary_category == "reasoning"

    def test_observation_rejects_wrong_shape(self):
        with pytest.raises(ValueError):
            Observation(
                state_vec=np.zeros(3),  # wrong size
                action_vec=np.zeros(N_CATEGORIES),
                next_state_vec=np.zeros(N_CATEGORIES),
            )

    def test_observation_coerces_lists_to_arrays(self):
        ob = Observation(
            state_vec=[0.1] * N_CATEGORIES,  # list, not ndarray
            action_vec=[0.1] * N_CATEGORIES,
            next_state_vec=[0.1] * N_CATEGORIES,
        )
        assert isinstance(ob.state_vec, np.ndarray)


# ══════════════════════════════════════════════════════════════════
# Math correctness: fit() recovers known (A, B)
# ══════════════════════════════════════════════════════════════════

class TestFitMathCorrectness:
    """The load-bearing math test. With noise-free, full-rank
    (gaussian) inputs, fit() must recover the true (A, B) to machine
    epsilon. If this fails, the entire dynamics module is broken.
    """

    def test_recovery_machine_epsilon_no_noise(self):
        A_true = _stable_random_A(seed=42, target_radius=0.7)
        B_true = _random_B(seed=43, scale=0.2)
        obs = synthetic_observations(
            n=100, A=A_true, B=B_true,
            noise_std=0.0, seed=123, distribution="gaussian",
        )
        dyn = CognitiveDynamics()
        result = dyn.fit(obs)
        A_err = float(np.max(np.abs(dyn.A - A_true)))
        B_err = float(np.max(np.abs(dyn.B - B_true)))
        assert A_err < 1e-10, f"A recovery failed: max err = {A_err}"
        assert B_err < 1e-10, f"B recovery failed: max err = {B_err}"
        assert result.r2 > 0.9999

    def test_recovery_robust_to_small_noise(self):
        A_true = _stable_random_A(seed=42, target_radius=0.7)
        B_true = _random_B(seed=43, scale=0.2)
        obs = synthetic_observations(
            n=2000, A=A_true, B=B_true,
            noise_std=0.02, seed=456, distribution="gaussian",
        )
        dyn = CognitiveDynamics()
        result = dyn.fit(obs)
        A_err = float(np.max(np.abs(dyn.A - A_true)))
        B_err = float(np.max(np.abs(dyn.B - B_true)))
        assert A_err < 0.05, f"A recovery degraded: {A_err}"
        assert B_err < 0.05, f"B recovery degraded: {B_err}"
        assert result.r2 > 0.99

    def test_dirichlet_predictions_exact_but_params_non_unique(self):
        """On simplex (Dirichlet) inputs, the regressor matrix is
        rank-deficient (each block has columns that sum to 1), so
        the parameter estimates are non-unique. Predictions on the
        training set are still exact (R^2 ≈ 1.0).
        """
        A_true = _stable_random_A(seed=42, target_radius=0.7)
        B_true = _random_B(seed=43, scale=0.2)
        obs = synthetic_observations(
            n=500, A=A_true, B=B_true,
            noise_std=0.0, seed=777, distribution="dirichlet",
        )
        dyn = CognitiveDynamics()
        result = dyn.fit(obs)
        # Predictions on training set should be exact
        assert result.r2 > 0.9999
        # Predict each tuple matches its training next_state to ~0
        for ob in obs[:20]:
            pred = (dyn.A @ ob.state_vec) + (dyn.B @ ob.action_vec)
            err = float(np.max(np.abs(pred - ob.next_state_vec)))
            assert err < 1e-9
        # But the parameters need NOT match the true values
        # (this is the rank-deficient identification regime)
        # We don't assert anything about A/B element error here.

    def test_fit_records_spectral_radius(self):
        A_true = _stable_random_A(seed=42, target_radius=0.5)
        B_true = _random_B(seed=43)
        obs = synthetic_observations(50, A_true, B_true, noise_std=0.0, seed=1)
        dyn = CognitiveDynamics()
        result = dyn.fit(obs)
        # Spectral radius of recovered A should equal target ~= 0.5
        assert abs(result.spectral_radius_A - 0.5) < 1e-3

    def test_is_stable_when_spectral_radius_below_one(self):
        A_true = _stable_random_A(seed=42, target_radius=0.5)
        B_true = _random_B(seed=43)
        obs = synthetic_observations(100, A_true, B_true, noise_std=0.0, seed=2)
        dyn = CognitiveDynamics()
        result = dyn.fit(obs)
        assert result.is_stable()

    def test_fit_rejects_empty_observations(self):
        dyn = CognitiveDynamics()
        with pytest.raises(ValueError):
            dyn.fit([])


# ══════════════════════════════════════════════════════════════════
# Predict
# ══════════════════════════════════════════════════════════════════

class TestPredict:

    def setup_method(self):
        A = _stable_random_A(seed=42, target_radius=0.7)
        B = _random_B(seed=43)
        self.A_true = A
        self.B_true = B
        obs = synthetic_observations(100, A, B, noise_std=0.0, seed=10)
        self.dyn = CognitiveDynamics()
        self.dyn.fit(obs)
        self.training_obs = obs

    def test_predict_returns_thought(self):
        s = state_to_thought(np.array([0.2, 0.3, 0.1, 0.2, 0.1, 0.1]))
        a = state_to_thought(np.array([0.1, 0.5, 0.0, 0.2, 0.1, 0.1]))
        result = self.dyn.predict(s, a)
        assert isinstance(result, Thought)

    def test_predict_raw_consistency_with_fit(self):
        """For a noise-free fit, raw-space predict must exactly
        reproduce training-set next_state_vec values."""
        for ob in self.training_obs[:30]:
            raw_pred = (self.dyn.A @ ob.state_vec) + (self.dyn.B @ ob.action_vec)
            assert np.allclose(raw_pred, ob.next_state_vec, atol=1e-10)

    def test_predict_unfitted_raises(self):
        dyn = CognitiveDynamics()
        s = state_to_thought(np.zeros(N_CATEGORIES))
        with pytest.raises(RuntimeError):
            dyn.predict(s, s)


# ══════════════════════════════════════════════════════════════════
# Simulate
# ══════════════════════════════════════════════════════════════════

class TestSimulate:

    def setup_method(self):
        A = _stable_random_A(seed=42, target_radius=0.6)
        B = _random_B(seed=43, scale=0.1)
        obs = synthetic_observations(100, A, B, noise_std=0.0, seed=20)
        self.dyn = CognitiveDynamics()
        self.dyn.fit(obs)

    def test_simulate_returns_correct_length(self):
        initial = state_to_thought(np.array([0.2, 0.2, 0.2, 0.2, 0.1, 0.1]))
        actions = [
            Thought.target("reasoning"),
            Thought.target("creative"),
            Thought.target("retrieval"),
        ]
        traj = self.dyn.simulate(initial, actions)
        assert len(traj) == 4  # initial + 3 steps
        for state in traj:
            assert isinstance(state, Thought)

    def test_simulate_first_state_is_initial(self):
        initial = state_to_thought(np.array([0.5, 0.1, 0.1, 0.1, 0.1, 0.1]))
        actions = [Thought.target("reasoning")]
        traj = self.dyn.simulate(initial, actions)
        # First element should be the initial state
        assert traj[0].mean_probs() == initial.mean_probs()

    def test_simulate_unfitted_raises(self):
        dyn = CognitiveDynamics()
        with pytest.raises(RuntimeError):
            dyn.simulate(
                state_to_thought(np.zeros(N_CATEGORIES)),
                [Thought.target("reasoning")],
            )


# ══════════════════════════════════════════════════════════════════
# Suggest (controller)
# ══════════════════════════════════════════════════════════════════

class TestSuggestController:
    """The model-predictive controller. Given (current, target),
    suggest() returns the action that minimizes ||target - next||
    in the linear model. With a well-conditioned B, this should
    drive the next state to the target exactly (raw vector space).
    """

    def setup_method(self):
        A = _stable_random_A(seed=42, target_radius=0.7)
        B = _random_B(seed=43)
        obs = synthetic_observations(100, A, B, noise_std=0.0, seed=30)
        self.dyn = CognitiveDynamics()
        self.dyn.fit(obs)

    def test_suggest_returns_thought(self):
        c = state_to_thought(np.array([0.3, 0.2, 0.1, 0.2, 0.1, 0.1]))
        t = state_to_thought(np.array([0.1, 0.7, 0.0, 0.1, 0.0, 0.1]))
        a = self.dyn.suggest(c, t)
        assert isinstance(a, Thought)

    def test_suggest_drives_raw_state_to_target(self):
        """The math: solving B·a = target - A·current. The resulting
        action, applied via the linear model, should produce the
        target exactly in raw vector space."""
        current_v = np.array([0.5, 0.1, 0.1, 0.1, 0.1, 0.1])
        target_v = np.array([0.05, 0.85, 0.0, 0.05, 0.0, 0.05])
        # Compute the optimal action directly in vector space
        delta = target_v - self.dyn.A @ current_v
        a_solved, *_ = np.linalg.lstsq(self.dyn.B, delta, rcond=None)
        # Apply
        predicted = self.dyn.A @ current_v + self.dyn.B @ a_solved
        err = float(np.linalg.norm(predicted - target_v))
        assert err < 1e-9, f"controller failed to converge: {err}"

    def test_suggest_unfitted_raises(self):
        dyn = CognitiveDynamics()
        s = state_to_thought(np.zeros(N_CATEGORIES))
        with pytest.raises(RuntimeError):
            dyn.suggest(s, s)


# ══════════════════════════════════════════════════════════════════
# Forecast horizon
# ══════════════════════════════════════════════════════════════════

class TestForecastHorizon:

    def test_forecast_horizon_correct_length(self):
        A = _stable_random_A(seed=42, target_radius=0.5)
        B = _random_B(seed=43)
        obs = synthetic_observations(100, A, B, noise_std=0.0, seed=40)
        dyn = CognitiveDynamics()
        dyn.fit(obs)
        initial = state_to_thought(np.array([0.4, 0.1, 0.1, 0.1, 0.2, 0.1]))
        fh = dyn.forecast_horizon(initial, n_steps=10)
        assert len(fh) == 11  # initial + 10 steps

    def test_forecast_horizon_unfitted_raises(self):
        dyn = CognitiveDynamics()
        with pytest.raises(RuntimeError):
            dyn.forecast_horizon(state_to_thought(np.zeros(N_CATEGORIES)), 5)


# ══════════════════════════════════════════════════════════════════
# Residual on held-out
# ══════════════════════════════════════════════════════════════════

class TestResidual:

    def test_residual_is_zero_on_training_tuple_no_noise(self):
        A = _stable_random_A(seed=42)
        B = _random_B(seed=43)
        obs = synthetic_observations(50, A, B, noise_std=0.0, seed=50)
        dyn = CognitiveDynamics()
        dyn.fit(obs)
        for ob in obs[:5]:
            r = dyn.residual(ob)
            assert r < 1e-10, f"residual on training tuple should be ~0, got {r}"

    def test_residual_bounded_on_held_out_with_noise(self):
        A = _stable_random_A(seed=42)
        B = _random_B(seed=43)
        train = synthetic_observations(2000, A, B, noise_std=0.02, seed=60)
        held = synthetic_observations(100, A, B, noise_std=0.02, seed=61)
        dyn = CognitiveDynamics()
        dyn.fit(train)
        residuals = [dyn.residual(ob) for ob in held]
        mean_res = sum(residuals) / len(residuals)
        # noise std=0.02 across 6 dims => expected L2 ~ 0.02 * sqrt(6) ~ 0.05
        assert mean_res < 0.2

    def test_residual_unfitted_raises(self):
        dyn = CognitiveDynamics()
        ob = Observation(
            state_vec=np.zeros(N_CATEGORIES),
            action_vec=np.zeros(N_CATEGORIES),
            next_state_vec=np.zeros(N_CATEGORIES),
        )
        with pytest.raises(RuntimeError):
            dyn.residual(ob)


# ══════════════════════════════════════════════════════════════════
# Save / load (.cogdyn file format)
# ══════════════════════════════════════════════════════════════════

class TestCogdynFileFormat:

    def setup_method(self):
        A = _stable_random_A(seed=42, target_radius=0.6)
        B = _random_B(seed=43)
        obs = synthetic_observations(100, A, B, noise_std=0.0, seed=70)
        self.dyn = CognitiveDynamics()
        self.dyn.fit(obs)

    def test_save_creates_file(self, tmp_path):
        path = tmp_path / "out.cogdyn"
        result = self.dyn.save(path)
        assert path.exists()
        assert path.stat().st_size > 100

    def test_save_load_roundtrip_exact(self, tmp_path):
        path = tmp_path / "rt.cogdyn"
        self.dyn.save(path)
        loaded = CognitiveDynamics.load(path)
        assert np.allclose(loaded.A, self.dyn.A, atol=1e-15)
        assert np.allclose(loaded.B, self.dyn.B, atol=1e-15)
        assert loaded.is_fitted

    def test_save_no_bom(self, tmp_path):
        path = tmp_path / "nobom.cogdyn"
        self.dyn.save(path)
        head = path.read_bytes()[:3]
        assert head != b"\xef\xbb\xbf"

    def test_canonical_dict_has_required_keys(self):
        d = self.dyn.as_dict()
        for key in ("cogdyn_format", "cogdyn_version", "dynamics_id",
                    "schema", "model", "fit"):
            assert key in d
        assert d["cogdyn_format"] == COGDYN_FORMAT
        assert d["cogdyn_version"] == COGDYN_VERSION
        assert d["model"]["A"] is not None
        assert d["model"]["B"] is not None

    def test_load_unfitted_save_raises(self):
        dyn = CognitiveDynamics()
        with pytest.raises(RuntimeError):
            dyn.save("/tmp/nope.cogdyn")

    def test_load_missing_file_raises(self, tmp_path):
        with pytest.raises(FileNotFoundError):
            CognitiveDynamics.load(tmp_path / "nope.cogdyn")

    def test_load_wrong_format_raises(self, tmp_path):
        path = tmp_path / "bad.cogdyn"
        path.write_text(json.dumps({"cogdyn_format": "not_dynamics"}))
        with pytest.raises(ValueError):
            CognitiveDynamics.load(path)

    def test_load_unknown_version_raises(self, tmp_path):
        path = tmp_path / "future.cogdyn"
        path.write_text(json.dumps({
            "cogdyn_format": COGDYN_FORMAT,
            "cogdyn_version": "99.99",
        }))
        with pytest.raises(ValueError):
            CognitiveDynamics.load(path)


# ══════════════════════════════════════════════════════════════════
# Public API
# ══════════════════════════════════════════════════════════════════

class TestPublicAPI:

    def test_dynamics_classes_exposed(self):
        assert hasattr(styxx, "CognitiveDynamics")
        assert hasattr(styxx, "Observation")
        assert hasattr(styxx, "FitResult")
        assert hasattr(styxx, "synthetic_observations")
        assert hasattr(styxx, "thought_to_state")
        assert hasattr(styxx, "state_to_thought")
        assert hasattr(styxx, "COGDYN_FORMAT")
        assert hasattr(styxx, "COGDYN_VERSION")

    def test_version_is_3_x(self):
        assert styxx.__version__.startswith(
            ("3.", "4.", "5.", "6.", "7.", "8.", "9.")
        )

    def test_end_to_end_via_public_api(self):
        A = _stable_random_A(seed=42)
        B = _random_B(seed=43)
        obs = styxx.synthetic_observations(50, A, B, seed=80)
        dyn = styxx.CognitiveDynamics()
        result = dyn.fit(obs)
        assert isinstance(result, styxx.FitResult)
        assert result.r2 > 0.9999
        # predict + simulate + suggest end-to-end
        s = styxx.state_to_thought(np.array([0.3, 0.2, 0.1, 0.2, 0.1, 0.1]))
        a = styxx.state_to_thought(np.array([0.1, 0.5, 0.0, 0.2, 0.1, 0.1]))
        pred = dyn.predict(s, a)
        assert isinstance(pred, styxx.Thought)
        traj = dyn.simulate(s, [a, a, a])
        assert len(traj) == 4
        action = dyn.suggest(s, styxx.Thought.target("reasoning"))
        assert isinstance(action, styxx.Thought)
