import cv2
import numpy as np
import logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s [%(levelname)s] %(message)s', handlers=[logging.StreamHandler()])

# Age
from mivolo.predictor import Predictor
import argparse
# Nudity
from nudenet import NudeDetector


class CV():
    def __init__(self):
        args = argparse.ArgumentParser()
        args.add_argument("--device", type=str, default="cpu")
        args.add_argument("--checkpoint", default="models/model_imdb_cross_person_4.22_99.46.pth.tar")
        args.add_argument("--detector_weights", default="models/yolov8x_person_face.pt")
        args.add_argument("--with-persons", action="store_true", default=False, help="If set model will run with persons, if available")
        args.add_argument("--disable-faces", action="store_true", default=False, help="If set model will use only persons if available")
        args.add_argument("--draw", action="store_true", default=False, help="If set, resulted images will be drawn")
        args = args.parse_args([])
        # Initialize
        self.predictor_age = Predictor(args)

        # Initialize
        self.nude_detector = NudeDetector(model_path="models/640m.onnx", inference_resolution=640)
        # detector = NudeDetector(model_path="downloaded_640m.onnx path", inference_resolution=640)
        # https://github.com/notAI-tech/NudeNet?tab=readme-ov-file#available-models

        # All labels list
        self.nudity_all_labels = [
            "FEMALE_GENITALIA_COVERED",
            "FACE_FEMALE",
            "BUTTOCKS_EXPOSED",
            "FEMALE_BREAST_EXPOSED",
            "FEMALE_GENITALIA_EXPOSED",
            "MALE_BREAST_EXPOSED",
            "ANUS_EXPOSED",
            "FEET_EXPOSED",
            "BELLY_COVERED",
            "FEET_COVERED",
            "ARMPITS_COVERED",
            "ARMPITS_EXPOSED",
            "FACE_MALE",
            "BELLY_EXPOSED",
            "MALE_GENITALIA_EXPOSED",
            "ANUS_COVERED",
            "FEMALE_BREAST_COVERED",
            "BUTTOCKS_COVERED",
        ]
        # Classes of interest
        self.nudity_classes_of_interest = ["BUTTOCKS_EXPOSED", "FEMALE_BREAST_EXPOSED", "FEMALE_GENITALIA_EXPOSED", "ANUS_EXPOSED", "MALE_GENITALIA_EXPOSED"]

    def _censor(self, image_bgr, detections):
        # Copy original image
        image_bgr_censored = image_bgr.copy()

        for detection in detections:
            box = detection["box"]
            x, y, w, h = box[0], box[1], box[2], box[3]
            # Change these pixels to pure black
            image_bgr_censored[y : y + h, x : x + w] = (0, 0, 0)

        return image_bgr_censored

    def process_image(self, image_bgr):
        ###################################################################
        # Predict
        detected_objects, out_img = self.predictor_age.recognize(image_bgr)
        logging.debug("#persons: {}, #faces: {}".format(detected_objects.n_persons, detected_objects.n_faces))

        # Num faces and persons detected
        detected_objects.n_faces, detected_objects.n_persons
        # Association
        detected_objects.associate_faces_with_persons()

        # detected_objects.face_to_person_map
        # {2: 1, 3: 0}
        # detected_objects.ages
        # [None, None, 27.18, 23.77]
        age_predictions = [e for e in detected_objects.ages if e is not None]

        # Crops of faces & persons
        # crops = detected_objects.collect_crops(img)
        any_minor_present = any([ a < 18 for a in detected_objects.ages if a is not None ])
        ###################################################################

        ###################################################################
        # Predict
        nude_detections = self.nude_detector.detect(image_bgr)
        logging.debug("Nude detections: {}".format(nude_detections))
        # Filter by classes of interest
        nude_detections = [ detection for detection in nude_detections if detection["class"] in self.nudity_classes_of_interest ]
        # Nude detections present?
        any_nude_detection = len(nude_detections) > 0
        ###################################################################

        ###################################################################
        # Censor image
        censored_img_bgr = self._censor(image_bgr, nude_detections)
        # Plot age predictions on censored image
        output_image = detected_objects.plot(img=censored_img_bgr)
        ###################################################################

        results = {
            "any_minor_present": any_minor_present,
            "any_nude_detection": any_nude_detection,
            "nudity_detections": nude_detections,
            "age_predictions": age_predictions,
            "image": output_image,
        }
        return results

def process(img_bgr):
    try:
        logging.info("Processing image")
        # Process
        results = CV().process_image(img_bgr)
        logging.info("Returning results")
        return results
        
    except Exception as e:
        logging.warning("Error processing image: {}".format(str(e)))
        return {}