import torch
from jaxtyping import Float

from .render import render, _make_tgt
from ..data import Subject


class DRR(torch.nn.Module):
    def __init__(
        self,
        k_inv: Float[torch.Tensor, "B 3 3"],
        sdd: Float[torch.Tensor, "B"],
        height: int,
        width: int,
    ):
        super().__init__()
        self._intrinsic_params = {"k_inv", "sdd", "height", "width"}
        self.k_inv = k_inv
        self.sdd = sdd
        self.height = height
        self.width = width
        self._compute_tgt()

    def forward(
        self,
        subject: Subject,
        rt_inv: Float[torch.Tensor, "B 4 4"],
        n_samples: int = 500,
        align_corners: bool = True,
    ):
        return render(
            subject,
            self.k_inv,
            rt_inv,
            self.sdd,
            self.height,
            self.width,
            n_samples,
            align_corners,
            self.tgt,
        )

    def __setattr__(self, name, value):
        super().__setattr__(name, value)
        if hasattr(self, "_intrinsic_params") and name in self._intrinsic_params:
            self._compute_tgt()

    def _compute_tgt(self):
        """Compute target points based on current intrinsic parameters."""
        self.tgt = _make_tgt(self.k_inv, self.sdd, self.height, self.width, self.k_inv.device, self.k_inv.dtype)
