# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

from typing import List, Optional, cast

import numpy as np
import torch

import theseus as th

from .util import random_small_quaternion


class Camera:
    def __init__(
        self,
        pose: th.SE3,
        focal_length: th.Vector,
        calib_k1: th.Vector,
        calib_k2: th.Vector,
    ):
        self.pose = pose
        self.focal_length = focal_length
        self.calib_k1 = calib_k1
        self.calib_k2 = calib_k2

    def to_params(self) -> List[float]:
        batch_size = self.pose.tensor.shape[0]
        assert batch_size == 1
        R = th.SO3(tensor=self.pose.tensor[:, :, :3]).log_map().squeeze(0)
        t = self.pose.tensor[:, :, 3].squeeze(0)
        return [
            *R.numpy(),
            *t.numpy(),
            float(self.focal_length[0, 0]),
            float(self.calib_k1[0, 0]),
            float(self.calib_k2[0, 0]),
        ]

    @staticmethod
    def from_params(params: List[float], name: str = "Cam") -> "Camera":
        r = th.SO3.exp_map(torch.tensor(params[:3], dtype=torch.float64).unsqueeze(0))
        t = torch.tensor([params[3:6]], dtype=torch.float64).unsqueeze(2)
        pose = th.SE3(tensor=torch.cat([r.tensor, t], dim=2), name=name + "_pose")
        focal_length = th.Vector(
            tensor=torch.tensor([params[6:7]], dtype=torch.float64),
            name=name + "_focal_length",
        )
        calib_k1 = th.Vector(
            tensor=torch.tensor([params[7:8]], dtype=torch.float64),
            name=name + "_calib_k1",
        )
        calib_k2 = th.Vector(
            tensor=torch.tensor([params[8:9]], dtype=torch.float64),
            name=name + "_calib_k2",
        )
        return Camera(pose, focal_length, calib_k1, calib_k2)

    def project_point(self, point: th.Point3) -> torch.Tensor:
        point_cam = self.pose.transform_from(point)
        proj = -point_cam.tensor[:, :2] / point_cam.tensor[:, 2:3]
        proj_sqn = (proj * proj).sum(dim=1).unsqueeze(1)
        proj_factor = self.focal_length.tensor * (
            1.0 + proj_sqn * (self.calib_k1.tensor + proj_sqn * self.calib_k2.tensor)
        )
        return proj * proj_factor

    def perturbed(self, rot_random: float = 0.1, pos_random: float = 0.2):
        batch_size = self.pose.shape[0]
        pert_rot = torch.cat(
            [
                random_small_quaternion(max_degrees=rot_random).unsqueeze(0)
                for _ in range(batch_size)
            ]
        )
        pert_tr = (
            torch.rand((batch_size, 3), dtype=torch.float64) * 2
            + torch.tensor([-1, -1, -1], dtype=torch.float64)
        ) * pos_random
        pert_data = torch.cat([pert_tr, pert_rot], dim=1)
        pert = th.SE3(pert_data)
        pert_pose = cast(th.SE3, pert.compose(self.pose))
        return Camera(
            pert_pose.copy(new_name=self.pose.name + "_copy"),
            self.focal_length.copy(new_name=self.focal_length.name + "_copy"),
            self.calib_k1.copy(new_name=self.calib_k1.name + "_copy"),
            self.calib_k2.copy(new_name=self.calib_k2.name + "_copy"),
        )

    @staticmethod
    def generate_synthetic(
        batch_size: int = 1,
        rot_random: float = 20.0,
        pos_random: float = 1.0,
        pos_base: Optional[torch.Tensor] = None,
        fl_random: float = 100.0,
        fl_base: float = 1000.0,
        k1_random: float = 0.1,
        k1_base: float = 0.0,
        k2_random: float = 0.05,
        k2_base: float = 0.0,
        name: str = "Cam",
    ):
        if pos_base is None:
            pos_base = torch.zeros(3, dtype=torch.float64)
        cam_rot = torch.cat(
            [
                random_small_quaternion(max_degrees=rot_random).unsqueeze(0)
                for _ in range(batch_size)
            ]
        )
        cam_tr = (
            torch.rand((batch_size, 3), dtype=torch.float64) * 2
            + torch.tensor([-1, -1, -1], dtype=torch.float64)
        ) * pos_random + pos_base
        cam_pose_data = torch.cat([cam_tr, cam_rot], dim=1)
        cam_pose = th.SE3(cam_pose_data, name=name + "_pose")

        focal_length = th.Vector(
            tensor=(torch.rand((batch_size, 1), dtype=torch.float64) * 2 - 1.0)
            * fl_random
            + fl_base,
            name=name + "_focal_length",
        )
        calib_k1 = th.Vector(
            tensor=(torch.rand((batch_size, 1), dtype=torch.float64) * 2 - 1.0)
            * k1_random
            + k1_base,
            name=name + "_calib_k1",
        )
        calib_k2 = th.Vector(
            tensor=(torch.rand((batch_size, 1), dtype=torch.float64) * 2 - 1.0)
            * k2_random
            + k2_base,
            name=name + "_calib_k2",
        )
        return Camera(cam_pose, focal_length, calib_k1, calib_k2)


