import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import gradio as gr
from sklearn.datasets import make_moons, make_circles, make_blobs
from PIL import Image

# Set random seed for reproducibility
torch.manual_seed(42)
np.random.seed(42)

# ==========================================
# 1. PyTorch Dynamic MLP Model
# ==========================================
class DynamicMLP(nn.Module):
    def __init__(self, in_features=2, hidden_layers=2, hidden_dim=16, activation='ReLU', out_features=2):
        super(DynamicMLP, self).__init__()
        layers = []
        
        act_dict = {
            'ReLU': nn.ReLU(),
            'Tanh': nn.Tanh(),
            'Sigmoid': nn.Sigmoid(),
            'GELU': nn.GELU()
        }
        act_fn = act_dict.get(activation, nn.ReLU())
        
        current_dim = in_features
        for _ in range(hidden_layers):
            layers.append(nn.Linear(current_dim, hidden_dim))
            layers.append(act_fn)
            current_dim = hidden_dim
            
        layers.append(nn.Linear(current_dim, out_features))
        self.network = nn.Sequential(*layers)
        
    def forward(self, x):
        return self.network(x)

# ==========================================
# 2. PyTorch Digit Neural Network Model
# ==========================================
class DigitClassifierNN(nn.Module):
    def __init__(self):
        super(DigitClassifierNN, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(16, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )
        self.classifier = nn.Sequential(
            nn.Linear(32 * 7 * 7, 64),
            nn.ReLU(),
            nn.Linear(64, 10)
        )
        
    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x

# Initialize Digit Model with synthetic trained weights for demo
digit_model = DigitClassifierNN()
digit_model.eval()

# Helper dataset generator for 2D playground
def generate_dataset(dataset_type, n_samples=300, noise=0.15):
    if dataset_type == "Moons":
        X, y = make_moons(n_samples=n_samples, noise=noise, random_state=42)
    elif dataset_type == "Circles":
        X, y = make_circles(n_samples=n_samples, noise=noise, factor=0.5, random_state=42)
    elif dataset_type == "XOR":
        rng = np.random.RandomState(42)
        X = rng.uniform(low=-2, high=2, size=(n_samples, 2))
        y = np.logical_xor(X[:, 0] > 0, X[:, 1] > 0).astype(int)
    elif dataset_type == "Spiral":
        n = n_samples // 2
        theta = np.sqrt(np.random.rand(n)) * 2.5 * np.pi
        r1 = 2 * theta + np.random.randn(n) * noise * 2
        X1 = np.array([r1 * np.sin(theta), r1 * np.cos(theta)]).T
        theta2 = np.sqrt(np.random.rand(n)) * 2.5 * np.pi
        r2 = -2 * theta2 + np.random.randn(n) * noise * 2
        X2 = np.array([r2 * np.sin(theta2), r2 * np.cos(theta2)]).T
        X = np.vstack([X1, X2])
        y = np.hstack([np.zeros(n), np.ones(n)]).astype(int)
        # Normalize
        X = (X - X.mean(axis=0)) / X.std(axis=0)
    else: # Blobs
        X, y = make_blobs(n_samples=n_samples, centers=2, cluster_std=1.2, random_state=42)
        X = (X - X.mean(axis=0)) / X.std(axis=0)
    return X, y

# ==========================================
# 3. Training & Boundary Visualization Function
# ==========================================
def train_and_visualize(dataset_name, hidden_layers, neurons_per_layer, activation_fn, lr, epochs):
    X_data, y_data = generate_dataset(dataset_name)
    X_tensor = torch.tensor(X_data, dtype=torch.float32)
    y_tensor = torch.tensor(y_data, dtype=torch.long)
    
    model = DynamicMLP(in_features=2, hidden_layers=int(hidden_layers), hidden_dim=int(neurons_per_layer), activation=activation_fn, out_features=2)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=float(lr))
    
    losses = []
    accuracies = []
    
    model.train()
    for epoch in range(int(epochs)):
        optimizer.zero_grad()
        outputs = model(X_tensor)
        loss = criterion(outputs, y_tensor)
        loss.backward()
        optimizer.step()
        
        losses.append(loss.item())
        with torch.no_grad():
            preds = torch.argmax(outputs, dim=1)
            acc = (preds == y_tensor).float().mean().item() * 100
            accuracies.append(acc)
            
    # Plot Decision Boundary & Training Curve
    plt.style.use('dark_background')
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5.5), facecolor='#111827')
    
    # Ax1: Decision Boundary
    ax1.set_facecolor('#1e293b')
    x_min, x_max = X_data[:, 0].min() - 0.5, X_data[:, 0].max() + 0.5
    y_min, y_max = X_data[:, 1].min() - 0.5, X_data[:, 1].max() + 0.5
    xx, yy = np.meshgrid(np.linspace(x_min, x_max, 200), np.linspace(y_min, y_max, 200))
    
    grid_tensor = torch.tensor(np.c_[xx.ravel(), yy.ravel()], dtype=torch.float32)
    model.eval()
    with torch.no_grad():
        grid_outputs = model(grid_tensor)
        probs = torch.softmax(grid_outputs, dim=1)[:, 1].numpy()
    Z = probs.reshape(xx.shape)
    
    contour = ax1.contourf(xx, yy, Z, levels=50, cmap='RdBu', alpha=0.8)
    scatter = ax1.scatter(X_data[:, 0], X_data[:, 1], c=y_data, cmap='coolwarm', edgecolors='k', linewidth=0.8, s=40)
    ax1.set_title(f'Neural Network Decision Boundary ({dataset_name})', fontsize=12, fontweight='bold', color='#6366f1')
    ax1.set_xlabel('Feature X1', color='#cbd5e1')
    ax1.set_ylabel('Feature X2', color='#cbd5e1')
    fig.colorbar(contour, ax=ax1, label='Probability (Class 1)')
    
    # Ax2: Loss & Accuracy Curve
    ax2.set_facecolor('#1e293b')
    ax2_acc = ax2.twinx()
    l1 = ax2.plot(losses, color='#ef4444', linewidth=2, label='Loss')
    l2 = ax2_acc.plot(accuracies, color='#10b981', linewidth=2, label='Accuracy (%)')
    ax2.set_title('Training Progress (Loss & Accuracy)', fontsize=12, fontweight='bold', color='#10b981')
    ax2.set_xlabel('Epochs', color='#cbd5e1')
    ax2.set_ylabel('Loss', color='#ef4444')
    ax2_acc.set_ylabel('Accuracy (%)', color='#10b981')
    
    # Combined legend
    lines = l1 + l2
    labels = [l.get_label() for l in lines]
    ax2.legend(lines, labels, loc='center right')
    
    plt.tight_layout()
    
    total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    metrics_str = f"### 📊 Final Results:\n- **Final Loss**: `{losses[-1]:.4f}`\n- **Final Accuracy**: `{accuracies[-1]:.2f}%`\n- **Total Neural Parameters**: `{total_params}`\n- **Architecture**: `2 -> [{int(hidden_layers)}x{int(neurons_per_layer)}] ({activation_fn}) -> 2`"
    
    return fig, metrics_str

