import torch
from einops import rearrange, repeat
from typing import Literal, Optional
from jaxtyping import Float, UInt8
from torch import Tensor
import json
import argparse
import os
from glob import glob
from tqdm import tqdm
from collections import OrderedDict
from jaxtyping import install_import_hook


# Configure beartype and jaxtyping.
with install_import_hook(
    ("src",),
    ("beartype", "beartype"),
):
    from src.dataset.view_sampler.view_sampler_bounded_v2 import farthest_point_sample


def convert_poses(
    poses: Float[Tensor, "batch 18"],
) -> tuple[
    Float[Tensor, "batch 4 4"],  # extrinsics
    Float[Tensor, "batch 3 3"],  # intrinsics
]:
    b, _ = poses.shape

    # Convert the intrinsics to a 3x3 normalized K matrix.
    intrinsics = torch.eye(3, dtype=torch.float32)
    intrinsics = repeat(intrinsics, "h w -> b h w", b=b).clone()
    fx, fy, cx, cy = poses[:, :4].T
    intrinsics[:, 0, 0] = fx
    intrinsics[:, 1, 1] = fy
    intrinsics[:, 0, 2] = cx
    intrinsics[:, 1, 2] = cy

    # Convert the extrinsics to a 4x4 OpenCV-style W2C matrix.
    w2c = repeat(torch.eye(4, dtype=torch.float32), "h w -> b h w", b=b).clone()
    w2c[:, :3] = rearrange(poses[:, 6:], "b (h w) -> b h w", h=3, w=4)
    return w2c.inverse(), intrinsics


def partition_list(lst, n_bins):
    if n_bins <= 0:
        raise ValueError("Number of bins must be greater than 0")
    if len(lst) < n_bins:
        raise ValueError("Number of bins cannot exceed the length of the list")

    bin_size = len(lst) // n_bins
    borders = [lst[0]]  # First border is always the first index
    for i in range(1, n_bins):
        border_index = min(
            i * bin_size, len(lst) - 1
        )  # Ensure last bin doesn't exceed list length
        borders.append(lst[border_index])
    borders.append(lst[-1])  # Last border is always the last index
    return borders


def find_train_and_test_index(chunk_path, scene_name=None, num_context_views=5,
                              num_target_skip=1, num_target_views=28,
                              start_frame=None,
                              frame_distance=None,
                              render_video=False,
                              ):
    chunk = torch.load(chunk_path)
    out_dict = OrderedDict()
    for example in chunk:
        cur_scene_name = example["key"]
        if scene_name is not None and cur_scene_name != scene_name:
            continue

        extrinsics, intrinsics = convert_poses(example["cameras"])

        # bounded evaluation to make the task easier
        if start_frame is not None:
            # assert frame_distance is not None
            if frame_distance is None:
                frame_distance = extrinsics.shape[0] - start_frame
            end_frame = start_frame + frame_distance

            extrinsics = extrinsics[start_frame:end_frame]

        n_views = extrinsics.shape[0]

        if render_video:
            if num_context_views > n_views:
                index_context = list(range(n_views))
            else:
                index_context = sorted(farthest_point_sample(
                    extrinsics[:, :3, -1].unsqueeze(0), num_context_views
                ).squeeze(0).tolist())
            index_target = list(range(n_views))
        else:
            all_indices = list(range(n_views))
            index_target_select = partition_list(all_indices, num_target_views)
            index_target = index_target_select[:num_target_views]

            context_candidate_indices = [i for i in all_indices if i not in index_target]
            context_candidate_extrinsics = extrinsics[context_candidate_indices]

            if num_context_views > len(context_candidate_indices):
                sampled_indices_in_candidates = list(range(len(context_candidate_indices)))
            else:
                sampled_indices_in_candidates = farthest_point_sample(
                    context_candidate_extrinsics[:, :3, -1].unsqueeze(0), num_context_views
                ).squeeze(0).tolist()
            
            index_context = sorted([context_candidate_indices[i] for i in sampled_indices_in_candidates])

        if start_frame is not None:
            # the original index in the full sequence
            index_context = [idx + start_frame for idx in index_context]
            if render_video:
                index_target = list(range(start_frame, end_frame))
            else:
                index_target = [idx + start_frame for idx in index_target]

        if not render_video:
            assert (
                len(index_target) == num_target_views
            ), f"double check {cur_scene_name} at {chunk_path}: target len: {len(index_target)} from {n_views}"

        out_dict[cur_scene_name] = {"context": index_context, "target": index_target}

    return out_dict


def generate_index_file(args):
    n_ctx_list = args.num_context_views
    n_tgt = args.num_target_views

    out_dir = f"assets"
    os.makedirs(out_dir, exist_ok=True)
    data_dir = "/users/bye/code/NoPoSplat_new/datasets/scannetpp/iphone/test"
    # data_dir = "/users/bye/code/NoPoSplat_new/datasets/scannetpp/dslr/test"
    chunk_paths = sorted(glob(os.path.join(data_dir, "*.torch")))
    for n_ctx in n_ctx_list:
        out_dict_all = OrderedDict()
        for chunk_path in tqdm(chunk_paths, desc=f"Processing for {n_ctx} context views"):
            out_dict = find_train_and_test_index(
                chunk_path, scene_name=None, num_context_views=n_ctx,
                num_target_views=n_tgt,
                start_frame=args.start_frame,
                frame_distance=args.frame_distance,
                render_video=args.render_video,
            )
            out_dict_all.update(out_dict)

        if args.start_frame is not None:
            if args.render_video:
                save_file = f"scannetpp_start_{args.start_frame}_distance_{args.frame_distance}_ctx_{n_ctx}v_video.json"
            else:
                save_file = f"scannetpp_start_{args.start_frame}_distance_{args.frame_distance}_ctx_{n_ctx}v_tgt_{n_tgt}v.json"
        elif args.render_video:
            save_file = f"scannetpp_ctx_{n_ctx}v_video.json"
        else:
            save_file = f"scannetpp_fixTgt_ctx_{n_ctx}v_tgt_{n_tgt}v.json"

        out_path = os.path.join(out_dir, save_file)

        with open(out_path, "w") as f:
            json.dump(out_dict_all, f)
        
        print(f"Generated: {out_path}")

    print("Done")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--num_target_views", type=int, default=16, help="test skip")
    parser.add_argument("--num_context_views", type=int, nargs='+', default=[32,64,128,320,640,1000], help="List of context views counts")
    parser.add_argument('--render_video', action='store_true')

    # bounded evaluation to make the task easier
    parser.add_argument('--start_frame', default=None, type=int)
    parser.add_argument('--frame_distance', default=None, type=int)

    args = parser.parse_args()

    generate_index_file(args)
