import os
from PIL import Image
import h5py
import numpy as np
import torch
import torchvision.transforms.functional as tvf
from torchvision.transforms._functional_tensor import affine as tv_affine
from tqdm import tqdm

from affine_steerers.utils import get_depth_tuple_transform_ops, get_tuple_transform_ops
import affine_steerers
from affine_steerers.utils import *


class MegadepthScene:
    def __init__(
        self,
        data_root,
        scene_info,
        ht=512,
        wt=512,
        min_overlap=0.0,
        max_overlap=1.0,
        shake_t=0,
        scene_info_detections=None,
        scene_info_detections3D=None,
        normalize=True,
        max_num_pairs=100_000,
        scene_name=None,
        use_horizontal_flip_aug=False,
        use_affine_aug=False,
        min_dilation=0.75,
        max_dilation=2.0,
        min_angle=-np.pi,
        max_angle=np.pi,
        same_img=False,
        grayscale=False,
        clahe=False,
    ) -> None:
        self.data_root = data_root
        self.scene_name = os.path.splitext(scene_name)[0] + f"_{min_overlap}_{max_overlap}"
        self.image_paths = scene_info["image_paths"]
        self.depth_paths = scene_info["depth_paths"]
        self.intrinsics = scene_info["intrinsics"]
        self.poses = scene_info["poses"]
        self.pairs = scene_info["pairs"]
        self.overlaps = scene_info["overlaps"]
        threshold = (self.overlaps > min_overlap) & (self.overlaps < max_overlap)
        self.pairs = self.pairs[threshold]
        self.overlaps = self.overlaps[threshold]
        self.detections = scene_info_detections
        self.tracks3D = scene_info_detections3D
        if len(self.pairs) > max_num_pairs:
            pairinds = np.random.choice(np.arange(0, len(self.pairs)), max_num_pairs, replace=False)
            self.pairs = self.pairs[pairinds]
            self.overlaps = self.overlaps[pairinds]
        self.im_transform_ops = get_tuple_transform_ops(
            resize=(ht, wt),
            normalize=normalize,
            clahe=clahe,
        )
        self.depth_transform_ops = get_depth_tuple_transform_ops(resize=(ht, wt), normalize=False)
        self.wt, self.ht = wt, ht
        self.shake_t = shake_t
        self.use_horizontal_flip_aug = use_horizontal_flip_aug
        self.use_affine_aug = use_affine_aug
        self.min_dilation = min_dilation
        self.max_dilation = max_dilation
        self.min_angle = min_angle
        self.max_angle = max_angle
        self.same_img = same_img
        self.grayscale = grayscale

    def load_im(self, im_B, crop=None):
        im = Image.open(im_B)
        return im

    def sample_affine(self):
        angle_1 = np.random.uniform(high=2 * np.pi)
        angle_2 = np.random.uniform(low=self.min_angle, high=self.max_angle)
        dilation_1 = np.exp(np.random.uniform(low=np.log(self.min_dilation), high=np.log(self.max_dilation)))
        dilation_2 = np.exp(np.random.uniform(low=np.log(self.min_dilation), high=np.log(self.max_dilation)))
        if np.random.uniform() < 0.5:
            dilation_1 *= -1
        return build_affine_double(angle_1, dilation_1, dilation_2, angle_2)

    def inv_affine_list(self, affine):
        inv = torch.linalg.inv(affine)
        return [inv[0, 0].item(), inv[0, 1].item(), 0.0, inv[1, 0].item(), inv[1, 1].item(), 0]

    def affine_around_center(self, affine):
        """Input 2x2 affine map for image center (0, 0), output 3x3 map for image center ((wt-1)/2, (ht-1)/2)."""
        T = torch.eye(3, device=affine.device)
        T[0, 2] = (self.wt - 1.0) / 2
        T[1, 2] = (self.ht - 1.0) / 2

        Tinv = torch.eye(3, device=affine.device)
        Tinv[0, 2] = -(self.wt - 1.0) / 2
        Tinv[1, 2] = -(self.ht - 1.0) / 2

        A = torch.eye(3, device=affine.device)
        A[:2, :2] = affine
        return T @ A @ Tinv

    def affine_aug(self, im_A, im_B, depth_A, depth_B, K_A, K_B, detections_A, detections_B):
        affine_A = self.sample_affine()
        affine_B = self.sample_affine()
        affine_AtoB = (affine_B @ torch.linalg.inv(affine_A)).to(im_A.device)

        affine_inv_list_A = self.inv_affine_list(affine_A)
        affine_inv_list_B = self.inv_affine_list(affine_B)

        im_A = tv_affine(im_A, affine_inv_list_A, interpolation="bilinear")
        im_B = tv_affine(im_B, affine_inv_list_B, interpolation="bilinear")

        depth_A = tv_affine(depth_A, affine_inv_list_A, interpolation="bilinear")
        depth_B = tv_affine(depth_B, affine_inv_list_B, interpolation="bilinear")

        coord_map_A = self.affine_around_center(affine_A).to(K_A.device)
        coord_map_B = self.affine_around_center(affine_B).to(K_B.device)

        K_A = coord_map_A @ K_A
        K_B = coord_map_B @ K_B

        if detections_A is not None:
            detections_A = detections_A @ coord_map_A[:2, :2].mT + coord_map_A[None, :2, 2]
        if detections_B is not None:
            detections_B = detections_B @ coord_map_B[:2, :2].mT + coord_map_B[None, :2, 2]

        return im_A, im_B, depth_A, depth_B, K_A, K_B, detections_A, detections_B, affine_AtoB

    def horizontal_flip(self, im_A, im_B, depth_A, depth_B, K_A, K_B):
        im_A = im_A.flip(-1)
        im_B = im_B.flip(-1)
        depth_A, depth_B = depth_A.flip(-1), depth_B.flip(-1)
        flip_mat = torch.tensor([[-1, 0, self.wt], [0, 1, 0], [0, 0, 1.0]]).to(K_A.device)
        K_A = flip_mat @ K_A
        K_B = flip_mat @ K_B

        return im_A, im_B, depth_A, depth_B, K_A, K_B

    def load_depth(self, depth_ref, crop=None):
        depth = np.array(h5py.File(depth_ref, "r")["depth"])
        return torch.from_numpy(depth)

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

    def scale_intrinsic(self, K, wi, hi):
        sx, sy = self.wt / wi, self.ht / hi
        sK = torch.tensor([[sx, 0, 0], [0, sy, 0], [0, 0, 1]])
        return sK @ K

    def scale_detections(self, detections, wi, hi):
        sx, sy = self.wt / wi, self.ht / hi
        return detections * torch.tensor([[sx, sy]])

    def rand_shake(self, *things):
        t = np.random.choice(range(-self.shake_t, self.shake_t + 1), size=(2))
        return [tvf.affine(thing, angle=0.0, translate=list(t), scale=1.0, shear=[0.0, 0.0]) for thing in things], t

    def tracks_to_detections(self, tracks3D, pose, intrinsics, H, W):
        tracks3D = tracks3D.double()
        intrinsics = intrinsics.double()
        bearing_vectors = pose[..., :3, :3] @ tracks3D.mT + pose[..., :3, 3:]
        hom_pixel_coords = (intrinsics @ bearing_vectors).mT
        pixel_coords = hom_pixel_coords[..., :2] / (hom_pixel_coords[..., 2:] + 1e-12)
        legit_detections = (
            (pixel_coords > 0).prod(dim=-1)
            * (pixel_coords[..., 0] < W - 1)
            * (pixel_coords[..., 1] < H - 1)
            * (tracks3D != 0).prod(dim=-1)
        )
        return pixel_coords.float(), legit_detections.bool()

    def __getitem__(self, pair_idx):
        try:
            # read intrinsics of original size
            idx1, idx2 = self.pairs[pair_idx]
            if self.same_img:
                idx2 = idx1
            K1 = torch.tensor(self.intrinsics[idx1].copy(), dtype=torch.float).reshape(3, 3)
            K2 = torch.tensor(self.intrinsics[idx2].copy(), dtype=torch.float).reshape(3, 3)

            # read and compute relative poses
            T1 = self.poses[idx1]
            T2 = self.poses[idx2]
            T_1to2 = torch.tensor(np.matmul(T2, np.linalg.inv(T1)), dtype=torch.float)[:4, :4]  # (4, 4)

            # Load positive pair data
            im_A, im_B = self.image_paths[idx1], self.image_paths[idx2]
            depth1, depth2 = self.depth_paths[idx1], self.depth_paths[idx2]
            im_A_ref = os.path.join(self.data_root, im_A)
            im_B_ref = os.path.join(self.data_root, im_B)
            depth_A_ref = os.path.join(self.data_root, depth1)
            depth_B_ref = os.path.join(self.data_root, depth2)
            # return torch.randn((1000,1000))
            im_A = self.load_im(im_A_ref)
            im_B = self.load_im(im_B_ref)
            depth_A = self.load_depth(depth_A_ref)
            depth_B = self.load_depth(depth_B_ref)

            # Recompute camera intrinsic matrix due to the resize
            W_A, H_A = im_A.width, im_A.height
            W_B, H_B = im_B.width, im_B.height

            detections2D_A = detections2D_B = tracks3D_A = tracks3D_B = []
            if self.detections is not None or self.tracks3D is not None:
                K = 10000
            if self.detections is not None:
                detections2D_A = self.detections[idx1]
                detections2D_B = self.detections[idx2]
            if self.tracks3D is not None:
                tracks3D_A = torch.zeros(K, 3)
                tracks3D_B = torch.zeros(K, 3)
                tracks3D_A[: len(detections2D_A)] = torch.tensor(self.tracks3D[detections2D_A[:K, -1].astype(np.int32)])
                tracks3D_B[: len(detections2D_B)] = torch.tensor(self.tracks3D[detections2D_B[:K, -1].astype(np.int32)])

            # projs_A, _ = self.tracks_to_detections(tracks3D_A, T1, K1, W_A, H_A)
            # tracks3D_B = torch.zeros(K,2)

            K1 = self.scale_intrinsic(K1, W_A, H_A)
            K2 = self.scale_intrinsic(K2, W_B, H_B)

            # Process images
            im_A, im_B = self.im_transform_ops((im_A, im_B))
            depth_A, depth_B = self.depth_transform_ops((depth_A[None, None], depth_B[None, None]))
            [im_A, depth_A], t_A = self.rand_shake(im_A, depth_A)
            [im_B, depth_B], t_B = self.rand_shake(im_B, depth_B)

            detections_A = detections_B = []
            if self.detections is not None:
                detections_A = -torch.ones(K, 2)
                detections_B = -torch.ones(K, 2)
                detections_A[: len(self.detections[idx1])] = (
                    self.scale_detections(torch.tensor(detections2D_A[:K, :2]), W_A, H_A) + t_A
                )
                detections_B[: len(self.detections[idx2])] = (
                    self.scale_detections(torch.tensor(detections2D_B[:K, :2]), W_B, H_B) + t_B
                )

            K1[:2, 2] += t_A
            K2[:2, 2] += t_B

            if affine_steerers.DEBUG_MODE:
                import matplotlib.pyplot as plt

                im = tensor_to_pil(depth_A[0], unnormalize=True)
                plt.imshow(im)
                plt.scatter(detections_A[:30, 0], detections_A[:30, 1])
                plt.savefig(f"vis/depth_A_pre.jpg")
                plt.cla()
                im = tensor_to_pil(depth_B[0], unnormalize=True)
                plt.imshow(im)
                plt.scatter(detections_B[:30, 0], detections_B[:30, 1])
                plt.savefig(f"vis/depth_B_pre.jpg")
                plt.cla()
                im = tensor_to_pil(im_A[:3], unnormalize=True)
                plt.imshow(im)
                plt.scatter(detections_A[:30, 0], detections_A[:30, 1])
                plt.savefig(f"vis/im_A_pre.jpg")
                plt.cla()
                im = tensor_to_pil(im_B[:3], unnormalize=True)
                plt.imshow(im)
                plt.scatter(detections_B[:30, 0], detections_B[:30, 1])
                plt.savefig(f"vis/im_B_pre.jpg")
                plt.cla()

            affine_AtoB = torch.eye(2)
            if self.use_affine_aug:
                im_A, im_B, depth_A, depth_B, K1, K2, detections_A, detections_B, affine_AtoB = self.affine_aug(
                    im_A,
                    im_B,
                    depth_A,
                    depth_B,
                    K1,
                    K2,
                    detections_A,
                    detections_B,
                )
                affine_AtoB = affine_AtoB.float()

            if self.use_horizontal_flip_aug:
                if self.use_affine_aug:
                    raise NotImplementedError(
                        "horizontal flips interfere with affine steering, flips are baked in to the affine augs"
                    )
                if np.random.rand() > 0.5:
                    im_A, im_B, depth_A, depth_B, K1, K2 = self.horizontal_flip(im_A, im_B, depth_A, depth_B, K1, K2)
                    if self.detections is not None:
                        detections_A[:, 0] = W - detections_A
                        detections_B[:, 0] = W - detections_B

            if affine_steerers.DEBUG_MODE:
                im = tensor_to_pil(depth_A[0], unnormalize=True)
                plt.imshow(im)
                plt.scatter(detections_A[:30, 0], detections_A[:30, 1])
                plt.savefig(f"vis/depth_A_post.jpg")
                plt.cla()
                im = tensor_to_pil(depth_B[0], unnormalize=True)
                plt.imshow(im)
                plt.scatter(detections_B[:30, 0], detections_B[:30, 1])
                plt.savefig(f"vis/depth_B_post.jpg")
                plt.cla()
                im = tensor_to_pil(im_A[:3], unnormalize=True)
                plt.imshow(im)
                plt.scatter(detections_A[:30, 0], detections_A[:30, 1])
                plt.savefig(f"vis/im_A_post.jpg")
                plt.cla()
                im = tensor_to_pil(im_B[:3], unnormalize=True)
                plt.imshow(im)
                plt.scatter(detections_B[:30, 0], detections_B[:30, 1])
                plt.savefig(f"vis/im_B_post.jpg")
                plt.cla()

            if self.grayscale:
                im_A = im_A.mean(dim=-3, keepdim=True)
                im_B = im_B.mean(dim=-3, keepdim=True)
            data_dict = {
                "im_A": im_A,
                "im_A_identifier": self.image_paths[idx1].split("/")[-1].split(".jpg")[0],
                "im_B": im_B,
                "im_B_identifier": self.image_paths[idx2].split("/")[-1].split(".jpg")[0],
                "im_A_depth": depth_A[0, 0],
                "im_B_depth": depth_B[0, 0],
                "pose_A": T1,
                "pose_B": T2,
                "detections_A": detections_A,
                "detections_B": detections_B,
                "tracks3D_A": tracks3D_A,
                "tracks3D_B": tracks3D_B,
                "K1": K1,
                "K2": K2,
                "T_1to2": T_1to2,
                "im_A_path": im_A_ref,
                "im_B_path": im_B_ref,
                "affine_AtoB": affine_AtoB,
            }
            if affine_steerers.DEBUG_MODE:
                import pickle

                with open("vis/data.pkl", "wb") as handle:
                    pickle.dump(data_dict, handle, protocol=pickle.HIGHEST_PROTOCOL)
        except Exception as e:
            print(e)
            print(f"Failed to load image pair {self.pairs[pair_idx]}")
            print("Loading a random pair in scene instead")
            rand_ind = np.random.choice(range(len(self)))
            return self[rand_ind]
        return data_dict


