#
# Copyright (C) 2023, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use
# under the terms of the LICENSE.md file.
#
# For inquiries contact  george.drettakis@inria.fr
#

import os
import sys
from PIL import Image
from typing import NamedTuple
from scene.colmap_loader import (
    read_extrinsics_text,
    read_intrinsics_text,
    qvec2rotmat,
    read_extrinsics_binary,
    read_intrinsics_binary,
    read_points3D_binary,
    read_points3D_text,
)
from utils.graphics_utils import getWorld2View2, focal2fov, fov2focal
from utils.camera_utils import get_uniform_poses
import numpy as np
import json
from pathlib import Path
from plyfile import PlyData, PlyElement
from utils.sh_utils import SH2RGB
from scene.gaussian_model import BasicPointCloud
from scene.cameras import Camera
import torch
import rembg
import mcubes
import trimesh


class CameraInfo(NamedTuple):
    uid: int
    R: np.array
    T: np.array
    FovY: np.array
    FovX: np.array
    image: np.array
    image_path: str
    image_name: str
    width: int
    height: int


class SceneInfo(NamedTuple):
    point_cloud: BasicPointCloud
    train_cameras: list
    test_cameras: list
    nerf_normalization: dict
    ply_path: str


def getNerfppNorm(cam_info):
    def get_center_and_diag(cam_centers):
        cam_centers = np.hstack(cam_centers)
        avg_cam_center = np.mean(cam_centers, axis=1, keepdims=True)
        center = avg_cam_center
        dist = np.linalg.norm(cam_centers - center, axis=0, keepdims=True)
        diagonal = np.max(dist)
        return center.flatten(), diagonal

    cam_centers = []

    for cam in cam_info:
        W2C = getWorld2View2(cam.R, cam.T)
        C2W = np.linalg.inv(W2C)
        cam_centers.append(C2W[:3, 3:4])

    center, diagonal = get_center_and_diag(cam_centers)
    radius = diagonal * 1.1

    translate = -center

    return {"translate": translate, "radius": radius}


def readColmapCameras(cam_extrinsics, cam_intrinsics, images_folder):
    cam_infos = []
    for idx, key in enumerate(cam_extrinsics):
        sys.stdout.write("\r")
        # the exact output you're looking for:
        sys.stdout.write("Reading camera {}/{}".format(idx + 1, len(cam_extrinsics)))
        sys.stdout.flush()

        extr = cam_extrinsics[key]
        intr = cam_intrinsics[extr.camera_id]
        height = intr.height
        width = intr.width

        uid = intr.id
        R = np.transpose(qvec2rotmat(extr.qvec))
        T = np.array(extr.tvec)

        if intr.model == "SIMPLE_PINHOLE":
            focal_length_x = intr.params[0]
            FovY = focal2fov(focal_length_x, height)
            FovX = focal2fov(focal_length_x, width)
        elif intr.model == "PINHOLE":
            focal_length_x = intr.params[0]
            focal_length_y = intr.params[1]
            FovY = focal2fov(focal_length_y, height)
            FovX = focal2fov(focal_length_x, width)
        else:
            assert (
                False
            ), "Colmap camera model not handled: only undistorted datasets (PINHOLE or SIMPLE_PINHOLE cameras) supported!"

        image_path = os.path.join(images_folder, os.path.basename(extr.name))
        image_name = os.path.basename(image_path).split(".")[0]
        image = Image.open(image_path)

        cam_info = CameraInfo(
            uid=uid,
            R=R,
            T=T,
            FovY=FovY,
            FovX=FovX,
            image=image,
            image_path=image_path,
            image_name=image_name,
            width=width,
            height=height,
        )
        cam_infos.append(cam_info)
    sys.stdout.write("\n")
    return cam_infos


def fetchPly(path):
    plydata = PlyData.read(path)
    vertices = plydata["vertex"]
    positions = np.vstack([vertices["x"], vertices["y"], vertices["z"]]).T
    colors = np.vstack([vertices["red"], vertices["green"], vertices["blue"]]).T / 255.0
    normals = np.vstack([vertices["nx"], vertices["ny"], vertices["nz"]]).T
    return BasicPointCloud(points=positions, colors=colors, normals=normals)


