import os
import numpy as np
from PIL import Image, ImageDraw
from modules import shared, processing
from modules.face_restoration import FaceRestoration
from modules import devices, processing_class


class YoLoResult:
    def __init__(self, score: float, box: list[int], mask: Image.Image = None, size: float = 0, width = 0, height = 0, args = {}):
        self.score = score
        self.box = box
        self.mask = mask
        self.size = size
        self.width = width
        self.height = height
        self.args = args


class FaceRestorerYolo(FaceRestoration):
    def name(self):
        return "Face HiRes"

    def __init__(self):
        from modules import paths
        self.model = None
        self.model_dir = os.path.join(paths.models_path, 'yolo')
        self.model_name = 'yolov8n-face.pt'
        self.model_url = 'https://github.com/akanametov/yolov8-face/releases/download/v0.0.0/yolov8n-face.pt'
        # self.model_name = 'yolov9-c-face.pt'
        # self.model_url = 'https://github.com/akanametov/yolov9-face/releases/download/1.0/yolov9-c-face.pt'

    def dependencies(self):
        import installer
        installer.install('ultralytics', ignore=True)

    def predict(
            self,
            image: Image.Image,
            imgsz: int = 640,
            half: bool = True,
            device = devices.device,
            augment: bool = True,
            agnostic: bool = False,
            retina: bool = False,
            mask: bool = True,
            offload: bool = shared.opts.face_restoration_unload,
        ) -> list[YoLoResult]:

        args = {
            'conf': shared.opts.facehires_conf,
            'iou': shared.opts.facehires_iou,
            'max_det': shared.opts.facehires_max,
        }
        self.model.to(device)
        predictions = self.model.predict(
            source=[image],
            stream=False,
            verbose=False,
            imgsz=imgsz,
            half=half,
            device=device,
            augment=augment,
            agnostic_nms=agnostic,
            retina_masks=retina,
            **args
        )
        if offload:
            self.model.to('cpu')
        result = []
        for prediction in predictions:
            boxes = prediction.boxes.xyxy.detach().int().cpu().numpy() if prediction.boxes is not None else []
            scores = prediction.boxes.conf.detach().float().cpu().numpy() if prediction.boxes is not None else []
            for score, box in zip(scores, boxes):
                box = box.tolist()
                mask_image = None
                w, h = box[2] - box[0], box[3] - box[1]
                size = w * h / (image.width * image.height)
                if (min(w, h) > shared.opts.facehires_min_size if shared.opts.facehires_min_size > 0 else True) and (max(w, h) < shared.opts.facehires_max_size if shared.opts.facehires_max_size > 0 else True):
                    if mask:
                        mask_image = image.copy()
                        mask_image = Image.new('L', image.size, 0)
                        draw = ImageDraw.Draw(mask_image)
                        draw.rectangle(box, fill="white", outline=None, width=0)
                    result.append(YoLoResult(score=round(score, 2), box=box, mask=mask_image, size=size, width=w, height=h, args=args))
        return result

    def load(self):
        from modules import modelloader
        self.dependencies()
        if self.model is None:
            model_file = modelloader.load_file_from_url(url=self.model_url, model_dir=self.model_dir, file_name=self.model_name)
            if model_file is not None:
                shared.log.info(f'Loading: type=FaceHires model={model_file}')
                from ultralytics import YOLO # pylint: disable=import-outside-toplevel
                self.model = YOLO(model_file)

    def restore(self, np_image, p: processing.StableDiffusionProcessing = None):
        if hasattr(p, 'recursion'):
            return
        if not hasattr(p, 'facehires'):
            p.facehires = 0
        if np_image is None or p.facehires >= p.batch_size * p.n_iter:
            return np_image
        self.load()
        if self.model is None:
            shared.log.error(f"Model load: type=FaceHires model='{self.model_name}' dir={self.model_dir} url={self.model_url}")
            return np_image
        image = Image.fromarray(np_image)
        faces = self.predict(image)
        if len(faces) == 0:
            return np_image

        # create backups
        orig_apply_overlay = shared.opts.mask_apply_overlay
        orig_p = p.__dict__.copy()
        orig_cls = p.__class__

        pp = None
        shared.opts.data['mask_apply_overlay'] = True
        resolution = 512 if shared.sd_model_type in ['none', 'sd', 'lcm', 'unknown'] else 1024
        args = {
            'batch_size': 1,
            'n_iter': 1,
            'inpaint_full_res': True,
            'inpainting_mask_invert': 0,
            'inpainting_fill': 1, # no fill
            'sampler_name': orig_p.get('hr_sampler_name', 'default'),
            'steps': orig_p.get('hr_second_pass_steps', 0),
            'negative_prompt': orig_p.get('refiner_negative', ''),
            'denoising_strength': shared.opts.facehires_strength if shared.opts.facehires_strength > 0 else orig_p.get('denoising_strength', 0.3),
            'styles': [],
            'prompt': orig_p.get('refiner_prompt', ''),
            'mask_blur': 10,
            'inpaint_full_res_padding': shared.opts.facehires_padding,
            'restore_faces': True,
            'width': resolution,
            'height': resolution,
        }
        if getattr(p, 'is_control', False):
            from modules.control import run
            run.restore_pipeline()

        p = processing_class.switch_class(p, processing.StableDiffusionProcessingImg2Img, args)
        p.facehires += 1 # set flag to avoid recursion

        if p.steps < 1:
            p.steps = orig_p.get('steps', 0)
        if len(p.prompt) == 0:
            p.prompt = orig_p.get('all_prompts', [''])[0]
        if len(p.negative_prompt) == 0:
            p.negative_prompt = orig_p.get('all_negative_prompts', [''])[0]

        report = [{'score': f.score, 'size': f'{f.width}x{f.height}' } for f in faces]
        shared.log.debug(f'Face HiRes: faces={report} args={faces[0].args} denoise={p.denoising_strength} blur={p.mask_blur} resolution={p.width}x{p.height} padding={p.inpaint_full_res_padding}')

        mask_all = []
        for face in faces:
            if face.mask is None:
                continue
            p.init_images = [image]
            p.image_mask = [face.mask]
            # mask_all.append(face.mask)
            p.recursion = True
            pp = processing.process_images_inner(p)
            del p.recursion
            p.overlay_images = None # skip applying overlay twice
            if pp is not None and pp.images is not None and len(pp.images) > 0:
                image = pp.images[0] # update image to be reused for next face
                if len(pp.images) > 1:
                    mask_all.append(pp.images[1])

        # restore pipeline
        p = processing_class.switch_class(p, orig_cls, orig_p)
        p.init_images = getattr(orig_p, 'init_images', None)
        p.image_mask = getattr(orig_p, 'image_mask', None)
        shared.opts.data['mask_apply_overlay'] = orig_apply_overlay
        np_image = np.array(image)

        if len(mask_all) > 0 and shared.opts.include_mask:
            from modules.control.util import blend
            p.image_mask = blend([np.array(m) for m in mask_all])
            # combined = blend([np_image, p.image_mask])
            # combined = Image.fromarray(combined)
            # combined.save('/tmp/face.png')
            p.image_mask = Image.fromarray(p.image_mask)
        return np_image


yolo = FaceRestorerYolo()
shared.face_restorers.append(yolo)