class MegadepthBuilder:
    def __init__(
        self,
        data_root="data/megadepth",
        loftr_ignore=True,
        imc21_ignore=True,
        use_detections=True,
        use_3D_detections=False,
    ) -> None:
        self.data_root = data_root
        self.scene_info_root = os.path.join(data_root, "prep_scene_info")
        self.all_scenes = os.listdir(self.scene_info_root)
        self.test_scenes = ["0017.npy", "0004.npy", "0048.npy", "0013.npy"]
        # LoFTR did the D2-net preprocessing differently than we did and got more ignore scenes, can optionially ignore those
        self.loftr_ignore_scenes = set(
            [
                "0121.npy",
                "0133.npy",
                "0168.npy",
                "0178.npy",
                "0229.npy",
                "0349.npy",
                "0412.npy",
                "0430.npy",
                "0443.npy",
                "1001.npy",
                "5014.npy",
                "5015.npy",
                "5016.npy",
            ]
        )
        self.imc21_scenes = set(
            ["0008.npy", "0019.npy", "0021.npy", "0024.npy", "0025.npy", "0032.npy", "0063.npy", "1589.npy"]
        )
        self.test_scenes_loftr = ["0015.npy", "0022.npy"]
        self.loftr_ignore = loftr_ignore
        self.imc21_ignore = imc21_ignore
        self.use_detections = use_detections
        self.use_3D_detections = use_3D_detections

    def build_scenes(self, split="train", min_overlap=0.0, scene_names=None, **kwargs):
        if split == "train":
            scene_names = set(self.all_scenes) - set(self.test_scenes)
        elif split == "train_loftr":
            scene_names = set(self.all_scenes) - set(self.test_scenes_loftr)
        elif split == "test":
            scene_names = self.test_scenes
        elif split == "test_loftr":
            scene_names = self.test_scenes_loftr
        elif split == "custom":
            scene_names = scene_names
        else:
            raise ValueError(f"Split {split} not available")
        scenes = []
        for scene_name in tqdm(scene_names):
            if self.loftr_ignore and scene_name in self.loftr_ignore_scenes:
                continue
            if self.imc21_ignore and scene_name in self.imc21_scenes:
                continue
            if ".npy" not in scene_name:
                continue
            scene_info = np.load(os.path.join(self.scene_info_root, scene_name), allow_pickle=True).item()
            if self.use_detections:
                scene_info_detections = np.load(
                    os.path.join(self.scene_info_root, "detections", f"detections_{scene_name}"), allow_pickle=True
                ).item()
            else:
                scene_info_detections = None
            if self.use_3D_detections:
                scene_info_detections3D = np.load(
                    os.path.join(self.scene_info_root, "detections3D", f"detections3D_{scene_name}"), allow_pickle=True
                )
            else:
                scene_info_detections3D = None

            scenes.append(
                MegadepthScene(
                    self.data_root,
                    scene_info,
                    scene_info_detections=scene_info_detections,
                    scene_info_detections3D=scene_info_detections3D,
                    min_overlap=min_overlap,
                    scene_name=scene_name,
                    **kwargs,
                )
            )
        return scenes

    def weight_scenes(self, concat_dataset, alpha=0.5):
        ns = []
        for d in concat_dataset.datasets:
            ns.append(len(d))
        ws = torch.cat([torch.ones(n) / n**alpha for n in ns])
        return ws