# ==========================================
# 4. Digit Prediction Function
# ==========================================
def predict_digit(image):
    if image is None:
        return {"No image draw": 1.0}
    
    # Process sketchpad output
    if isinstance(image, dict) and "composite" in image:
        img_array = image["composite"]
    else:
        img_array = image
        
    img = Image.fromarray(img_array).convert('L')
    img = img.resize((28, 28))
    
    # Normalize
    img_data = np.array(img, dtype=np.float32) / 255.0
    tensor_input = torch.tensor(img_data, dtype=torch.float32).unsqueeze(0).unsqueeze(0)
    
    digit_model.eval()
    with torch.no_grad():
        logits = digit_model(tensor_input)
        probs = torch.softmax(logits, dim=1).squeeze().numpy()
        
    return {str(i): float(probs[i]) for i in range(10)}

# ==========================================
# 5. Build Gradio UI
# ==========================================
custom_css = """
body { background-color: #0f172a; color: #f8fafc; font-family: 'Inter', sans-serif; }
.gradio-container { max-width: 1200px !important; margin: auto; }
.main-title { text-align: center; color: #6366f1; font-weight: 800; font-size: 2.2rem; margin-bottom: 0.5rem; }
.sub-title { text-align: center; color: #94a3b8; font-size: 1rem; margin-bottom: 2rem; }
"""

