import torch
from jaxtyping import Float

from ..data import Subject
from ..drr import render
from ..geometry import convert, transform_point


class Registration(torch.nn.Module):
    def __init__(
        self,
        subject: Subject,
        rt_inv: Float[torch.Tensor, "1 4 4"],
        k_inv: Float[torch.Tensor, "1 3 3"],
        sdd: Float[torch.Tensor, "1"],
        height: int,
        width: int,
        eps: float = 1e-8,
    ):
        super().__init__()
        self.subject = subject
        self.rt_inv = rt_inv
        self.k_inv = k_inv
        self.sdd = sdd
        self.height = height
        self.width = width

        # Rotation pivot: isocenter projected into camera frame
        c = transform_point(rt_inv.inverse(), subject.isocenter[None])
        self.pivot = torch.eye(4, device=c.device, dtype=c.dtype)[None]
        self.pivot_inv = torch.eye(4, device=c.device, dtype=c.dtype)[None]
        self.pivot[:, :3, 3] = c
        self.pivot_inv[:, :3, 3] = -c

        # Parameterization of the perturbation
        self._rot = torch.nn.Parameter(eps * torch.randn(1, 3, device=c.device))
        self._xyz = torch.nn.Parameter(eps * torch.randn(1, 3, device=c.device))

    def forward(self) -> Float[torch.Tensor, "1 C H W"]:
        return render(
            self.subject,
            self.k_inv,
            self.rt_inv @ self.pose,
            self.sdd,
            self.height,
            self.width,
        )

    @property
    def pose(self) -> Float[torch.Tensor, "1 4 4"]:
        T = convert(self._rot, self._xyz, "so3_log")
        return self.pivot @ T @ self.pivot_inv
