"""Multi-Scale Flow-Warp-Mask U-Net: predicts flow at multiple resolutions."""
import torch
import torch.nn as nn
import torch.nn.functional as F


class ResConvBlock(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1)
        self.gn1 = nn.GroupNorm(min(8, out_ch), out_ch)
        self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1)
        self.gn2 = nn.GroupNorm(min(8, out_ch), out_ch)
        self.proj = nn.Conv2d(in_ch, out_ch, 1) if in_ch != out_ch else nn.Identity()

    def forward(self, x):
        residual = self.proj(x)
        x = F.silu(self.gn1(self.conv1(x)))
        x = F.silu(self.gn2(self.conv2(x)))
        return x + residual


class MultiScaleFlowUNet(nn.Module):
    def __init__(self, in_channels=12, channels=[64, 128, 256]):
        super().__init__()
        # Encoder
        self.encoders = nn.ModuleList()
        self.pools = nn.ModuleList()
        prev_ch = in_channels
        for ch in channels:
            self.encoders.append(ResConvBlock(prev_ch, ch))
            self.pools.append(nn.MaxPool2d(2))
            prev_ch = ch

        # Bottleneck
        self.bottleneck = ResConvBlock(channels[-1], channels[-1] * 2)

        # Decoder
        self.upconvs = nn.ModuleList()
        self.decoders = nn.ModuleList()
        dec_channels = list(reversed(channels))
        prev_ch = channels[-1] * 2
        for ch in dec_channels:
            self.upconvs.append(nn.ConvTranspose2d(prev_ch, ch, 2, stride=2))
            self.decoders.append(ResConvBlock(ch * 2, ch))
            prev_ch = ch

        # Multi-scale flow heads at each decoder level
        # dec_channels = [256, 128, 64] (coarsest to finest)
        # Level 0 (coarsest, 8x8): flow refinement
        # Level 1 (16x16): flow refinement
        # Level 2 (finest, 64x64): flow refinement + mask + gen_frame
        self.flow_heads = nn.ModuleList()
        for ch in dec_channels:
            head = nn.Conv2d(ch, 2, 1)
            nn.init.zeros_(head.weight)
            nn.init.zeros_(head.bias)
            self.flow_heads.append(head)

        # Mask and generation heads only at finest level (level 2, 64x64)
        self.mask_head = nn.Conv2d(dec_channels[-1], 1, 1)
        nn.init.zeros_(self.mask_head.weight)
        nn.init.zeros_(self.mask_head.bias)

        self.gen_head = nn.Conv2d(dec_channels[-1], 3, 1)

    def forward(self, x):
        skips = []
        for enc, pool in zip(self.encoders, self.pools):
            x = enc(x)
            skips.append(x)
            x = pool(x)

        x = self.bottleneck(x)

        flows = []  # flow at each level, from coarsest to finest
        for i, (upconv, dec, skip) in enumerate(zip(self.upconvs, self.decoders, reversed(skips))):
            x = upconv(x)
            x = torch.cat([x, skip], dim=1)
            x = dec(x)

            # Predict flow refinement at this level
            flow_refine = self.flow_heads[i](x)

            if i == 0:
                # Coarsest level: just the flow refinement
                flow = flow_refine
            else:
                # Upsample previous flow and add refinement
                prev_flow_up = F.interpolate(flows[-1], scale_factor=2, mode='bilinear', align_corners=True)
                # Scale flow values by 2 since coordinates double
                prev_flow_up = prev_flow_up * 2
                flow = prev_flow_up + flow_refine

            flows.append(flow)

        # Final level outputs
        mask = torch.sigmoid(self.mask_head(x))
        gen_frame = self.gen_head(x)

        # flows[-1] is the finest (64x64) flow
        return flows, mask, gen_frame