def storePly(path, xyz, rgb):
    # Define the dtype for the structured array
    dtype = [
        ("x", "f4"),
        ("y", "f4"),
        ("z", "f4"),
        ("nx", "f4"),
        ("ny", "f4"),
        ("nz", "f4"),
        ("red", "u1"),
        ("green", "u1"),
        ("blue", "u1"),
    ]

    normals = np.zeros_like(xyz)

    elements = np.empty(xyz.shape[0], dtype=dtype)
    attributes = np.concatenate((xyz, normals, rgb), axis=1)
    elements[:] = list(map(tuple, attributes))

    # Create the PlyData object and write to file
    vertex_element = PlyElement.describe(elements, "vertex")
    ply_data = PlyData([vertex_element])
    ply_data.write(path)


def readColmapSceneInfo(path, images, eval, llffhold=8):
    try:
        cameras_extrinsic_file = os.path.join(path, "sparse/0", "images.bin")
        cameras_intrinsic_file = os.path.join(path, "sparse/0", "cameras.bin")
        cam_extrinsics = read_extrinsics_binary(cameras_extrinsic_file)
        cam_intrinsics = read_intrinsics_binary(cameras_intrinsic_file)
    except:
        cameras_extrinsic_file = os.path.join(path, "sparse/0", "images.txt")
        cameras_intrinsic_file = os.path.join(path, "sparse/0", "cameras.txt")
        cam_extrinsics = read_extrinsics_text(cameras_extrinsic_file)
        cam_intrinsics = read_intrinsics_text(cameras_intrinsic_file)

    reading_dir = "images" if images == None else images
    cam_infos_unsorted = readColmapCameras(
        cam_extrinsics=cam_extrinsics,
        cam_intrinsics=cam_intrinsics,
        images_folder=os.path.join(path, reading_dir),
    )
    cam_infos = sorted(cam_infos_unsorted.copy(), key=lambda x: x.image_name)

    if eval:
        train_cam_infos = [c for idx, c in enumerate(cam_infos) if idx % llffhold != 0]
        test_cam_infos = [c for idx, c in enumerate(cam_infos) if idx % llffhold == 0]
    else:
        train_cam_infos = cam_infos
        test_cam_infos = []

    nerf_normalization = getNerfppNorm(train_cam_infos)

    ply_path = os.path.join(path, "sparse/0/points3D.ply")
    bin_path = os.path.join(path, "sparse/0/points3D.bin")
    txt_path = os.path.join(path, "sparse/0/points3D.txt")
    if not os.path.exists(ply_path):
        print(
            "Converting point3d.bin to .ply, will happen only the first time you open the scene."
        )
        try:
            xyz, rgb, _ = read_points3D_binary(bin_path)
        except:
            xyz, rgb, _ = read_points3D_text(txt_path)
        storePly(ply_path, xyz, rgb)
    try:
        pcd = fetchPly(ply_path)
    except:
        pcd = None

    scene_info = SceneInfo(
        point_cloud=pcd,
        train_cameras=train_cam_infos,
        test_cameras=test_cam_infos,
        nerf_normalization=nerf_normalization,
        ply_path=ply_path,
    )
    return scene_info


def readCamerasFromTransforms(path, transformsfile, white_background, extension=".png"):
    cam_infos = []

    with open(os.path.join(path, transformsfile)) as json_file:
        contents = json.load(json_file)
        fovx = contents["camera_angle_x"]

        frames = contents["frames"]
        for idx, frame in enumerate(frames):
            cam_name = os.path.join(path, frame["file_path"] + extension)

            # NeRF 'transform_matrix' is a camera-to-world transform
            c2w = np.array(frame["transform_matrix"])
            # change from OpenGL/Blender camera axes (Y up, Z back) to COLMAP (Y down, Z forward)
            c2w[:3, 1:3] *= -1

            # get the world-to-camera transform and set R, T
            w2c = np.linalg.inv(c2w)
            R = np.transpose(
                w2c[:3, :3]
            )  # R is stored transposed due to 'glm' in CUDA code
            T = w2c[:3, 3]

            image_path = os.path.join(path, cam_name)
            image_name = Path(cam_name).stem
            image = Image.open(image_path)

            im_data = np.array(image.convert("RGBA"))

            bg = np.array([1, 1, 1]) if white_background else np.array([0, 0, 0])

            norm_data = im_data / 255.0
            if norm_data.shape[-1] != 3:
                arr = norm_data[:, :, :3] * norm_data[:, :, 3:4] + bg * (
                    1 - norm_data[:, :, 3:4]
                )
            image = Image.fromarray(np.array(arr * 255.0, dtype=np.byte), "RGB")

            fovy = focal2fov(fov2focal(fovx, image.size[0]), image.size[1])
            FovY = fovy
            FovX = fovx

            cam_infos.append(
                CameraInfo(
                    uid=idx,
                    R=R,
                    T=T,
                    FovY=FovY,
                    FovX=FovX,
                    image=image,
                    image_path=image_path,
                    image_name=image_name,
                    width=image.size[0],
                    height=image.size[1],
                )
            )

    return cam_infos


