import torch
import gradio as gr
import cv2
import numpy as np
from PIL import Image

from model import AdvancedCrimeDetectionModel
from feature_extractor import extract_features

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Load model
print("Loading model...")
try:
    checkpoint = torch.load(
        "crime_model_advanced.pth",
        map_location=device,
        weights_only=False
    )
    threshold = checkpoint.get("threshold", 0.5)
    model = AdvancedCrimeDetectionModel().to(device)
    model.load_state_dict(checkpoint["model_state_dict"])
    model.eval()
    print(f"Model loaded! Threshold: {threshold:.3f}")
except Exception as e:
    print(f"Error loading model: {e}")
    threshold = 0.5


def extract_frames_optimized(video_path, max_frames=72, target_fps=4):
    """Efficiently extract frames from video using adaptive sampling."""
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise ValueError("Could not open video file")

    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = cap.get(cv2.CAP_PROP_FPS)
    duration = total_frames / fps if fps > 0 else 0

    print(f"Video: {total_frames} frames, {fps:.1f} FPS, {duration:.1f}s")

    frame_interval = max(1, int(fps / target_fps)) if duration > 0 else 1

    estimated_frames = total_frames // frame_interval
    if estimated_frames > max_frames:
        frame_interval = total_frames // max_frames

    frames = []
    tmp_paths = []
    frame_count = 0

    while len(frames) < max_frames:
        ret, frame = cap.read()
        if not ret:
            break

        if frame_count % frame_interval == 0:
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frames.append(frame_rgb)

            # Save to tmp for feature_extractor
            path = f"/tmp/frame_{len(frames) - 1}.png"
            cv2.imwrite(path, frame)  # feature_extractor may handle BGR
            tmp_paths.append(path)

        frame_count += 1

    cap.release()
    print(f"Extracted {len(frames)} frames (every {frame_interval} frames)")
    return frames, tmp_paths


def get_risk_level(prob):
    """Categorize risk level based on probability."""
    if prob < 0.3:
        return "🟢 Low Risk"
    elif prob < 0.5:
        return "🟡 Medium Risk"
    elif prob < 0.7:
        return "🟠 Elevated Risk"
    else:
        return "🔴 High Risk"


def get_interpretation(max_prob, avg_prob, threshold):
    """Provide interpretation of results."""
    if max_prob < threshold:
        return "The video appears to show normal activity. No suspicious behavior detected."
    else:
        severity = "significant" if max_prob > 0.7 else "potential"
        consistency = "throughout the video" if avg_prob > threshold else "in specific segments"
        return f"The model detected {severity} anomalous behavior {consistency}. Manual review recommended."


def format_output(result):
    """Format result dictionary into readable markdown."""
    if "error" in result:
        return f"❌ {result['error']}"

    output = f"""
# {result['prediction']}

**Confidence:** {result['confidence']}
**Risk Level:** {result['risk_level']}

---

### Analysis Details
- **Max Anomaly Score:** {result['max_anomaly_score']}
- **Average Anomaly Score:** {result['avg_anomaly_score']}
- **Detection Threshold:** {result['threshold']}
- **Frames Analyzed:** {result['frames_analyzed']}
- **Segments Analyzed:** {result['segments_analyzed']}

---

### Interpretation
{result['details']}

---

**Note:** This is an AI-based analysis tool. Results should be used as a supplementary assessment and verified by trained personnel.
"""
    return output


def predict_video(video_path, progress=gr.Progress()):
    """Main prediction function using model.py and feature_extractor.py."""
    try:
        progress(0, desc="Extracting frames...")

        frames, tmp_paths = extract_frames_optimized(video_path, max_frames=72, target_fps=4)

        if len(frames) < 24:
            return format_output({
                "error": f"Video too short. Extracted only {len(frames)} frames (need at least 24)."
            })

        progress(0.3, desc="Extracting features...")

        # Sliding window over frame tmp_paths
        window_size = 24
        stride = 12
        all_probs = []
        num_windows = max(1, (len(tmp_paths) - window_size) // stride + 1)

        for i in range(0, len(tmp_paths) - window_size + 1, stride):
            window_paths = tmp_paths[i:i + window_size]
            feats = extract_features(window_paths).unsqueeze(0).to(device)

            with torch.no_grad():
                prob = torch.sigmoid(model(feats)).item()
            all_probs.append(prob)

            progress(
                0.3 + 0.6 * ((i // stride + 1) / num_windows),
                desc=f"Analyzing segments... ({i // stride + 1}/{num_windows})"
            )

        progress(0.9, desc="Finalizing results...")

        max_prob = max(all_probs)
        avg_prob = float(np.mean(all_probs))
        is_anomaly = max_prob > threshold

        result = {
            "prediction": "🚨 ANOMALY DETECTED" if is_anomaly else "✅ NORMAL ACTIVITY",
            "confidence": f"{max_prob * 100:.1f}%",
            "max_anomaly_score": f"{max_prob:.3f}",
            "avg_anomaly_score": f"{avg_prob:.3f}",
            "threshold": f"{threshold:.3f}",
            "frames_analyzed": len(frames),
            "segments_analyzed": len(all_probs),
            "risk_level": get_risk_level(max_prob),
            "details": get_interpretation(max_prob, avg_prob, threshold),
        }

        progress(1.0, desc="Complete!")
        return format_output(result)

    except Exception as e:
        return format_output({"error": f"Error processing video: {str(e)}"})


# Gradio Interface
with gr.Blocks(theme=gr.themes.Soft(), title="Crime Detection AI") as demo:
    gr.Markdown("""
# 🔍 Advanced Crime Detection System

Upload a video to analyze for potential anomalous or criminal behavior.
The model uses dual-backbone feature extraction with Transformer + GRU architecture.

**Supported formats:** MP4, AVI, MOV, MKV
**Optimized for:** Videos from 10 seconds to several minutes
""")

    with gr.Row():
        with gr.Column():
            video_input = gr.Video(label="Upload Video", height=400)
            analyze_btn = gr.Button("🔍 Analyze Video", variant="primary", size="lg")

            gr.Markdown("""
### Tips:
- Videos are processed efficiently using adaptive frame sampling
- Longer videos are analyzed in overlapping segments
- Analysis typically takes 30-60 seconds
""")

        with gr.Column():
            output_text = gr.Markdown(label="Analysis Results")

    gr.Markdown("""
---
### About This Model

This system uses a deep learning model trained on the UCF-Crime dataset to detect anomalous activities in videos.

**Architecture:**
- Dual backbone feature extraction (Transformer + GRU)
- Transformer encoder for long-range dependencies
- Bidirectional GRU for temporal modeling
- Multi-head attention mechanism

**Detection Capabilities:**
- Identifies unusual patterns and behaviors
- Analyzes motion and spatial features
- Provides confidence scores and risk levels

""")

    analyze_btn.click(
        fn=predict_video,
        inputs=video_input,
        outputs=output_text
    )

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