import torch


class RF:
    def __init__(self, model, ln=True):
        self.model = model
        self.ln = ln

    def forward(self, x0, x1):
        b = x0.size(0)
        if self.ln:
            nt = torch.randn((b,)).to(x0.device)
            t = torch.sigmoid(nt)
        else:
            t = torch.rand((b,)).to(x0.device)
        texp = t.view([b, *([1] * len(x0.shape[1:]))])
        zt = (1 - texp) * x0 + texp * x1
        vtheta = self.model(zt, t)
        batchwise_mse = ((x1 - x0 - vtheta) ** 2).mean(
            dim=list(range(1, len(x0.shape)))
        )
        tlist = batchwise_mse.detach().cpu().reshape(-1).tolist()
        ttloss = [(tv, tloss) for tv, tloss in zip(t, tlist)]
        return batchwise_mse.mean(), ttloss

    @torch.no_grad()
    def sample(self, x0, sample_steps=50):
        b = x0.size(0)
        dt = 1.0 / sample_steps
        dt = torch.tensor([dt] * b).to(x0.device).view([b, *([1] * len(x0.shape[1:]))])
        images = [x0]
        z = x0
        for i in range(sample_steps):
            t = i / sample_steps
            t = torch.tensor([t] * b).to(z.device)

            vc = self.model(z, t)

            z = z + dt * vc
            images.append(z)
        return images