def uniform_surface_sampling_from_vertices_and_faces(
    vertices, faces, num_points: int
) -> torch.Tensor:
    """
    Uniformly sample points from the surface of a mesh.

    Args:
        vertices (torch.Tensor): Vertices of the mesh.
        faces (torch.Tensor): Faces of the mesh.
        num_points (int): Number of points to sample.

    Returns:
        torch.Tensor: Points sampled from the surface of the mesh.
    """
    mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
    n = num_points
    points = []
    while n > 0:
        p, _ = trimesh.sample.sample_surface_even(mesh, n)
        n -= p.shape[0]
        if n >= 0:
            points.append(p)
        else:
            points.append(p[:n])

    if len(points) > 1:
        points = np.concatenate(points, axis=0)
    else:
        points = points[0]

    points = torch.from_numpy(points.astype(np.float32))

    return points, torch.rand_like(points)


def occ_from_sparse_initialize(poses, images, cameras, grid_reso, num_points):
    # fov is in degrees
    this_session = rembg.new_session()

    imgs = [rembg.remove(im, session=this_session) for im in images]

    reso = grid_reso
    occ_grid = torch.ones((reso, reso, reso), dtype=torch.bool, device="cuda")

    c2ws = poses
    center = c2ws[..., :3, 3].mean(axis=0)
    radius = np.linalg.norm(c2ws[..., :3, 3] - center, axis=-1).mean()
    xx, yy, zz = torch.meshgrid(
        torch.linspace(-radius, radius, reso, device="cuda"),
        torch.linspace(-radius, radius, reso, device="cuda"),
        torch.linspace(-radius, radius, reso, device="cuda"),
        indexing="ij",
    )
    print("radius", radius)

    # xyz_grid = torch.stack((xx.flatten(), yy.flatten(), zz.flatten()), dim=-1)
    ww = torch.ones((reso, reso, reso), dtype=torch.float32, device="cuda")
    xyzw_grid = torch.stack((xx, yy, zz, ww), dim=-1)
    xyzw_grid[..., :3] += torch.from_numpy(center).cuda()

    c2ws = torch.tensor(c2ws, dtype=torch.float32)

    for c2w, camera, img in zip(c2ws, cameras, imgs):
        img = np.asarray(img)
        alpha = img[..., 3].astype(np.float32) / 255.0
        is_foreground = alpha > 0.05
        is_foreground = torch.from_numpy(is_foreground).cuda()

        full_proj_mtx = Camera(
            colmap_id=camera.uid,
            R=camera.R,
            T=camera.T,
            FoVx=camera.FovX,
            FoVy=camera.FovY,
            image=torch.randn(3, 10, 10),
            gt_alpha_mask=None,
            image_name="no",
            uid=0,
            data_device="cuda",
        ).full_proj_transform
        # check the scale

        ij = xyzw_grid @ full_proj_mtx
        ij = (ij + 1) / 2.0
        h, w = img.shape[:2]
        ij = ij[..., :2] * torch.tensor([w, h], dtype=torch.float32, device="cuda")
        ij = (
            ij.clamp(
                min=torch.tensor([0.0, 0.0], device="cuda"),
                max=torch.tensor([w - 1, h - 1], dtype=torch.float32, device="cuda"),
            )
            .to(torch.long)
            .cuda()
        )

        occ_grid = torch.logical_and(occ_grid, is_foreground[ij[..., 1], ij[..., 0]])

    # To mesh
    occ_grid = occ_grid.to(torch.float32).cpu().numpy()
    vertices, triangles = mcubes.marching_cubes(occ_grid, 0.5)

    # vertices = (vertices / reso - 0.5) * radius * 2 + center
    # vertices = (vertices / (reso - 1.0) - 0.5) * radius * 2 * 2 + center
    vertices = vertices / (grid_reso - 1) * 2 - 1
    vertices = vertices * radius + center
    # mcubes.export_obj(vertices, triangles, "./tmp/occ_voxel.obj")

    xyz, rgb = uniform_surface_sampling_from_vertices_and_faces(
        vertices, triangles, num_points
    )

    return xyz


