import gradio as gr
import numpy as np
import torch
import torch.nn as nn
from torch.nn.utils import spectral_norm
from PIL import Image

device = torch.device('cpu')
print(f"Device: {device}")

# ── SR Model Architecture ─────────────────────────────────
class ResidualBlock(nn.Module):
    def __init__(self, ch):
        super().__init__()
        self.conv1 = nn.Conv2d(ch, ch, 3, padding=1)
        self.bn1   = nn.BatchNorm2d(ch)
        self.conv2 = nn.Conv2d(ch, ch, 3, padding=1)
        self.bn2   = nn.BatchNorm2d(ch)
        self.relu  = nn.PReLU()
    def forward(self, x):
        return x + self.bn2(self.conv2(self.relu(self.bn1(self.conv1(x)))))

class SRNet(nn.Module):
    def __init__(self, num_res_blocks=8, channels=64):
        super().__init__()
        self.entry      = nn.Sequential(nn.Conv2d(1, channels, 9, padding=4), nn.PReLU())
        self.res_blocks = nn.Sequential(*[ResidualBlock(channels) for _ in range(num_res_blocks)])
        self.mid_conv   = nn.Sequential(nn.Conv2d(channels, channels, 3, padding=1), nn.BatchNorm2d(channels))
        self.upsample   = nn.Sequential(nn.Conv2d(channels, channels*4, 3, padding=1), nn.PixelShuffle(2), nn.PReLU())
        self.exit       = nn.Conv2d(channels, 1, 9, padding=4)
        self.sigmoid    = nn.Sigmoid()
    def forward(self, x):
        e = self.entry(x)
        return self.sigmoid(self.exit(self.upsample(self.mid_conv(self.res_blocks(e)) + e)))

# ── Colorization Model Architecture ──────────────────────
def sn_conv(i, o, k, s=1, p=0):
    return spectral_norm(nn.Conv2d(i, o, k, s, p))

class DownBlock(nn.Module):
    def __init__(self, i, o, norm=True):
        super().__init__()
        layers = [sn_conv(i, o, 4, s=2, p=1)]
        if norm: layers.append(nn.InstanceNorm2d(o))
        layers.append(nn.LeakyReLU(0.2, True))
        self.block = nn.Sequential(*layers)
    def forward(self, x): return self.block(x)

class UpBlock(nn.Module):
    def __init__(self, i, o, drop=False):
        super().__init__()
        self.up   = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False)
        self.conv = nn.Sequential(sn_conv(i, o, 3, s=1, p=1), nn.InstanceNorm2d(o), nn.ReLU(True))
        self.drop = nn.Dropout(0.5) if drop else nn.Identity()
    def forward(self, x): return self.drop(self.conv(self.up(x)))

class UNetGenerator(nn.Module):
    def __init__(self):
        super().__init__()
        self.e1 = nn.Sequential(sn_conv(1, 64, 4, s=2, p=1), nn.LeakyReLU(0.2, True))
        self.e2 = DownBlock(64,  128); self.e3 = DownBlock(128, 256)
        self.e4 = DownBlock(256, 512); self.e5 = DownBlock(512, 512)
        self.e6 = DownBlock(512, 512); self.e7 = DownBlock(512, 512)
        self.e8 = DownBlock(512, 512, norm=False)
        self.d1 = UpBlock(512,  512, drop=True); self.d2 = UpBlock(1024, 512, drop=True)
        self.d3 = UpBlock(1024, 512, drop=True); self.d4 = UpBlock(1024, 512)
        self.d5 = UpBlock(1024, 256);             self.d6 = UpBlock(512,  128)
        self.d7 = UpBlock(256,  64)
        self.final = nn.Sequential(
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
            sn_conv(128, 3, 3, s=1, p=1), nn.Tanh()
        )
    def forward(self, x):
        e1=self.e1(x); e2=self.e2(e1); e3=self.e3(e2); e4=self.e4(e3)
        e5=self.e5(e4); e6=self.e6(e5); e7=self.e7(e6); e8=self.e8(e7)
        d1=self.d1(e8)
        d2=self.d2(torch.cat([d1,e7],1)); d3=self.d3(torch.cat([d2,e6],1))
        d4=self.d4(torch.cat([d3,e5],1)); d5=self.d5(torch.cat([d4,e4],1))
        d6=self.d6(torch.cat([d5,e3],1)); d7=self.d7(torch.cat([d6,e2],1))
        return self.final(torch.cat([d7,e1],1))

# ── Load models ───────────────────────────────────────────
print("Loading SR model...")
sr_model = SRNet().to(device)
sr_model.load_state_dict(torch.load('sr_model_best_psnr.pth', map_location=device))
sr_model.eval()
print("SR model loaded ✓")

print("Loading colorization model...")
G_color = UNetGenerator().to(device)
G_color.load_state_dict(torch.load('color_G_best.pth', map_location=device))
G_color.eval()
print("Colorization model loaded ✓")

# ── Inference function ────────────────────────────────────
def colorize(npy_file):
    if npy_file is None:
        return None

    # Load the .npy file
    tir = np.load(npy_file.name).astype(np.float32)
    print(f"Input shape: {tir.shape}, min={tir.min():.4f}, max={tir.max():.4f}")

    # Validate shape
    if tir.ndim == 2:
        tir = tir[np.newaxis]  # add channel dim if missing
    if tir.shape != (1, 256, 256):
        return None  # wrong shape

    # Normalize to [0, 1]
    lo, hi = tir.min(), tir.max()
    tir_norm = (tir - lo) / (hi - lo + 1e-8)

    with torch.no_grad():
        # Stage A: Super-Resolution (256x256 → 512x512)
        tir_sr = sr_model(
            torch.from_numpy(tir_norm).unsqueeze(0).to(device)
        ).squeeze(0).cpu().numpy()

        # Stage B: Colorization (TIR → RGB)
        rgb_fake = G_color(
            torch.from_numpy(tir_sr).unsqueeze(0).to(device)
        ).squeeze(0).cpu().numpy()

    # Convert from [-1,1] to [0,255] uint8
    rgb = np.clip((rgb_fake + 1) / 2, 0, 1)
    rgb_uint8 = (rgb * 255).astype(np.uint8).transpose(1, 2, 0)  # CHW → HWC

    return Image.fromarray(rgb_uint8, 'RGB')

# ── Gradio UI ─────────────────────────────────────────────
demo = gr.Interface(
    fn=colorize,
    inputs=gr.File(label="Upload tir_200m.npy patch (shape: 1×256×256)"),
    outputs=gr.Image(label="Colorized RGB Output (512×512)", type="pil"),
    title="🛰️ IR2RGB-Net — Thermal Infrared to RGB Colorization",
    description="""
**Team VoidBit | BAH2026 PS10**

Upload a Landsat 9 thermal infrared patch (`.npy` file, shape `1×256×256`) 
and the model will:
1. **Super-resolve** it from 200m → 100m resolution (256×256 → 512×512)
2. **Colorize** it into a synthetic RGB representation

The output is a 512×512 RGB image showing an interpretable color version of the thermal data.
    """,
    flagging_mode="never"
)

demo.launch()