"""
Tests for the NFXP (Nested Fixed Point) estimation workflow.

This module reproduces the NFXP logic to verify that the Estimator 
can correctly recover parameters from synthetic data generated by 
a known Structural Model.
"""

import logging
from typing import Dict, Any

import jax.numpy as jnp
import numpy as np
import pytest

from econox import (
    Model,
    ParameterSpace,
    ValueIterationSolver,
    Estimator,
    MaximumLikelihood,
    LinearUtility,
    GumbelDistribution,
    EstimationResult
)

# =============================================================================
# Configuration
# =============================================================================

TOLERANCE_PARAMS: float = 0.15  # Tolerance for parameter recovery
TOLERANCE_LOSS: float = 1e-4

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# =============================================================================
# Fixtures
# =============================================================================

@pytest.fixture(scope="module")
def nfxp_environment() -> Dict[str, Any]:
    """
    Defines the structural environment (State Space & Transitions).
    
    Returns:
        Dict containing state_variables, transition_matrix, num_states, num_actions.
    """
    # State variables tensor (S=4, A=2, K=3)
    state_variables_tensor = jnp.array([
        [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]],  # State 0
        [[0.0, 0.5, 0.0], [1.0, 0.0, 0.5]],  # State 1
        [[0.0, 1.0, 0.0], [1.0, 0.0, 1.0]],  # State 2
        [[0.0, 1.5, 0.0], [1.0, 0.0, 1.5]],  # State 3
    ], dtype=jnp.float64)

    # Transition probabilities (S, A, S')
    transition_tensor = jnp.array([
        [[0.4, 0.3, 0.2, 0.1], [0.1, 0.2, 0.3, 0.4]],
        [[0.2, 0.3, 0.3, 0.2], [0.3, 0.2, 0.2, 0.3]],
        [[0.1, 0.2, 0.4, 0.3], [0.25, 0.25, 0.25, 0.25]],
        [[0.3, 0.2, 0.1, 0.4], [0.2, 0.4, 0.2, 0.2]],
    ], dtype=jnp.float64)
    
    num_states = 4
    num_actions = 2
    transition_matrix = transition_tensor.reshape(num_states * num_actions, num_states)

    return {
        "features": state_variables_tensor,
        "transitions": transition_matrix,
        "num_states": num_states,
        "num_actions": num_actions
    }


@pytest.fixture(scope="module")
def nfxp_model(nfxp_environment) -> Model:
    """Creates the Econox Model instance."""
    env = nfxp_environment
    return Model.from_data(
        num_states=env["num_states"],
        num_actions=env["num_actions"],
        data={"features": env["features"]},
        transitions=env["transitions"],
        num_periods=np.inf
    )


@pytest.fixture(scope="module")
def true_parameters() -> Dict[str, float]:
    """Returns the ground truth parameters used for data generation."""
    return {"p0": 0.5, "p1": 1.0, "p2": -1.0}


@pytest.fixture(scope="module")
def synthetic_data(nfxp_model, true_parameters) -> Dict[str, Any]:
    """
    Generates synthetic choice data based on the true parameters.
    
    Returns:
        Dict containing 'observations' for the estimator and 'true_probs' for verification.
    """
    # 1. Setup Logic
    utility = LinearUtility(param_keys=("p0", "p1", "p2"), feature_key="features")
    dist = GumbelDistribution()
    solver = ValueIterationSolver(discount_factor=0.9)

    # 2. Solve for True Policy
    result_true = solver.solve(true_parameters, nfxp_model, utility, dist)
    choice_prob_true = result_true.profile  # (S, A)

    # 3. Generate Samples
    n_observations = 100000
    np.random.seed(42)
    
    states = np.random.randint(0, nfxp_model.num_states, size=n_observations)
    choices = np.array([
        np.random.choice([0, 1], p=np.array(choice_prob_true[s]))
        for s in states
    ], dtype=int)

    observations = {
        "state_indices": jnp.array(states),
        "choice_indices": jnp.array(choices)
    }

    return {
        "observations": observations,
        "true_probs": choice_prob_true
    }


