from typing import Optional, Sequence

import torch
from torch import Tensor, nn

from .diffusion import (
    ADPM2Sampler,
    Diffusion,
    DiffusionSampler,
    Distribution,
    KarrasSchedule,
    LogNormalDistribution,
    Sampler,
    Schedule,
)
from .modules import UNet1d


class Model1d(nn.Module):
    def __init__(
        self,
        in_channels: int,
        channels: int,
        patch_size: int,
        kernel_sizes_init: Sequence[int],
        out_means: int,
        multipliers: Sequence[int],
        factors: Sequence[int],
        num_blocks: Sequence[int],
        attentions: Sequence[bool],
        attention_heads: int,
        attention_features: int,
        attention_multiplier: int,
        use_attention_bottleneck: bool,
        resnet_groups: int,
        kernel_multiplier_downsample: int,
        use_nearest_upsample: bool,
        use_skip_scale: bool,
        diffusion_sigma_distribution: Distribution,
        diffusion_sigma_data: int,
        diffusion_dynamic_threshold: float,
        out_channels: Optional[int] = None,
        context_channels: Optional[Sequence[int]] = None,
    ):
        super().__init__()

        self.unet = UNet1d(
            in_channels=in_channels,
            channels=channels,
            patch_size=patch_size,
            resnet_groups=resnet_groups,
            kernel_multiplier_downsample=kernel_multiplier_downsample,
            kernel_sizes_init=kernel_sizes_init,
            out_means=out_means,
            multipliers=multipliers,
            factors=factors,
            num_blocks=num_blocks,
            attentions=attentions,
            attention_heads=attention_heads,
            attention_features=attention_features,
            attention_multiplier=attention_multiplier,
            use_nearest_upsample=use_nearest_upsample,
            use_skip_scale=use_skip_scale,
            use_attention_bottleneck=use_attention_bottleneck,
            out_channels=out_channels,
            context_channels=context_channels,
        )

        self.diffusion = Diffusion(
            net=self.unet,
            sigma_distribution=diffusion_sigma_distribution,
            sigma_data=diffusion_sigma_data,
            dynamic_threshold=diffusion_dynamic_threshold,
        )

    def forward(self, x: Tensor, **kwargs) -> Tensor:
        return self.diffusion(x, **kwargs)

    def sample(
        self,
        noise: Tensor,
        num_steps: int,
        sigma_schedule: Schedule,
        sampler: Sampler,
        **kwargs
    ) -> Tensor:
        diffusion_sampler = DiffusionSampler(
            diffusion=self.diffusion,
            sampler=sampler,
            sigma_schedule=sigma_schedule,
            num_steps=num_steps,
        )
        return diffusion_sampler(noise, **kwargs)


class AudioDiffusionModel(Model1d):
    def __init__(self, *args, **kwargs):
        default_kwargs = dict(
            in_channels=1,
            channels=128,
            patch_size=16,
            kernel_sizes_init=[1, 3, 7],
            out_means=4,
            multipliers=[1, 2, 4, 4, 4, 4, 4],
            factors=[4, 4, 4, 2, 2, 2],
            num_blocks=[2, 2, 2, 2, 2, 2],
            attentions=[False, False, False, True, True, True],
            attention_heads=8,
            attention_features=64,
            attention_multiplier=2,
            use_attention_bottleneck=True,
            resnet_groups=8,
            kernel_multiplier_downsample=2,
            use_nearest_upsample=False,
            use_skip_scale=True,
            diffusion_sigma_distribution=LogNormalDistribution(mean=-3.0, std=1.0),
            diffusion_sigma_data=0.1,
            diffusion_dynamic_threshold=0.0,
        )

        super().__init__(*args, **{**default_kwargs, **kwargs})

    def sample(self, *args, **kwargs):
        default_kwargs = dict(
            sigma_schedule=KarrasSchedule(sigma_min=0.0001, sigma_max=3.0, rho=9.0),
            sampler=ADPM2Sampler(rho=1.0),
        )
        return super().sample(*args, **{**default_kwargs, **kwargs})


class AudioDiffusionUpsampler(Model1d):
    def __init__(self, factor: int, in_channels: int = 1, *args, **kwargs):
        self.factor = factor

        default_kwargs = dict(
            in_channels=in_channels,
            channels=128,
            patch_size=16,
            out_means=4,
            kernel_sizes_init=[1, 3, 7],
            multipliers=[1, 2, 4, 4, 4, 4, 4],
            factors=[4, 4, 4, 2, 2, 2],
            num_blocks=[2, 2, 2, 2, 2, 2],
            attentions=[False, False, False, True, True, True],
            attention_heads=8,
            attention_features=64,
            attention_multiplier=2,
            use_attention_bottleneck=True,
            resnet_groups=8,
            kernel_multiplier_downsample=2,
            use_nearest_upsample=False,
            use_skip_scale=True,
            diffusion_sigma_distribution=LogNormalDistribution(mean=-3.0, std=1.0),
            diffusion_sigma_data=0.1,
            diffusion_dynamic_threshold=0.0,
            context_channels=[in_channels],
        )

        super().__init__(*args, **{**default_kwargs, **kwargs})  # type: ignore

    def forward(self, x: Tensor, **kwargs) -> Tensor:
        # Downsample by picking every `factor` item
        downsampled = x[:, :, :: self.factor]
        # Upsample by interleaving to get context
        context = torch.repeat_interleave(downsampled, repeats=self.factor, dim=2)
        return self.diffusion(x, context=[context], **kwargs)

    def sample(self, start: Tensor, *args, **kwargs):  # type: ignore
        context = torch.repeat_interleave(start, repeats=self.factor, dim=2)
        noise = torch.randn_like(context)
        default_kwargs = dict(
            context=[context],
            sigma_schedule=KarrasSchedule(sigma_min=0.0001, sigma_max=3.0, rho=9.0),
            sampler=ADPM2Sampler(rho=1.0),
        )
        return super().sample(noise, *args, **{**default_kwargs, **kwargs})  # type: ignore # noqa
