import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt

class WaterfallControl(nn.Module):
    """
    A simple model that we will train using the 'Waterfall' 
    Control Theory approach.
    """
    def __init__(self, input_dim=10, output_dim=1):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, 32),
            nn.Tanh(),
            nn.Linear(32, 16),
            nn.Tanh(),
            nn.Linear(16, output_dim)
        )

    def forward(self, x):
        return self.net(x)

class PIDLossController:
    """
    Control Theory implementation to ensure 'Waterfall' loss reduction.
    It monitors the slope of the loss and adjusts 'Pressure' (lr_mult).
    """
    def __init__(self, target_slope=-0.1, kp=0.1, ki=0.01, kd=0.05):
        self.target_slope = target_slope # The 'Steepness' we want
        self.kp = kp  # Proportional gain
        self.ki = ki  # Integral gain
        self.kd = kd  # Derivative gain
        
        self.prev_loss = None
        self.prev_slope = 0
        self.integral = 0

    def compute_pressure(self, current_loss):
        if self.prev_loss is None:
            self.prev_loss = current_loss
            return 1.0

        # 1. Calculate current slope (The 'Flow Rate')
        slope = current_loss - self.prev_loss
        
        # 2. Error = How far are we from our 'Waterfall' slope?
        error = slope - self.target_slope
        
        # 3. PID computation
        self.integral += error
        derivative = slope - self.prev_slope
        
        # Pressure is the multiplier for the learning rate
        # If slope is too flat (error > 0), pressure increases.
        pressure = 1.0 + (self.kp * error) + (self.ki * self.integral) + (self.kd * derivative)
        
        # Clamp pressure to prevent explosion (The 'Dam' limit)
        pressure = max(0.1, min(pressure, 10.0))
        
        self.prev_loss = current_loss
        self.prev_slope = slope
        return pressure

def train_waterfall():
    # Hyperparameters
    input_dim = 10
    batch_size = 32
    epochs = 200
    
    # Data: Simple target function y = sum(x^2)
    X = torch.randn(1000, input_dim)
    Y = torch.sum(X**2, dim=1, keepdim=True)

    model = WaterfallControl()
    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=1e-3)
    
    # The Control Engine
    controller = PIDLossController(target_slope=-0.05)
    
    loss_history = []
    
    print("Initiating Waterfall Training...")
    
    for epoch in range(epochs):
        # Batching
        permutation = torch.randperm(X.size(0))
        epoch_losses = []
        
        for i in range(0, X.size(0), batch_size):
            indices = permutation[i:i+batch_size]
            batch_x, batch_y = X[indices], Y[indices]
            
            optimizer.zero_grad()
            outputs = model(batch_x)
            loss = criterion(outputs, batch_y)
            
            # --- Control Theory Step ---
            # Calculate the 'Pressure' based on loss trajectory
            pressure = controller.compute_pressure(loss.item())
            
            # Scale the gradient by pressure to force the 'Fall'
            loss.backward()
            
            # We manually scale the gradients to simulate increased pressure
            for param in model.parameters():
                if param.grad is not None:
                    param.grad.data.mul_(pressure)
            
            optimizer.step()
            epoch_losses.append(loss.item())
            
        avg_loss = np.mean(epoch_losses)
        loss_history.append(avg_loss)
        
        if epoch % 20 == 0:
            print(f"Epoch {epoch} | Avg Loss: {avg_loss:.4f} | Pressure: {pressure:.2f}")

    return loss_history

# Run and Plot
history = train_waterfall()

plt.figure(figsize=(10, 6))
plt.plot(history, color='blue', linewidth=2)
plt.title("Loss History: Waterfall Trajectory")
plt.xlabel("Epoch")
plt.ylabel("MSE Loss")
plt.grid(True, alpha=0.3)
plt.yscale('log') # Log scale to visualize the steep fall
plt.show()