"""LiquidInbetween: a small learned in-betweener for photos.

The architecture is RIFE's IFNet (Huang et al., ECCV 2022, MIT licence,
github.com/hzwer/ECCV2022-RIFE), extended with a timestep input so one pass
gives the frame at any t in (0, 1), not only the middle. The extra channel
starts at zero, so training begins exactly at RIFE's behaviour at t = 0.5.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F


def conv(cin, cout, k=3, s=1, p=1):
    return nn.Sequential(nn.Conv2d(cin, cout, k, s, p, bias=True), nn.PReLU(cout))


def base_grid(n, h, w, device, dtype):
    gx = torch.linspace(-1.0, 1.0, w, device=device, dtype=dtype).view(1, 1, 1, w).expand(n, -1, h, -1)
    gy = torch.linspace(-1.0, 1.0, h, device=device, dtype=dtype).view(1, 1, h, 1).expand(n, -1, -1, w)
    return torch.cat([gx, gy], 1)


def warp(img, flow):
    """Backward warp: each output pixel reads `img` at its position + `flow` (pixels)."""
    n, _, h, w = flow.shape
    f = torch.cat([flow[:, 0:1] / ((w - 1.0) / 2.0), flow[:, 1:2] / ((h - 1.0) / 2.0)], 1)
    # Clamping the grid gives border padding on every backend (MPS and Core ML too).
    g = (base_grid(n, h, w, flow.device, flow.dtype) + f).clamp(-1.0, 1.0).permute(0, 2, 3, 1)
    return F.grid_sample(img, g, mode="bilinear", padding_mode="zeros", align_corners=True)


class IFBlock(nn.Module):
    def __init__(self, in_planes, c=90):
        super().__init__()
        self.conv0 = nn.Sequential(conv(in_planes, c // 2, 3, 2, 1), conv(c // 2, c, 3, 2, 1))
        self.convblock0 = nn.Sequential(conv(c, c), conv(c, c))
        self.convblock1 = nn.Sequential(conv(c, c), conv(c, c))
        self.convblock2 = nn.Sequential(conv(c, c), conv(c, c))
        self.convblock3 = nn.Sequential(conv(c, c), conv(c, c))
        self.conv1 = nn.Sequential(nn.ConvTranspose2d(c, c // 2, 4, 2, 1), nn.PReLU(c // 2), nn.ConvTranspose2d(c // 2, 4, 4, 2, 1))
        self.conv2 = nn.Sequential(nn.ConvTranspose2d(c, c // 2, 4, 2, 1), nn.PReLU(c // 2), nn.ConvTranspose2d(c // 2, 1, 4, 2, 1))

    def forward(self, x, flow, scale):
        x = F.interpolate(x, scale_factor=1.0 / scale, mode="bilinear", align_corners=False)
        flow = F.interpolate(flow, scale_factor=1.0 / scale, mode="bilinear", align_corners=False) / scale
        feat = self.conv0(torch.cat((x, flow), 1))
        feat = self.convblock0(feat) + feat
        feat = self.convblock1(feat) + feat
        feat = self.convblock2(feat) + feat
        feat = self.convblock3(feat) + feat
        flow = F.interpolate(self.conv1(feat), scale_factor=scale, mode="bilinear", align_corners=False) * scale
        mask = F.interpolate(self.conv2(feat), scale_factor=scale, mode="bilinear", align_corners=False)
        return flow, mask


class IFNet(nn.Module):
    """RIFE's IFNet. With `timestep`, every block also sees t (and the
    mirrored pass 1 − t), so the frame at any t comes out in one pass."""

    def __init__(self, timestep=False, c=90):
        super().__init__()
        self.timestep = timestep
        extra = 1 if timestep else 0
        self.block0 = IFBlock(7 + 4 + extra, c)
        self.block1 = IFBlock(7 + 4 + extra, c)
        self.block2 = IFBlock(7 + 4 + extra, c)

    def forward(self, img0, img1, t=None, all_scales=False):
        n, _, h, w = img0.shape
        flow = torch.zeros(n, 4, h, w, device=img0.device, dtype=img0.dtype)
        mask = torch.zeros(n, 1, h, w, device=img0.device, dtype=img0.dtype)
        if self.timestep:
            tt = (t if t is not None else torch.full((n, 1, 1, 1), 0.5, device=img0.device, dtype=img0.dtype)).expand(n, 1, h, w)
        w0, w1 = img0, img1
        outputs = []
        for block, scale in ((self.block0, 4), (self.block1, 2), (self.block2, 1)):
            a = [w0, w1, mask] + ([tt] if self.timestep else [])
            b = [w1, w0, -mask] + ([1 - tt] if self.timestep else [])
            f0, m0 = block(torch.cat(a, 1), flow, scale)
            f1, m1 = block(torch.cat(b, 1), torch.cat((flow[:, 2:4], flow[:, :2]), 1), scale)
            flow = flow + (f0 + torch.cat((f1[:, 2:4], f1[:, :2]), 1)) / 2
            mask = mask + (m0 - m1) / 2
            w0 = warp(img0, flow[:, :2])
            w1 = warp(img1, flow[:, 2:4])
            if all_scales:
                m = torch.sigmoid(mask)
                outputs.append(w0 * m + w1 * (1 - m))
        m = torch.sigmoid(mask)
        out = w0 * m + w1 * (1 - m)
        return outputs if all_scales else out


def load_teacher(path):
    """RIFE's ECCV 2022 weights (safetensors) into IFNet()."""
    from safetensors.torch import load_file
    net = IFNet(timestep=False)
    state = load_file(path)
    state = {k.replace("module.", ""): v for k, v in state.items() if not k.startswith("block_tea")}
    missing, unexpected = net.load_state_dict(state, strict=False)
    assert not missing, missing
    return net


def student_from_teacher(teacher):
    """IFNet(timestep=True) starting exactly as the teacher: the new t channel's
    weights are zero."""
    student = IFNet(timestep=True)
    sd = student.state_dict()
    for k, v in teacher.state_dict().items():
        if k.endswith("conv0.0.0.weight"):
            new = torch.zeros_like(sd[k])
            # Inputs: warped0 (3), warped1 (3), mask (1) | t (1) | flow (4)
            new[:, :7] = v[:, :7]
            new[:, 8:] = v[:, 7:]
            sd[k] = new
        else:
            sd[k] = v
    student.load_state_dict(sd)
    return student