import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np

# Set deterministic seeds for ML reproducibility
torch.manual_seed(369)
np.random.seed(369)

class AQARIONWorldModel(nn.Module):
    """
    A foundational world model integrating the AQARION Defect Regularizer.
    """
    def __init__(self, obs_dim: int, latent_dim: int, num_clusters: int):
        super().__init__()
        # Encoder: Observation -> Continuous Latent Space
        self.encoder = nn.Sequential(
            nn.Linear(obs_dim, 64),
            nn.ReLU(),
            nn.Linear(64, latent_dim)
        )
        
        # AQARION Defect Loss Module (from the theoretical foundation)
        self.defect_regularizer = AQARIONDefectLoss(num_clusters, latent_dim, tau=0.5)
        
        # Decoder: Continuous Latent Space -> Reconstructed Observation
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, 64),
            nn.ReLU(),
            nn.Linear(64, obs_dim)
        )

    def forward(self, obs_t, obs_t_next):
        # 1. Encode into continuous latent representations
        z_t = self.encoder(obs_t)
        z_t_next = self.encoder(obs_t_next)
        
        # 2. Evaluate Differentiable Defect Loss
        loss_defect, active_clusters = self.defect_regularizer(z_t, z_t_next)
        
        # 3. Standard Reconstruction
        recon_t = self.decoder(z_t)
        
        return recon_t, loss_defect, active_clusters

def run_ml_reproducibility_benchmark():
    print("Initializing AQARION Track 4 ML Benchmark...")
    
    # Mock high-dimensional continuous data (e.g., noisy sensors tracking the torus)
    batch_size, obs_dim, latent_dim, num_clusters = 256, 128, 16, 6
    obs_t = torch.randn(batch_size, obs_dim)
    
    # Simulate a noisy transition
    transition_noise = torch.randn(batch_size, obs_dim) * 0.1
    obs_t_next = obs_t + transition_noise 
    
    model = AQARIONWorldModel(obs_dim, latent_dim, num_clusters)
    optimizer = optim.Adam(model.parameters(), lr=1e-3)
    
    # Weight of the AQARION regularizer
    beta = 10.0 
    
    model.train()
    optimizer.zero_grad()
    
    recon_t, loss_defect, clusters = model(obs_t, obs_t_next)
    
    # Standard MSE + AQARION Defect Regularization
    loss_recon = nn.MSELoss()(recon_t, obs_t)
    total_loss = loss_recon + beta * loss_defect
    
    total_loss.backward()
    optimizer.step()
    
    print(f"✅ Epoch 1 - Total Loss: {total_loss.item():.4f}")
    print(f"🔍 Defect Penalty (||D_P||_F^2): {loss_defect.item():.6f}")
    print(f"📊 Active Latent Clusters Utilized: {len(torch.unique(clusters))}/{num_clusters}")
    
    return {
        "benchmark": "AQ-ML-TRACK4-DEFECT-LOSS",
        "initial_defect_loss": float(loss_defect.item()),
        "status": "PASS"
    }

if __name__ == "__main__":
    run_ml_reproducibility_benchmark()