import os
import torch
from cleanfid import fid as FID
from PIL import Image
from torch.utils.data import Dataset
from torchmetrics.image import StructuralSimilarityIndexMeasure
from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
from torchvision import transforms
from tqdm import tqdm

from prettytable import PrettyTable

def scan_files_in_dir(directory, postfix: set[str] = None, progress_bar: tqdm = None) -> list:
    file_list = []
    progress_bar = tqdm(total=0, desc=f"Scanning", ncols=100) if progress_bar is None else progress_bar
    for entry in os.scandir(directory):
        if entry.is_file():
            if postfix is None or os.path.splitext(entry.path)[1] in postfix:
                file_list.append(entry)
                progress_bar.total += 1
                progress_bar.update(1)
        elif entry.is_dir():
            file_list += scan_files_in_dir(entry.path, postfix=postfix, progress_bar=progress_bar)
    return file_list

class EvalDataset(Dataset):
    def __init__(self, gt_folder, pred_folder, height=1024):
        self.gt_folder = gt_folder
        self.pred_folder = pred_folder
        self.height = height
        self.data = self.prepare_data()
        self.to_tensor = transforms.ToTensor()
       
        self.clip_transforms = transforms.Compose([
            transforms.Resize((224, 224),interpolation=2),  # 调整图像大小以匹配模型输入
            #transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
        """
        self.clip_transforms = transforms.Compose([
            transforms.Resize((224, 224)), 
            transforms.Normalize((0.48145466, 0.4578275, 0.40821073),
                                (0.26862954, 0.26130258, 0.27577711))
        ])
        """
        
    def extract_id_from_filename(self, filename):
        # find first number in filename
        start_i = None
        for i, c in enumerate(filename):
            if c.isdigit():
                start_i = i
                break
        if start_i is None:
            assert False, f"Cannot find number in filename {filename}"
        return filename[start_i:start_i+8]
    
    def prepare_data(self):
        gt_files = scan_files_in_dir(self.gt_folder, postfix={'.jpg', '.png'})
        gt_dict = {self.extract_id_from_filename(file.name): file for file in gt_files}
        pred_files = scan_files_in_dir(self.pred_folder, postfix={'.jpg', '.png'})
        
        tuples = []
        for pred_file in pred_files:
            pred_id = self.extract_id_from_filename(pred_file.name)
            if pred_id not in gt_dict:
                print(f"Cannot find gt file for {pred_file}")
            else:
                tuples.append((gt_dict[pred_id].path, pred_file.path))
        return tuples
        
    def resize(self, img):
        w, h = img.size
        new_w = int(w * self.height / h)
        return img.resize((new_w, self.height), Image.LANCZOS)
        #return transforms.Resize((new_w, self.height), interpolation=2)(img)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        gt_path, pred_path = self.data[idx]
        gt, pred = self.resize(Image.open(gt_path)), self.resize(Image.open(pred_path))
        if gt.height != self.height:
            gt = self.resize(gt)
        if pred.height != self.height:
            pred = self.resize(pred)
        gt = self.to_tensor(gt)
        pred = self.to_tensor(pred)
        clip_gt = self.clip_transforms(gt)
        clip_pred = self.clip_transforms(pred)
        return gt, pred, clip_gt, clip_pred


def copy_resize_gt(gt_folder, height):
    new_folder = f"{gt_folder}_{height}"
    if not os.path.exists(new_folder):
        os.makedirs(new_folder, exist_ok=True)
    for file in tqdm(os.listdir(gt_folder)):
        if os.path.exists(os.path.join(new_folder, file)):
            continue
        img = Image.open(os.path.join(gt_folder, file))
        w, h = img.size
        new_w = int(w * height / h)
        img = img.resize((new_w, height), Image.LANCZOS)
        #img = transforms.Resize((new_w, height), interpolation=2)(img) 
        #img = img.resize((new_w, height))

        img.save(os.path.join(new_folder, file))
    return new_folder

#cuda = 7
@torch.no_grad()
def ssim(dataloader):
    device = torch.device(f"cuda" if torch.cuda.is_available() else "cpu")
    #device = "cuda:2"
    ssim_score = 0
    ssim = StructuralSimilarityIndexMeasure(data_range=1.0).to(device)
    for gt, pred,_,_ in tqdm(dataloader, desc="Calculating SSIM"):
        batch_size = gt.size(0)
        gt, pred = gt.to(device), pred.to(device)
        ssim_score += ssim(pred, gt) * batch_size
    return ssim_score / len(dataloader.dataset)


@torch.no_grad()
def lpips(dataloader):
    device = torch.device(f"cuda" if torch.cuda.is_available() else "cpu")
    #device = "cuda:2"
    lpips_score = LearnedPerceptualImagePatchSimilarity(net_type='squeeze').to(device)
    score = 0
    for gt, pred,_,_ in tqdm(dataloader, desc="Calculating LPIPS"):
        batch_size = gt.size(0)
        pred = pred.to(device)
        gt = gt.to(device)
        # LPIPS needs the images to be in the [-1, 1] range.
        gt = (gt * 2) - 1
        pred = (pred * 2) - 1
        score += lpips_score(gt, pred) * batch_size
    return score / len(dataloader.dataset)


@torch.no_grad()
def clip_score(dataloader):
    # 初始化 CLIPScore
    from transformers import CLIPModel, CLIPProcessor
    import open_clip

    #model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32") # 98
    model = CLIPModel.from_pretrained("openai/clip-vit-large-patch14")
    #model, _, _ = open_clip.create_model_and_transforms('ViT-H-14', pretrained='laion2b_s32b_b79k')
    #processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14")
    device = torch.device(f"cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)

    total_score = 0.0

    for _,_,gt, pred in tqdm(dataloader, desc="Calculating CLIP Score"):
        batch_size = gt.size(0)
        gt = gt.to(device)
        pred = pred.to(device)

        # 提取图像特征
        gt_features = model.get_image_features(gt)
        pred_features = model.get_image_features(pred)
        #gt_features = model.encode_image(gt)
        #pred_features = model.encode_image(pred)

        # 计算余弦相似度
        similarity = torch.nn.functional.cosine_similarity(gt_features, pred_features)
        total_score += similarity.sum().item()

    return total_score / len(dataloader.dataset)



@torch.no_grad()
def dino_score(dataloader, model_name="dinov2_vitl14"):
    """
    计算 DINO 分数
    :param dataloader: 数据加载器，返回 (gt, pred) 图像对
    :param model_name: DINOv2 模型名称
    :param cuda: GPU 设备索引
    """
    # 加载预训练的 DINOv2 模型  95
    #model = torch.hub.load('facebookresearch/dinov2', model_name)
    model = torch.hub.load('facebookresearch/dinov2', model_name)

    device = torch.device(f"cuda" if torch.cuda.is_available() else "cpu")
    #device = "cuda:2"
    model.to(device)

    # 定义图像预处理
    

    total_score = 0.0

    for _, _, gt, pred in tqdm(dataloader, desc="Calculating DINO Score"):
        gt = gt.to(device)
        pred = pred.to(device)

        # 提取图像特征
        gt_features = model(gt)
        pred_features = model(pred)

        # 计算余弦相似度
        similarity = torch.nn.functional.cosine_similarity(gt_features, pred_features)
        total_score += similarity.sum().item()

    return total_score / len(dataloader.dataset)



def eval(args):
    # Check gt_folder has images with target height, resize if not
    pred_sample = os.listdir(args.pred_folder)[0]
    gt_sample = os.listdir(args.gt_folder)[0]
    img = Image.open(os.path.join(args.pred_folder, pred_sample))
    gt_img = Image.open(os.path.join(args.gt_folder, gt_sample))
    if img.height != gt_img.height:
        title = "--"*30 + f"Resizing GT Images to height {img.height}" + "--"*30
        print(title)
        args.gt_folder = copy_resize_gt(args.gt_folder, img.height)
        print("-"*len(title))
    
    # Form dataset
    dataset = EvalDataset(args.gt_folder, args.pred_folder, img.height)
    dataloader = torch.utils.data.DataLoader(
        dataset, batch_size=args.batch_size, num_workers=args.num_workers, shuffle=False, drop_last=False
    )
    
    # Calculate Metrics
    header = []
    row = []
    
    if args.paired:
        header += ['dino_score',"CLIP_Score","SSIM", "LPIPS"]
        dino_ = dino_score(dataloader)
        clip_ = 0#clip_score(dataloader)
        ssim_ = ssim(dataloader).item()
        lpips_ = lpips(dataloader).item()
        row += [dino_,clip_ ,ssim_, lpips_]

    header += ["FID", "KID"]
    #os.environ["FID_WEIGHTS_PATH"] = "~/.cache/torch/hub/checkpoints"
    fid_ = FID.compute_fid(args.gt_folder, args.pred_folder)
    kid_ = FID.compute_kid(args.gt_folder, args.pred_folder) * 1000
    #kid_ = 0
    row += [fid_, kid_]
    
    # Print Results
    print("GT Folder  : ", args.gt_folder)
    print("Pred Folder: ", args.pred_folder)
    table = PrettyTable()
    table.field_names = header
    table.add_row(row)
    print(table)
    
         
if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--gt_folder", type=str, required=True)
    parser.add_argument("--pred_folder", type=str, required=True)
    parser.add_argument("--paired", action="store_true")
    parser.add_argument("--batch_size", type=int, default=16)
    parser.add_argument("--num_workers", type=int, default=4)
    args = parser.parse_args()
    
    eval(args)