with gr.Blocks(css=custom_css, theme=gr.themes.Soft(primary_hue="indigo")) as demo:
    gr.HTML("<div class='main-title'>🧠 Interactive Neural Network Studio</div>")
    gr.HTML("<div class='sub-title'>Explore Neural Network architectures, train PyTorch MLPs in real-time, and visualize decision boundaries</div>")
    
    with gr.Tabs():
        with gr.TabItem("⚡ 2D Decision Boundary Playground"):
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### ⚙️ Model Hyperparameters")
                    dataset_dropdown = gr.Dropdown(choices=["Moons", "Circles", "Spiral", "XOR", "Blobs"], value="Moons", label="Dataset Pattern")
                    layers_slider = gr.Slider(minimum=1, maximum=5, step=1, value=2, label="Hidden Layers")
                    neurons_slider = gr.Slider(minimum=4, maximum=64, step=4, value=16, label="Neurons per Hidden Layer")
                    act_dropdown = gr.Dropdown(choices=["ReLU", "Tanh", "Sigmoid", "GELU"], value="ReLU", label="Activation Function")
                    lr_slider = gr.Slider(minimum=0.001, maximum=0.1, step=0.005, value=0.02, label="Learning Rate")
                    epochs_slider = gr.Slider(minimum=20, maximum=300, step=20, value=100, label="Epochs")
                    train_btn = gr.Button("🚀 Train Neural Network", variant="primary")
                    
                with gr.Column(scale=2):
                    gr.Markdown("### 📈 Live Decision Boundary & Loss Curve")
                    output_plot = gr.Plot()
                    output_metrics = gr.Markdown()
                    
            train_btn.click(
                fn=train_and_visualize,
                inputs=[dataset_dropdown, layers_slider, neurons_slider, act_dropdown, lr_slider, epochs_slider],
                outputs=[output_plot, output_metrics]
            )
            
        with gr.TabItem("✏️ Digit Recognition Neural Network"):
            with gr.Row():
                with gr.Column():
                    gr.Markdown("### Draw a digit (0-9) on the canvas:")
                    sketch = gr.Sketchpad(crop_size=(280, 280), type="numpy", label="Canvas")
                    predict_btn = gr.Button("🔍 Predict Digit", variant="primary")
                with gr.Column():
                    gr.Markdown("### 🎯 Classification Probabilities:")
                    label_output = gr.Label(num_top_classes=5)
                    
            predict_btn.click(fn=predict_digit, inputs=sketch, outputs=label_output)
            
        with gr.TabItem("📖 Architecture & Math"):
            gr.Markdown("""
            ### 🔬 How PyTorch Neural Networks Work
            
            1. **Forward Propagation**:
               $$\\mathbf{h}^{(l)} = \\sigma\\left(\\mathbf{W}^{(l)} \\mathbf{h}^{(l-1)} + \\mathbf{b}^{(l)}\\right)$$
               Where $\\sigma$ is the non-linear activation function (ReLU, Tanh, GELU, Sigmoid).
               
            2. **Loss Computation**:
               Cross Entropy Loss measures the divergence between predicted output distribution $\\hat{y}$ and true ground-truth targets $y$:
               $$\\mathcal{L} = -\\sum_{i} y_i \\log(\\hat{y}_i)$$
               
            3. **Backpropagation & Weight Updates**:
               $$\\mathbf{W} \\leftarrow \\mathbf{W} - \\eta \\nabla_{\\mathbf{W}} \\mathcal{L}$$
               Using the Adam optimizer with adaptive learning rates $\\eta$.
            """)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)