"""
Test script for Bayesian Linear Regression functionality in DeepLatent.
"""

import jax.numpy as jnp
from jax import random
import sys
sys.path.append("src2")
from models import DeepLatent

# Set random seed
seed = 42
key = random.PRNGKey(seed)

print("=" * 80)
print("Testing Bayesian Linear Regression in DeepLatent")
print("=" * 80)

# Generate synthetic data
print("\n1. Generating synthetic data...")
N = 200  # Number of samples
latent_dim = 3
embedding_dim = 50

# Generate latent factors (true ideology)
key, subkey = random.split(key)
true_z = random.normal(subkey, (N, latent_dim))

# Generate text embeddings from latent factors
key, subkey = random.split(key)
W_text = random.normal(subkey, (latent_dim, embedding_dim))
text_embeddings = true_z @ W_text + random.normal(subkey, (N, embedding_dim)) * 0.1

# Generate outcome with known coefficients
key, subkey = random.split(key)
true_beta = jnp.array([2.0, -1.5, 0.8])  # True coefficients for z
true_intercept = 0.5
outcomes = true_z @ true_beta + true_intercept + random.normal(subkey, (N,)) * 0.5

print(f"   - N = {N} samples")
print(f"   - Latent dim = {latent_dim}")
print(f"   - Embedding dim = {embedding_dim}")
print(f"   - True coefficients: {true_beta}")
print(f"   - True intercept: {true_intercept}")

# Setup model with linear regression predictor
print("\n2. Setting up DeepLatent with Bayesian Linear Regression...")

modality_specs = {
    "text": {"type": "gaussian", "dim": embedding_dim}
}

predictor_specs = {
    "policy_outcome": {
        "type": "linear_regression",
        "outcome_type": "gaussian",
        "output_dim": 1,
        "beta_prior_scale": 2.0,
        "sigma_prior_scale": 1.0,
        "intercept_prior_scale": 2.0,
        "feature_names": ["ideology_1", "ideology_2", "ideology_3"]
    }
}

model = DeepLatent(
    modality_specs=modality_specs,
    predictor_specs=predictor_specs,
    latent_dim=latent_dim,
    encoder_hidden_dims=(64, 32),
    decoder_hidden_dims=(32, 64),
)

print("   - Model created successfully")

# Prepare data
data = {"text": text_embeddings}
outcomes_dict = {"policy_outcome": outcomes.reshape(-1, 1)}

# Train model
print("\n3. Training model...")
model.fit(
    data=data,
    outcomes=outcomes_dict,
    num_steps=2000,
    learning_rate=0.01,
    seed=seed
)
print("   - Training completed")

# Extract coefficients
print("\n4. Extracting coefficient estimates with credible intervals...")
coef_results = model.get_linear_coefficients(
    data=data,
    outcomes=outcomes_dict,
    num_samples=1000,
    ci_level=0.95,
    seed=seed
)
print("   - Coefficient extraction completed")

# Print results
print("\n5. Coefficient Results:")
model.print_linear_coefficients(coef_results)

# Compare with true values
print("\n6. Comparison with True Values:")
print("-" * 80)
beta_mean = coef_results["policy_outcome"]["beta_mean"]
intercept_mean = coef_results["policy_outcome"]["intercept_mean"]

print(f"{'Parameter':<20} {'True':>10} {'Estimated':>12} {'Difference':>12}")
print("-" * 80)
for i in range(latent_dim):
    true_val = true_beta[i]
    est_val = beta_mean[i, 0]
    diff = est_val - true_val
    print(f"ideology_{i+1:<14} {true_val:10.3f} {est_val:12.3f} {diff:12.3f}")

print(f"{'Intercept':<20} {true_intercept:10.3f} {float(intercept_mean):12.3f} {float(intercept_mean) - true_intercept:12.3f}")

print("\n" + "=" * 80)
print("Test completed successfully!")
print("=" * 80)
