import torch
import torch.nn as nn
import torch.nn.functional as F

from copy import deepcopy
from torchvision.transforms.functional import rgb_to_grayscale
import segmentation_models_pytorch as smp
from diffusers import StableDiffusionInpaintPipeline
from diffusers.utils.torch_utils import randn_tensor
from transformers import PretrainedConfig, PreTrainedModel

class SDGrayInpaintConfig(PretrainedConfig):
    model_type = "sd_gray_inpaint"
    def __init__(
        self,
        base_model="stabilityai/stable-diffusion-2-inpainting",
        height=512,
        width=512,
        **kwargs
    ):
        self.base_model=base_model
        self.height=height
        self.width=width
        super().__init__(**kwargs)

class SDGrayInpaintModel(PreTrainedModel):
    config_class = SDGrayInpaintConfig
    def __init__(self, config):
        super().__init__(config)
        pipe = StableDiffusionInpaintPipeline.from_pretrained(config.base_model)
        self.mask_predictor = smp.Unet(
            encoder_name="mit_b4",        
            encoder_weights="imagenet",     
            in_channels=3,                  
            classes=1,                      
        )
        self.image_processor = pipe.image_processor
        self.scheduler = pipe.scheduler
        self.unet = pipe.unet
        self.vae = pipe.vae
        self.prompt_embeds = nn.Parameter(torch.randn(1,77,1024))
        self.height=config.height
        self.width=config.width

    def forward(
        self, 
        images_gray_masked, 
        masks=None, 
        num_inference_steps=250,
        seed=42,
        input_type='pil', 
        output_type='pil'
    ):
        generator = torch.Generator()
        generator.manual_seed(seed)
        if input_type=='pil':
            images_gray_masked = self.image_processor.preprocess(images_gray_masked, height=self.height, width=self.width).float()
        elif input_type=='pt':
            images_gray_masked=images_gray_masked
        else:
            raise ValueError('unsupported input_type')
        images_gray_masked = images_gray_masked.to(self.vae.device)
        if masks is None:
            masks_logits = self.mask_predictor(images_gray_masked)
            masks = (torch.sigmoid(masks_logits)>0.5)*1.
        masks = masks.float().to(self.vae.device)
        images_gray_masked = (1-masks) * images_gray_masked
        B, C, H, W = images_gray_masked.shape
        prompt_embeds = self.prompt_embeds.repeat(B,1,1)

        scheduler = deepcopy(self.scheduler)
        scheduler.set_timesteps(num_inference_steps=num_inference_steps, device=self.vae.device)
        masked_image_latents = self.vae.encode(images_gray_masked).latent_dist.mode() * self.vae.config.scaling_factor
        mask_latents = F.interpolate(masks, size=(self.unet.config.sample_size, self.unet.config.sample_size))
        latents = randn_tensor(masked_image_latents.shape, generator=generator).to(self.device) * self.scheduler.init_noise_sigma
        for t in scheduler.timesteps:
            latents = scheduler.scale_model_input(latents, t)
            latent_model_input = torch.cat([latents, mask_latents, masked_image_latents], dim=1)
            noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=prompt_embeds)[0]
            latents = scheduler.step(noise_pred, t, latents)[0]
        latents = latents / self.vae.config.scaling_factor
        images_gray_restored = self.vae.decode(latents.detach())[0]
        images_gray_restored = images_gray_masked * (1-masks) + images_gray_restored.detach() * masks
        images_gray_restored = rgb_to_grayscale(images_gray_restored)
        
        if output_type=='pil':
            images_gray_restored = self.image_processor.postprocess(images_gray_restored)
        elif output_type=='np':
            images_gray_restored = self.image_processor.postprocess(images_gray_restored, 'np')
        elif output_type=='pt':
            images_gray_restored = self.image_processor.postprocess(images_gray_restored, 'pt')
        elif output_type=='none':
            images_gray_restored = images_gray_restored
        else:
            raise ValueError('unsupported output_type')

        return images_gray_restored