import torch

from os import path
import random
import numpy as np

from . import LCM_Sampler,LCM_LoRa_Sampler,LCM_img2img_Sampler,LCM_LoRa_img2img_Sampler,LCM_onnx_Sampler,LCM_onnx_img2img_Sampler

from PIL import Image,ImageOps

MAX_SEED = np.iinfo(np.int32).max

def make_seed(seed: int, random_seed:bool) -> int:
    if random_seed:
        seed = random.randint(0,MAX_SEED)
    return seed

class fasteasySD:
    """ LCM model pipeline control class.

    Create and manage pipeline objects for LCM models,
    It has functions that process the pipeline input and output values as its main methods.

    """

    def __init__(self, device:str='cpu', use_fp16:bool=False,mode:str='LCM',img2img:bool=False):

        """ Class constructors.

        device : device to use (ex: 'cpu' or 'cuda').

        use_fp16 : Enable fp16 mode. (bool)

        """

        self.user_device = device
        self.user_fp16 = use_fp16
        self.__makeSampler(mode,img2img)
        
    def __get_sampler(self):
        
        if not self.mode_img2img :
            if self.mode == 'LCM' :
                return LCM_Sampler.LCM_Sampler()
            elif self.mode in ["SD","SDXL","SSD-1B"]:
                return LCM_LoRa_Sampler.LCM_LoRa_Sampler()
            elif self.mode == 'ONNX':
                return LCM_onnx_Sampler.LCM_onnx_Sampler()
            # elif mode == ''
        elif self.mode_img2img :
            if self.mode == 'LCM' :
                return LCM_img2img_Sampler.LCM_img2img_Sampler()
            elif self.mode in ["SD","SDXL","SSD-1B"]:
                return LCM_LoRa_img2img_Sampler.LCM_LoRa_img2img_Sampler()
            elif self.mode == 'ONNX':
                return LCM_onnx_img2img_Sampler.LCM_onnx_img2img_Sampler()
            # elif mode == ''
    

    def __makeSampler(self,mode:str,img2img:bool):
        """ Create a txt2img, img2img sampler object. (automatic load with init)

        Create sampler objects for LCM model use.

        """
        
        if mode in ['LCM','SD','SDXL','SSD-1B','ONNX'] :
            self.mode = mode
        else :
            raise Exception(f'%{mode} is not support. ')
        
        if img2img :
            self.mode_img2img = True
        else :
            self.mode_img2img = False
        
        self.pipeline = self.__get_sampler()
    
    def make_seed(self,seed: int, random_seed:bool) -> int:

        """ Automatically generate seed value (random number)

        Automatically generate seed value (random number)

        seed : user input seed value (int)

        random_seed : True, False for use random_seed
        
        """

        if random_seed:
            seed = random.randint(0,MAX_SEED)
        return seed
    
    def __load_img(self,img_dir,width,height):

        """ Load image file for img2img input

        Load the file specified in img_dir and return it to the form available in img2img sampler.

        img_dir : path for input img.

        """

        i = Image.open(img_dir)
        i = i.resize((width,height))
        i = ImageOps.exif_transpose(i)
        image = i.convert("RGB")
        image = np.array(image).astype(np.float32) / 255.0
        image = torch.from_numpy(image)[None,]

        return image
    
    def save_PIL(self,pils,save_name):

        """ PIL image list storage function.

        Store a list of PIL images generated by the LCM model.

        pils : list of PIL images

        save_name : Set image save filename. (ex: {save_name}_1.png)

        """

        counter = 0 
        for img in pils:
            img.save(save_name + f"_{counter}.png",compress_level=4)
            counter += 1
    
    def return_PIL(self,images):

        """ Converts LCM Model Tensor results to a PIL list.

        Converts the Tensor list generated through the LCM model to a PIL list.

        images : LCM pipeline output tensor

        """

        images = images[0]

        if images.ndim == 3:
            images = images.numpy()[None, ...]
        else :
            images = images.numpy()
        images = (images * 255).round().astype("uint8")
        if images.shape[-1] == 1:
            # special case for grayscale (single channel) images
            pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images]
        else:
            pil_images = [Image.fromarray(image) for image in images]

        return pil_images
    
    def i2i_batch_save(self,images_list,base_name):

        """ Save img for img2img result values.

        First clean up the Tensor list generated by img2img function.

        and save img2img result.

        images_list : LCM img2img pipeline output tensor list.

        base_name : base name for saving img. (ex : {base_name}_{save_name}_1.png)

        """

        counter = 0
        for images in images_list:
            images = self.return_PIL(images=images)
            self.save_PIL(images,base_name + f"_{counter}")

    def make(self,model_path:str,lora_path:str=None,lora_name:str=None,**kwargs):
        
        """ Process user input and forward it to the LCM pipeline.

        Forward variable values for image creation to the LCM pipeline and return the corresponding result values
        
        model_path : path for model (huggingface repo id or path str)
        
        model_type : type of model ("LCM","SD","SDXL","SSD-1B")
        
        lora_path : path of lora file (ex : "./path/for/lora")
        
        lora_name : name for lora file (ex : "test.safetensor")
        
        input_image_dir : (only for img2img) input image dir

        output_image_dir : output image dir (it will not make dir)

        prompt : model input prompt (ex : "masterpeice, best quality, anime style" )
        
        n_prompt : model negative input prompt (ex : "nsfw,nude,bad quality" )

        seed : seed for LCM model (input 0 will make random seed)

        steps : steps for LCM model (recommend 2~4)

        cfg : cfg for LCM model (recommend 6~8)

        height , width : setting height and width for img (** if you are using img2img w and h should be the same as the input image. **)

        num_images : How many images will you create for this input

        prompt_strength : (only for img2img) How Strong will the prompts be applied in the img2img feature

        """
        if "input_image_dir" in kwargs:
            input_image_dir = kwargs.get("input_image_dir")
        else:
            input_image_dir = "./input.jpg"
        
        if "prompt" in kwargs:
            prompt = kwargs.get("prompt")
        else :
            prompt = "masterpiece"
        if "n_prompt" in kwargs:
            n_prompt = kwargs.get("n_prompt")
        else :
            n_prompt = "nsfw, nude, worst quality, low quality"
        if "seed" in kwargs:
            seed = kwargs.get("seed")
        else : 
            seed = 0
        if "steps" in kwargs:
            steps = kwargs.get("steps")
        else :
            steps = 4
        if "cfg" in kwargs:
            cfg = kwargs.get("cfg")
        else :
            cfg = 2
        if "width" in kwargs:
            width = kwargs.get("width")
        else :
            width = 512
        if "height" in kwargs:
            height = kwargs.get("height")
        else :
            height = 512
        if "num_images" in kwargs:
            num_images = kwargs.get("num_images")
        else :
            num_images = 1
        if "prompt_strength" in kwargs:
            prompt_strength = kwargs.get("prompt_strength")
        else :
            prompt_strength = 0.5
        
        if seed == 0 :
            seed = self.make_seed(0,True)
                
        if not self.mode_img2img :
            images = self.pipeline.sample(model_path=model_path, lora_path=lora_path, lora_name=lora_name, model_type=self.mode,seed=seed,steps=steps,cfg=cfg, prompt_strength=prompt_strength,
                            positive_prompt=prompt, negative_prompt=n_prompt, height=height,width=width,num_images=num_images,use_fp16=self.user_fp16,device=self.user_device)
        elif self.mode_img2img :
            images = self.pipeline.sample(model_path=model_path, lora_path=lora_path, lora_name=lora_name, model_type=self.mode, seed=seed,steps=steps,cfg=cfg, prompt_strength=prompt_strength, 
                            images=self.__load_img(input_image_dir,width,height),positive_prompt=prompt, negative_prompt=n_prompt, height=height,width=width,num_images=num_images,use_fp16=self.user_fp16,device=self.user_device)
        
        return images


    def make_image(self,model_path:str=None,lora_path:str=None,lora_name:str=None,output_image_dir:str=".",**kwargs):
        
        """ Most Simplified Image Generation Function

        Save the image generated by the txt2img, img2img function as a separate file based on user input.

        the output img will be save like output_image_dir/fesd_0.png(txt2img) or output_image_dir/fesd_i2i_0_0.png(img2img)
        
        model_path : path for model (huggingface repo id or path str)
        
        model_type : type of model ("LCM","SD","SDXL","SSD-1B")
        
        lora_path : path of lora file (ex : "./path/for/lora")
        
        lora_name : name for lora file (ex : "test.safetensor")
        
        input_image_dir : (only for img2img) input image dir

        output_image_dir : output image dir (it will not make dir)

        prompt : model input prompt (ex : "masterpeice, best quality, anime style" )
        
        n_prompt : model negative input prompt (ex : "nsfw,nude,bad quality" )

        seed : seed for LCM model (input 0 will make random seed)

        steps : steps for LCM model (recommend 2~4)

        cfg : cfg for LCM model (recommend 6~8)

        height , width : setting height and width for img (** if you are using img2img w and h should be the same as the input image. **)

        num_images : How many images will you create for this input

        prompt_strength : (only for img2img) How Strong will the prompts be applied in the img2img feature

        """
        
        images = self.make(model_path=model_path,lora_path=lora_path,lora_name=lora_name,**kwargs)

        if not self.mode_img2img and images is not None:
            
            pil_images = self.return_PIL(images)

            self.save_PIL(pils=pil_images,save_name=output_image_dir + "/fesd")

            return True

        elif self.mode_img2img and images is not None:
            
            self.i2i_batch_save(images_list=images,base_name=output_image_dir + "/fesd_i2i")

            return True
        
        else :
            return False
