import torch
import gradio as gr
from torchvision import transforms
from huggingface_hub import hf_hub_download
import torch.nn as nn
import math


# window_partition
def window_partition(x, win):
    B, H, W, C = x.shape
    x = x.view(B,H//win,win,W//win,win,C)
    x = x.permute(0,1,3,2,4,5)
    x = x.reshape(-1,win,win,C)
    return x

# window_reverse
def window_reverse(windows, win, H, W):
    B = windows.shape[0]//(H//win * W//win)
    x = windows.view(B,H//win,W//win, win,win,-1)
    x = x.permute(0,1,3,2,4,5)
    x=x.reshape(B,H,W,-1)
    return x


# WindowAttention

class WindowAttention(nn.Module):
    def __init__(self, dim, num_heads, win):
        super().__init__()
        
        self.dim = dim
        self.num_heads = num_heads
        self.head_dim = self.dim//self.num_heads
        self.scale = self.head_dim**-0.5
        self.win = win
        
        self.q = nn.Linear(dim,dim)
        self.k = nn.Linear(dim,dim)
        self.v = nn.Linear(dim,dim)
        
        self.proj = nn.Linear(dim,dim)
        
        coords = torch.stack(torch.meshgrid(torch.arange(win), torch.arange(win), indexing = "ij"))
        coords_flat = coords.flatten(1)
        
        rel = coords_flat[:,:,None] - coords_flat[:,None,:]
        rel = rel.permute(1,2,0)
        
        rel[:,:,0] =rel[:,:,0] + (win-1)
        rel[:,:,1] =rel[:,:,1] + (win-1)
        rel[:,:,0] =rel[:,:,0] * (2*win-1)
        index = rel.sum(-1)
        
        self.register_buffer("pos_index", index)
        self.rel_bias = nn.Parameter(torch.zeros((2*win-1) * (2*win-1), num_heads))
        
    
    def forward(self,x,mask=None):
        B_,N,C =x.shape
        q = self.q(x)
        k = self.k(x)
        v = self.v(x)
        
        q = q.view(B_, N, self.num_heads, C // self.num_heads).transpose(1, 2)
        k = k.view(B_, N, self.num_heads, C // self.num_heads).transpose(1, 2)
        v = v.view(B_, N, self.num_heads, C // self.num_heads).transpose(1, 2)
        
        q = q * self.scale
        attn = q @ k.transpose(-2,-1)
        
        rb = self.rel_bias[self.pos_index.view(-1)]
        rb = rb.view(N,N, self.num_heads).permute(2,0,1)
        attn = attn + rb.unsqueeze(0)
        
        
        if mask is not None:
            nw = mask.shape[0]
            attn = attn.view(B_//nw,nw,self.num_heads,N,N) + mask.unsqueeze(1).unsqueeze(0)
            attn = attn.view(-1,self.num_heads,N,N)

        attn = attn.softmax(dim=-1)
        out = (attn @ v).transpose(1,2).reshape(B_,N,C)
        out = self.proj(out)
        return out
    


# SwinBlock

class SwinBlock(nn.Module):
    def __init__(self, dim, res, heads, win, shift):
        super().__init__()
        
        self.dim = dim
        self.res = res
        self.win = win
        self.shift = shift

        self.norm1 = nn.LayerNorm(dim)
        self.attn = WindowAttention(dim, heads, win)
        self.norm2 = nn.LayerNorm(dim)
        self.mlp = nn.Sequential(
          nn.Linear(dim, 4 * dim),
          nn.GELU(),
          nn.Linear(4 * dim, dim),
      )

        H, W = res
        if shift > 0:
            self.mask = self.create_mask(H, W, win, shift)
        else:
            self.mask = None
    
    def create_mask(self, H, W, win, shift):
        img_mask = torch.zeros((1, H, W, 1))
        count = 0

        for h in (slice(0,-win), slice(-win,-shift), slice(-shift, None)):
            for w in (slice(0,-win), slice(-win, -shift), slice(-shift, None)):
                img_mask[:,h,w,:] = count
                count += 1


        mask = window_partition(img_mask, win)
        mask = mask.view(-1, win*win)
        mask = mask.unsqueeze(1) - mask.unsqueeze(2)
        mask = mask.masked_fill(mask!=0, -10000.0)
        return mask
    
    def forward(self, x):
        B, L, C = x.shape
        H, W = self.res
        residual = x
        x = self.norm1(x)
        x = x.view(B,H,W,C)

        if self.shift > 0:
            x = torch.roll(x, shifts=(-self.shift, -self.shift), dims = (1,2))
            
        win_x = window_partition(x, self.win).view(-1, self.win*self.win, C)
            
        attn_out = self.attn(win_x, mask = self.mask.to(x.device) if self.mask is not None else None)
            
        x = window_reverse(attn_out, self.win, H, W)

        if self.shift > 0:
            x = torch.roll(x, shifts=(self.shift, self.shift), dims = (1,2))

        x = residual + x.view(B, L, C)
        residual2 = x
        x = self.norm2(x)
        x = self.mlp(x)
        return residual2 + x
    



# PatchMerging


class PatchMerging(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

        self.reduction = nn.Linear(4*dim, 2*dim, bias = False)
        self.norm = nn.LayerNorm(4*dim)

    def forward(self, x, H, W):
        B, L, C = x.shape
        x = x.view(B, H, W, C)

        x0 = x[:,0::2,0::2,:]
        x1 = x[:,1::2,0::2,:]
        x2 = x[:,0::2,1::2,:]
        x3 = x[:,1::2,1::2,:]

        x0 = x0.reshape(B, -1, C)
        x1 = x1.reshape(B, -1, C)
        x2 = x2.reshape(B, -1, C)
        x3 = x3.reshape(B, -1, C)

        x = torch.cat([x0,x1,x2,x3], -1)

        x = self.norm(x)
        x= self.reduction(x)
        return x, H//2, W//2
    


# SwinTransformer

class SwinTransformer(nn.Module):
      def __init__(self, embed_dim = 96, win = 7 , num_classes = 45):
        super().__init__()
        self.patch_embed = nn.Conv2d(3, embed_dim, kernel_size= 4, stride = 4)
        self.patch_norm = nn.LayerNorm(embed_dim)       
        
        
        dims  = [embed_dim, embed_dim*2, embed_dim*4, embed_dim*8]
        heads = [3, 6, 12, 24]
        shift = win//2

        #stage_1

        self.stage1_blocks =  nn.Sequential(
            SwinBlock(dims[0], (56,56), heads = heads[0], win = win, shift = 0),
            SwinBlock(dims[0], (56,56), heads = heads[0], win = win, shift = shift)
        )

        self.patch_merge1 = PatchMerging(dims[0])

        #stage_2
        self.stage2_blocks = nn.Sequential(
            SwinBlock(dims[1], (28,28), heads = heads[1], win = win, shift = 0),
            SwinBlock(dims[1], (28,28), heads = heads[1], win = win, shift = shift)
        )
        
        self.patch_merge2 = PatchMerging(dims[1])
        
        #stage_3
        self.stage3_blocks = nn.Sequential(*[
            SwinBlock(dims[2], (14,14), heads = heads[2], win = win, shift = 0 if i%2 == 0 else shift)
            for i in range(6)
        ])
        
        self.patch_merge3 = PatchMerging(dims[2])
        
        #stage_4
        self.stage4_blocks = nn.Sequential(
            SwinBlock(dims[3], (7,7), heads = heads[3], win = win, shift = 0),
            SwinBlock(dims[3], (7,7), heads = heads[3], win = win, shift = 0)
        )
        
        self.norm =nn.LayerNorm(dims[3])
        self.fc = nn.Linear(dims[3], num_classes)

  
      def forward(self, x):
        x = self.patch_embed(x)                  # (B, 96, 56, 56)
        B, C, H, W = x.shape
        x = self.patch_norm(x.flatten(2).transpose(1,2))  # (B, 3136, 96)

        x = self.stage1_blocks(x)
        x, H, W = self.patch_merge1(x, H, W)    # (B, 784, 192)

        x = self.stage2_blocks(x)
        x, H, W = self.patch_merge2(x, H, W)    # (B, 196, 384)

        x = self.stage3_blocks(x)
        x, H, W = self.patch_merge3(x, H, W)    # (B, 49, 768)

        x = self.stage4_blocks(x)
        x = self.norm(x).mean(dim=1)
        return self.fc(x)
      


MODEL_REPO = "Sathya77/swin-transformer-satellite"

checkpoint = torch.load(
    hf_hub_download(MODEL_REPO, "swin_resisc45.pth"),
    map_location='cpu'
)

model = SwinTransformer(
    embed_dim   = 96,
    num_classes = checkpoint['num_classes']
)
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()

classes = checkpoint['classes']


transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406],
                         [0.229, 0.224, 0.225]),
])


def predict(image):
    x = transform(image).unsqueeze(0)
    with torch.no_grad():
        out  = model(x)
        prob = torch.softmax(out, dim=1)
    top3_probs, top3_idx = prob.topk(3, dim=1)
    return {
        classes[top3_idx[0][i].item()]: float(top3_probs[0][i])
        for i in range(3)
    }



#gradio for hugging_face

import gradio as gr

with gr.Blocks(theme=gr.themes.Soft(), title="Swin Transformer") as demo:

    gr.Markdown("""
    # 🛰️ Swin Transformer — Satellite Image Classification
    Trained on NWPU-RESISC45 | 45 land use categories | 82% test accuracy
    
    > Upload a satellite image or click an example below to classify it.
    """)

    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(
                type="pil",
                label="Input Satellite Image",
                height=500
            )
            with gr.Row():
                clear_btn  = gr.Button("Clear",  variant="secondary")
                submit_btn = gr.Button("Classify ✈️", variant="primary")

        with gr.Column(scale=1):
            label_output = gr.Label(
                num_top_classes=5,      # show top 5 instead of 3
                label="Predictions"
            )

    gr.Markdown("### 📷 Example Images")
    gr.Examples(
        examples=[
            "samples/airplane_062.jpg",
            "samples/airport_048.jpg",
            "samples/baseball_diamond_047.jpg",
            "samples/bridge_047.jpg",
            "samples/chaparral_045.jpg",
            "samples/church_031.jpg",
            "samples/desert_062.jpg",
            "samples/golf_course_062.jpg",
            "samples/intersection_033.jpg",
            "samples/medium_residential_073.jpg",
            "samples/palace_090.jpg",
            "samples/ship_036.jpg",
        ],
        inputs=image_input,
        cache_examples=False
    )

    gr.Markdown("""
    ---
    ### 📊 About this model
    | Property | Value |
    |---|---|
    | Architecture | Swin Transformer (full 4 stages) |
    | Dataset | NWPU-RESISC45 |
    | Classes | 45 land use categories |
    | Test Accuracy | 81% |
    | Framework | PyTorch |
    
    ### 📄 References
    - [Swin Transformer Paper](https://arxiv.org/abs/2103.14030) — Liu et al. 2021
    - [NWPU-RESISC45 Dataset](https://www.tensorflow.org/datasets/catalog/resisc45)
    """)

    # button actions
    submit_btn.click(fn=predict, inputs=image_input, outputs=label_output)
    clear_btn.click(fn=lambda: (None, None), outputs=[image_input, label_output])

demo.launch()