import glob
import os
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from torchvision import transforms
from tqdm import tqdm
import joblib
from huggingface_hub import hf_hub_download

from .load_u2net import load_cloth_segm_model
from .utils import NormalizeImage, naive_cutout, resize_by_bigger_index, image_resize


def segment_garment(inputs_dir, outputs_dir, cls="all"):
    os.makedirs(outputs_dir, exist_ok=True)

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    transform_fn = transforms.Compose(
        [transforms.ToTensor(),
         NormalizeImage(0.5, 0.5)]
    )

    # load model from huggingface
    file_path = hf_hub_download(repo_id="tryonlabs/u2net-cloth-segmentation", filename="u2net_cloth_segm.pth")
    
    print("model loaded from huggingface:", file_path)
    
    net = load_cloth_segm_model(device, file_path, in_ch=3, out_ch=4)

    images_list = sorted(os.listdir(inputs_dir))
    pbar = tqdm(total=len(images_list))

    for image_name in images_list:
        img = Image.open(os.path.join(inputs_dir, image_name)).convert('RGB')
        img_size = img.size
        img = img.resize((768, 768), Image.BICUBIC)
        image_tensor = transform_fn(img)
        image_tensor = torch.unsqueeze(image_tensor, 0)

        with torch.no_grad():
            output_tensor = net(image_tensor.to(device))
            output_tensor = F.log_softmax(output_tensor[0], dim=1)
            output_tensor = torch.max(output_tensor, dim=1, keepdim=True)[1]
            output_tensor = torch.squeeze(output_tensor, dim=0)
            output_arr = output_tensor.cpu().numpy()

        if cls == "all":
            classes_to_save = []

            # Check which classes are present in the image
            for cls in range(1, 4):  # Exclude background class (0)
                if np.any(output_arr == cls):
                    classes_to_save.append(cls)
        elif cls == "upper":
            classes_to_save = [1]
        elif cls == "lower":
            classes_to_save = [2]
        elif cls == "dress":
            classes_to_save = [3]
        else:
            raise ValueError(f"Unknown cls: {cls}")

        for cls1 in classes_to_save:
            alpha_mask = (output_arr == cls1).astype(np.uint8) * 255
            alpha_mask = alpha_mask[0]  # Selecting the first channel to make it 2D
            alpha_mask_img = Image.fromarray(alpha_mask, mode='L')
            alpha_mask_img = alpha_mask_img.resize(img_size, Image.BICUBIC)
            alpha_mask_img.save(os.path.join(outputs_dir, f'{image_name.split(".")[0]}_{cls1}.jpg'))

        pbar.update(1)

    pbar.close()


def extract_garment(inputs_dir, outputs_dir, cls="all", resize_to_width=None):
    os.makedirs(outputs_dir, exist_ok=True)
    cloth_mask_dir = os.path.join(Path(outputs_dir).parent.absolute(), "cloth-mask")
    os.makedirs(cloth_mask_dir, exist_ok=True)

    segment_garment(inputs_dir, os.path.join(Path(outputs_dir).parent.absolute(), "cloth-mask"), cls=cls)

    images_path = sorted(glob.glob(os.path.join(inputs_dir, "*")))
    masks_path = sorted(glob.glob(os.path.join(cloth_mask_dir, "*")))
    img = Image.open(images_path[0])

    for mask_path in masks_path:
        mask = Image.open(mask_path)

        cutout = np.array(naive_cutout(img, mask))
        cutout = resize_by_bigger_index(cutout)

        canvas = np.ones((1024, 768, 3), np.uint8) * 255
        y1, y2 = (canvas.shape[0] - cutout.shape[0]) // 2, (canvas.shape[0] + cutout.shape[0]) // 2
        x1, x2 = (canvas.shape[1] - cutout.shape[1]) // 2, (canvas.shape[1] + cutout.shape[1]) // 2

        alpha_s = cutout[:, :, 3] / 255.0
        alpha_l = 1.0 - alpha_s

        for c in range(0, 3):
            canvas[y1:y2, x1:x2, c] = (alpha_s * cutout[:, :, c] +
                                       alpha_l * canvas[y1:y2, x1:x2, c])

        # resize image before saving
        if resize_to_width:
            canvas = image_resize(canvas, width=resize_to_width)

        canvas = Image.fromarray(canvas)

        canvas.save(os.path.join(outputs_dir, f"{os.path.basename(mask_path).split('.')[0]}.jpg"), format='JPEG')