import os
import yaml
import torch
import torchvision
from tqdm import tqdm
from inference.utils import *
from train import ControlNetCore, WurstCoreB
import warnings
warnings.filterwarnings("ignore")
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")



class Upscale_CaseCade:
    def __init__(self) -> None:
        self.config_file = './configs/inference/controlnet_c_3b_sr.yaml'
        # SETUP STAGE C
        with open(self.config_file, "r", encoding="utf-8") as file:
            loaded_config = yaml.safe_load(file)
        self.core = ControlNetCore(config_dict=loaded_config, device=device, training=False)
        # SETUP STAGE B
        self.config_file_b = './configs/inference/stage_b_3b.yaml'
        with open(self.config_file_b, "r", encoding="utf-8") as file:
            self.config_file_b = yaml.safe_load(file)
        self.core_b = WurstCoreB(config_dict=self.config_file_b, device=device, training=False)
        self.extras = self.core.setup_extras_pre()
        self.models = self.core.setup_models(self.extras)
        self.models.generator.eval().requires_grad_(False)
        print("CONTROLNET READY")
        self.extras_b = self.core_b.setup_extras_pre()
        self.models_b = self.core_b.setup_models(self.extras_b, skip_clip=True)
        self.models_b = WurstCoreB.Models(
        **{**self.models_b.to_dict(), 'tokenizer': self.models.tokenizer, 'text_model': self.models.text_model}
        )
        self.models_b.generator.eval().requires_grad_(False)
        print("STAGE B READY")


    def upscale_image(self,image_pil,scale_fator):
        batch_size = 1
        cnet_override = None
        images = resize_image(image_pil).unsqueeze(0).expand(batch_size, -1, -1, -1)
        
        batch = {'images': images}

        with torch.no_grad(), torch.cuda.amp.autocast(dtype=torch.bfloat16):
            effnet_latents = self.core.encode_latents(batch, self.models, self.extras)
            effnet_latents_up = torch.nn.functional.interpolate(effnet_latents, scale_factor=scale_fator, mode="nearest")
            cnet = self.models.controlnet(effnet_latents_up)
            cnet_uncond = cnet
            cnet_input = torch.nn.functional.interpolate(images, scale_factor=scale_fator, mode="nearest")
            # cnet, cnet_input = core.get_cnet(batch, models, extras)
            # cnet_uncond = cnet
        og=show_images(batch['images'],return_images=True)
        upsclae=show_images(cnet_input,return_images=True)
        return og,upsclae