from dataclasses import dataclass

import torch
from einops import rearrange
from jaxtyping import Float
from lpips import LPIPS
from torch import Tensor

from ..dataset.types import BatchedExample
from ..model.decoder.decoder import DecoderOutput
from ..model.types import Gaussians
from .loss import Loss


@dataclass
class LossLpipsCfg:
    weight: float
    apply_after_step: int


@dataclass
class LossLpipsCfgWrapper:
    lpips: LossLpipsCfg


class LossLpips(Loss[LossLpipsCfg, LossLpipsCfgWrapper]):
    lpips: LPIPS

    def __init__(self, cfg: LossLpipsCfgWrapper) -> None:
        super().__init__(cfg)

        lpips = LPIPS(net="vgg").eval()
        lpips.requires_grad_(False)
        self.__dict__["lpips"] = lpips

    def lpips_on_device(self, device: torch.device) -> LPIPS:
        lpips = self.lpips
        param = next(iter(lpips.parameters()))
        if param.device != device or param.dtype != torch.float32:
            lpips.to(device=device, dtype=torch.float32)
        return lpips

    def forward(
        self,
        prediction: DecoderOutput,
        batch: BatchedExample,
        gaussians: Gaussians,
        global_step: int,
        valid_depth_mask: Tensor | None,
        reduction: str = "mean",
    ) -> Float[Tensor, "*batch"]:
        image = batch["target"]["image"]

        # Before the specified step, don't apply the loss.
        if global_step < self.cfg.apply_after_step:
            return torch.tensor(0, dtype=torch.float32, device=image.device)

        prediction_color = prediction.color
        if valid_depth_mask is not None and valid_depth_mask.max() > 0.5:
            prediction_color = prediction_color.clone()
            image = image.clone()
            prediction_color[valid_depth_mask] = 0
            image[valid_depth_mask] = 0

        with torch.autocast(device_type=image.device.type, enabled=False):
            loss = self.lpips_on_device(image.device).forward(
                rearrange(prediction_color.float(), "b v c h w -> (b v) c h w"),
                rearrange(image.float(), "b v c h w -> (b v) c h w"),
                normalize=True,
            )
        batch_size, num_views = image.shape[:2]
        sample_loss = loss.reshape(batch_size, num_views, -1).mean(dim=(1, 2))

        if reduction == "mean":
            return self.cfg.weight * sample_loss.mean()
        if reduction == "none":
            return self.cfg.weight * sample_loss

        raise ValueError(f"Unsupported reduction={reduction}")