import gradio as gr
import numpy as np
import torch
import torch.nn as nn
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

class SentinelNoiseSchedule:
    def __init__(self, timesteps=1000, z=2.0):
        self.timesteps = timesteps
        self.z = z
        self.betas = self._sentinel_schedule()
        self.alphas = 1.0 - self.betas
        self.alpha_bars = torch.cumprod(self.alphas, dim=0)
    
    def _sentinel_schedule(self):
        n = torch.arange(1, self.timesteps + 1, dtype=torch.float64)
        t_norm = n / self.timesteps
        beta = torch.zeros_like(n)
        for i in range(self.timesteps):
            t = t_norm[i].item()
            if t < 0.5:
                beta[i] = 0.0001 + 0.01 * (2 * t) ** (1 / (2 * t + 0.01))
            else:
                beta[i] = 0.01 + 0.02 * ((2 * t - 1) ** (2 * t - 1))
        return torch.clamp(beta, 0.0001, 0.999).float()
    
    def add_noise(self, x, t):
        sqrt_alpha_bar = torch.sqrt(self.alpha_bars[t])
        sqrt_one_minus = torch.sqrt(1.0 - self.alpha_bars[t])
        noise = torch.randn_like(x)
        return sqrt_alpha_bar.view(-1,1,1,1) * x + sqrt_one_minus.view(-1,1,1,1) * noise, noise

def visualize_schedule(timesteps, z):
    """Visualize Sentinel noise schedule."""
    schedule = SentinelNoiseSchedule(timesteps, z)
    
    fig, axes = plt.subplots(1, 3, figsize=(15, 4))
    
    t = np.arange(timesteps)
    axes[0].plot(t, schedule.betas.numpy(), linewidth=2, color='purple')
    axes[0].set_title('Sentinel β Schedule (Super-Exponential)')
    axes[0].set_xlabel('Timestep')
    axes[0].set_ylabel('β')
    axes[0].grid(True, alpha=0.3)
    
    axes[1].plot(t, schedule.alpha_bars.numpy(), linewidth=2, color='blue')
    axes[1].set_title('ᾱ (Cumulative Product)')
    axes[1].set_xlabel('Timestep')
    axes[1].set_ylabel('ᾱ')
    axes[1].grid(True, alpha=0.3)
    
    # Compare with cosine schedule
    cos_betas = np.cos(np.linspace(0, np.pi/2, timesteps)) ** 2 * 0.02
    axes[2].plot(t, schedule.betas.numpy(), label='Sentinel', linewidth=2, color='purple')
    axes[2].plot(t, cos_betas, label='Cosine', linewidth=2, color='orange', linestyle='--')
    axes[2].set_title('Schedule Comparison')
    axes[2].set_xlabel('Timestep')
    axes[2].set_ylabel('β')
    axes[2].legend()
    axes[2].grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('/tmp/diffusion_sched.png', dpi=150)
    plt.close()
    
    return '/tmp/diffusion_sched.png'

def add_noise_demo(image_size, timesteps, step, z):
    """Demo noise addition on synthetic image."""
    schedule = SentinelNoiseSchedule(timesteps, z)
    
    # Create synthetic image (colored pattern)
    img = torch.zeros(1, 3, image_size, image_size)
    for c in range(3):
        for i in range(image_size):
            for j in range(image_size):
                img[0, c, i, j] = np.sin(i * 0.3 + c) * np.cos(j * 0.3 + c) * 0.5 + 0.5
    
    t = torch.tensor([step])
    noisy_img, noise = schedule.add_noise(img, t)
    
    fig, axes = plt.subplots(1, 3, figsize=(15, 5))
    
    def show_tensor(ax, tensor, title):
        arr = tensor[0].permute(1, 2, 0).numpy()
        arr = np.clip(arr, 0, 1)
        ax.imshow(arr)
        ax.set_title(title)
        ax.axis('off')
    
    show_tensor(axes[0], img, 'Original Image')
    show_tensor(axes[1], noisy_img, f'Noisy (t={step}, β={schedule.betas[step]:.4f})')
    show_tensor(axes[2], noise * 0.3 + 0.5, 'Noise (scaled)')
    
    plt.tight_layout()
    plt.savefig('/tmp/diffusion_noise.png', dpi=150)
    plt.close()
    
    info = f"""
## Sentinel Diffusion Noise Addition

| Property | Value |
|----------|-------|
| Timestep | {step}/{timesteps} |
| β (noise level) | {schedule.betas[step]:.6f} |
| ᾱ (signal retained) | {schedule.alpha_bars[step]:.6f} |
| Schedule type | **Super-exponential** |

### Key Innovation
Sentinel noise schedule uses **super-exponential growth** of β:
- Early steps: small noise (preserve structure)
- Late steps: rapid increase (destroy structure)
- Sharper transitions than cosine/linear schedules
"""
    return '/tmp/diffusion_noise.png', info

with gr.Blocks(title="Sentinel Diffusion Model") as demo:
    gr.Markdown("""
    # 🎨 Sentinel Diffusion Model
    
    **Super-exponential noise schedule for sharper transitions.**
    
    The Sentinel partition function F(z) = Σ zⁿ/nⁿ inspires a noise schedule
    with super-exponential β growth — potentially requiring fewer steps.
    """)
    
    with gr.Tab("Noise Schedule"):
        with gr.Row():
            ts_sched = gr.Slider(100, 2000, value=1000, step=100, label="Timesteps")
            z_sched = gr.Slider(0.5, 5.0, value=2.0, label="z Parameter")
        btn_sched = gr.Button("Visualize Schedule", variant="primary")
        img_sched = gr.Image()
        btn_sched.click(visualize_schedule, [ts_sched, z_sched], img_sched)
    
    with gr.Tab("Noise Addition Demo"):
        with gr.Row():
            img_size = gr.Slider(16, 128, value=64, step=16, label="Image Size")
            ts_noise = gr.Slider(100, 2000, value=1000, step=100, label="Total Timesteps")
            step_noise = gr.Slider(0, 999, value=500, label="Current Step")
            z_noise = gr.Slider(0.5, 5.0, value=2.0, label="z Parameter")
        btn_noise = gr.Button("Add Noise", variant="primary")
        img_noise = gr.Image()
        info_noise = gr.Markdown()
        btn_noise.click(add_noise_demo, [img_size, ts_noise, step_noise, z_noise], [img_noise, info_noise])
    
    gr.Markdown("""
    ## About Sentinel Diffusion
    
    - **Noise schedule**: Super-exponential β growth (from partition function)
    - **Transition**: Sharper than cosine/linear (phase-like)
    - **Structure preservation**: Strong early, weak late
    - **Potential**: Fewer diffusion steps needed
    
    [Model Repo](https://huggingface.co/5dimension/sentinel-diffusion)
    """)

if __name__ == "__main__":
    demo.launch()