import importlib.util
import os
import sys
import cv2
import gradio as gr
import numpy as np
import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from safetensors.torch import load_file
from transformers import ViTImageProcessor

# -----------------------------
# CONFIGURATION
# -----------------------------

MODEL_REPO = "SARVM/Refined_ViT"
HF_TOKEN = os.getenv("HF_TOKEN")
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

print("=" * 60)
print("FORESIGHT - MODEL INITIALIZATION")
print("=" * 60)

processor = ViTImageProcessor(
    do_resize=True,
    size={"height": 224, "width": 224},
    do_rescale=True,
    rescale_factor=1 / 255,
    do_normalize=True,
    image_mean=[0.485, 0.456, 0.406],
    image_std=[0.229, 0.224, 0.225],
)

config_path = hf_hub_download(repo_id=MODEL_REPO, filename="configuration_vit_refiner.py", token=HF_TOKEN)
model_path = hf_hub_download(repo_id=MODEL_REPO, filename="modeling_vit_refiner.py", token=HF_TOKEN)

config_spec = importlib.util.spec_from_file_location("configuration_vit_refiner", config_path)
config_module = importlib.util.module_from_spec(config_spec)
sys.modules["configuration_vit_refiner"] = config_module
config_spec.loader.exec_module(config_module)

model_spec = importlib.util.spec_from_file_location("modeling_vit_refiner", model_path)
model_module = importlib.util.module_from_spec(model_spec)
sys.modules["modeling_vit_refiner"] = model_module
model_spec.loader.exec_module(model_module)

ViTRefinerConfig = config_module.ViTRefinerConfig
ViTRefinerModel = model_module.ViTRefinerModel

config = ViTRefinerConfig.from_pretrained(MODEL_REPO, token=HF_TOKEN)
model = ViTRefinerModel(config)

# Force 'eager' attention mode to enable explicit attention weights output
if hasattr(model, "vit"):
    if hasattr(model.vit, "config"):
        model.vit.config._attn_implementation = "eager"
        model.vit.config.output_attentions = True
    if hasattr(model.vit, "vit") and hasattr(model.vit.vit, "config"):
        model.vit.vit.config._attn_implementation = "eager"
        model.vit.vit.config.output_attentions = True

weights_path = hf_hub_download(repo_id=MODEL_REPO, filename="model.safetensors", token=HF_TOKEN)
state_dict = load_file(weights_path)
model.load_state_dict(state_dict, strict=False)

model.to(DEVICE)
model.eval()

model.config.id2label = {0: "FAKE", 1: "REAL"}
model.config.label2id = {"FAKE": 0, "REAL": 1}

# -----------------------------
# ATTENTION HOOK EXTRACTION
# -----------------------------

captured_attentions = []

def attention_hook(module, input, output):
    if isinstance(output, tuple):
        if len(output) > 1 and output[1] is not None:
            captured_attentions.append(output[1].detach().cpu())
        elif isinstance(output[0], torch.Tensor) and output[0].ndim == 4:
            captured_attentions.append(output[0].detach().cpu())

# Register forward hooks on all ViTSelfAttention blocks
for module in model.modules():
    if module.__class__.__name__ == "ViTSelfAttention":
        module.register_forward_hook(attention_hook)

# -----------------------------
# ATTENTION ROLLOUT
# -----------------------------

def compute_attention_rollout(attentions):
    if not attentions:
        return None

    att_mat = torch.stack(attentions)
    if att_mat.dim() == 5:
        att_mat = att_mat[:, 0]

    att_mat = att_mat.mean(dim=1)

    residual_att = torch.eye(att_mat.size(-1), device=att_mat.device)
    aug_att_mat = att_mat + residual_att
    aug_att_mat = aug_att_mat / aug_att_mat.sum(dim=-1, keepdim=True)

    joint_attentions = torch.zeros_like(aug_att_mat)
    joint_attentions[0] = aug_att_mat[0]

    for n in range(1, aug_att_mat.size(0)):
        joint_attentions[n] = aug_att_mat[n] @ joint_attentions[n - 1]

    return joint_attentions[-1]

