import gradio as gr
import torch
import torch.nn as nn
import torch.nn.functional as F
import pickle
import numpy as np
from PIL import Image
from torchvision import models, transforms

# ==============================================================================
# EXACT MODEL ARCHITECTURE — copied 1-to-1 from training code
# ==============================================================================

class Encoder(nn.Module):
    def __init__(self, feature_size=2048, hidden_size=512, dropout=0.5):
        super(Encoder, self).__init__()
        self.fc      = nn.Linear(feature_size, hidden_size)
        self.dropout = nn.Dropout(dropout)
        self.bn      = nn.BatchNorm1d(hidden_size)

    def forward(self, features):
        encoded = self.fc(features)
        encoded = self.bn(encoded)
        encoded = F.relu(encoded)
        encoded = self.dropout(encoded)
        return encoded


class Attention(nn.Module):
    def __init__(self, encoder_dim, decoder_dim, attention_dim):
        super(Attention, self).__init__()
        self.encoder_att = nn.Linear(encoder_dim, attention_dim)
        self.decoder_att = nn.Linear(decoder_dim, attention_dim)
        self.full_att    = nn.Linear(attention_dim, 1)
        self.relu        = nn.ReLU()
        self.softmax     = nn.Softmax(dim=1)

    def forward(self, encoder_out, decoder_hidden):
        att1  = self.encoder_att(encoder_out)
        att2  = self.decoder_att(decoder_hidden)
        att   = self.full_att(self.relu(att1 + att2)).squeeze(-1)
        alpha = self.softmax(att.unsqueeze(1))
        attention_weighted_encoding = encoder_out * alpha
        return attention_weighted_encoding, alpha


class Decoder(nn.Module):
    def __init__(self, embed_size, hidden_size, vocab_size, num_layers=1, dropout=0.5):
        super(Decoder, self).__init__()
        self.embed_size  = embed_size
        self.hidden_size = hidden_size
        self.vocab_size  = vocab_size
        self.num_layers  = num_layers

        self.embedding = nn.Embedding(vocab_size, embed_size)
        self.dropout   = nn.Dropout(dropout)
        self.lstm      = nn.LSTM(
            embed_size + hidden_size,
            hidden_size,
            num_layers,
            batch_first=True,
            dropout=dropout if num_layers > 1 else 0
        )
        self.attention = Attention(hidden_size, hidden_size, hidden_size)
        self.fc        = nn.Linear(hidden_size, vocab_size)
        self._init_weights()

    def _init_weights(self):
        self.embedding.weight.data.uniform_(-0.1, 0.1)
        self.fc.bias.data.fill_(0)
        self.fc.weight.data.uniform_(-0.1, 0.1)

    def generate(self, features, max_len=50, start_token=1, end_token=2):
        """Greedy search"""
        h = features.unsqueeze(0).repeat(self.num_layers, 1, 1)
        c = torch.zeros_like(h)
        input_token = torch.LongTensor([start_token]).to(features.device)
        generated   = []

        for _ in range(max_len):
            embed       = self.embedding(input_token).unsqueeze(1)
            hiddens     = h[-1]
            context, _  = self.attention(features, hiddens)
            lstm_input  = torch.cat([embed.squeeze(1), context], dim=1).unsqueeze(1)
            _, (h, c)   = self.lstm(lstm_input, (h, c))
            output      = self.fc(h[-1])
            predicted   = output.argmax(1)
            generated.append(predicted.item())
            if predicted.item() == end_token:
                break
            input_token = predicted

        return generated

    def beam_search(self, features, beam_width=5, max_len=50, start_token=1, end_token=2):
        """Beam search"""
        device = features.device
        h = features.unsqueeze(0).repeat(self.num_layers, 1, 1)
        c = torch.zeros_like(h)

        beams           = [([start_token], 0.0, h, c)]
        completed_beams = []

        for _ in range(max_len):
            all_candidates = []
            for seq, score, h_state, c_state in beams:
                if seq[-1] == end_token:
                    completed_beams.append((seq, score))
                    continue
                input_token = torch.LongTensor([seq[-1]]).to(device)
                embed       = self.embedding(input_token).unsqueeze(1)
                hiddens     = h_state[-1]
                context, _  = self.attention(features, hiddens)
                lstm_input  = torch.cat([embed.squeeze(1), context], dim=1).unsqueeze(1)
                _, (h_new, c_new) = self.lstm(lstm_input, (h_state, c_state))
                log_probs   = F.log_softmax(self.fc(h_new[-1]), dim=1)
                topk_probs, topk_idx = log_probs.topk(beam_width, dim=1)
                for i in range(beam_width):
                    all_candidates.append((
                        seq + [topk_idx[0, i].item()],
                        score + topk_probs[0, i].item(),
                        h_new, c_new
                    ))
            beams = sorted(all_candidates, key=lambda x: x[1], reverse=True)[:beam_width]
            if not beams:
                break

        completed_beams.extend([(seq, score) for seq, score, _, _ in beams])
        if completed_beams:
            best_seq, _ = max(completed_beams, key=lambda x: x[1])
            return best_seq
        return beams[0][0] if beams else [start_token, end_token]


