import torch
import cv2
import numpy as np
import torchvision.transforms as T
from collections import OrderedDict
import base64

from model import DeepfakeModel
from cam import GradCAM, overlay_heatmap

# CONFIG

IMG_SIZE = 240
NUM_FRAMES = 10

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

# LOAD MODEL

model = DeepfakeModel()

state_dict = torch.load(
    "best_model.pth",
    map_location="cpu"
)

new_state = OrderedDict()

for k, v in state_dict.items():

    name = k.replace(
        "module.",
        ""
    )

    new_state[name] = v

model.load_state_dict(
    new_state
)

model = model.to(device)

model.eval()

print("Model loaded successfully")

# GRADCAM

target_layer = model.cnn.blocks[-1]

grad_cam = GradCAM(
    model,
    target_layer
)

# FACE DETECTOR

face_detector = cv2.CascadeClassifier(
    cv2.data.haarcascades +
    "haarcascade_frontalface_default.xml"
)

# CACHE

LAST_FRAMES = []

# TRANSFORM

transform = T.Compose([

    T.ToPILImage(),

    T.Resize(
        (IMG_SIZE, IMG_SIZE)
    ),

    T.ToTensor(),

    T.Normalize(
        mean=[0.5, 0.5, 0.5],
        std=[0.5, 0.5, 0.5]
    )
])

# FRAME EXTRACTION

def extract_and_crop(
    video_path,
    num_frames=NUM_FRAMES
):

    cap = cv2.VideoCapture(
        video_path
    )

    total_frames = int(
        cap.get(
            cv2.CAP_PROP_FRAME_COUNT
        )
    )

    if total_frames <= 0:

        cap.release()

        return []

    idx = np.linspace(
        0,
        total_frames - 1,
        num_frames,
        dtype=int
    )

    frames = []

    for i in idx:

        cap.set(
            cv2.CAP_PROP_POS_FRAMES,
            int(i)
        )

        ret, frame = cap.read()

        if not ret:
            continue

        gray = cv2.cvtColor(
            frame,
            cv2.COLOR_BGR2GRAY
        )

        faces = face_detector.detectMultiScale(
            gray,
            scaleFactor=1.1,
            minNeighbors=5,
            minSize=(80, 80)
        )

        if len(faces) > 0:

            x, y, w, h = max(
                faces,
                key=lambda f: f[2] * f[3]
            )

            pad = int(
                0.15 * max(w, h)
            )

            x1 = max(0, x - pad)
            y1 = max(0, y - pad)

            x2 = min(
                frame.shape[1],
                x + w + pad
            )

            y2 = min(
                frame.shape[0],
                y + h + pad
            )

            face = frame[
                y1:y2,
                x1:x2
            ]

        else:

            face = frame

        face = cv2.resize(
            face,
            (IMG_SIZE, IMG_SIZE)
        )

        face = cv2.cvtColor(
            face,
            cv2.COLOR_BGR2RGB
        )

        frames.append(
            face
        )

    cap.release()

    if len(frames) == 0:

        return []

    while len(frames) < num_frames:

        frames.append(
            frames[-1]
        )

    return frames

# INFERENCE

def run_inference(video_path):

    global LAST_FRAMES

    frames = extract_and_crop(
        video_path
    )

    LAST_FRAMES = frames

    if len(frames) == 0:

        return {
            "label": "Video tidak terbaca",
            "confidence": 0,
            "frames": []
        }

    imgs = []

    for frame in frames:

        imgs.append(
            transform(frame)
        )

    imgs = torch.stack(
        imgs
    )

    imgs = imgs.unsqueeze(0)

    imgs = imgs.to(device)

    with torch.no_grad():

        outputs = model(
            imgs
        )

        probs = torch.softmax(
            outputs,
            dim=1
        )[0]

        pred = torch.argmax(
            probs
        ).item()

        confidence = float(
            probs[pred].item() * 100
        )

    label = (
        "Real"
        if pred == 0
        else "Fake"
    )

    encoded_frames = []

    for frame in frames:

        _, buffer = cv2.imencode(
            ".jpg",
            cv2.cvtColor(
                frame,
                cv2.COLOR_RGB2BGR
            )
        )

        encoded_frames.append(
            base64.b64encode(
                buffer
            ).decode("utf-8")
        )

    return {
        "label": label,
        "confidence": confidence,
        "frames": encoded_frames
    }

# REGION ANALYSIS

def compute_regions(cam):

    regions = {
        "Forehead": cam[0:60, :].mean(),
        "Eyes": cam[60:110, :].mean(),
        "Cheeks": cam[110:170, :].mean(),
        "Mouth": cam[170:220, :].mean(),
        "Chin": cam[220:240, :].mean()
    }

    total = (
        sum(regions.values()) +
        1e-8
    )

    result = []

    for k, v in regions.items():

        result.append({
            "name": k,
            "value": float(v / total)
        })

    return result

# HEATMAP

def generate_heatmap(frame_index):

    global LAST_FRAMES

    if frame_index >= len(LAST_FRAMES):

        return None, None

    frame = LAST_FRAMES[
        frame_index
    ]

    img = transform(
        frame
    )

    seq = torch.stack(
        [img] * NUM_FRAMES
    )

    seq = seq.unsqueeze(0)

    seq = seq.to(device)

    cam = grad_cam.generate(
        seq
    )

    regions = compute_regions(
        cam
    )

    heatmap = overlay_heatmap(
        cv2.cvtColor(
            frame,
            cv2.COLOR_RGB2BGR
        ),
        cam
    )

    return heatmap, regions