import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import cv2
import numpy as np
import gradio as gr

# Force CPU inference for Hugging Face free tier
device = torch.device("cpu")

# 1. Exact Architecture matching the trained weights
class SigNetTwinEncoder(nn.Module):
    def __init__(self):
        super(SigNetTwinEncoder, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2),

            nn.Conv2d(32, 64, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2),

            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2)
        )
        self.fc = nn.Sequential(
            nn.Linear(128 * 27 * 19, 256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.3),
            nn.Linear(256, 128)
        )

    def forward_once(self, x):
        features = self.conv(x)
        features = features.view(features.size(0), -1)
        embeddings = self.fc(features)
        return F.normalize(embeddings, p=2, dim=1)

    def forward(self, input1, input2):
        out1 = self.forward_once(input1)
        out2 = self.forward_once(input2)
        return out1, out2

# 2. Instantiate and load model weights safely
model = SigNetTwinEncoder().to(device)
model_path = "signet_twin_model.pth"

if os.path.exists(model_path):
    state_dict = torch.load(model_path, map_location=device)
    model.load_state_dict(state_dict)
    model.eval()
    print("Model weights loaded successfully.")
else:
    print(f"Warning: {model_path} not found in the repository root.")

OPTIMAL_THRESHOLD = 0.5581

# 3. Preprocessing function
def preprocess_for_inference(image):
    if image is None:
        return None
    if len(image.shape) == 3:
        gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    else:
        gray = image
        
    _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
    resized = cv2.resize(binary, (155, 220))
    normalized = resized.astype(np.float32) / 255.0
    tensor = torch.tensor(normalized, dtype=torch.float32).unsqueeze(0).unsqueeze(0)
    return tensor.to(device)

# 4. Verification Logic
def verify_signatures(ref_img, query_img):
    if ref_img is None or query_img is None:
        return "Please upload both reference and query signatures."
    
    t_ref = preprocess_for_inference(ref_img)
    t_query = preprocess_for_inference(query_img)
    
    with torch.no_grad():
        out_ref, out_query = model(t_ref, t_query)
        distance = F.pairwise_distance(out_ref, out_query).item()
        
    if distance <= OPTIMAL_THRESHOLD:
        verdict = "GENUINE MATCH"
        symbol = "✅"
    else:
        verdict = "FORGERY DETECTED"
        symbol = "❌"
        
    return f"{symbol} {verdict}\n\n• Euclidean Distance: {distance:.4f}\n• Decision Threshold: {OPTIMAL_THRESHOLD:.4f}"

# 5. Gradio Interface
interface = gr.Interface(
    fn=verify_signatures,
    inputs=[
        gr.Image(label="Reference Signature (Known Genuine)"),
        gr.Image(label="Query Signature (To be Verified)")
    ],
    outputs=gr.Textbox(label="Biometric Verification Verdict"),
    title="SigNet-Verify: Offline Signature Forgery Detection",
    description="Siamese Neural Network (Twin CNN) for writer-independent biometric signature verification."
)

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