"""
Human Face to Comic Face Generator - Gradio Web App
Converts human faces to comic-style illustrations using a trained Pix2Pix conditional GAN.
"""

import os
import functools
import torch
import torch.nn as nn
import torch.nn.functional as F
import gradio as gr
from PIL import Image
import numpy as np
from torchvision import transforms
import math

# ─────────────────────────────────────────────────────────────────────
# MODEL ARCHITECTURE (Same as in training notebook)
# ─────────────────────────────────────────────────────────────────────

def norm(c):
    """GroupNorm works at any spatial size"""
    for g in [32, 16, 8, 4, 2, 1]:
        if c % g == 0:
            return nn.GroupNorm(g, c)

SN = nn.utils.spectral_norm


class ChannelAttention(nn.Module):
    def __init__(self, c, r=16):
        super().__init__()
        mid = max(c // r, 4)
        self.avg = nn.AdaptiveAvgPool2d(1)
        self.mx = nn.AdaptiveMaxPool2d(1)
        self.mlp = nn.Sequential(
            nn.Conv2d(c, mid, 1, bias=False), nn.ReLU(),
            nn.Conv2d(mid, c, 1, bias=False))

    def forward(self, x):
        return x * torch.sigmoid(self.mlp(self.avg(x)) + self.mlp(self.mx(x)))


class SpatialAttention(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(2, 1, 7, padding=3, bias=False)

    def forward(self, x):
        return x * torch.sigmoid(
            self.conv(torch.cat([x.mean(1, keepdim=True),
                                 x.max(1, keepdim=True)[0]], 1)))


class CBAM(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.ca = ChannelAttention(c)
        self.sa = SpatialAttention()

    def forward(self, x):
        return self.sa(self.ca(x))


class SelfAttention(nn.Module):
    """fp32 softmax + clamped logits — numerically stable under AMP."""
    def __init__(self, c):
        super().__init__()
        mid = max(c // 8, 1)
        self.q = SN(nn.Conv2d(c, mid, 1))
        self.k = SN(nn.Conv2d(c, mid, 1))
        self.v = SN(nn.Conv2d(c, c, 1))
        self.gamma = nn.Parameter(torch.zeros(1))

    def forward(self, x):
        B, C, H, W = x.shape
        q = self.q(x).view(B, -1, H*W).permute(0, 2, 1).float()
        k = self.k(x).view(B, -1, H*W).float()
        v = self.v(x).view(B, -1, H*W).float()
        attn = F.softmax(
            (torch.bmm(q, k) * (q.shape[-1] ** -0.5)).clamp(-10, 10), dim=-1)
        out = torch.bmm(v, attn.permute(0, 2, 1)).view(B, C, H, W)
        return (x + self.gamma * out.to(x.dtype)).contiguous()


class ResBlock(nn.Module):
    def __init__(self, c, dropout=0.0):
        super().__init__()
        layers = [norm(c), nn.ReLU(),
                  nn.Conv2d(c, c, 3, padding=1, bias=False),
                  norm(c), nn.ReLU()]
        if dropout > 0:
            layers.append(nn.Dropout(dropout))
        layers.append(nn.Conv2d(c, c, 3, padding=1, bias=False))
        self.block = nn.Sequential(*layers)

    def forward(self, x):
        return x + self.block(x)


class DownBlock(nn.Module):
    def __init__(self, ic, oc, outermost=False):
        super().__init__()
        seq = []
        if not outermost:
            seq.append(nn.LeakyReLU(0.2))
        seq.append(nn.Conv2d(ic, oc, 4, stride=2, padding=1, bias=outermost))
        if not outermost:
            seq.append(norm(oc))
        self.block = nn.Sequential(*seq)

    def forward(self, x):
        return self.block(x)


class UpBlock(nn.Module):
    def __init__(self, in_ch, out_ch, outermost=False, dropout=0.0):
        super().__init__()
        seq = [nn.ReLU(),
               nn.ConvTranspose2d(in_ch, out_ch, 4, stride=2, padding=1)]
        if outermost:
            seq.append(nn.Tanh())
        else:
            seq.append(norm(out_ch))
            if dropout > 0:
                seq.append(nn.Dropout(dropout))
        self.block = nn.Sequential(*seq)

    def forward(self, x):
        return self.block(x)


class HumanToComicGenerator(nn.Module):
    def __init__(self, input_nc=3, output_nc=3, ngf=64, dropout=0.3):
        super().__init__()
        nf = ngf
        self._nf = nf
        self._dropout = dropout

        # Encoder
        self.enc1 = DownBlock(input_nc, nf, outermost=True)
        self.enc2 = DownBlock(nf, nf*2)
        self.enc3 = DownBlock(nf*2, nf*4)
        self.enc4 = DownBlock(nf*4, nf*8)
        self.enc5 = DownBlock(nf*8, nf*8)
        self.enc6 = DownBlock(nf*8, nf*8)

        # SA at enc3 level
        self.sa_enc3 = SelfAttention(nf*4)

        # Bottleneck placeholder
        self.bottleneck = nn.Identity()

        # CBAM skip gates
        self.cbam1 = CBAM(nf)
        self.cbam2 = CBAM(nf*2)
        self.cbam3 = CBAM(nf*4)
        self.cbam4 = CBAM(nf*8)
        self.cbam5 = CBAM(nf*8)

        # Decoder blocks
        self.dec6_ns = UpBlock(nf*8, nf*8, dropout=dropout)
        self.dec5_ns = UpBlock(nf*8, nf*8, dropout=dropout)
        self.dec4_ns = UpBlock(nf*8, nf*4)
        self.dec5 = UpBlock(nf*8 + nf*8, nf*8, dropout=dropout)
        self.dec4 = UpBlock(nf*8 + nf*8, nf*4)
        self.dec3 = UpBlock(nf*4 + nf*4, nf*2)
        self.dec2 = UpBlock(nf*2 + nf*2, nf)
        self.dec1 = UpBlock(nf + nf, output_nc, outermost=True)

        self.depth = 6

    def set_depth(self, resolution: int, n_res_blocks: int = 6):
        nf = self._nf
        self.depth = int(math.log2(resolution)) - 2
        bot_sz = resolution // (2 ** self.depth)
        res_blocks = [ResBlock(nf*8, self._dropout) for _ in range(n_res_blocks)]
        res_blocks.append(SelfAttention(nf*8))
        self.bottleneck = nn.Sequential(*res_blocks).to(
            next(self.parameters()).device)

    def forward(self, x):
        e1 = self.enc1(x)
        e2 = self.enc2(e1)
        e3 = self.sa_enc3(self.enc3(e2))
        e4 = self.enc4(e3)
        e5 = self.enc5(e4) if self.depth >= 5 else None
        e6 = self.enc6(e5) if self.depth >= 6 else None

        bot_in = {6: e6, 5: e5, 4: e4}[self.depth]
        b = self.bottleneck(bot_in)

        if self.depth == 6:
            d = self.dec6_ns(b)
            d = self.dec5(torch.cat([d, self.cbam5(e5)], 1).contiguous())
            d = self.dec4(torch.cat([d, self.cbam4(e4)], 1).contiguous())
        elif self.depth == 5:
            d = self.dec5_ns(b)
            d = self.dec4(torch.cat([d, self.cbam4(e4)], 1).contiguous())
        else:
            d = self.dec4_ns(b)

        d = self.dec3(torch.cat([d, self.cbam3(e3)], 1).contiguous())
        d = self.dec2(torch.cat([d, self.cbam2(e2)], 1).contiguous())
        d = self.dec1(torch.cat([d, self.cbam1(e1)], 1).contiguous())
        return d


# ─────────────────────────────────────────────────────────────────────
# MODEL LOADING & INFERENCE
# ─────────────────────────────────────────────────────────────────────

DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
BASE_DIR = os.path.dirname(os.path.abspath(__file__))


def _candidate_checkpoint_paths():
    """Build robust checkpoint candidates across local/dev/docker layouts."""
    # Optional explicit override.
    env_path = os.getenv("GENERATOR_CHECKPOINT", "").strip()
    if env_path:
        return (env_path,)

    search_roots = [
        BASE_DIR,
        os.getcwd(),
        "/workspace",
        "/workspace/HumanFace2ComicFaceGenerator",
        "/home/engmatix/HumanFace2ComicFaceGenerator",
    ]
    # Keep order stable and deduplicated.
    seen = set()
    ordered_roots = []
    for root in search_roots:
        if root and root not in seen:
            seen.add(root)
            ordered_roots.append(root)

    filenames = ("generator_best.pth",)

    candidates = []
    for root in ordered_roots:
        for name in filenames:
            candidates.append(os.path.join(root, "model_weights", name))
            candidates.append(os.path.join(root, "results", name))
    return tuple(candidates)


def _extract_state_dict(checkpoint):
    """Handle common checkpoint wrappers from training scripts."""
    if not isinstance(checkpoint, dict):
        return checkpoint

    for key in ("state_dict", "model_state_dict", "model", "generator", "G"):
        if key in checkpoint and isinstance(checkpoint[key], dict):
            checkpoint = checkpoint[key]
            break

    # Remove optional DistributedDataParallel prefix.
    if any(k.startswith("module.") for k in checkpoint.keys()):
        checkpoint = {k.replace("module.", "", 1): v for k, v in checkpoint.items()}

    return checkpoint


@functools.lru_cache(maxsize=1)
def load_model(resolution=256):
    """Load the trained generator model."""
    model = HumanToComicGenerator(ngf=64, dropout=0.3)
    # Trained production checkpoint is 256px depth (stage-3 progressive setup).
    model.set_depth(256, n_res_blocks=6)
    
    model_candidates = _candidate_checkpoint_paths()
    model_path = next((p for p in model_candidates if os.path.exists(p)), None)
    if model_path is None:
        searched = "\n".join(model_candidates)
        raise FileNotFoundError(f"Checkpoint not found. Searched:\n{searched}")

    checkpoint = torch.load(model_path, map_location=DEVICE)
    state_dict = _extract_state_dict(checkpoint)
    missing, unexpected = model.load_state_dict(state_dict, strict=False)
    if missing or unexpected:
        raise RuntimeError(
            "Checkpoint mismatch detected. "
            f"missing_keys={len(missing)}, unexpected_keys={len(unexpected)}. "
            "This prevents random/invalid generations from partial weight loading."
        )
    
    model.to(DEVICE).eval()
    return model


def preprocess_image(image, size=256):
    """Convert PIL Image to normalized tensor."""
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    
    # Resize
    image = image.resize((size, size), Image.Resampling.LANCZOS)
    
    # To tensor and normalize
    img_tensor = transforms.ToTensor()(image)
    img_tensor = transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))(img_tensor)
    
    return img_tensor.unsqueeze(0).to(DEVICE)


def postprocess_image(tensor):
    """Convert model output tensor to PIL Image."""
    with torch.no_grad():
        tensor = tensor.squeeze(0)
        tensor = (tensor.cpu() * 0.5 + 0.5).clamp(0, 1)
        tensor = tensor.permute(1, 2, 0)
        img = (tensor.numpy() * 255).astype(np.uint8)
    
    return Image.fromarray(img)


def convert_face(input_image, resolution=256):
    """Main inference function."""
    try:
        if input_image is None:
            return None, "❌ Error: Please upload an input face image."

        # Load model
        # Keep UI compatible, but always run with trained 256px generator depth.
        model = load_model(256)
        
        # Preprocess
        img_tensor = preprocess_image(input_image, size=256)
        
        # Inference
        with torch.no_grad():
            output = model(img_tensor)
        
        # Postprocess
        result = postprocess_image(output)
        
        return result, "✅ Conversion successful!"
    
    except Exception as e:
        return None, f"❌ Error: {str(e)}"


# ─────────────────────────────────────────────────────────────────────
# GRADIO INTERFACE
# ─────────────────────────────────────────────────────────────────────

def create_interface():
    """Create Gradio interface."""
    
    with gr.Blocks(title="Human Face to Comic Generator", theme=gr.themes.Soft()) as app:
        gr.Markdown("""
        # 🎨 Human Face to Comic Face Generator
        
        Convert your human face photos to comic-style illustrations using AI!
        
        **Model:** Pix2Pix Conditional GAN trained on human↔comic face pairs
        
        **Features:**
        - High-quality 256×256 output resolution
        - Trained on 10,000+ paired face images
        - Preserves facial identity while adding comic-style effects
        """)
        
        with gr.Row():
            with gr.Column():
                input_img = gr.Image(
                    label="📸 Upload Your Face",
                    type="pil",
                    scale=1
                )
                
                resolution = gr.Slider(
                    label="🎯 Resolution",
                    minimum=128,
                    maximum=256,
                    step=128,
                    value=256,
                    info="Higher resolution = better quality but slower"
                )
                
                convert_btn = gr.Button("🚀 Convert to Comic", size="lg", variant="primary")
            
            with gr.Column():
                output_img = gr.Image(
                    label="🎭 Comic Face Output",
                    type="pil"
                )
                status_msg = gr.Textbox(
                    label="Status",
                    interactive=False,
                    value="Upload an image and click Convert!"
                )
        
        # Event handler
        convert_btn.click(
            fn=convert_face,
            inputs=[input_img, resolution],
            outputs=[output_img, status_msg]
        )
        
        # Example images
        gr.Examples(
            examples=[
                ["./examples/example1.jpg", 256],
                ["./examples/example2.jpg", 256],
            ] if os.path.exists("./examples") else [],
            inputs=[input_img, resolution],
            outputs=[output_img, status_msg],
            fn=convert_face,
            cache_examples=False,
        )
        
        gr.Markdown("""
        ---
        
        **Tips:**
        - Upload a clear, well-lit photo of a face
        - Square aspect ratio works best
        - The model works best with frontal face views
        
        **Model Details:**
        - Architecture: U-Net with attention mechanisms
        - Trained with progressive growing (64→128→256px)
        - Loss functions: L1, Perceptual, Feature Matching, Frequency
        - Training time: ~100 GPU hours on A100
        """)
    
    return app


if __name__ == "__main__":
    interface = create_interface()
    interface.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
    )