import os
import glob
import os.path as osp
import numpy as np
import torch
import torch.nn.functional as F
from torch.utils.data import Dataset
from loguru import logger
import random
from src.utils.dataset import read_pretrain_gray

class PretrainDataset(Dataset):
    def __init__(self,
                 root_dir,
                 mode='train',
                 img_resize=None,
                 df=None,
                 img_padding=False,
                 frame_gap=2,
                 **kwargs):
        """
        Manage image pairs of KAIST Multispectral Pedestrian Detection Benchmark Dataset.
        
        Args:
            root_dir (str): KAIST Multispectral Pedestrian  root directory that has `phoenix`.
            mode (str): options are ['train', 'val']
            img_resize (int, optional): the longer edge of resized images. None for no resize. 640 is recommended.
                                        This is useful during training with batches and testing with memory intensive algorithms.
            df (int, optional): image size division factor. NOTE: this will change the final image size after img_resize.
            img_padding (bool): If set to 'True', zero-pad the image to squared size. This is useful during training.
            augment_fn (callable, optional): augments images with pre-defined visual effects.
        """
        super().__init__()
        self.root_dir = root_dir
        self.mode = mode

        # specify which part of the data is used for trainng and testing
        if mode == 'train':
            assert img_resize is not None and img_padding 
            self.start_ratio = 0.0
            self.end_ratio = 0.9
        elif mode == 'val':
            assert img_resize is not None and img_padding 
            self.start_ratio = 0.9
            self.end_ratio = 1.0
        else:
            raise NotImplementedError()
        
        # parameters for image resizing, padding 
        self.img_resize = img_resize
        self.df = df
        self.img_padding = img_padding

        # for training XoFTR
        self.coarse_scale = getattr(kwargs, 'coarse_scale', 0.125)

        self.pair_paths = self.generate_kaist_pairs(root_dir, frame_gap=frame_gap, second_frame_range=0)

    def get_kaist_image_paths(self, root_dir):
        vis_img_paths = []
        lwir_img_paths = []
        img_num_per_folder = []

        # Recursively search for folders named "image"
        for folder, subfolders, filenames in os.walk(root_dir):
            if "visible" in subfolders and "lwir" in subfolders:
                vis_img_folder = osp.join(folder, "visible")
                lwir_img_folder = osp.join(folder, "lwir")
                # Use glob to find image files (you can add more extensions if needed)
                vis_imgs_i = glob.glob(osp.join(vis_img_folder, '*.jpg'))
                vis_imgs_i.sort()
                lwir_imgs_i = glob.glob(osp.join(lwir_img_folder, '*.jpg'))      
                lwir_imgs_i.sort()     
                vis_img_paths.append(vis_imgs_i)
                lwir_img_paths.append(lwir_imgs_i)
                img_num_per_folder.append(len(vis_imgs_i))
                assert len(vis_imgs_i) == len(lwir_imgs_i), f"Image numbers do not match in {folder}, {len(vis_imgs_i)} != {len(lwir_imgs_i)}"
                # Add more image file extensions as necessary
        return vis_img_paths, lwir_img_paths, img_num_per_folder
    
    def generate_kaist_pairs(self, root_dir, frame_gap, second_frame_range):
        """ generate image pairs (Vis-TIR) from KAIST Pedestrian dataset
        Args:
            root_dir: root directory for the dataset
            frame_gap (int): the frame gap between consecutive images 
            second_frame_range (int): the range for second image i.e. for the first ind i, second ind j element of [i-10, i+10]
        Returns:
            pair_paths (list)
        """
        vis_img_paths, lwir_img_paths, img_num_per_folder = self.get_kaist_image_paths(root_dir)
        pair_paths = []
        for i in range(len(img_num_per_folder)):
            num_img = img_num_per_folder[i]
            inds_vis = torch.arange(int(self.start_ratio * num_img),
                                    int(self.end_ratio * num_img),
                                    frame_gap, dtype=int)
            if second_frame_range > 0:
                inds_lwir = inds_vis + torch.randint(-second_frame_range, second_frame_range, (inds_vis.shape[0],))
                inds_lwir[inds_lwir<int(self.start_ratio * num_img)] = int(self.start_ratio * num_img)
                inds_lwir[inds_lwir>int(self.end_ratio * num_img)-1] = int(self.end_ratio * num_img)-1
            else:
                inds_lwir = inds_vis
            for j, k in zip(inds_vis, inds_lwir):
                img_name0 = os.path.relpath(vis_img_paths[i][j], root_dir)
                img_name1 = os.path.relpath(lwir_img_paths[i][k], root_dir)

                if torch.rand(1) > 0.5:
                    img_name0, img_name1 = img_name1, img_name0

                pair_paths.append([img_name0, img_name1])
        
        random.shuffle(pair_paths)
        return pair_paths

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

    def __getitem__(self, idx):
        # read grayscale and normalized image, and mask. (1, h, w) and (h, w)
        img_name0 = osp.join(self.root_dir, self.pair_paths[idx][0])
        img_name1 = osp.join(self.root_dir, self.pair_paths[idx][1])

        if self.mode == "train" and torch.rand(1) > 0.5:
            img_name0, img_name1 = img_name1, img_name0

        image0, image0_norm, mask0, scale0, image0_mean, image0_std = read_pretrain_gray(
            img_name0, self.img_resize, self.df, self.img_padding, None)
        image1, image1_norm, mask1, scale1, image1_mean, image1_std = read_pretrain_gray(
            img_name1, self.img_resize, self.df, self.img_padding, None)

        data = {
            'image0': image0,  # (1, h, w)
            'image1': image1,
            'image0_norm': image0_norm,
            'image1_norm': image1_norm,
            'scale0': scale0,  # [scale_w, scale_h]
            'scale1': scale1,
            "image0_mean": image0_mean,
            "image0_std": image0_std,
            "image1_mean": image1_mean,
            "image1_std": image1_std,
            'dataset_name': 'PreTrain',
            'pair_id': idx,
            'pair_names': (self.pair_paths[idx][0], self.pair_paths[idx][1]),
        }

        # for XoFTR training
        if mask0 is not None:  # img_padding is True
            if self.coarse_scale:
                [ts_mask_0, ts_mask_1] = F.interpolate(torch.stack([mask0, mask1], dim=0)[None].float(),
                                                       scale_factor=self.coarse_scale,
                                                       mode='nearest',
                                                       recompute_scale_factor=False)[0].bool()
            data.update({'mask0': ts_mask_0, 'mask1': ts_mask_1})

        return data