import torch
import torch.nn as nn
import torch.nn.functional as F

from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.models.modeling_utils import ModelMixin


def make_group_norm(
    channels: int, max_groups: int = 32, eps: float = 1e-6
) -> nn.GroupNorm:
    groups = min(max_groups, channels)
    while channels % groups != 0 and groups > 1:
        groups -= 1
    return nn.GroupNorm(groups, channels, eps=eps)


# ---- chunked inference execution -------------------------------------------
# At large latents a block's transient temporaries (SiLU copies, full-res
# FiLM scale/shift maps, conv workspaces) dwarf its input/output tensors.
# The chunked path caps the live set at ~3 activation-sized tensors per
# block: convs write row strips (1px halo for 3x3) into preallocated
# outputs, FiLM applies in strips, and GN runs as the native full-tensor op
# (its output just fills a live-set slot the convs already require, so
# keeping it native costs no memory and keeps the chunked path bit-identical
# to the standard one). Inference-only; callers gate on
# torch.is_grad_enabled().

CHUNK_MIN_PIXELS = 512 * 512     # engage the chunked path at/above this H*W
CHUNK_STRIP_PIXELS = 128 * 1024  # target H*W of one conv row strip


def _strip_rows(height: int, width: int) -> int:
    return max(8, min(height, CHUNK_STRIP_PIXELS // max(1, width)))


def _use_chunked(x: torch.Tensor) -> bool:
    return not torch.is_grad_enabled() and x.shape[-2] * x.shape[-1] >= CHUNK_MIN_PIXELS


def chunked_conv(
    x: torch.Tensor,
    weight: torch.Tensor,
    bias: torch.Tensor | None = None,
    out: torch.Tensor | None = None,
    accumulate: bool = False,
) -> torch.Tensor:
    """Stride-1 odd-kernel conv2d in row strips.

    With ``accumulate`` the strips are added into ``out``, which lets the last
    conv of a block write straight into its residual buffer.
    """
    pad = weight.shape[-1] // 2
    batch, _, height, width = x.shape
    if out is None:
        out = x.new_empty(batch, weight.shape[0], height, width)
    rows = _strip_rows(height, width)
    for r0 in range(0, height, rows):
        r1 = min(r0 + rows, height)
        a = max(r0 - pad, 0)
        y = F.conv2d(x[:, :, a : min(r1 + pad, height)], weight, bias, padding=pad)
        y = y[:, :, r0 - a : r0 - a + r1 - r0]
        if accumulate:
            out[:, :, r0:r1] += y
        else:
            out[:, :, r0:r1] = y
    return out


def chunked_film_(film: "FilmCond2D", x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
    """FilmCond2D applied in place in row strips.

    Avoids materializing the full-res 2x-channel scale/shift map (0.74GiB at
    an 8K-UHD latent for film_cond_1).
    """
    silu, proj = film.cond_proj[0], film.cond_proj[1]
    height, width = x.shape[-2:]
    rows = _strip_rows(height, width)
    for r0 in range(0, height, rows):
        r1 = min(r0 + rows, height)
        scale_shift = proj(silu(cond[:, :, r0:r1]))
        scale, shift = scale_shift.chunk(2, dim=1)
        x[:, :, r0:r1].mul_(1 + scale).add_(shift)
    return x


class SinusoidalTimeEmbedding(nn.Module):
    def __init__(self, dim: int = 128, max_period: int = 10000):
        super().__init__()
        self.dim = dim
        self.max_period = max_period

    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
        half = self.dim // 2

        freqs = torch.exp(
            -torch.log(torch.tensor(float(self.max_period), device=timesteps.device))
            * torch.arange(half, device=timesteps.device, dtype=timesteps.dtype)
            / half
        )
        args = timesteps[:, None] * freqs[None]
        emb = torch.cat([torch.sin(args), torch.cos(args)], dim=-1)

        if self.dim % 2 == 1:
            emb = F.pad(emb, (0, 1))

        return emb


class ConditioningEncoder(nn.Module):
    def __init__(self, time_dim: int = 128, cond_dim: int = 256):
        super().__init__()
        self.time_embed = SinusoidalTimeEmbedding(time_dim)

        self.time_proj = nn.Sequential(
            nn.Linear(time_dim, cond_dim),
            nn.SiLU(),
            nn.Linear(cond_dim, cond_dim),
        )

    def forward(self, timestep: torch.Tensor) -> torch.Tensor:
        # Sinusoidal embedding stays fp32; cast to the projection weight
        # dtype so the model can run in bf16.
        emb = self.time_embed(timestep).to(self.time_proj[0].weight.dtype)
        return self.time_proj(emb)


class ConditionedResidualBlock(nn.Module):
    """
    SDXL-style residual block:
      GN -> SiLU -> Conv
      + condition (scale/shift)
      GN -> SiLU -> Dropout -> Conv
      + skip connection
    """

    def __init__(
        self,
        input_channels: int,
        output_channels: int,
        cond_dim: int = 256,
        dropout: float = 0.0,
    ):
        super().__init__()

        self.norm1 = make_group_norm(input_channels)
        self.conv1 = nn.Conv2d(
            input_channels, output_channels, kernel_size=3, padding=1
        )

        self.cond_proj = nn.Sequential(
            nn.SiLU(),
            nn.Linear(cond_dim, 2 * output_channels),
        )

        self.norm2 = make_group_norm(output_channels)
        self.dropout = nn.Dropout(dropout)
        self.conv2 = nn.Conv2d(
            output_channels, output_channels, kernel_size=3, padding=1
        )

        if input_channels != output_channels:
            self.skip = nn.Conv2d(
                input_channels, output_channels, kernel_size=1, bias=False
            )
        else:
            self.skip = nn.Identity()

    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
        if _use_chunked(x):
            return self._forward_chunked(x, cond)
        # In-place ops need the grad guard: silu/mul_ backward reads the
        # pre-op values. At inference they drop several activation-sized
        # temporaries per block (~0.5GiB peak at an 8K-UHD latent).
        inplace = not torch.is_grad_enabled()
        residual = self.skip(x)

        h = self.norm1(x)
        h = F.silu(h, inplace=inplace)
        h = self.conv1(h)

        scale_shift = self.cond_proj(cond)
        scale, shift = scale_shift.chunk(2, dim=1)
        scale = scale[:, :, None, None]
        shift = shift[:, :, None, None]

        h = self.norm2(h)
        if inplace:
            h.mul_(1 + scale).add_(shift)
        else:
            h = h * (1 + scale) + shift
        h = F.silu(h, inplace=inplace)
        h = self.dropout(h)
        h = self.conv2(h)

        return h.add_(residual) if inplace else h + residual

    def _forward_chunked(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
        # Large-latent inference: native out-of-place GN keeps x pristine so
        # it doubles as the residual; convs run in row strips. Peak live set
        # is 3 activation-sized tensors, and the output is bit-identical to
        # the standard path. Dropout is identity at inference and is skipped.
        residual = x if isinstance(self.skip, nn.Identity) else self.skip(x)
        h = self.norm1(x)
        F.silu(h, inplace=True)
        h = chunked_conv(h, self.conv1.weight, self.conv1.bias)

        scale_shift = self.cond_proj(cond)
        scale, shift = scale_shift.chunk(2, dim=1)
        h = self.norm2(h)
        h.mul_(1 + scale[:, :, None, None]).add_(shift[:, :, None, None])
        F.silu(h, inplace=True)
        out = chunked_conv(h, self.conv2.weight, self.conv2.bias)
        del h
        return out.add_(residual)


class DownStage(nn.Module):
    def __init__(
        self,
        input_channels: int,
        output_channels: int,
        cond_dim: int = 256,
        dropout: float = 0.0,
        num_blocks: int = 1,
        downsample_first: bool = False,
    ):
        super().__init__()
        self.downsample_first = downsample_first

        self.blocks = nn.ModuleList()
        for i in range(num_blocks):
            in_ch = input_channels if i == 0 else output_channels
            self.blocks.append(
                ConditionedResidualBlock(
                    input_channels=in_ch,
                    output_channels=output_channels,
                    cond_dim=cond_dim,
                    dropout=dropout,
                )
            )

        self.downsample = nn.Conv2d(
            output_channels, output_channels, kernel_size=3, stride=2, padding=1
        )

    def forward(self, x: torch.Tensor, cond: torch.Tensor):

        if self.downsample_first:
            x = self.downsample(x)

        for block in self.blocks:
            x = block(x, cond)
        skip = x

        if not self.downsample_first:
            x = self.downsample(x)

        return x, skip


class UpStage(nn.Module):
    def __init__(
        self,
        input_channels: int,
        skip_channels: int,
        output_channels: int,
        cond_dim: int = 256,
        dropout: float = 0.0,
        num_blocks: int = 1,
    ):
        super().__init__()

        self.upsample = nn.Upsample(
            scale_factor=2, mode="bilinear", align_corners=False
        )

        self.blocks = nn.ModuleList()
        for i in range(num_blocks):
            in_ch = (input_channels + skip_channels) if i == 0 else output_channels
            self.blocks.append(
                ConditionedResidualBlock(
                    input_channels=in_ch,
                    output_channels=output_channels,
                    cond_dim=cond_dim,
                    dropout=dropout,
                )
            )

    def forward(
        self, x: torch.Tensor, skip: torch.Tensor, cond: torch.Tensor
    ) -> torch.Tensor:
        x = self.upsample(x)

        if x.shape[-2:] != skip.shape[-2:]:
            x = F.interpolate(
                x, size=skip.shape[-2:], mode="bilinear", align_corners=False
            )

        if torch.is_grad_enabled() or not self._can_split(x, skip):
            x = torch.cat([x, skip], dim=1)
            for block in self.blocks:
                x = block(x, cond)
            return x

        if _use_chunked(x):
            # Hand the tensors over in a list so this frame stops pinning
            # them — the method frees each one as soon as it is consumed.
            tensors = [x, skip]
            del x, skip
            x = self._split_concat_block0_chunked(tensors, cond)
        else:
            x = self._split_concat_block0(x, skip, cond)
        for block in self.blocks[1:]:
            x = block(x, cond)
        return x

    def _can_split(self, x: torch.Tensor, skip: torch.Tensor) -> bool:
        b0 = self.blocks[0]
        ca, cb = x.shape[1], skip.shape[1]
        groups = b0.norm1.num_groups
        ga = groups * ca // (ca + cb)
        return (
            isinstance(b0.skip, nn.Conv2d)
            and groups * ca % (ca + cb) == 0
            and 0 < ga < groups
            and ca % ga == 0
            and cb % (groups - ga) == 0
        )

    def _split_concat_block0(
        self, x: torch.Tensor, skip: torch.Tensor, cond: torch.Tensor
    ) -> torch.Tensor:
        """Run blocks[0] without materializing the (ca+cb)-channel concat.

        Valid when the GroupNorm group boundaries align with the concat
        boundary (checked by _can_split): GN/SiLU apply per half, and the
        skip-conv/conv1 outputs are the sums of per-half convolutions. Only
        the fp reduction order differs from the concat path — this is the
        peak-memory point of the whole UNet at large latents (the concat is
        full-res at double width), worth ~1.1GiB at an 8K-UHD output.
        Inference-only (in-place ops), guarded by the grad check in forward.
        """
        b0 = self.blocks[0]
        ca = x.shape[1]
        cb = skip.shape[1]
        n1 = b0.norm1
        ga = n1.num_groups * ca // (ca + cb)

        w = b0.skip.weight
        residual = F.conv2d(x, w[:, :ca]).add_(F.conv2d(skip, w[:, ca:]))

        h = F.silu(
            F.group_norm(x, ga, n1.weight[:ca], n1.bias[:ca], n1.eps), inplace=True
        )
        del x
        w1 = b0.conv1.weight
        out = F.conv2d(h, w1[:, :ca], b0.conv1.bias, padding=b0.conv1.padding)
        del h
        h = F.silu(
            F.group_norm(
                skip, n1.num_groups - ga, n1.weight[ca:], n1.bias[ca:], n1.eps
            ),
            inplace=True,
        )
        del skip
        out.add_(F.conv2d(h, w1[:, ca:], padding=b0.conv1.padding))
        del h

        scale_shift = b0.cond_proj(cond)
        scale, shift = scale_shift.chunk(2, dim=1)
        out = b0.norm2(out)
        out.mul_(1 + scale[:, :, None, None]).add_(shift[:, :, None, None])
        out = F.silu(out, inplace=True)
        out = b0.conv2(b0.dropout(out))
        return out.add_(residual)

    def _split_concat_block0_chunked(
        self, tensors: list, cond: torch.Tensor
    ) -> torch.Tensor:
        """_split_concat_block0 with chunked convs and native per-half GN;
        conv2 accumulates into the residual buffer. Output is bit-identical
        to _split_concat_block0.

        ``tensors`` is [x, skip]; it is emptied so each half can be freed the
        moment its contributions are computed.
        """
        b0 = self.blocks[0]
        x, skip = tensors
        tensors.clear()
        ca = x.shape[1]
        n1 = b0.norm1
        ga = n1.num_groups * ca // (ca + skip.shape[1])

        w = b0.skip.weight
        residual = chunked_conv(x, w[:, :ca])
        chunked_conv(skip, w[:, ca:], out=residual, accumulate=True)

        h = F.group_norm(x, ga, n1.weight[:ca], n1.bias[:ca], n1.eps)
        del x
        F.silu(h, inplace=True)
        w1 = b0.conv1.weight
        out = chunked_conv(h, w1[:, :ca], b0.conv1.bias)
        del h
        h = F.group_norm(skip, n1.num_groups - ga, n1.weight[ca:], n1.bias[ca:], n1.eps)
        del skip
        F.silu(h, inplace=True)
        chunked_conv(h, w1[:, ca:], out=out, accumulate=True)
        del h

        scale_shift = b0.cond_proj(cond)
        scale, shift = scale_shift.chunk(2, dim=1)
        h = F.group_norm(out, b0.norm2.num_groups, b0.norm2.weight, b0.norm2.bias, b0.norm2.eps)
        del out
        h.mul_(1 + scale[:, :, None, None]).add_(shift[:, :, None, None])
        F.silu(h, inplace=True)
        return chunked_conv(h, b0.conv2.weight, b0.conv2.bias, out=residual, accumulate=True)


class LowResEncoder(nn.Module):
    def __init__(
        self,
        sample_channels: int = 32,
        base_channels: int = 128,
        cond_dim: int = 1024,
        dropout: float = 0.0,
    ):
        super().__init__()

        self.in_conv = nn.Conv2d(
            sample_channels, base_channels, kernel_size=1, padding=0
        )

        self.block_1 = ConditionedResidualBlock(
            input_channels=base_channels,
            output_channels=base_channels,
            cond_dim=cond_dim,
            dropout=dropout,
        )

        self.block_2 = DownStage(
            input_channels=base_channels,
            output_channels=base_channels,
            cond_dim=cond_dim,
            dropout=dropout,
            num_blocks=1,
            downsample_first=True,
        )

        self.block_3 = DownStage(
            input_channels=base_channels,
            output_channels=base_channels,
            cond_dim=cond_dim,
            dropout=dropout,
            num_blocks=1,
            downsample_first=True,
        )

    def forward(self, latents_small, cond):
        x = self.in_conv(latents_small)
        block_1_out = self.block_1(x, cond)
        block_2_out, _ = self.block_2(block_1_out, cond)
        block_3_out, _ = self.block_3(block_2_out, cond)

        return block_1_out, block_2_out, block_3_out


class FilmCond2D(nn.Module):
    def __init__(self, base_channels: int = 256, cond_channels: int = 256):
        super().__init__()

        self.cond_proj = nn.Sequential(
            nn.SiLU(),
            nn.Conv2d(cond_channels, base_channels * 2, kernel_size=1),
        )

    def forward(self, x, cond):
        if _use_chunked(x):
            return chunked_film_(self, x, cond)

        scale_shift = self.cond_proj(cond)
        scale, shift = scale_shift.chunk(2, dim=1)

        if torch.is_grad_enabled():
            return x * (1 + scale) + shift
        # Callers rebind x, so modifying it in place at inference is safe and
        # avoids two full-res temporaries.
        return x.mul_(1 + scale).add_(shift)


class UpscalerUNet(ModelMixin, ConfigMixin):
    """FlowUpscaler velocity-prediction UNet (diffusers `ModelMixin`).

    Attention-free U-Net with SDXL-style residual blocks conditioned on the
    timestep (FiLM from a sinusoidal embedding) and on the low-resolution
    latents (multi-scale FiLM from `LowResEncoder`). The module tree is
    identical to the original standalone `nn.Module`, so the state dict is
    interchangeable with the flat `flow_upscaler.safetensors` checkpoint.
    """

    @register_to_config
    def __init__(
        self,
        sample_channels: int = 32,
        base_channels: int = 384,
        time_dim: int = 512,
        cond_dim: int = 1024,
        dropout: float = 0.01,
    ):
        super().__init__()

        self.conditioning = ConditioningEncoder(
            time_dim=time_dim,
            cond_dim=cond_dim,
        )

        self.in_conv = nn.Conv2d(
            sample_channels, base_channels, kernel_size=1, padding=0
        )

        self.low_res_encoder = LowResEncoder(base_channels=base_channels)

        self.film_cond_1 = FilmCond2D(
            base_channels=base_channels, cond_channels=base_channels
        )
        self.film_cond_2 = FilmCond2D(
            base_channels=base_channels, cond_channels=base_channels
        )
        self.film_cond_3 = FilmCond2D(
            base_channels=base_channels, cond_channels=base_channels
        )

        self.down_stages = nn.ModuleList(
            [
                DownStage(
                    input_channels=base_channels,
                    output_channels=base_channels,
                    cond_dim=cond_dim,
                    dropout=dropout,
                    num_blocks=3,
                ),
                DownStage(
                    input_channels=base_channels,
                    output_channels=base_channels,
                    cond_dim=cond_dim,
                    dropout=dropout,
                    num_blocks=2,
                ),
            ]
        )

        self.mid_stages = nn.ModuleList(
            [
                ConditionedResidualBlock(
                    input_channels=base_channels,
                    output_channels=base_channels,
                    cond_dim=cond_dim,
                    dropout=dropout,
                )
                for i in range(1)
            ]
        )

        self.up_stages = nn.ModuleList(
            [
                UpStage(
                    input_channels=base_channels,
                    skip_channels=base_channels,
                    output_channels=base_channels,
                    cond_dim=cond_dim,
                    dropout=dropout,
                    num_blocks=2,
                ),
                UpStage(
                    input_channels=base_channels,
                    skip_channels=base_channels,
                    output_channels=base_channels,
                    cond_dim=cond_dim,
                    dropout=dropout,
                    num_blocks=3,
                ),
            ]
        )

        self.out_conv = nn.Conv2d(
            base_channels, sample_channels, kernel_size=1, padding=0
        )

    def forward(
        self, sample: torch.Tensor, timestep: torch.Tensor, latents_small: torch.Tensor
    ) -> torch.Tensor:
        cond = self.conditioning(timestep)

        B, C, H, W = sample.shape

        lr_cond_1, lr_cond_2, lr_cond_3 = self.low_res_encoder(latents_small, cond)

        lr_cond_1 = torch.nn.functional.interpolate(lr_cond_1, (H, W), mode="bilinear")
        lr_cond_2 = torch.nn.functional.interpolate(
            lr_cond_2, (H // 2, W // 2), mode="bilinear"
        )
        lr_cond_3 = torch.nn.functional.interpolate(
            lr_cond_3, (H // 4, W // 4), mode="bilinear"
        )

        x = self.in_conv(sample)
        x = self.film_cond_1(x, lr_cond_1)
        # Each lr_cond is dead after its film layer; without the dels the
        # full-res lr_cond_1 (~0.4GiB at an 8K-UHD latent) stays pinned for
        # the whole forward.
        del lr_cond_1

        skips = []

        x, skip = self.down_stages[0](x, cond)
        skips.append(skip)

        x = self.film_cond_2(x, lr_cond_2)
        del lr_cond_2

        x, skip = self.down_stages[1](x, cond)
        skips.append(skip)

        x = self.film_cond_3(x, lr_cond_3)
        del lr_cond_3

        for mid in self.mid_stages:
            x = mid(x, cond)

        for up in self.up_stages:
            x = up(x, skips.pop(), cond)

        x = self.out_conv(x)
        return x