def readNerfSyntheticInfo(path, white_background, eval, extension=".png"):
    print("Reading Training Transforms")
    train_cam_infos = readCamerasFromTransforms(
        path, "transforms_train.json", white_background, extension
    )
    print("Reading Test Transforms")
    test_cam_infos = readCamerasFromTransforms(
        path, "transforms_test.json", white_background, extension
    )

    if not eval:
        train_cam_infos.extend(test_cam_infos)
        test_cam_infos = []

    nerf_normalization = getNerfppNorm(train_cam_infos)

    ply_path = os.path.join(path, "points3d.ply")
    if not os.path.exists(ply_path):
        # Since this data set has no colmap data, we start with random points
        num_pts = 100_000
        print(f"Generating random point cloud ({num_pts})...")

        # We create random points inside the bounds of the synthetic Blender scenes
        xyz = np.random.random((num_pts, 3)) * 2.6 - 1.3
        shs = np.random.random((num_pts, 3)) / 255.0
        pcd = BasicPointCloud(
            points=xyz, colors=SH2RGB(shs), normals=np.zeros((num_pts, 3))
        )

        storePly(ply_path, xyz, SH2RGB(shs) * 255)
    try:
        pcd = fetchPly(ply_path)
    except:
        pcd = None

    scene_info = SceneInfo(
        point_cloud=pcd,
        train_cameras=train_cam_infos,
        test_cameras=test_cam_infos,
        nerf_normalization=nerf_normalization,
        ply_path=ply_path,
    )
    return scene_info


def constructVideoNVSInfo(
    num_frames,
    radius,
    elevation,
    fov,
    reso,
    images,
    masks,
    num_pts=100_000,
    train=True,
):
    poses = get_uniform_poses(num_frames, radius, elevation)
    w2cs = np.linalg.inv(poses)
    train_cam_infos = []

    for idx, pose in enumerate(w2cs):
        train_cam_infos.append(
            CameraInfo(
                uid=idx,
                R=np.transpose(pose[:3, :3]),
                T=pose[:3, 3],
                FovY=np.deg2rad(fov),
                FovX=np.deg2rad(fov),
                image=images[idx],
                image_path=None,
                image_name=idx,
                width=reso,
                height=reso,
            )
        )

    nerf_normalization = getNerfppNorm(train_cam_infos)
    # xyz = np.random.random((num_pts, 3)) * radius / 3 - radius / 3
    xyz = np.random.randn(num_pts, 3) * radius / 16
    # if len(poses) <= 24:
    #     xyz = occ_from_sparse_initialize(poses, images, train_cam_infos, 256, num_pts)
    #     num_pts = xyz.shape[0]
    # else:
    #     xyz = np.random.randn(num_pts, 3) * radius / 16
    xyz = np.random.randn(num_pts, 3) * radius / 16
    # shs = np.random.random((num_pts, 3)) / 255.0
    shs = np.ones((num_pts, 3)) * 0.2
    pcd = BasicPointCloud(
        points=xyz, colors=SH2RGB(shs), normals=np.zeros((num_pts, 3))
    )

    ply_path = "./tmp/points3d.ply"
    storePly(ply_path, xyz, SH2RGB(shs) * 255)
    pcd = fetchPly(ply_path)

    scene_info = SceneInfo(
        point_cloud=pcd,
        train_cameras=train_cam_infos,
        test_cameras=[],
        nerf_normalization=nerf_normalization,
        ply_path="./tmp/points3d.ply",
    )

    return scene_info


sceneLoadTypeCallbacks = {
    "Colmap": readColmapSceneInfo,
    "Blender": readNerfSyntheticInfo,
    "VideoNVS": constructVideoNVSInfo,
}