import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import torchvision.datasets as datasets
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np

# Check for torchdiffeq, install if needed
try:
    from torchdiffeq import odeint
    USE_TORCHDIFFEQ = True
except ImportError:
    import subprocess
    subprocess.run(['pip', 'install', 'torchdiffeq', '-q'])
    from torchdiffeq import odeint
    USE_TORCHDIFFEQ = True


class PixelODEFunc(nn.Module):
    """
    ODE function where EACH PIXEL follows its own differential equation.
    
    Input X: (batch, channels, height, width) - e.g., (B, 1, 28, 28)
    Each pixel x_ij evolves according to: dx_ij/dt = f(x_ij, t, params)
    
    The hidden state tracks pixel values across time.
    """
    
    def __init__(self, hidden_dim=64):
        super().__init__()
        
        # ODE function that computes pixel dynamics
        # Each pixel value changes based on its neighbors and time
        self.fc1 = nn.Linear(2, hidden_dim)  # input: (pixel_value, time)
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        self.fc3 = nn.Linear(hidden_dim, 1)   # output: dx/dt for pixel
        
        self.norm = nn.LayerNorm(hidden_dim)
        
    def forward(self, t, state):
        """
        Args:
            t: current time (scalar)
            state: (batch, num_pixels) pixel values
        Returns:
            dx/dt for each pixel
        """
        batch_size = state.shape[0]
        
        # Combine pixel value with time
        pixel_time = torch.stack([state, t.expand(batch_size, -1)], dim=-1)
        
        # Compute derivative
        h = torch.relu(self.fc1(pixel_time))
        h = self.norm(h)
        h = torch.relu(self.fc2(h))
        dxdt = self.fc3(h)
        
        return dxdt.squeeze(-1)


class ImageODEClassifier(nn.Module):
    """
    Classifier that treats image pixels as an ODE system.
    
    Process:
    1. Flatten image to sequence of pixel values
    2. Evolve pixels through ODE from t=0 to t=T
    3. Use final states + aggregation for classification
    """
    
    def __init__(self, img_size=28, num_classes=10, t_span=(0, 3)):
        super().__init__()
        
        self.img_size = img_size
        self.num_pixels = img_size * img_size
        self.num_classes = num_classes
        self.t_span = t_span
        
        # ODE function for pixel evolution
        self.ode_func = PixelODEFunc(hidden_dim=64)
        
        # Classifier head after ODE evolution
        self.classifier = nn.Sequential(
            nn.Linear(self.num_pixels + 50, 256),  # final states + integration features
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(128, num_classes)
        )
        
        # Learnable time embedding
        self.time_embed = nn.Parameter(torch.randn(50))
        
        # Statistics collector for aggregation
        self.stats_proj = nn.Linear(50, 50)
        
    def compute_trajectory_stats(self, trajectory):
        """
        Compute statistics along the trajectory for additional features.
        trajectory: (batch, num_steps, num_pixels)
        """
        # Mean over time and pixels
        mean_t = trajectory.mean(dim=1)  # (batch, num_pixels)
        
        # Standard deviation over time
        std_t = trajectory.std(dim=1)    # (batch, num_pixels)
        
        # Max activation over time
        max_t = trajectory.max(dim=1)[0]  # (batch, num_pixels)
        
        return mean_t, std_t, max_t
    
    def forward(self, x):
        """
        Args:
            x: (batch, 1, 28, 28) MNIST images
        Returns:
            logits: (batch, num_classes)
        """
        batch_size = x.shape[0]
        
        # Flatten image to pixel sequence
        x_flat = x.view(batch_size, -1)  # (batch, 784)
        
        # Normalize to [0, 1]
        x_flat = x_flat / 255.0
        
        # Time points for ODE solver
        t_points = torch.linspace(self.t_span[0], self.t_span[1], 50, device=x.device)
        
        # Solve ODE for each pixel
        trajectory = odeint(
            self.ode_func,
            x_flat,           # initial state
            t_points,
            method='rk4'      # 4th-order Runge-Kutta
        )  # (num_steps, batch, num_pixels)
        
        # Take final state
        final_state = trajectory[-1]  # (batch, num_pixels)
        
        # Compute trajectory statistics
        mean_state, std_state, max_state = self.compute_trajectory_stats(trajectory)
        
        # Concatenate all features
        combined = torch.cat([
            final_state,
            self.time_embed.unsqueeze(0).expand(batch_size, -1)
        ], dim=-1)
        
        # Classify
        logits = self.classifier(combined)
        
        return logits


