"""Compact but faithful conditional-regression and residual-EDM CorrDiff."""

import math

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


class ConvBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        groups = min(8, out_channels)
        self.block = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 3, padding=1),
            nn.GroupNorm(groups, out_channels),
            nn.SiLU(),
            nn.Conv2d(out_channels, out_channels, 3, padding=1),
            nn.GroupNorm(groups, out_channels),
            nn.SiLU(),
        )
        self.skip = nn.Conv2d(in_channels, out_channels, 1)

    def forward(self, x):
        return self.block(x) + self.skip(x)


class RegressionUNet(nn.Module):
    """Deterministic conditional mean, computed cheaply before 448 px recovery."""

    def __init__(self, in_channels=12, out_channels=4, base_channels=16, feature_size=56,
                 output_size=448):
        super().__init__()
        self.feature_size = feature_size
        self.output_size = output_size
        self.net = nn.Sequential(
            ConvBlock(in_channels, base_channels),
            ConvBlock(base_channels, base_channels * 2),
            ConvBlock(base_channels * 2, base_channels),
            nn.Conv2d(base_channels, out_channels, 1),
        )

    def forward(self, coarse):
        x = F.interpolate(coarse, (self.feature_size, self.feature_size), mode="bilinear",
                          align_corners=False)
        x = self.net(x)
        return F.interpolate(x, (self.output_size, self.output_size), mode="bilinear",
                             align_corners=False)


class ResidualDenoiser(nn.Module):
    """Noise-conditional network used inside EDM preconditioning."""

    def __init__(self, condition_channels=16, out_channels=4, base_channels=16,
                 feature_size=56):
        super().__init__()
        self.feature_size = feature_size
        self.noise_mlp = nn.Sequential(
            nn.Linear(1, base_channels), nn.SiLU(), nn.Linear(base_channels, base_channels)
        )
        self.input = ConvBlock(condition_channels, base_channels)
        self.body = nn.Sequential(
            ConvBlock(base_channels, base_channels * 2),
            ConvBlock(base_channels * 2, base_channels),
            nn.Conv2d(base_channels, out_channels, 1),
        )

    def forward(self, noisy, condition, c_noise):
        size = noisy.shape[-2:]
        x = torch.cat((noisy, condition), dim=1)
        x = F.interpolate(x, (self.feature_size, self.feature_size), mode="bilinear",
                          align_corners=False)
        x = self.input(x)
        x = x + self.noise_mlp(c_noise[:, None].float())[:, :, None, None].to(x.dtype)
        x = self.body(x)
        return F.interpolate(x, size, mode="bilinear", align_corners=False)


class CorrDiff(nn.Module):
    def __init__(self, in_channels=12, out_channels=4, base_channels=16, feature_size=56,
                 output_size=448, sigma_data=0.5):
        super().__init__()
        self.sigma_data = sigma_data
        self.output_size = output_size
        self.regression = RegressionUNet(in_channels, out_channels, base_channels,
                                         feature_size, output_size)
        self.diffusion = ResidualDenoiser(in_channels + 2 * out_channels, out_channels,
                                          base_channels, feature_size)

    def mean(self, coarse):
        return self.regression(coarse)

    def condition(self, coarse, mean):
        coarse = F.interpolate(coarse, mean.shape[-2:], mode="bilinear", align_corners=False)
        return torch.cat((coarse, mean), dim=1)

    def denoise(self, noisy_residual, coarse, mean, sigma):
        sigma = sigma.reshape(-1, 1, 1, 1).to(noisy_residual.dtype)
        sigma_data = self.sigma_data
        c_skip = sigma_data**2 / (sigma.square() + sigma_data**2)
        c_out = sigma * sigma_data / (sigma.square() + sigma_data**2).sqrt()
        c_in = (sigma.square() + sigma_data**2).rsqrt()
        c_noise = sigma.flatten().log() / 4
        network = self.diffusion(c_in * noisy_residual, self.condition(coarse, mean), c_noise)
        return c_skip * noisy_residual + c_out * network

    @staticmethod
    def karras_schedule(steps, sigma_min, sigma_max, rho, device):
        ramp = torch.linspace(0, 1, steps, device=device)
        maximum = sigma_max ** (1 / rho)
        minimum = sigma_min ** (1 / rho)
        sigmas = (maximum + ramp * (minimum - maximum)) ** rho
        return torch.cat((sigmas, sigmas.new_zeros(1)))

    def sample(self, coarse, steps=4, sigma_min=0.002, sigma_max=5.0, rho=7.0,
               solver="heun"):
        mean = self.mean(coarse)
        sigmas = self.karras_schedule(steps, sigma_min, sigma_max, rho, coarse.device)
        x = torch.randn_like(mean) * sigmas[0]
        for index, (current, following) in enumerate(zip(sigmas[:-1], sigmas[1:])):
            sigma = current.expand(coarse.shape[0])
            denoised = self.denoise(x, coarse, mean, sigma)
            derivative = (x - denoised) / current
            proposal = x + (following - current) * derivative
            if solver == "heun" and index < len(sigmas) - 2:
                next_sigma = following.expand(coarse.shape[0])
                next_denoised = self.denoise(proposal, coarse, mean, next_sigma)
                next_derivative = (proposal - next_denoised) / following
                x = x + (following - current) * (derivative + next_derivative) / 2
            else:
                x = proposal
        return mean + x

    def forward(self, coarse, mode="sample", mean=None, noisy=None, sigma=None,
                **sample_options):
        if mode == "mean":
            return self.mean(coarse)
        if mode == "denoise":
            return self.denoise(noisy, coarse, mean, sigma)
        return self.sample(coarse, **sample_options)


__all__ = ["CorrDiff", "RegressionUNet", "ResidualDenoiser"]