import sys
import os.path as osp

sys.path.append(osp.dirname(osp.dirname(osp.abspath(__file__))))

from tqdm import tqdm

from utils.io_utils import *
from utils.cv import *
from utils.visualize import *


def deduplate_execlist(exec_list):
    uniquedict = {}
    for p in exec_list:
        k = osp.basename(p).split('NONE')[0]
        if k not in uniquedict:
            uniquedict[k] = {'flist': [], 'ids': []}
        uniquedict[k]['flist'].append(p) 
        uniquedict[k]['ids'].append(exec_list.index(p))   

    unique_list = []
    for k, v in uniquedict.items():
        flist = uniquedict[k]['flist']
        unique_list.append(flist[len(flist) // 2])

    return unique_list


def vis_parts(srcd: str, tag_list, nmax_samples=12, cols=4):
    partsd = osp.join(srcd, 'parts')
    rst_list = []
    nparts = 0
    for tag in tag_list:
        p = osp.join(partsd, tag + '_vis.png')
        if not osp.exists(p):
            continue
        img = Image.open(p)
        pil_draw_text(img, tag, point=(0, 0), font_size=128, stroke_width=12)
        rst_list.append(img)
        if len(rst_list) >= nmax_samples:
            vis = imglist2imgrid(rst_list, cols=cols)
            Image.fromarray(vis).save(osp.join(srcd, f'part_vis{nparts}.jpg'), q=97)
            rst_list = []
            nparts += 1

    if len(rst_list) > 0:
        vis = imglist2imgrid(rst_list, cols=cols)
        Image.fromarray(vis).save(osp.join(srcd, f'part_vis{nparts}.jpg'), q=97)




src_list = ['workspace/datasets/l2d_eval_oa/l2d_eval2_output0', 'workspace/datasets/l2d_eval2_output_woattn']
src_list = ['workspace/datasets/testcaseall_output', 'workspace/datasets/testcaseall_output_woattn']


src_list = ['workspace/datasets/l2d_eval_oa/l2d_eval2_output0', 'workspace/datasets/l2deval_sam3_ouput']
src_list = ['workspace/datasets/testcaseall_output', 'workspace/datasets/l2deval_sam3_ouput']


dedupliacte = False


save_dir = 'tmp/cmp_part_extr'

src = src_list[0]

exec_list = [osp.join(src, d) for d in os.listdir(src)]

# if osp.isfile(src):
#     exec_list = load_exec_list('workspace/datasets/eval_chunk3.txt')
# else:
#     exec_list = find_all_imgs_recursive(src)

if dedupliacte:
    exec_list = deduplate_execlist(exec_list)

sz = (448, 448)

for srcp in tqdm(exec_list):
    src_name = osp.basename(srcp)
    
    flist = find_all_imgs(exec_list[0])
    sd = osp.join(save_dir, src_name)
    os.makedirs(sd, exist_ok=True)
    for filename in flist:
        if '_depth' in filename:
            continue
        row = []
        # if 'reconstruction' in filename:
        p = osp.join(src_list[0], src_name, 'src_img.png')
        img = pil_ensure_rgb(p)
        img = img.resize(sz, resample=Image.Resampling.LANCZOS)
        img = np.array(img)
        row.append(img)
        for srcd in src_list:
            p = osp.join(srcd, src_name, filename)
            if osp.exists(p):
                img = pil_ensure_rgb(p)
                img = img.resize(sz, resample=Image.Resampling.LANCZOS)
                img = np.array(img)
            else:
                img = np.full((sz[1], sz[0], 3), 255, np.uint8)
            row.append(img)
        row = np.concatenate(row, axis=1)

        savep = osp.join(sd, filename)
        save_tmp_img(row, savep)
        

pass