class MNISTODENet(nn.Module):
    """
    Alternative architecture: Conv features -> ODE -> Classify
    Processes spatial features through ODE.
    """
    
    def __init__(self, num_classes=10, t_span=(0, 2)):
        super().__init__()
        
        self.t_span = t_span
        
        # Convolutional feature extractor
        self.encoder = nn.Sequential(
            nn.Conv2d(1, 32, 3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2),           # 28 -> 14
            nn.Conv2d(32, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2),           # 14 -> 7
            nn.Conv2d(64, 128, 3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
        )
        
        # ODE function on feature map
        self.ode_hidden = 128
        self.ode_net = nn.Sequential(
            nn.Linear(128 + 1, 256),  # features + time
            nn.ReLU(),
            nn.Linear(256, 128),
        )
        
        # Classifier
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(128 * 7 * 7 + 50, 256),
            nn.ReLU(),
            nn.Dropout(0.4),
            nn.Linear(256, num_classes)
        )
        
        self.time_embed = nn.Parameter(torch.randn(50))
        
    def ode_func(self, t, h):
        return self.ode_net(torch.cat([h, t.expand(h.shape[0], 1)], dim=-1))
    
    def forward(self, x):
        batch_size = x.shape[0]
        
        # Extract features
        features = self.encoder(x)  # (batch, 128, 7, 7)
        
        # Flatten spatial dimensions
        h = features.view(batch_size, -1)  # (batch, 128*7*7 = 6272)
        
        # Take subset for ODE processing
        h_ode = h[:, :128]  # Use first 128 features
        
        # Time points
        t_points = torch.linspace(self.t_span[0], self.t_span[1], 50, device=x.device)
        
        # Solve ODE
        trajectory = odeint(self.ode_func, h_ode, t_points, method='rk4')
        final_h = trajectory[-1]
        
        # Combine ODE output with remaining features
        combined = torch.cat([
            final_h,
            h[:, 128:6272],  # remaining features
            self.time_embed.unsqueeze(0).expand(batch_size, -1)
        ], dim=-1)
        
        return self.classifier(combined)


def train_model(model, train_loader, test_loader, epochs=5, device='cuda'):
    """Train the ODE classifier."""
    
    model = model.to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=2, gamma=0.5)
    
    train_losses = []
    train_accs = []
    test_accs = []
    
    for epoch in range(epochs):
        model.train()
        running_loss = 0.0
        correct = 0
        total = 0
        
        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.to(device), target.to(device)
            
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
            
            running_loss += loss.item()
            _, predicted = output.max(1)
            total += target.size(0)
            correct += predicted.eq(target).sum().item()
            
            if batch_idx % 100 == 0:
                print(f'  Batch {batch_idx}/{len(train_loader)} | Loss: {loss.item():.4f}')
        
        scheduler.step()
        
        # Evaluate
        train_acc = 100. * correct / total
        test_acc = evaluate(model, test_loader, device)
        
        avg_loss = running_loss / len(train_loader)
        train_losses.append(avg_loss)
        train_accs.append(train_acc)
        test_accs.append(test_acc)
        
        print(f'Epoch {epoch+1}/{epochs} | Loss: {avg_loss:.4f} | Train Acc: {train_acc:.2f}% | Test Acc: {test_acc:.2f}%')
        print('-' * 60)
    
    return train_losses, train_accs, test_accs


def evaluate(model, test_loader, device='cuda'):
    """Evaluate model on test set."""
    model.eval()
    correct = 0
    total = 0
    
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            _, predicted = output.max(1)
            total += target.size(0)
            correct += predicted.eq(target).sum().item()
    
    return 100. * correct / total