class ImageCaptioningModel(nn.Module):
    def __init__(self, encoder, decoder):
        super(ImageCaptioningModel, self).__init__()
        self.encoder = encoder
        self.decoder = decoder

    def forward(self, features, captions, lengths):
        return self.decoder(self.encoder(features), captions, lengths)


# ==============================================================================
# LOAD MODELS & VOCAB  (runs once at startup)
# ==============================================================================

DEVICE = torch.device("cpu")   # HF free-tier has no GPU

# ── Hyperparameters — must match training exactly ──────────────────────────────
FEATURE_SIZE = 2048
HIDDEN_SIZE  = 512
EMBED_SIZE   = 512
NUM_LAYERS   = 2
DROPOUT      = 0.5

def load_vocab(path="vocab.pkl"):
    with open(path, "rb") as f:
        return pickle.load(f)

def load_model(vocab, path="best_model.pth"):
    encoder = Encoder(feature_size=FEATURE_SIZE, hidden_size=HIDDEN_SIZE, dropout=DROPOUT)
    decoder = Decoder(
        embed_size=EMBED_SIZE,
        hidden_size=HIDDEN_SIZE,
        vocab_size=len(vocab),
        num_layers=NUM_LAYERS,
        dropout=DROPOUT
    )
    model = ImageCaptioningModel(encoder, decoder).to(DEVICE)
    checkpoint = torch.load(path, map_location=DEVICE, weights_only=False)
    model.load_state_dict(checkpoint["model_state_dict"])
    model.eval()
    return model

def load_resnet():
    resnet = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
    resnet = nn.Sequential(*list(resnet.children())[:-1])   # drop classifier
    resnet = resnet.to(DEVICE)
    resnet.eval()
    return resnet

print("Loading vocab...")
vocab  = load_vocab("vocab.pkl")
print(f"  Vocab size: {len(vocab)} words")

print("Loading captioning model...")
model  = load_model(vocab, "best_model.pth")
print("  Model ready.")

print("Loading ResNet50 feature extractor...")
resnet = load_resnet()
print("  ResNet50 ready.")

# ── Image pre-processing — same as training ────────────────────────────────────
IMG_TRANSFORM = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])

# ==============================================================================
# INFERENCE HELPERS
# ==============================================================================

def extract_features(pil_image: Image.Image) -> torch.Tensor:
    """Run image through ResNet50 → (1, 2048) feature vector."""
    img_tensor = IMG_TRANSFORM(pil_image).unsqueeze(0).to(DEVICE)
    with torch.no_grad():
        features = resnet(img_tensor).view(1, -1)   # (1, 2048)
    return features


def ids_to_caption(ids, vocab) -> str:
    """Convert list of token ids → clean sentence string."""
    words = []
    for idx in ids:
        if idx == vocab.stoi["<end>"]:
            break
        if idx not in (vocab.stoi["<start>"], vocab.stoi["<pad>"]):
            words.append(vocab.itos[idx])
    return " ".join(words).capitalize() if words else "Could not generate a caption."


