#!/usr/bin/env python3
"""
World Models for Eden
Learn physics dynamics, predict future states, understand object interactions
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from tqdm import tqdm
import matplotlib.pyplot as plt
import random

# =============================================================================
# PHYSICS SIMULATION ENVIRONMENT
# =============================================================================

class SimplePhysicsWorld:
    """Simple 2D physics environment"""
    
    def __init__(self, width=100, height=100, gravity=0.5, friction=0.99):
        self.width = width
        self.height = height
        self.gravity = gravity
        self.friction = friction
        self.objects = []
    
    def add_object(self, x, y, vx, vy, mass=1.0, radius=5.0):
        """Add an object to the world"""
        self.objects.append({
            'x': x, 'y': y,
            'vx': vx, 'vy': vy,
            'mass': mass,
            'radius': radius
        })
    
    def step(self):
        """Simulate one time step"""
        for obj in self.objects:
            # Apply gravity
            obj['vy'] += self.gravity
            
            # Apply friction
            obj['vx'] *= self.friction
            obj['vy'] *= self.friction
            
            # Update position
            obj['x'] += obj['vx']
            obj['y'] += obj['vy']
            
            # Bounce off walls
            if obj['x'] < obj['radius']:
                obj['x'] = obj['radius']
                obj['vx'] = -obj['vx'] * 0.8
            elif obj['x'] > self.width - obj['radius']:
                obj['x'] = self.width - obj['radius']
                obj['vx'] = -obj['vx'] * 0.8
            
            if obj['y'] < obj['radius']:
                obj['y'] = obj['radius']
                obj['vy'] = -obj['vy'] * 0.8
            elif obj['y'] > self.height - obj['radius']:
                obj['y'] = self.height - obj['radius']
                obj['vy'] = -obj['vy'] * 0.8
        
        # Check collisions between objects
        for i in range(len(self.objects)):
            for j in range(i + 1, len(self.objects)):
                self._handle_collision(self.objects[i], self.objects[j])
    
    def _handle_collision(self, obj1, obj2):
        """Handle collision between two objects"""
        dx = obj2['x'] - obj1['x']
        dy = obj2['y'] - obj1['y']
        dist = np.sqrt(dx**2 + dy**2)
        
        min_dist = obj1['radius'] + obj2['radius']
        
        if dist < min_dist and dist > 0:
            # Normalize
            nx = dx / dist
            ny = dy / dist
            
            # Relative velocity
            dvx = obj2['vx'] - obj1['vx']
            dvy = obj2['vy'] - obj1['vy']
            
            # Relative velocity in collision normal direction
            dvn = dvx * nx + dvy * ny
            
            # Do not resolve if velocities are separating
            if dvn < 0:
                return
            
            # Collision impulse
            impulse = 2 * dvn / (obj1['mass'] + obj2['mass'])
            
            # Apply impulse
            obj1['vx'] += impulse * obj2['mass'] * nx
            obj1['vy'] += impulse * obj2['mass'] * ny
            obj2['vx'] -= impulse * obj1['mass'] * nx
            obj2['vy'] -= impulse * obj1['mass'] * ny
            
            # Separate objects
            overlap = min_dist - dist
            obj1['x'] -= overlap * 0.5 * nx
            obj1['y'] -= overlap * 0.5 * ny
            obj2['x'] += overlap * 0.5 * nx
            obj2['y'] += overlap * 0.5 * ny
    
    def get_state(self):
        """Get current state as array"""
        state = []
        for obj in self.objects:
            state.extend([obj['x'], obj['y'], obj['vx'], obj['vy']])
        return np.array(state, dtype=np.float32)
    
    def set_state(self, state):
        """Set state from array"""
        for i, obj in enumerate(self.objects):
            obj['x'] = state[i * 4]
            obj['y'] = state[i * 4 + 1]
            obj['vx'] = state[i * 4 + 2]
            obj['vy'] = state[i * 4 + 3]

# =============================================================================
# WORLD MODEL NEURAL NETWORK
# =============================================================================

class WorldModel(nn.Module):
    """
    Neural network that learns to predict physics
    
    Input: current state
    Output: next state (after one time step)
    """
    
    def __init__(self, state_dim=8, hidden_dim=256):
        super().__init__()
        
        # State encoder
        self.encoder = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.ReLU()
        )
        
        # Dynamics model (RNN to capture temporal dependencies)
        self.dynamics = nn.GRU(hidden_dim, hidden_dim, num_layers=2, batch_first=True)
        
        # State decoder
        self.decoder = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, state_dim)
        )
    
    def forward(self, state):
        """Predict next state"""
        # Encode
        encoded = self.encoder(state)
        
        # Add sequence dimension
        if len(encoded.shape) == 2:
            encoded = encoded.unsqueeze(1)
        
        # Dynamics
        dynamics_out, _ = self.dynamics(encoded)
        
        # Decode
        next_state = self.decoder(dynamics_out.squeeze(1))
        
        return next_state
    
    def rollout(self, initial_state, steps=10):
        """Predict multiple steps into the future"""
        predictions = [initial_state]
        current_state = initial_state
        
        for _ in range(steps):
            next_state = self.forward(current_state)
            predictions.append(next_state)
            current_state = next_state
        
        return torch.stack(predictions)

# =============================================================================
# DATA GENERATION
# =============================================================================

def generate_physics_data(n_trajectories=1000, n_steps=50):
    """Generate training data from physics simulations"""
    print("Generating physics simulation data...")
    
    trajectories = []
    
    for _ in tqdm(range(n_trajectories)):
        # Create random initial conditions
        world = SimplePhysicsWorld()
        
        # Add 2 objects with random positions and velocities
        for _ in range(2):
            x = np.random.uniform(20, 80)
            y = np.random.uniform(20, 80)
            vx = np.random.uniform(-5, 5)
            vy = np.random.uniform(-5, 5)
            world.add_object(x, y, vx, vy)
        
        # Simulate and record trajectory
        trajectory = []
        for _ in range(n_steps):
            state = world.get_state()
            trajectory.append(state)
            world.step()
        
        trajectories.append(np.array(trajectory))
    
    return trajectories

# =============================================================================
# TRAINING
# =============================================================================

def train_world_model(epochs=100, n_trajectories=1000):
    """Train world model on physics data"""
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"Device: {device}\n")
    
    # Generate data
    trajectories = generate_physics_data(n_trajectories=n_trajectories, n_steps=50)
    
    print(f"\nGenerated {len(trajectories)} trajectories")
    print(f"State dimension: {trajectories[0].shape[1]}\n")
    
    # Split train/test
    split = int(0.8 * len(trajectories))
    train_traj = trajectories[:split]
    test_traj = trajectories[split:]
    
    # Create model
    state_dim = trajectories[0].shape[1]
    model = WorldModel(state_dim=state_dim).to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=10)
    
    print(f"Training for {epochs} epochs...\n")
    
    best_loss = float('inf')
    
    for epoch in range(epochs):
        model.train()
        train_loss = 0
        
        # Shuffle trajectories
        random.shuffle(train_traj)
        
        for traj in tqdm(train_traj, desc=f"Epoch {epoch+1}/{epochs}"):
            # Convert to tensor
            traj_tensor = torch.tensor(traj, dtype=torch.float32).to(device)
            
            # One-step predictions
            for t in range(len(traj) - 1):
                current_state = traj_tensor[t].unsqueeze(0)
                next_state_true = traj_tensor[t + 1].unsqueeze(0)
                
                optimizer.zero_grad()
                
                # Predict next state
                next_state_pred = model(current_state)
                
                # Loss
                loss = F.mse_loss(next_state_pred, next_state_true)
                
                loss.backward()
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                optimizer.step()
                
                train_loss += loss.item()
        
        avg_train_loss = train_loss / (len(train_traj) * 49)
        
        # Test
        model.eval()
        test_loss = 0
        
        with torch.no_grad():
            for traj in test_traj:
                traj_tensor = torch.tensor(traj, dtype=torch.float32).to(device)
                
                for t in range(len(traj) - 1):
                    current_state = traj_tensor[t].unsqueeze(0)
                    next_state_true = traj_tensor[t + 1].unsqueeze(0)
                    next_state_pred = model(current_state)
                    loss = F.mse_loss(next_state_pred, next_state_true)
                    test_loss += loss.item()
        
        avg_test_loss = test_loss / (len(test_traj) * 49)
        
        scheduler.step(avg_test_loss)
        
        if (epoch + 1) % 10 == 0:
            print(f"Epoch {epoch+1}: Train Loss: {avg_train_loss:.6f}, Test Loss: {avg_test_loss:.6f}")
        
        if avg_test_loss < best_loss:
            best_loss = avg_test_loss
            torch.save(model.state_dict(), 'world_model.pth')
    
    print(f"\n✅ Best Test Loss: {best_loss:.6f}")
    return model

# =============================================================================
# TESTING
# =============================================================================

def test_world_model():
    """Test world model predictions"""
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    
    print("\n" + "="*70)
    print("TESTING WORLD MODEL")
    print("="*70)
    
    # Load model
    state_dim = 8  # 2 objects, 4 values each
    model = WorldModel(state_dim=state_dim).to(device)
    model.load_state_dict(torch.load('world_model.pth'))
    model.eval()
    
    # Create test scenario
    world = SimplePhysicsWorld()
    world.add_object(30, 30, 3, 0)
    world.add_object(70, 30, -3, 0)
    
    print("\nTest: Two objects moving toward each other")
    print("Will they collide? Will the model predict it?\n")
    
    # Get initial state
    initial_state = world.get_state()
    initial_tensor = torch.tensor(initial_state, dtype=torch.float32).unsqueeze(0).to(device)
    
    # Ground truth simulation
    true_states = [initial_state]
    for _ in range(20):
        world.step()
        true_states.append(world.get_state())
    
    # Model prediction
    with torch.no_grad():
        predicted_states = model.rollout(initial_tensor, steps=20)
        predicted_states = predicted_states.squeeze(1).cpu().numpy()
    
    # Calculate error
    errors = []
    for i in range(len(true_states)):
        error = np.mean((true_states[i] - predicted_states[i])**2)
        errors.append(error)
    
    avg_error = np.mean(errors)
    
    print(f"Average prediction error (MSE): {avg_error:.6f}")
    
    # Check if model captures collision
    # At collision, velocities should change sign
    true_vx_before = true_states[5][2]  # vx of object 1 at step 5
    true_vx_after = true_states[15][2]  # vx of object 1 at step 15
    
    pred_vx_before = predicted_states[5][2]
    pred_vx_after = predicted_states[15][2]
    
    true_collision = (true_vx_before * true_vx_after < 0)  # Sign changed
    pred_collision = (pred_vx_before * pred_vx_after < 0)
    
    print(f"\nTrue physics: {'Collision detected' if true_collision else 'No collision'}")
    print(f"Model prediction: {'Collision predicted' if pred_collision else 'No collision predicted'}")
    
    if avg_error < 10.0:
        print("\n✅ LOW ERROR - Model accurately predicts physics!")
        success = True
    elif avg_error < 50.0:
        print("\n⚠️ MEDIUM ERROR - Model captures general dynamics")
        success = True
    else:
        print("\n❌ HIGH ERROR - Model needs more training")
        success = False
    
    return success

def main():
    # Train model
    model = train_world_model(epochs=100, n_trajectories=1000)
    
    # Test model
    success = test_world_model()
    
    print("\n" + "="*70)
    print("WORLD MODELS STATUS")
    print("="*70)
    
    if success:
        print("\n✅ World Models capability: WORKING")
        print("\nEden can now:")
        print("  1. Learn physics dynamics from simulation")
        print("  2. Predict object motion")
        print("  3. Anticipate collisions")
        print("  4. Rollout future states")
        print("\n✅ CAPABILITY #5 COMPLETE")
    else:
        print("\n⚠️ Needs more training or tuning")

if __name__ == "__main__":
    main()