def visualize_results(train_losses, train_accs, test_accs, model, test_loader, device):
    """Visualize training and model predictions."""
    
    fig, axes = plt.subplots(2, 2, figsize=(14, 10))
    
    # Training loss
    axes[0, 0].plot(train_losses, 'b-', linewidth=2)
    axes[0, 0].set_title('Training Loss', fontsize=14)
    axes[0, 0].set_xlabel('Epoch')
    axes[0, 0].set_ylabel('Loss')
    axes[0, 0].grid(True)
    
    # Accuracy
    axes[0, 1].plot(train_accs, 'b-', label='Train', linewidth=2)
    axes[0, 1].plot(test_accs, 'r-', label='Test', linewidth=2)
    axes[0, 1].set_title('Accuracy', fontsize=14)
    axes[0, 1].set_xlabel('Epoch')
    axes[0, 1].set_ylabel('Accuracy (%)')
    axes[0, 1].legend()
    axes[0, 1].grid(True)
    
    # Show sample predictions
    model.eval()
    samples = next(iter(test_loader))
    images, labels = samples[0][:10].to(device), samples[1][:10]
    
    with torch.no_grad():
        outputs = model(images)
        _, predictions = outputs.max(1)
    
    for i in range(10):
        axes[1, 0].subplot(2, 5, i + 1)
        axes[1, 0].imshow(images[i].cpu().squeeze(), cmap='gray')
        axes[1, 0].set_title(f'True: {labels[i].item()}\nPred: {predictions[i].item()}')
        axes[1, 0].axis('off')
    
    axes[1, 0].set_title('Sample Predictions')
    
    # ODE evolution visualization (sample image)
    sample_img = images[0:1] / 255.0
    t_points = torch.linspace(0, 3, 50, device=device)
    
    # Get trajectory for first pixel
    model.eval()
    with torch.no_grad():
        # Trace the evolution manually
        state = sample_img.view(1, -1)
        trajectory = [state.cpu().numpy()]
        for t in t_points[1:]:
            # Simple Euler step for visualization
            dxdt = model.ode_func(t.item(), state)
            state = state + dxdt * (3/50)
            trajectory.append(state.cpu().numpy())
    
    trajectory = np.array(trajectory).squeeze()
    
    # Plot pixel evolution
    for i in [100, 200, 300, 400]:
        axes[1, 1].plot(t_points.cpu().numpy(), trajectory[:, i], label=f'Pixel {i}')
    axes[1, 1].set_title('Pixel ODE Evolution (sample pixels)')
    axes[1, 1].set_xlabel('Time t')
    axes[1, 1].set_ylabel('Pixel Value')
    axes[1, 1].legend()
    axes[1, 1].grid(True)
    
    plt.tight_layout()
    plt.savefig('ode_mnist_results.png', dpi=150)
    plt.show()


def main():
    # Configuration
    DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
    BATCH_SIZE = 64
    EPOCHS = 5
    
    print(f"Using device: {DEVICE}")
    print("=" * 60)
    
    # Load MNIST
    transform = transforms.Compose([
        transforms.ToTensor(),
    ])
    
    train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
    test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
    
    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)
    test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=2)
    
    print(f"Train samples: {len(train_dataset)}, Test samples: {len(test_dataset)}")
    print("=" * 60)
    
    # Create model - choose architecture
    # Option 1: Pure pixel ODE (memory intensive)
    # model = ImageODEClassifier(img_size=28, num_classes=10, t_span=(0, 3))
    
    # Option 2: Conv + ODE (more practical)
    model = MNISTODENet(num_classes=10, t_span=(0, 2))
    
    print(f"\nModel architecture:")
    print(model)
    print("=" * 60)
    
    # Train
    print("\nStarting training...\n")
    train_losses, train_accs, test_accs = train_model(
        model, train_loader, test_loader, 
        epochs=EPOCHS, device=DEVICE
    )
    
    # Final evaluation
    print("\n" + "=" * 60)
    print("FINAL RESULTS:")
    final_acc = evaluate(model, test_loader, DEVICE)
    print(f"Test Accuracy: {final_acc:.2f}%")
    print("=" * 60)
    
    # Visualize
    visualize_results(train_losses, train_accs, test_accs, model, test_loader, DEVICE)
    
    return model, train_losses, train_accs, test_accs


if __name__ == '__main__':
    model, losses, train_acc, test_acc = main()