import sys, os

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(__dir__)
sys.path.append(os.path.abspath(os.path.join(__dir__, '..')))

from typing import Dict, List, Any
import uuid
import torch
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

import shlex
import subprocess
import base64
import os, shutil, cv2
import numpy as np
from io import BytesIO
from PIL import Image
import albumentations

from deepfake_detection.model.network import Recce
from deepfake_detection.model.common import freeze_weights
from albumentations.pytorch.transforms import ToTensorV2

from deepfake_detection.helper import ImageDataset

class DeepfakeDetection:
    def __init__(self, batch_size=16, device=None, gpu_id=0, weights='weights/model_params_ffpp_c23.pickle'):
        if device:
            self.device = device
        elif torch.cuda.is_available() and gpu_id != -1:
            self.device = 'cuda:{}'.format(gpu_id)
        else:
            self.device = 'cpu'

        self.batch_size = batch_size

        self.device = torch.device(self.device)
        self.model = eval("Recce")(num_classes=1)
        weights = torch.load(weights, map_location=self.device)

        self.model.load_state_dict(weights)
        self.model = self.model.to(self.device)
        freeze_weights(self.model)
        self.model.eval()
        self.transform = self.get_transform()
    
    def preprocess(self, images):
        images = [cv2.cvtColor(img, cv2.COLOR_BGR2RGB) for img in images]
        return images

    def get_transform(self, img_size=299):
        return albumentations.Compose([
            albumentations.Resize(height=img_size, width=img_size),
            albumentations.Normalize(mean=[0.5] * 3, std=[0.5] * 3),
            ToTensorV2()])

    def get_dataloader(self, images):
        # images = self.preprocess(images)
        dataset = ImageDataset(images, transform=self.transform)
        dataloader = DataLoader(dataset, batch_size=self.batch_size, shuffle=False)

        return dataloader
    def classify_fake(self, dataloader, model, device):
        probs = []
        for batch in dataloader:
            with torch.no_grad():
                batch = batch.to(self.device)

                # Process predictions
                outputs = model(batch)
                if len(outputs.shape) == 1:
                    outputs = outputs.unsqueeze(0)
                outputs = [prob[0].sigmoid().cpu().tolist() for prob in outputs]
                probs += (outputs)

        return probs

    def __call__(self, images, confidence_threshold=0.6):
        dataloader = self.get_dataloader(images)
        fake_prob = self.classify_fake(dataloader, model=self.model, device=self.device)
        labels = ['fake' if prob >= confidence_threshold else 'real' for prob in fake_prob]
            
        return labels, fake_prob

if __name__ == '__main__':
    fake_img = cv2.imread('sample_files/face_fake.png')
    real_img = cv2.imread('sample_files/face_real.png')
    detector = DeepfakeDetection()
    result = detector([fake_img])
    print(result)