# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Pose and Intrinsics Optimizers
"""

from __future__ import annotations

import functools
from dataclasses import dataclass, field
from typing import Literal, Optional, Type, Union

import numpy
import torch
import tyro
from jaxtyping import Float, Int
from torch import Tensor, nn
from typing_extensions import assert_never

from nerfstudio.cameras.cameras import Cameras
from nerfstudio.cameras.lie_groups import exp_map_SE3, exp_map_SO3xR3
from nerfstudio.cameras.rays import RayBundle
from nerfstudio.configs.base_config import InstantiateConfig
from nerfstudio.engine.optimizers import OptimizerConfig
from nerfstudio.engine.schedulers import SchedulerConfig
from nerfstudio.utils import poses as pose_utils


@dataclass
class CameraOptimizerConfig(InstantiateConfig):
    """Configuration of optimization for camera poses."""

    _target: Type = field(default_factory=lambda: CameraOptimizer)

    mode: Literal["off", "SO3xR3", "SE3"] = "off"
    """Pose optimization strategy to use. If enabled, we recommend SO3xR3."""
    
    refine_intrinsics: bool = False
    """Whether to refine the camera intrinsic parameters or not."""
        
    refine_focal_length: bool = True
    """Whether to refine the focal length parameters. Only used if refine_intrinsics is True."""
    
    refine_camera_center: bool = True
    """Whether to refine the camera center parameters. Only used if refine_intrinsics is True."""
    
    refine_distortion_params: bool = False
    """Whether to refine the distortion parameters. Only used if refine_intrinsics is True."""

    camp_translation_scale: float = 0.7
    """Camera preconditioning scale for translation parameters."""

    camp_orientation_scale: float = 0.8
    """Camera preconditioning scale for orientation parameters."""
    
    camp_focal_length_scale: float = 55.0
    """Camera preconditioning scale for focal length parameters."""

    camp_camera_center_scale: float = 10.0
    """Camera preconditioning scale for camera center parameters."""
    
    camp_distortion_scale: float = 0.01
    """Camera preconditioning scale for distortion parameters (WARNING untested)."""

    #trans_l2_penalty: float = 1e-2
    trans_l2_penalty: float = 0.0
    """L2 penalty on translation parameters."""

    #rot_l2_penalty: float = 1e-3
    rot_l2_penalty: float = 0.0
    """L2 penalty on rotation parameters."""

    # tyro.conf.Suppress prevents us from creating CLI arguments for these fields.
    optimizer: tyro.conf.Suppress[Optional[OptimizerConfig]] = field(default=None)
    """Deprecated, now specified inside the optimizers dict"""

    scheduler: tyro.conf.Suppress[Optional[SchedulerConfig]] = field(default=None)
    """Deprecated, now specified inside the optimizers dict"""

    def __post_init__(self):
        if self.optimizer is not None:
            import warnings

            from nerfstudio.utils.rich_utils import CONSOLE

            CONSOLE.print(
                "\noptimizer is no longer specified in the CameraOptimizerConfig, it is now defined with the rest of the param groups inside the config file under the name 'camera_opt'\n",
                style="bold yellow",
            )
            warnings.warn("above message coming from", FutureWarning, stacklevel=3)

        if self.scheduler is not None:
            import warnings

            from nerfstudio.utils.rich_utils import CONSOLE

            CONSOLE.print(
                "\nscheduler is no longer specified in the CameraOptimizerConfig, it is now defined with the rest of the param groups inside the config file under the name 'camera_opt'\n",
                style="bold yellow",
            )
            warnings.warn("above message coming from", FutureWarning, stacklevel=3)


class CameraOptimizer(nn.Module):
    """Layer that modifies camera poses to be optimized as well as the field during training."""

    config: CameraOptimizerConfig

    def __init__(
        self,
        config: CameraOptimizerConfig,
        num_cameras: int,
        device: Union[torch.device, str],
        non_trainable_camera_indices: Optional[Int[Tensor, "num_non_trainable_cameras"]] = None,
        **kwargs,
    ) -> None:
        super().__init__()
        self.config = config
        self.num_cameras = num_cameras
        self.device = device
        self.non_trainable_camera_indices = non_trainable_camera_indices

        # Initialize learnable parameters.
        if self.config.mode == "off":
            pass
        elif self.config.mode in ("SO3xR3", "SE3"):
            self.pose_adjustment = torch.nn.Parameter(torch.zeros((num_cameras, 6), device=device))
        else:
            assert_never(self.config.mode)
        
        if self.config.refine_intrinsics:
            if self.config.refine_focal_length:
                self.focal_length_adjustment = torch.nn.Parameter(torch.zeros((num_cameras, 2), device=device))
            if self.config.refine_camera_center:
                self.camera_center_adjustment = torch.nn.Parameter(torch.zeros((num_cameras, 2), device=device))
            if self.config.refine_distortion_params:
                self.distortion_adjustment = torch.nn.Parameter(torch.zeros((num_cameras, 6), device=device))

    def forward(
        self,
        indices: Int[Tensor, "camera_indices"],
    ) -> Float[Tensor, "camera_indices 3 4"]:
        """Indexing into camera adjustments.
        Args:
            indices: indices of Cameras to optimize.
        Returns:
            Transformation matrices from optimized camera coordinates
            to given camera coordinates.
        """

        outputs = []

        # Apply learned transformation delta.
        if self.config.mode == "off":
            pass
        else:
            pose_adj = self.pose_adjustment[indices, :].clone()
            pose_adj[:, :3] *= self.config.camp_translation_scale
            pose_adj[:, 3:] *= self.config.camp_orientation_scale
            if self.config.mode == "SO3xR3":
                outputs.append(exp_map_SO3xR3(pose_adj))
            elif self.config.mode == "SE3":
                outputs.append(exp_map_SE3(pose_adj))
            else:
                assert_never(self.config.mode)
        # Detach non-trainable indices by setting to identity transform
        if self.non_trainable_camera_indices is not None:
            if self.non_trainable_camera_indices.device != self.pose_adjustment.device:
                self.non_trainable_camera_indices = self.non_trainable_camera_indices.to(self.pose_adjustment.device)
            outputs[0][self.non_trainable_camera_indices] = torch.eye(4, device=self.pose_adjustment.device)[:3, :4]

        # Return: identity if no transforms are needed, otherwise multiply transforms together.
        if len(outputs) == 0:
            # Note that using repeat() instead of tile() here would result in unnecessary copies.
            return torch.eye(4, device=self.device)[None, :3, :4].tile(indices.shape[0], 1, 1)
        return functools.reduce(pose_utils.multiply, outputs)

    def apply_to_raybundle(self, raybundle: RayBundle) -> None:
        raise NotImplementedError()

        """Apply the pose correction to the raybundle"""
        if self.config.mode != "off":
            correction_matrices = self(raybundle.camera_indices.squeeze())  # type: ignore
            raybundle.origins = raybundle.origins + correction_matrices[:, :3, 3]
            raybundle.directions = torch.bmm(correction_matrices[:, :3, :3], raybundle.directions[..., None]).squeeze()
    
    '''
    def apply_to_camera(self, camera: Cameras) -> torch.Tensor:
        """Apply the pose correction to the world-to-camera matrix in a Camera object"""
        if self.config.mode == "off":
            return camera.camera_to_worlds

        if camera.metadata is None or "cam_idx" not in camera.metadata:
            # Viser cameras
            return camera.camera_to_worlds

        camera_idx = camera.metadata["cam_idx"]
        adj = self(torch.tensor([camera_idx], dtype=torch.long)).to(camera.device)  # type: ignore

        return torch.cat(
            [
                # Apply rotation to directions in world coordinates, without touching the origin.
                # Equivalent to: directions -> correction[:3,:3] @ directions
                torch.bmm(adj[..., :3, :3], camera.camera_to_worlds[..., :3, :3]),
                # Apply translation in world coordinate, independently of rotation.
                # Equivalent to: origins -> origins + correction[:3,3]
                camera.camera_to_worlds[..., :3, 3:] + adj[..., :3, 3:],
            ],
            dim=-1,
        )
    '''
    
    def apply_to_camera(self, camera: Cameras) -> Cameras:
        """Apply the pose correction to a copy of the camera object"""

        camera_refined = camera.clone()

        if camera.metadata is None or "cam_idx" not in camera.metadata:
            assert False
            ## Viser cameras
            #return camera.camera_to_worlds
            
        camera_idx = camera.metadata["cam_idx"] if torch.is_tensor(camera.metadata["cam_idx"]) else torch.tensor([camera.metadata["cam_idx"]], dtype=torch.long)

        if self.config.mode != "off":
            pose_adj = self(camera_idx).to(camera.device)  # type: ignore
            camera_refined.camera_to_worlds = torch.cat(
                [
                    # Apply rotation to directions in world coordinates, without touching the origin.
                    # Equivalent to: directions -> correction[:3,:3] @ directions
                    torch.bmm(pose_adj[..., :3, :3], camera.camera_to_worlds[..., :3, :3]),
                    # Apply translation in world coordinate, independently of rotation.
                    # Equivalent to: origins -> origins + correction[:3,3]
                    camera.camera_to_worlds[..., :3, 3:] + pose_adj[..., :3, 3:],
                ],
                dim=-1,
            )
        
        if self.config.refine_intrinsics:
            if self.config.refine_focal_length:
                fxfy_adj = self.config.camp_focal_length_scale * self.focal_length_adjustment[camera_idx, :]
                camera_refined.fx += fxfy_adj[:, 0]
                camera_refined.fy += fxfy_adj[:, 1]

            if self.config.refine_camera_center:
                cxcy_adj = self.config.camp_camera_center_scale * self.camera_center_adjustment[camera_idx, :]
                camera_refined.cx += cxcy_adj[:, 0]
                camera_refined.cy += cxcy_adj[:, 1]
                
            if self.config.refine_distortion_params and camera_refined.distortion_params is not None:
                d_adj = self.config.camp_distortion_scale * self.distortion_adjustment[camera_idx, :]
                camera_refined.distortion_params += d_adj
        
        return camera_refined

    def get_loss_dict(self, loss_dict: dict) -> None:
        """Add regularization"""
        if self.config.mode != "off":
            loss_dict["camera_opt_regularizer"] = (
                self.pose_adjustment[:, :3].norm(dim=-1).mean() * self.config.trans_l2_penalty
                + self.pose_adjustment[:, 3:].norm(dim=-1).mean() * self.config.rot_l2_penalty
            )
    
    def get_correction_matrices(self):
        """Get optimized pose correction matrices"""
        return self(torch.arange(0, self.num_cameras).long())

    def get_metrics_dict(self, metrics_dict: dict) -> None:
        """Get camera optimizer metrics"""
        if self.config.mode != "off":
            trans = self.config.camp_translation_scale * self.pose_adjustment[:, :3].detach().norm(dim=-1)
            rot = self.config.camp_orientation_scale * self.pose_adjustment[:, 3:].detach().norm(dim=-1)
            metrics_dict["camera_opt_translation_diff_max"] = trans.max()
            metrics_dict["camera_opt_translation_diff_mean"] = trans.mean()
            metrics_dict["camera_opt_rotation_mean_diff_degrees"] = numpy.rad2deg(rot.mean().cpu())
            metrics_dict["camera_opt_rotation_max_diff_degrees"] = numpy.rad2deg(rot.max().cpu())

        if self.config.refine_intrinsics:
            if self.config.refine_focal_length:
                f = self.config.camp_focal_length_scale * self.focal_length_adjustment.norm(dim=-1)
                metrics_dict["camera_opt_focal_length_diff_max_pixels"] = f.max()
                metrics_dict["camera_opt_focal_length_diff_mean_pixels"] = f.mean()
            if self.config.refine_camera_center:
                c = self.config.camp_camera_center_scale * self.camera_center_adjustment.norm(dim=-1)
                metrics_dict["camera_opt_camera_center_diff_max_pixels"] = c.max()
                metrics_dict["camera_opt_camera_center_diff_mean_pixels"] = c.mean()
            if self.config.refine_distortion_params:
                d = self.config.camp_distortion_scale * self.distortion_adjustment.norm(dim=-1)
                metrics_dict["camera_opt_distortion_diff_max"] = d.max()
                metrics_dict["camera_opt_distortion_diff_mean"] = d.mean()

    def get_param_groups(self, param_groups: dict) -> None:
        """Get camera optimizer parameters"""
        camera_opt_params = list(self.parameters())
        if self.config.mode != "off":
            assert len(camera_opt_params) > 0
            param_groups["camera_opt"] = camera_opt_params
        else:
            assert len(camera_opt_params) == 0
