import torch
from torchvision import transforms
from diffusers import DDPMScheduler, AutoencoderKL
from transformers import (
    CLIPImageProcessor,
    CLIPVisionModelWithProjection,
    CLIPTextModel,
    CLIPTextModelWithProjection,
    AutoTokenizer,
)
from src.tryon.src.tryon_pipeline import (
    StableDiffusionXLInpaintPipeline as TryonPipeline,
)
from src.tryon.src.unet_hacked_garmnet import (
    UNet2DConditionModel as UNet2DConditionModel_ref,
)
from src.tryon.src.unet_hacked_tryon import UNet2DConditionModel
from src.tryon.preprocess.humanparsing.run_parsing import Parsing
from src.tryon.preprocess.openpose.run_openpose import OpenPose
from src.tryon.utils_mask import get_mask_location
from torchvision.transforms.functional import to_pil_image
from src.tryon.detectron2.data.detection_utils import (
    convert_PIL_to_numpy,
    _apply_exif_orientation,
)
import torch
from PIL import Image
from src.tryon import apply_net


pipe = None
UNet_Encoder = None
parsing_model = None
openpose_model = None
tensor_transform = None


def load_model():
    global pipe, UNet_Encoder, parsing_model, openpose_model, tensor_transform

    if pipe is not None:
        return {
            "pipe": pipe,
            "unet_encoder": UNet_Encoder,
            "parsing_model": parsing_model,
            "openpose_model": openpose_model,
            "tensor_transform": tensor_transform,
        }

    device = "cuda" if torch.cuda.is_available() else "cpu"
    base_path = "yisol/IDM-VTON"

    # Load sub-models
    unet = (
        UNet2DConditionModel.from_pretrained(
            base_path, subfolder="unet", torch_dtype=torch.float16
        )
        .eval()
        .requires_grad_(False)
    )

    tokenizer_one = AutoTokenizer.from_pretrained(
        base_path, subfolder="tokenizer", use_fast=False
    )
    tokenizer_two = AutoTokenizer.from_pretrained(
        base_path, subfolder="tokenizer_2", use_fast=False
    )

    noise_scheduler = DDPMScheduler.from_pretrained(base_path, subfolder="scheduler")

    text_encoder_one = CLIPTextModel.from_pretrained(
        base_path, subfolder="text_encoder", torch_dtype=torch.float16
    )
    text_encoder_two = CLIPTextModelWithProjection.from_pretrained(
        base_path, subfolder="text_encoder_2", torch_dtype=torch.float16
    )
    image_encoder = CLIPVisionModelWithProjection.from_pretrained(
        base_path, subfolder="image_encoder", torch_dtype=torch.float16
    )
    vae = AutoencoderKL.from_pretrained(
        base_path, subfolder="vae", torch_dtype=torch.float16
    )

    UNet_Encoder = UNet2DConditionModel_ref.from_pretrained(
        base_path, subfolder="unet_encoder", torch_dtype=torch.float16
    )
    UNet_Encoder.requires_grad_(False)

    parsing_model = Parsing(0)
    openpose_model = OpenPose(0)

    tensor_transform = transforms.Compose(
        [
            transforms.ToTensor(),
            transforms.Normalize([0.5], [0.5]),
        ]
    )

    pipe = TryonPipeline.from_pretrained(
        base_path,
        unet=unet,
        vae=vae,
        feature_extractor=CLIPImageProcessor(),
        text_encoder=text_encoder_one,
        text_encoder_2=text_encoder_two,
        tokenizer=tokenizer_one,
        tokenizer_2=tokenizer_two,
        scheduler=noise_scheduler,
        image_encoder=image_encoder,
        torch_dtype=torch.float16,
    )
    pipe.unet_encoder = UNet_Encoder

    pipe.to(device)

    return {
        "pipe": pipe,
        "unet_encoder": UNet_Encoder,
        "parsing_model": parsing_model,
        "openpose_model": openpose_model,
        "tensor_transform": tensor_transform,
    }