class Observation:
    def __init__(
        self, camera_index: int, point_index: int, image_feature_point: th.Point2
    ):
        self.camera_index = camera_index
        self.point_index = point_index
        self.image_feature_point = image_feature_point


class BundleAdjustmentDataset:
    def __init__(
        self,
        cameras: List[Camera],
        points: List[th.Point3],
        observations: List[Observation],
        gt_cameras: Optional[List[Camera]] = None,
        gt_points: Optional[List[th.Point3]] = None,
    ):
        self.cameras = cameras
        self.points = points
        self.observations = observations
        self.gt_cameras = gt_cameras
        self.gt_points = gt_points

    @staticmethod
    def load_bal_dataset(path: str):
        observations = []
        cameras = []
        points = []
        with open(path, "rt") as out:
            num_cameras, num_points, num_observations = [
                int(x) for x in out.readline().rstrip().split()
            ]
            for i in range(num_observations):
                fields = out.readline().rstrip().split()
                feat = th.Point2(
                    tensor=torch.tensor(
                        [float(fields[2]), float(fields[3])], dtype=torch.float64
                    ).unsqueeze(0),
                    name=f"Feat{i}",
                )
                observations.append(
                    Observation(
                        camera_index=int(fields[0]),
                        point_index=int(fields[1]),
                        image_feature_point=feat,
                    )
                )

            for i in range(num_cameras):
                params = []
                for _ in range(9):
                    params.append(float(out.readline().rstrip()))
                cameras.append(Camera.from_params(params, name=f"Cam{i}"))

            for i in range(num_points):
                params = []
                for _ in range(3):
                    params.append(float(out.readline().rstrip()))
                points.append(
                    th.Point3(
                        tensor=torch.tensor(params, dtype=torch.float64).unsqueeze(0),
                        name=f"Pt{i}",
                    )
                )
        return cameras, points, observations

    @staticmethod
    def save_bal_dataset(
        path: str,
        cameras: List[Camera],
        points: List[th.Point3],
        observations: List[Observation],
    ):
        with open(path, "wt") as out:
            print(f"{len(cameras)} {len(points)} {len(observations)}", file=out)
            for obs in observations:
                f = obs.image_feature_point.tensor.squeeze(0).numpy()
                print(f"{obs.camera_index} {obs.point_index} {f[0]} {f[1]}", file=out)
            for cam in cameras:
                params = cam.to_params()
                for p in params:
                    print(f"{p}", file=out)
            for pt in points:
                params = pt.tensor.squeeze(0).numpy()
                for p in params:
                    print(f"{p}", file=out)

    @staticmethod
    def load_from_file(path, gt_path: Optional[str] = None):
        cameras, points, observations = BundleAdjustmentDataset.load_bal_dataset(path)
        if gt_path is not None:
            gt_points, gt_cameras, _ = BundleAdjustmentDataset.load_bal_dataset(gt_path)
        else:
            gt_points, gt_cameras = None, None
        return BundleAdjustmentDataset(
            cameras=cameras,
            points=points,
            observations=observations,
            gt_cameras=gt_cameras,
            gt_points=gt_points,
        )

    def save_to_file(self, path, gt_path: Optional[str] = None):
        BundleAdjustmentDataset.save_bal_dataset(
            path, self.cameras, self.points, self.observations
        )
        if gt_path is not None:
            BundleAdjustmentDataset.save_bal_dataset(
                gt_path, self.gt_cameras, self.gt_points, self.observations
            )

    def histogram(self) -> str:
        buckets = np.zeros(11)
        for obs in self.observations:
            proj_pt = self.cameras[obs.camera_index].project_point(
                self.points[obs.point_index]
            )
            error = float((obs.image_feature_point.tensor - proj_pt).norm())
            idx = min(int(error), len(buckets) - 1)
            buckets[idx] += 1
        max_buckets = max(buckets)
        hist_str = ""
        for i in range(len(buckets)):
            bi = buckets[i]
            label = f"{i}-{i+1}" if i + 1 < len(buckets) else f"{i}+"
            barlen = round(bi * 80 / max_buckets)
            hist_str += f"{label}: {'#' * barlen} {bi}\n"
        return hist_str

    @staticmethod
    def generate_synthetic(
        num_cameras: int,
        num_points: int,
        average_track_length: int = 7,
        track_locality: float = 0.1,
        feat_random: float = 1.5,
        prob_feat_is_outlier: float = 0.02,
        outlier_feat_random: float = 70,
    ):

        # add cameras
        gt_cameras = [
            Camera.generate_synthetic(
                pos_base=torch.tensor(
                    [-i * 100.0 / (num_cameras - 1), 0, -100], dtype=torch.float64
                ),
                name=f"Cam{i}",
            )
            for i in range(num_cameras)
        ]
        cameras = [cam.perturbed() for cam in gt_cameras]

        # add points
        gt_points = [
            th.Point3(
                tensor=(torch.rand((1, 3), dtype=torch.float64) * 2 - 1) * 20
                + torch.tensor([i * 100.0 / (num_points), 0, 0], dtype=torch.float64),
                name=f"Pt{i}",
            )
            for i in range(num_points)
        ]
        points = [
            th.Point3(
                tensor=gt_points[i].tensor + (torch.rand((1, 3)) * 2 - 1) * 0.2,
                name=gt_points[i].name + "_copy",
            )
            for i in range(num_points)
        ]

        # add observations
        pts_per_cam = average_track_length * num_points // num_cameras
        observations = []
        obs_num = 0
        for i in range(num_cameras):
            span_size = min(pts_per_cam + int(track_locality * num_points), num_points)
            span_start = (num_points - span_size) * i // num_cameras
            obs_pts = torch.randperm(span_size)[:pts_per_cam] + span_start
            for j in obs_pts:
                feat = (
                    gt_cameras[i].project_point(gt_points[j])
                    + (torch.rand(1, 2) * 2 - 1) * feat_random
                )
                if torch.rand(()) < prob_feat_is_outlier:
                    feat = feat + (torch.rand(1, 2) * 2 - 1) * outlier_feat_random
                observations.append(
                    Observation(
                        camera_index=i,
                        point_index=j,
                        image_feature_point=th.Point2(
                            tensor=feat, name=f"Feat{obs_num}"
                        ),
                    )
                )
                obs_num += 1

        return BundleAdjustmentDataset(
            cameras, points, observations, gt_cameras, gt_points
        )


def ba_histogram(
    cameras: List[Camera],
    points: List[th.Point3],
    observations: List[Observation],
) -> str:
    ba = BundleAdjustmentDataset(cameras, points, observations)
    return ba.histogram()