def caption_image(pil_image: Image.Image, method: str, beam_width: int) -> tuple[str, str]:
    """
    Main inference function called by Gradio.
    Returns (caption_string, info_string).
    """
    if pil_image is None:
        return "Please upload an image.", ""

    # 1. Extract features
    features = extract_features(pil_image)

    # 2. Encode
    with torch.no_grad():
        encoded = model.encoder(features)   # (1, 512)

    # 3. Decode
    with torch.no_grad():
        if method == "Beam Search":
            ids = model.decoder.beam_search(
                encoded,
                beam_width=beam_width,
                max_len=50,
                start_token=vocab.stoi["<start>"],
                end_token=vocab.stoi["<end>"],
            )
        else:   # Greedy
            ids = model.decoder.generate(
                encoded,
                max_len=50,
                start_token=vocab.stoi["<start>"],
                end_token=vocab.stoi["<end>"],
            )

    caption = ids_to_caption(ids, vocab)
    info    = f"Method: {method}" + (f"  |  Beam width: {beam_width}" if method == "Beam Search" else "")
    return caption, info


# ==============================================================================
# GRADIO UI
# ==============================================================================

CSS = """
@import url('https://fonts.googleapis.com/css2?family=Playfair+Display:ital,wght@0,700;1,400&family=DM+Mono:wght@400;500&display=swap');

:root {
    --cream:   #F5F0E8;
    --ink:     #1A1208;
    --sienna:  #B85C2A;
    --gold:    #D4A843;
    --warm-gray: #8C7B6B;
    --border:  #D9CFC2;
}

body, .gradio-container {
    background: var(--cream) !important;
    font-family: 'DM Mono', monospace !important;
}

/* ── Header ── */
#header {
    text-align: center;
    padding: 2.5rem 1rem 1rem;
    border-bottom: 2px solid var(--ink);
    margin-bottom: 2rem;
}
#header h1 {
    font-family: 'Playfair Display', serif !important;
    font-size: clamp(2rem, 5vw, 3.5rem) !important;
    font-weight: 700 !important;
    color: var(--ink) !important;
    letter-spacing: -0.02em;
    margin: 0 !important;
    line-height: 1.1;
}
#header .subtitle {
    font-family: 'Playfair Display', serif;
    font-style: italic;
    font-size: 1.1rem;
    color: var(--sienna);
    margin-top: 0.4rem;
}
#header .tag-line {
    font-size: 0.72rem;
    letter-spacing: 0.18em;
    text-transform: uppercase;
    color: var(--warm-gray);
    margin-top: 0.6rem;
}

/* ── Main columns ── */
.main-col { padding: 0 0.5rem; }

/* ── Image upload box ── */
.image-upload .wrap { border: 2px dashed var(--border) !important; border-radius: 4px !important; background: white !important; }
.image-upload .wrap:hover { border-color: var(--sienna) !important; }

/* ── Caption output ── */
#caption-box textarea {
    font-family: 'Playfair Display', serif !important;
    font-style: italic !important;
    font-size: 1.35rem !important;
    color: var(--ink) !important;
    background: white !important;
    border: 2px solid var(--ink) !important;
    border-radius: 4px !important;
    padding: 1.2rem !important;
    line-height: 1.6 !important;
    min-height: 80px !important;
}

/* ── Info box ── */
#info-box textarea {
    font-family: 'DM Mono', monospace !important;
    font-size: 0.78rem !important;
    color: var(--warm-gray) !important;
    background: transparent !important;
    border: none !important;
    border-top: 1px solid var(--border) !important;
    padding: 0.6rem 0 0 !important;
}

/* ── Controls ── */
.gr-radio label, .gr-slider label {
    font-family: 'DM Mono', monospace !important;
    font-size: 0.8rem !important;
    letter-spacing: 0.05em !important;
    color: var(--ink) !important;
}

/* ── Generate button ── */
#generate-btn {
    background: var(--ink) !important;
    color: var(--cream) !important;
    border: none !important;
    border-radius: 2px !important;
    font-family: 'DM Mono', monospace !important;
    font-size: 0.85rem !important;
    letter-spacing: 0.12em !important;
    text-transform: uppercase !important;
    padding: 0.75rem 2rem !important;
    cursor: pointer !important;
    transition: background 0.2s, transform 0.1s !important;
    width: 100% !important;
}
#generate-btn:hover {
    background: var(--sienna) !important;
    transform: translateY(-1px) !important;
}

/* ── Section labels ── */
.section-label {
    font-size: 0.68rem;
    letter-spacing: 0.18em;
    text-transform: uppercase;
    color: var(--warm-gray);
    border-bottom: 1px solid var(--border);
    padding-bottom: 0.4rem;
    margin-bottom: 0.8rem;
}

/* ── Examples strip ── */
.examples-label {
    font-size: 0.68rem;
    letter-spacing: 0.18em;
    text-transform: uppercase;
    color: var(--warm-gray);
    margin-top: 1.5rem;
    margin-bottom: 0.5rem;
}

/* ── Footer ── */
#footer {
    text-align: center;
    padding: 1.5rem;
    border-top: 1px solid var(--border);
    margin-top: 2rem;
    font-size: 0.72rem;
    letter-spacing: 0.1em;
    color: var(--warm-gray);
}

/* hide default Gradio labels where we use custom ones */
.hide-label > label { display: none !important; }
"""