def unload_model():
    global pipe, UNet_Encoder, parsing_model, openpose_model, tensor_transform
    pipe = None
    UNet_Encoder = None
    parsing_model = None
    openpose_model = None
    tensor_transform = None
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def predict(
    # pipe,
    # UNet_Encoder,
    # parsing_model,
    # openpose_model,
    # tensor_transform,
    dict_img,
    garm_img,
    garment_des,
    is_checked,
    is_checked_crop,
    denoise_steps,
    seed,
    category,
):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    openpose_model.preprocessor.body_estimation.model.to(device)
    pipe.unet_encoder.to(device)
    pipe.to(device)

    # Resize and convert
    garm_img = garm_img.convert("RGB").resize((768, 1024))
    human_img_orig = dict_img["background"].convert("RGB")
    orig_size = human_img_orig.size

    if is_checked_crop:
        width, height = human_img_orig.size
        target_width = int(min(width, height * (3 / 4)))
        target_height = int(min(height, width * (4 / 3)))
        left = (width - target_width) / 2
        top = (height - target_height) / 2
        right = (width + target_width) / 2
        bottom = (height + target_height) / 2
        cropped_img = human_img_orig.crop((left, top, right, bottom))
        crop_size = cropped_img.size
        human_img = cropped_img.resize((768, 1024))
    else:
        human_img = human_img_orig.resize((768, 1024))

    # Mask generation
    if is_checked:
        keypoints = openpose_model(human_img.resize((384, 512)))
        model_parse, _ = parsing_model(human_img.resize((384, 512)))
        mask, _ = get_mask_location("hd", category, model_parse, keypoints)
        mask = mask.resize((768, 1024))
    else:
        if dict_img["layers"][0] is not None:
            mask = dict_img["layers"][0].convert("RGB").resize((768, 1024))
            mask = Image.fromarray(
                (torch.tensor(mask) > 0).numpy().astype("uint8") * 255
            )
        else:
            mask = Image.new("RGB", (768, 1024), color="white")

    # Mask gray
    mask_gray = (1 - transforms.ToTensor()(mask)) * tensor_transform(human_img)
    mask_gray = to_pil_image((mask_gray + 1.0) / 2.0)

    # Pose estimation
    human_np = convert_PIL_to_numpy(
        _apply_exif_orientation(human_img.resize((384, 512))), format="BGR"
    )
    args = apply_net.create_argument_parser().parse_args(
        (
            "show",
            "./tryon/configs/densepose_rcnn_R_50_FPN_s1x.yaml",
            "./tryon/ckpt/densepose/model_final_162be9.pkl",
            "dp_segm",
            "-v",
            "--opts",
            "MODEL.DEVICE",
            device,
        )
    )
    pose_img = args.func(args, human_np)
    pose_img = Image.fromarray(pose_img[:, :, ::-1]).resize((768, 1024))

    # Inference
    with torch.no_grad(), torch.cuda.amp.autocast():
        prompt = f"((best quality, masterpiece, ultra-detailed, high quality photography, photo realistic)), the model is wearing {garment_des}"
        neg_prompt = "monochrome, lowres, bad anatomy, worst quality, normal quality, low quality, blurry, jpeg artifacts, sketch"

        prompt_embeds, neg_embeds, pooled_prompt_embeds, neg_pooled_prompt_embeds = (
            pipe.encode_prompt(
                prompt,
                num_images_per_prompt=1,
                do_classifier_free_guidance=True,
                negative_prompt=neg_prompt,
            )
        )

        garment_prompt = f"((best quality, masterpiece, ultra-detailed, high quality photography, photo realistic)), a photo of {garment_des}"
        prompt_embeds_c, *_ = pipe.encode_prompt(
            [garment_prompt],
            num_images_per_prompt=1,
            do_classifier_free_guidance=False,
            negative_prompt=[neg_prompt],
        )

        pose_tensor = tensor_transform(pose_img).unsqueeze(0).to(device, torch.float16)
        garm_tensor = tensor_transform(garm_img).unsqueeze(0).to(device, torch.float16)
        generator = torch.Generator(device).manual_seed(seed) if seed != -1 else None

        images = pipe(
            prompt_embeds=prompt_embeds.to(device),
            negative_prompt_embeds=neg_embeds.to(device),
            pooled_prompt_embeds=pooled_prompt_embeds.to(device),
            negative_pooled_prompt_embeds=neg_pooled_prompt_embeds.to(device),
            num_inference_steps=denoise_steps,
            generator=generator,
            strength=1.0,
            pose_img=pose_tensor,
            text_embeds_cloth=prompt_embeds_c.to(device),
            cloth=garm_tensor,
            mask_image=mask,
            image=human_img,
            height=1024,
            width=768,
            ip_adapter_image=garm_img.resize((768, 1024)),
            guidance_scale=2.0,
        )[0]

    if is_checked_crop:
        return images[0].resize(crop_size), mask_gray.resize(crop_size)
    else:
        return images[0].resize(orig_size), mask_gray.resize(orig_size)


# from PIL import Image
# import torch
# from tryon_model import load_model, predict  # replace 'tryon_model' with your actual filename/module

# # Load the models
# pipe, UNet_Encoder, parsing_model, openpose_model, tensor_transform = load_model()

# # Load input images
# human_image = Image.open("path_to_human_image.jpg").convert("RGB")
# garment_image = Image.open("path_to_garment_image.jpg").convert("RGB")

# # Create mock dict_img
# dict_img = {
#     "background": human_image,
#     "layers": [None]  # or use a mask PIL.Image
# }

# # Set parameters
# garment_des = "a red floral dress with puff sleeves"
# is_checked = True             # Auto-generate mask
# is_checked_crop = True        # Crop input to 3:4 ratio
# denoise_steps = 30            # Diffusion steps
# seed = 42                     # Set -1 for random
# category = "upper"            # Use correct clothing type

# # Run prediction
# result_img, result_mask = predict(
#     pipe=pipe,
#     UNet_Encoder=UNet_Encoder,
#     parsing_model=parsing_model,
#     openpose_model=openpose_model,
#     tensor_transform=tensor_transform,
#     dict_img=dict_img,
#     garm_img=garment_image,
#     garment_des=garment_des,
#     is_checked=is_checked,
#     is_checked_crop=is_checked_crop,
#     denoise_steps=denoise_steps,
#     seed=seed,
#     category=category
# )

# # Save or show results
# result_img.save("tryon_output.png")
# result_mask.save("tryon_mask.png")
# print("Prediction complete. Saved output and mask.")