"""LaCT (Large-Chunk TTT, LVSM 계열) 를 abl kit 학습 루프에 끼우는 wrapper.

모델 코드는 git clone 한 LaCT repo (lact_nvs/model.py, lact_ttt.py) 를 수정 없이 그대로 쓴다.
kit 의 입력 (images, rays, camera_cond, num_cond_views) 을 LaCT 입력으로 바꿀 뿐이다:
  camera_cond["extrinsics"] = 정규화된 c2w 의 역행렬 (w2c, normalize_for_da3=False)
  camera_cond["intrinsics"] = 학습 해상도 (256x256) 기준 K
-> LaCT 의 c2w / fxfycxcy. kit 의 renderer 가 쓰는 Plucker ray 와 같은 카메라 정규화
   (첫 context 카메라 기준, 가장 먼 context 카메라가 1/1.35) 를 그대로 받는다.
출력은 kit 모델과 같이 [context 이미지, 렌더링한 target] 을 view 축으로 붙여 돌려준다.
"""

import os
import sys

import torch
import torch.nn as nn
import torch.utils.checkpoint
from omegaconf import OmegaConf


def _import_lact(lact_nvs_dir):
    # LaCT config 의 block type 이 "model.SelfAttention", "lact_ttt.FastWeightGluMLPMultihead"
    # 처럼 top-level 모듈 이름이라 lact_nvs 폴더를 sys.path 에 넣어야 한다.
    lact_nvs_dir = os.path.abspath(lact_nvs_dir)
    if lact_nvs_dir not in sys.path:
        sys.path.insert(0, lact_nvs_dir)
    import model as lact_model  # noqa: E402  (LaCT/lact_nvs/model.py)
    assert os.path.dirname(os.path.abspath(lact_model.__file__)) == lact_nvs_dir, lact_model.__file__
    return lact_model


class LaCTWrapper(nn.Module):
    def __init__(self, lact_nvs_dir, lact_config, ckpt_blocks=0, compile_blocks=False):
        super().__init__()
        lact_model = _import_lact(lact_nvs_dir)
        cfg = OmegaConf.to_container(OmegaConf.load(lact_config), resolve=True)
        self.lact = lact_model.LaCTLVSM(**cfg)
        # 학습 중 앞에서부터 ckpt_blocks 개 Block 을 activation checkpointing 한다 (0 = 끔).
        # Block.forward 만 바꾸므로 state_dict key 는 LaCTLVSM 그대로다.
        self.ckpt_blocks = int(ckpt_blocks)
        for i, blk in enumerate(self.lact.blocks):
            if i < self.ckpt_blocks:
                blk.forward = self._ckpt_forward(blk.forward)
        # LaCT README: TTT 모델은 compile 해야 제 속도가 난다. Block 을 하나씩 in-place 로 compile
        # 하고 checkpoint 는 compile 안에서 돈다 (2026-11-13 memtest, B300):
        #   LaCTLVSM 전체를 compile -> inductor 커널이 "misaligned address" 로 죽는다.
        #   checkpoint 바깥 + compile 안쪽 -> backward 재계산에서 CheckpointError.
        self.compile_blocks = bool(compile_blocks)
        if self.compile_blocks:
            # block 마다 guard (forward closure id) 가 달라 block 수만큼 재컴파일된다 (기본 한도 8).
            torch._dynamo.config.recompile_limit = 128
            torch._dynamo.config.accumulated_recompile_limit = 1024
            for blk in self.lact.blocks:
                blk.compile()

    def _ckpt_forward(self, fwd):
        def run(x, info):
            if self.training and torch.is_grad_enabled():
                return torch.utils.checkpoint.checkpoint(fwd, x, info, use_reentrant=False)
            return fwd(x, info)
        return run

    def forward(self, images, rays, camera_cond, num_cond_views, timeit=False):
        n = int(num_cond_views)
        w2c = camera_cond["extrinsics"].to(images.device).float()
        K = camera_cond["intrinsics"].to(images.device).float()
        c2w = torch.linalg.inv(w2c)
        fxfycxcy = torch.stack([K[..., 0, 0], K[..., 1, 1], K[..., 0, 2], K[..., 1, 2]], dim=-1)

        input_data = {"image": images[:, :n].float(), "fxfycxcy": fxfycxcy[:, :n], "c2w": c2w[:, :n]}
        target_data = {"fxfycxcy": fxfycxcy[:, n:], "c2w": c2w[:, n:]}
        rendered = self.lact(input_data, target_data)
        return torch.cat([images[:, :n], rendered.to(images.dtype)], dim=1)