# -----------------------------
# PREDICTION FUNCTION
# -----------------------------

def predict(image):
    global captured_attentions
    if image is None:
        return None, None, None

    if not isinstance(image, Image.Image):
        image = Image.fromarray(image)
    image = image.convert("RGB")

    inputs = processor(images=image, return_tensors="pt")
    inputs = {key: value.to(DEVICE) for key, value in inputs.items()}

    captured_attentions.clear()

    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs["logits"] if isinstance(outputs, dict) else outputs.logits

    probs = torch.nn.functional.softmax(logits, dim=-1)
    confidence, predicted_class_idx = torch.max(probs, dim=-1)

    prediction = model.config.id2label[predicted_class_idx.item()]
    confidence_pct = round(confidence.item() * 100, 2)

    rollout = compute_attention_rollout(captured_attentions)

    if rollout is not None:
        mask = rollout[0, 1:]
        size = int(mask.shape[0] ** 0.5)
        mask = mask.reshape(size, size).numpy()

        mask = cv2.resize(mask, image.size)
        mask = (mask - mask.min()) / (mask.max() - mask.min() + 1e-8)

        heatmap = cv2.applyColorMap(np.uint8(255 * mask), cv2.COLORMAP_JET)
        heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)

        overlay = cv2.addWeighted(np.array(image), 0.6, heatmap, 0.4, 0)
    else:
        print("Warning: No attention matrices captured.")
        overlay = np.array(image)

    return prediction, f"{confidence_pct}%", overlay

# -----------------------------
# UI DESIGN
# -----------------------------

custom_css = """
:root { --primary-blue: #2563eb; --hero-text: #0f172a; }
.dark { --hero-text: #f8fafc; }
body { background-color: var(--background-fill-primary); }
.hero { text-align: center; font-family: 'Inter', sans-serif; font-size: 48px; font-weight: 800; letter-spacing: -0.04em; margin-top: 50px; color: var(--hero-text) !important; }
.sub { text-align: center; opacity: 0.7; font-size: 14px; font-weight: 600; letter-spacing: 0.1em; text-transform: uppercase; margin-bottom: 40px; color: var(--body-text-color); }
.glass { background: var(--block-background-fill) !important; border: 1px solid var(--border-color-primary) !important; border-radius: 12px !important; padding: 24px !important; box-shadow: var(--block-shadow); transition: all 0.2s ease; }
.glass:hover { border-color: var(--primary-blue) !important; box-shadow: 0 4px 20px rgba(37, 99, 235, 0.1); }
button.primary { background: var(--primary-blue) !important; color: white !important; border: none !important; font-weight: 600 !important; padding: 12px 24px !important; border-radius: 8px !important; box-shadow: 0 4px 12px rgba(37, 99, 235, 0.2) !important; }
button.primary:hover { background: #1d4ed8 !important; transform: translateY(-1px); box-shadow: 0 6px 16px rgba(37, 99, 235, 0.3) !important; }
.gr-label { font-weight: 600 !important; font-size: 12px !important; text-transform: uppercase; color: var(--primary-blue) !important; }
"""

with gr.Blocks() as demo:
    gr.Markdown("<div class='hero'>FORESIGHT<span style='color:#3b82f6'>.</span></div>")
    gr.Markdown("<div class='sub'>Deep Intelligence Neural Analysis</div>")

    with gr.Row():
        with gr.Column():
            with gr.Group(elem_classes="glass"):
                input_image = gr.Image(type="pil", label="Source Input")
                run_btn = gr.Button("RUN DIAGNOSTIC", variant="primary")

        with gr.Column():
            with gr.Group(elem_classes="glass"):
                output_label = gr.Label(label="Classification Verdict")
                output_conf = gr.Textbox(label="Confidence Rating", interactive=False)
                heatmap_output = gr.Image(label="Vulnerability Visualization")

    run_btn.click(
        fn=predict,
        inputs=input_image,
        outputs=[output_label, output_conf, heatmap_output]
    )

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=int(os.getenv("PORT", "7860")),
        css=custom_css,
        theme=gr.themes.Soft(
            primary_hue="blue", 
            font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui"]
        )
    )