import os
import os.path as osp 
import json
import openpyxl
from openpyxl.drawing.image import Image
from openpyxl.utils import get_column_letter
from openpyxl.styles import Alignment
import random
import torch

import copy
from diffusers.utils import load_image

from customization_kontext.dataset import TestDataset
from customization_kontext.utils import * 


class Evaluator:
    def __init__(
            self,
            json_file,
            is_combine=False,
            image_size=512,    # only used for visualization
            gen_num=1,
            num_inference_steps=28,
        ):
        self.json_file = json_file
        self.test_set = None
        self.image_size = image_size
        self.gen_num = gen_num
        self.is_combine = is_combine
        self.pipe = None
        self.num_inference_steps = num_inference_steps

    def evaluate(self, pipe, output_dir):
        self.test_set = TestDataset(self.json_file)
        self.pipe = pipe
        res = []
        os.makedirs(output_dir, exist_ok=True)
        aspect_ratio_image = eval(f"ASPECT_RATIO_{self.image_size}")
        for idx, sample in enumerate(self.test_set):
            res_dict = copy.deepcopy(sample) 
            # 获取样本中的相关信息
            ref_image = sample['ref_images'][0]
            ref_image_path = ref_image["image_path"]
            input_image = load_image(ref_image_path)
            height = 512
            width = 512
            prompt = sample['prompt']
            res_dict["output_image"] = []
            # 从1-10000随机生成seed
            seed = [42]
            seed.extend([random.randint(1, 10000) for _ in range(self.gen_num-1)])
            for i in range(self.gen_num):
                torch.manual_seed(seed[i])
                output_name = generate_image_name(ref_image_path, prompt, i)
                output_path = os.path.join(output_dir, output_name)
                if not os.path.exists(output_path):
                    # 生成图像
                    result_img = self.pipe(
                        image=input_image,
                        prompt=prompt,
                        num_inference_steps=self.num_inference_steps
                    ).images[0]
                    result_img = result_img.resize((width, height))
                    if self.is_combine:
                        combined_img = PILImage.new("RGB", (width * 2, height))
                        input_image = input_image.resize((width, height))
                        combined_img.paste(input_image, (0, 0))
                        combined_img.paste(result_img, (width, 0))
                        combined_img.save(output_path)
                    else:
                        result_img.save(output_path)
                res_dict["output_image"].append(os.path.abspath(output_path))
            res.append(res_dict)
        
        with open(os.path.join(output_dir, "res.json"), "w") as f:
            json.dump(res, f, indent=4, ensure_ascii=False)