import gradio as gr
import torch
import torch.nn.functional as F
from facenet_pytorch import MTCNN, InceptionResnetV1
import cv2
from pytorch_grad_cam import GradCAM
from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
from pytorch_grad_cam.utils.image import show_cam_on_image
from PIL import Image
import numpy as np
import warnings
from typing import Tuple, List, Dict
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import io

warnings.filterwarnings("ignore")

# Device configuration
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'

# Load models
mtcnn = MTCNN(select_largest=False, post_process=False, device=DEVICE).to(DEVICE).eval()
model = InceptionResnetV1(pretrained="vggface2", classify=True, num_classes=1, device=DEVICE)

checkpoint = torch.load("df_model.pth", map_location=torch.device('cpu'))
model.load_state_dict(checkpoint['model_state_dict'])
model.to(DEVICE)
model.eval()

def predict_frame(frame: np.ndarray) -> Tuple[str, np.ndarray, Dict[str, float]]:
    """Predict whether the input frame contains a real or fake face"""
    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    frame_pil = Image.fromarray(frame)

    face = mtcnn(frame_pil)
    if face is None:
        return None, None, None  # No face detected

    # Preprocess the face
    face = F.interpolate(face.unsqueeze(0), size=(256, 256), mode='bilinear', align_corners=False)
    face = face.to(DEVICE, dtype=torch.float32) / 255.0

    # Predict
    with torch.no_grad():
        output = torch.sigmoid(model(face).squeeze(0))
        fake_confidence = output.item()
        real_confidence = 1 - fake_confidence
        prediction = "real" if real_confidence > fake_confidence else "fake"
        
        confidences = {
            'real': real_confidence,
            'fake': fake_confidence
        }

    # Visualize
    target_layers = [model.block8.branch1[-1]]
    cam = GradCAM(model=model, target_layers=target_layers, use_cuda=torch.cuda.is_available())
    targets = [ClassifierOutputTarget(0)]
    grayscale_cam = cam(input_tensor=face, targets=targets, eigen_smooth=True)
    grayscale_cam = grayscale_cam[0, :]
    face_np = face.squeeze(0).permute(1, 2, 0).cpu().numpy()
    visualization = show_cam_on_image(face_np, grayscale_cam, use_rgb=True)
    face_with_mask = cv2.addWeighted((face_np * 255).astype(np.uint8), 1, (visualization * 255).astype(np.uint8), 0.5, 0)

    return prediction, face_with_mask, confidences

def predict_video(input_video: str) -> Tuple[str, np.ndarray, Dict[str, List[float]], List[np.ndarray]]:
    cap = cv2.VideoCapture(input_video)

    frames = []
    predictions = []
    confidences_real = []
    confidences_fake = []
    frame_count = 0
    skip_frames = 20

    while True:
        ret, frame = cap.read()
        if not ret:
            break
        frame_count += 1
        if frame_count % skip_frames != 0:
            continue

        prediction, frame_with_mask, confidence = predict_frame(frame)
        if prediction is None:
            continue

        frames.append(frame_with_mask)
        predictions.append(prediction)
        confidences_real.append(confidence['real'])
        confidences_fake.append(confidence['fake'])

    cap.release()

    # Determine the final prediction based on the average confidence
    avg_real_confidence = sum(confidences_real) / len(confidences_real)
    avg_fake_confidence = sum(confidences_fake) / len(confidences_fake)
    final_prediction = 'real' if avg_real_confidence > avg_fake_confidence else 'fake'
    final_confidence = max(avg_real_confidence, avg_fake_confidence)

    # Create an animated summary plot
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 10))
    
    def animate(i):
        ax1.clear()
        ax2.clear()
        
        # Confidence over time
        ax1.plot(confidences_real[:i+1], label='Real', color='green')
        ax1.plot(confidences_fake[:i+1], label='Fake', color='red')
        ax1.set_title('Confidence Scores Over Time')
        ax1.set_xlabel('Frame')
        ax1.set_ylabel('Confidence')
        ax1.legend()
        ax1.grid(True)
        ax1.set_ylim(0, 1)

        # Prediction distribution
        labels, counts = np.unique(predictions[:i+1], return_counts=True)
        ax2.bar(labels, counts, color=['green', 'red'])
        ax2.set_title('Distribution of Predictions')
        ax2.set_xlabel('Prediction')
        ax2.set_ylabel('Count')
        ax2.set_ylim(0, len(predictions))
        
        plt.tight_layout()

    anim = FuncAnimation(fig, animate, frames=len(confidences_real), repeat=False)
    
    # Save the animation as a gif
    buf = io.BytesIO()
    anim.save(buf, writer='pillow', fps=5)
    buf.seek(0)
    summary_plot = Image.open(buf)

    return final_prediction, final_confidence, summary_plot, {
        'real': confidences_real,
        'fake': confidences_fake
    }, frames

# Custom CSS for a more appealing interface
custom_css = """
#output-container {
    display: flex;
    justify-content: center;
    align-items: center;
    flex-direction: column;
}
#confidence-label {
    font-size: 24px;
    font-weight: bold;
    margin-bottom: 10px;
}
#confidence-bar {
    width: 100%;
    height: 30px;
    background-color: #f0f0f0;
    border-radius: 15px;
    overflow: hidden;
}
#confidence-fill {
    height: 100%;
    background-color: #4CAF50;
    transition: width 0.5s ease-in-out;
}
"""

# Gradio Interface
with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🕵️‍♂️ DeepFake Video Detective 🎭")
    gr.Markdown("Upload a video to determine if it's real or a deepfake. Our AI will analyze it frame by frame!")
    
    with gr.Row():
        input_video = gr.Video(label="📹 Upload Your Video")
    
    with gr.Row():
        submit_btn = gr.Button("🔍 Analyze Video", variant="primary")
    
    with gr.Row():
        with gr.Column():
            output_label = gr.Label(label="🏷️ Prediction")
            confidence_output = gr.HTML(
                """
                <div id="output-container">
                    <div id="confidence-label">Confidence: 0%</div>
                    <div id="confidence-bar">
                        <div id="confidence-fill" style="width: 0%;"></div>
                    </div>
                </div>
                """
            )
        summary_plot = gr.Image(label="📊 Analysis Summary")
    
    with gr.Row():
        output_video = gr.Video(label="🎞️ Processed Video")
    
    def update_confidence(prediction, confidence):
        color = "#4CAF50" if prediction == "real" else "#FF5722"
        return f"""
        <div id="output-container">
            <div id="confidence-label">Confidence: {confidence:.2%}</div>
            <div id="confidence-bar">
                <div id="confidence-fill" style="width: {confidence:.2%}; background-color: {color};"></div>
            </div>
        </div>
        """
    
    def process_video(video):
        prediction, confidence, summary, _, frames = predict_video(video)
        processed_video = np.stack(frames, axis=0)
        confidence_html = update_confidence(prediction, confidence)
        return {output_label: prediction, confidence_output: confidence_html, summary_plot: summary, output_video: processed_video}
    
    submit_btn.click(
        process_video,
        inputs=[input_video],
        outputs=[output_label, confidence_output, summary_plot, output_video]
    )

demo.launch()