def run(image, method, beam_width):
    return caption_image(image, method, int(beam_width))


with gr.Blocks(css=CSS, title="Neural Storyteller") as demo:

    # ── Header ────────────────────────────────────────────────────────────────
    gr.HTML("""
    <div id="header">
        <div class="tag-line">ResNet50 · LSTM · Bahdanau Attention · Flickr30k</div>
        <h1>Neural Storyteller</h1>
        <div class="subtitle">Image Captioning with Seq2Seq</div>
    </div>
    """)

    # ── Main layout ───────────────────────────────────────────────────────────
    with gr.Row():

        # Left — image upload
        with gr.Column(scale=5, elem_classes="main-col"):
            gr.HTML('<div class="section-label">Upload Image</div>')
            image_input = gr.Image(
                type="pil",
                label="",
                elem_classes=["image-upload", "hide-label"],
                height=380,
            )

        # Right — controls + output
        with gr.Column(scale=5, elem_classes="main-col"):
            gr.HTML('<div class="section-label">Decoding Strategy</div>')

            method = gr.Radio(
                choices=["Beam Search", "Greedy"],
                value="Beam Search",
                label="",
                elem_classes="hide-label",
            )
            beam_width = gr.Slider(
                minimum=1, maximum=10, value=5, step=1,
                label="Beam Width  (only applies to Beam Search)",
            )

            gr.HTML('<div class="section-label" style="margin-top:1.2rem">Generated Caption</div>')
            caption_out = gr.Textbox(
                label="",
                elem_id="caption-box",
                elem_classes="hide-label",
                lines=3,
                interactive=False,
                placeholder="Your caption will appear here…",
            )
            info_out = gr.Textbox(
                label="",
                elem_id="info-box",
                elem_classes="hide-label",
                lines=1,
                interactive=False,
            )

            generate_btn = gr.Button("Generate Caption", elem_id="generate-btn")

    # ── Wire up ───────────────────────────────────────────────────────────────
    generate_btn.click(
        fn=run,
        inputs=[image_input, method, beam_width],
        outputs=[caption_out, info_out],
    )
    # Also trigger on image upload for instant feel
    image_input.upload(
        fn=run,
        inputs=[image_input, method, beam_width],
        outputs=[caption_out, info_out],
    )

    # ── Footer ────────────────────────────────────────────────────────────────
    gr.HTML("""
    <div id="footer">
        Neural Storyteller &nbsp;·&nbsp; Seq2Seq with Bahdanau Attention
        &nbsp;·&nbsp; Trained on Flickr30k &nbsp;·&nbsp; Built with PyTorch
    </div>
    """)


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