# =============================================================================
# Tests
# =============================================================================

def test_model_validity(nfxp_model):
    """Ensure the model is correctly initialized."""
    assert nfxp_model.num_states == 4
    assert nfxp_model.num_actions == 2
    assert nfxp_model.transitions.shape == (8, 4)


def test_solver_consistency(nfxp_model, true_parameters, synthetic_data):
    """
    Verify that the solver produces valid probabilities summing to 1.
    """
    true_probs = synthetic_data["true_probs"]
    
    # Check shape
    assert true_probs.shape == (nfxp_model.num_states, nfxp_model.num_actions)
    
    # Check probability sum constraint
    prob_sums = jnp.sum(true_probs, axis=1)
    assert jnp.allclose(prob_sums, 1.0, atol=1e-6)


def test_estimator_execution(nfxp_model, synthetic_data):
    """
    Verify that the Estimator runs successfully without errors.
    """
    initial_params = {"p0": 1.0, "p1": 2.0, "p2": -2.0}
    param_space = ParameterSpace.create(initial_params=initial_params)
    
    estimator = Estimator(
        model=nfxp_model,
        param_space=param_space,
        solver=ValueIterationSolver(discount_factor=0.9),
        utility=LinearUtility(param_keys=("p0", "p1", "p2"), feature_key="features"),
        dist=GumbelDistribution(),
        method=MaximumLikelihood(),
        verbose=False
    )

    observations = synthetic_data["observations"]
    result = estimator.fit(observations, sample_size=100000)

    assert result.success, "Estimator failed to converge."
    assert isinstance(result, EstimationResult)
    assert result.loss > 0


def test_parameter_recovery(nfxp_model, true_parameters, synthetic_data):
    """
    Verify that the Estimator recovers the true parameters within tolerance.
    """
    # Perturbed initial parameters
    initial_params = {"p0": 1.0, "p1": 2.0, "p2": -2.0}
    
    estimator = Estimator(
        model=nfxp_model,
        param_space=ParameterSpace.create(initial_params=initial_params),
        solver=ValueIterationSolver(discount_factor=0.9),
        utility=LinearUtility(param_keys=("p0", "p1", "p2"), feature_key="features"),
        dist=GumbelDistribution(),
        method=MaximumLikelihood(),
        verbose=True
    )

    observations = synthetic_data["observations"]
    result = estimator.fit(observations, sample_size=100000)

    estimated_params = result.params
    
    logger.info("True Params: %s", true_parameters)
    logger.info("Est Params:  %s", estimated_params)

    for key, true_val in true_parameters.items():
        est_val = estimated_params[key]
        diff = abs(true_val - est_val)
        assert diff < TOLERANCE_PARAMS, \
            f"Parameter {key} not recovered. True={true_val}, Est={est_val}, Diff={diff}"


def test_standard_errors(nfxp_model, synthetic_data):
    """
    Verify that standard errors are computed and are finite.
    """
    initial_params = {"p0": 0.5, "p1": 1.0, "p2": -1.0}  # Start at true for speed
    
    estimator = Estimator(
        model=nfxp_model,
        param_space=ParameterSpace.create(initial_params=initial_params),
        solver=ValueIterationSolver(discount_factor=0.9),
        utility=LinearUtility(param_keys=("p0", "p1", "p2"), feature_key="features"),
        dist=GumbelDistribution(),
        method=MaximumLikelihood(),
        verbose=False
    )

    observations = synthetic_data["observations"]
    result = estimator.fit(observations, sample_size=100000)

    assert result.std_errors is not None, "Standard errors were not computed."
    
    for key, se in result.std_errors.items():
        assert jnp.isfinite(se), f"Standard error for {key} is not finite: {se}"
        assert se > 0, f"Standard error for {key} should be positive: {se}"