import sys
import os.path as osp
import os
import argparse
sys.path.append(osp.dirname(osp.dirname(osp.abspath(__file__))))

from utils.cv import batch_load_masks
import os.path as osp
from utils.visualize import visualize_segs, visualize_segs_with_labels
from live2d.scrap_model import VALID_BODY_PARTS_V2
from utils.io_utils import find_all_imgs, imglist2imgrid
import numpy as np
from tqdm import tqdm
from PIL import Image

src_dir = 'workspace/datasets/testcase'
mask_dir = 'workspace/datasets/evalsam_iter0'

VALID_BODY_PARTS_V1 = [
    'hair', 'headwear', 'face', 'eyes', 'eyewear', 'ears', 'earwear', 'nose', 'beard', 'mouth', 'neck', 'neckwear',
    'skin', 'topwear', 'handwear',
    'bottomwear', 'legwear', 'footwear', 
    'tail', 'wings'
]

VALID_BODY_PARTS_V2 = [
    'hair', 'headwear', 'face', 'eyes', 'eyewear', 'ears', 'earwear', 'nose', 'mouth', 
    'neck', 'neckwear', 'topwear', 'handwear', 'bottomwear', 'legwear', 'footwear', 
    'tail', 'wings', 'objects'
]

# mask_valid = []

# reind_mask = not mask_dir.endswith('_iter2')

# for srcp in tqdm(find_all_imgs(src_dir, abs_path=True)):
#     maskp = osp.join(mask_dir, osp.basename(srcp) + '_masks.json')

#     masks = batch_load_masks(maskp)

#     if reind_mask:
#         mask_reind = []
#         for t in VALID_BODY_PARTS_V2:
#             if t in VALID_BODY_PARTS_V1:
#                 mask_reind.append(masks[VALID_BODY_PARTS_V1.index(t)])
#             else:
#                 mask_reind.append(np.zeros_like(masks[0]))
#         masks = mask_reind

#     masks[-4] = np.zeros_like(masks[-4])

#     img = np.array(Image.open(srcp).convert('RGB'))

#     Image.fromarray(visualize_segs(masks, img, image_weight=0.2, )).save(osp.join(mask_dir, osp.basename(srcp) + '.png'))



# compose_grid
src_list = ['workspace/datasets/evalsam_iter0', 'workspace/datasets/evalsam_iter1', 'workspace/datasets/evalsam_iter2']
save_dir = 'workspace/datasets/sam_visual_comparison'
os.makedirs(save_dir, exist_ok=True)

imgw = 768
for imgn in tqdm(find_all_imgs(src_list[0])):
    img_list = []
    for d in src_list:
        imgp = osp.join(d, imgn)
        img = Image.open(imgp).convert('RGB')
        w, h = img.width, img.height
        nh = int(round(imgw / w * h))
        nsize = (imgw, nh)
        img = img.resize(nsize, Image.Resampling.LANCZOS)
        img_list.append(np.array(img))
    img = np.concatenate(img_list, axis=1)
    savep = osp.join(save_dir, osp.splitext(imgn)[0] + '.png')
    Image.fromarray(img).save(savep)