from functools import partial
import torch
from torch import nn
from torch.nn import Module, ModuleList
import torch.nn.functional as F
from tools.uvr5.bs_roformer.attend import Attend
from torch.utils.checkpoint import checkpoint
from typing import Tuple, Optional, Callable
from tools.uvr5.rotary_embedding_torch import RotaryEmbedding
from einops import rearrange, pack, unpack
from einops.layers.torch import Rearrange

def exists(val):
    return val is not None

def default(v, d):
    return v if exists(v) else d

def pack_one(t, pattern):
    return pack([t], pattern)

def unpack_one(t, ps, pattern):
    return unpack(t, ps, pattern)[0]

def l2norm(t):
    return F.normalize(t, dim=-1, p=2)

class RMSNorm(Module):

    def __init__(self, dim):
        super().__init__()
        self.scale = dim ** 0.5
        self.gamma = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        return F.normalize(x, dim=-1) * self.scale * self.gamma

class FeedForward(Module):

    def __init__(self, dim, mult=4, dropout=0.0):
        super().__init__()
        dim_inner = int(dim * mult)
        self.net = nn.Sequential(RMSNorm(dim), nn.Linear(dim, dim_inner), nn.GELU(), nn.Dropout(dropout), nn.Linear(dim_inner, dim), nn.Dropout(dropout))

    def forward(self, x):
        return self.net(x)

class Attention(Module):

    def __init__(self, dim, heads=8, dim_head=64, dropout=0.0, rotary_embed=None, flash=True):
        super().__init__()
        self.heads = heads
        self.scale = dim_head ** (-0.5)
        dim_inner = heads * dim_head
        self.rotary_embed = rotary_embed
        self.attend = Attend(flash=flash, dropout=dropout)
        self.norm = RMSNorm(dim)
        self.to_qkv = nn.Linear(dim, dim_inner * 3, bias=False)
        self.to_gates = nn.Linear(dim, heads)
        self.to_out = nn.Sequential(nn.Linear(dim_inner, dim, bias=False), nn.Dropout(dropout))

    def forward(self, x):
        x = self.norm(x)
        (q, k, v) = rearrange(self.to_qkv(x), 'b n (qkv h d) -> qkv b h n d', qkv=3, h=self.heads)
        if exists(self.rotary_embed):
            q = self.rotary_embed.rotate_queries_or_keys(q)
            k = self.rotary_embed.rotate_queries_or_keys(k)
        out = self.attend(q, k, v)
        gates = self.to_gates(x)
        out = out * rearrange(gates, 'b n h -> b h n 1').sigmoid()
        out = rearrange(out, 'b h n d -> b n (h d)')
        return self.to_out(out)

class LinearAttention(Module):
    """
    this flavor of linear attention proposed in https://arxiv.org/abs/2106.09681 by El-Nouby et al.
    """

    def __init__(self, *, dim, dim_head=32, heads=8, scale=8, flash=False, dropout=0.0):
        super().__init__()
        dim_inner = dim_head * heads
        self.norm = RMSNorm(dim)
        self.to_qkv = nn.Sequential(nn.Linear(dim, dim_inner * 3, bias=False), Rearrange('b n (qkv h d) -> qkv b h d n', qkv=3, h=heads))
        self.temperature = nn.Parameter(torch.ones(heads, 1, 1))
        self.attend = Attend(scale=scale, dropout=dropout, flash=flash)
        self.to_out = nn.Sequential(Rearrange('b h d n -> b n (h d)'), nn.Linear(dim_inner, dim, bias=False))

    def forward(self, x):
        x = self.norm(x)
        (q, k, v) = self.to_qkv(x)
        (q, k) = map(l2norm, (q, k))
        q = q * self.temperature.exp()
        out = self.attend(q, k, v)
        return self.to_out(out)

class Transformer(Module):

    def __init__(self, *, dim, depth, dim_head=64, heads=8, attn_dropout=0.0, ff_dropout=0.0, ff_mult=4, norm_output=True, rotary_embed=None, flash_attn=True, linear_attn=False):
        super().__init__()
        self.layers = ModuleList([])
        for _ in range(depth):
            if linear_attn:
                attn = LinearAttention(dim=dim, dim_head=dim_head, heads=heads, dropout=attn_dropout, flash=flash_attn)
            else:
                attn = Attention(dim=dim, dim_head=dim_head, heads=heads, dropout=attn_dropout, rotary_embed=rotary_embed, flash=flash_attn)
            self.layers.append(ModuleList([attn, FeedForward(dim=dim, mult=ff_mult, dropout=ff_dropout)]))
        self.norm = RMSNorm(dim) if norm_output else nn.Identity()

    def forward(self, x):
        for (attn, ff) in self.layers:
            x = attn(x) + x
            x = ff(x) + x
        return self.norm(x)

class BandSplit(Module):

    def __init__(self, dim, dim_inputs):
        super().__init__()
        self.dim_inputs = dim_inputs
        self.to_features = ModuleList([])
        for dim_in in dim_inputs:
            net = nn.Sequential(RMSNorm(dim_in), nn.Linear(dim_in, dim))
            self.to_features.append(net)

    def forward(self, x):
        x = x.split(self.dim_inputs, dim=-1)
        outs = []
        for (split_input, to_feature) in zip(x, self.to_features):
            split_output = to_feature(split_input)
            outs.append(split_output)
        return torch.stack(outs, dim=-2)

def MLP(dim_in, dim_out, dim_hidden=None, depth=1, activation=nn.Tanh):
    dim_hidden = default(dim_hidden, dim_in)
    net = []
    dims = (dim_in, *(dim_hidden,) * (depth - 1), dim_out)
    for (ind, (layer_dim_in, layer_dim_out)) in enumerate(zip(dims[:-1], dims[1:])):
        is_last = ind == len(dims) - 2
        net.append(nn.Linear(layer_dim_in, layer_dim_out))
        if is_last:
            continue
        net.append(activation())
    return nn.Sequential(*net)

class MaskEstimator(Module):

    def __init__(self, dim, dim_inputs, depth, mlp_expansion_factor=4):
        super().__init__()
        self.dim_inputs = dim_inputs
        self.to_freqs = ModuleList([])
        dim_hidden = dim * mlp_expansion_factor
        for dim_in in dim_inputs:
            net = []
            mlp = nn.Sequential(MLP(dim, dim_in * 2, dim_hidden=dim_hidden, depth=depth), nn.GLU(dim=-1))
            self.to_freqs.append(mlp)

    def forward(self, x):
        x = x.unbind(dim=-2)
        outs = []
        for (band_features, mlp) in zip(x, self.to_freqs):
            freq_out = mlp(band_features)
            outs.append(freq_out)
        return torch.cat(outs, dim=-1)
DEFAULT_FREQS_PER_BANDS = (2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 12, 12, 12, 12, 12, 12, 12, 12, 24, 24, 24, 24, 24, 24, 24, 24, 48, 48, 48, 48, 48, 48, 48, 48, 128, 129)

class BSRoformer(Module):

    def __init__(self, dim, *, depth, stereo=False, num_stems=1, time_transformer_depth=2, freq_transformer_depth=2, linear_transformer_depth=0, freqs_per_bands=DEFAULT_FREQS_PER_BANDS, dim_head=64, heads=8, attn_dropout=0.0, ff_dropout=0.0, flash_attn=True, dim_freqs_in=1025, stft_n_fft=2048, stft_hop_length=512, stft_win_length=2048, stft_normalized=False, stft_window_fn=None, mask_estimator_depth=2, multi_stft_resolution_loss_weight=1.0, multi_stft_resolutions_window_sizes=(4096, 2048, 1024, 512, 256), multi_stft_hop_size=147, multi_stft_normalized=False, multi_stft_window_fn=torch.hann_window, mlp_expansion_factor=4, use_torch_checkpoint=False, skip_connection=False):
        super().__init__()
        self.stereo = stereo
        self.audio_channels = 2 if stereo else 1
        self.num_stems = num_stems
        self.use_torch_checkpoint = use_torch_checkpoint
        self.skip_connection = skip_connection
        self.layers = ModuleList([])
        transformer_kwargs = dict(dim=dim, heads=heads, dim_head=dim_head, attn_dropout=attn_dropout, ff_dropout=ff_dropout, flash_attn=flash_attn, norm_output=False)
        time_rotary_embed = RotaryEmbedding(dim=dim_head)
        freq_rotary_embed = RotaryEmbedding(dim=dim_head)
        for _ in range(depth):
            tran_modules = []
            if linear_transformer_depth > 0:
                tran_modules.append(Transformer(depth=linear_transformer_depth, linear_attn=True, **transformer_kwargs))
            tran_modules.append(Transformer(depth=time_transformer_depth, rotary_embed=time_rotary_embed, **transformer_kwargs))
            tran_modules.append(Transformer(depth=freq_transformer_depth, rotary_embed=freq_rotary_embed, **transformer_kwargs))
            self.layers.append(nn.ModuleList(tran_modules))
        self.final_norm = RMSNorm(dim)
        self.stft_kwargs = dict(n_fft=stft_n_fft, hop_length=stft_hop_length, win_length=stft_win_length, normalized=stft_normalized)
        self.stft_window_fn = partial(default(stft_window_fn, torch.hann_window), stft_win_length)
        self._stft_windows = {}
        freqs = torch.stft(torch.randn(1, 4096), **self.stft_kwargs, window=torch.ones(stft_win_length), return_complex=True).shape[1]
        assert len(freqs_per_bands) > 1
        assert sum(freqs_per_bands) == freqs, f'the number of freqs in the bands must equal {freqs} based on the STFT settings, but got {sum(freqs_per_bands)}'
        freqs_per_bands_with_complex = tuple((2 * f * self.audio_channels for f in freqs_per_bands))
        self.band_split = BandSplit(dim=dim, dim_inputs=freqs_per_bands_with_complex)
        self.mask_estimators = nn.ModuleList([])
        for _ in range(num_stems):
            mask_estimator = MaskEstimator(dim=dim, dim_inputs=freqs_per_bands_with_complex, depth=mask_estimator_depth, mlp_expansion_factor=mlp_expansion_factor)
            self.mask_estimators.append(mask_estimator)
        self.multi_stft_resolution_loss_weight = multi_stft_resolution_loss_weight
        self.multi_stft_resolutions_window_sizes = multi_stft_resolutions_window_sizes
        self.multi_stft_n_fft = stft_n_fft
        self.multi_stft_window_fn = multi_stft_window_fn
        self.multi_stft_kwargs = dict(hop_length=multi_stft_hop_size, normalized=multi_stft_normalized)

    def _get_stft_window(self, device):
        key = str(device)
        window = self._stft_windows.get(key)
        if window is None:
            window = self.stft_window_fn(device=device, dtype=torch.float32)
            self._stft_windows[key] = window
        return window

    def forward(self, raw_audio, target=None, return_loss_breakdown=False):
        """
        einops

        b - batch
        f - freq
        t - time
        s - audio channel (1 for mono, 2 for stereo)
        n - number of 'stems'
        c - complex (2)
        d - feature dimension
        """
        device = raw_audio.device
        x_is_dml = device.type == 'privateuseone'
        x_is_mps = True if device.type == 'mps' else False
        if raw_audio.ndim == 2:
            raw_audio = rearrange(raw_audio, 'b t -> b 1 t')
        channels = raw_audio.shape[1]
        assert not self.stereo and channels == 1 or (self.stereo and channels == 2), 'stereo needs to be set to True if passing in audio signal that is stereo (channel dimension of 2). also need to be False if mono (channel dimension of 1)'
        (raw_audio, batch_audio_channel_packed_shape) = pack_one(raw_audio, '* t')
        if x_is_dml:
            # DirectML has no complex/STFT kernels.  Keep only the spectral
            # boundary on CPU and move its real representation to DirectML.
            stft_window = self._get_stft_window('cpu')
            stft_complex = torch.stft(
                raw_audio.cpu(),
                **self.stft_kwargs,
                window=stft_window,
                return_complex=True,
            )
            stft_repr_cpu = torch.view_as_real(stft_complex)
            stft_repr_cpu = unpack_one(
                stft_repr_cpu, batch_audio_channel_packed_shape, '* f t c'
            )
            stft_repr_cpu = rearrange(
                stft_repr_cpu, 'b s f t c -> b (f s) t c'
            )
            stft_repr = stft_repr_cpu.to(device)
        else:
            stft_window = self._get_stft_window(device)
            try:
                stft_repr = torch.stft(raw_audio, **self.stft_kwargs, window=stft_window, return_complex=True)
            except:
                stft_repr = torch.stft(raw_audio.cpu() if x_is_mps else raw_audio, **self.stft_kwargs, window=stft_window.cpu() if x_is_mps else stft_window, return_complex=True).to(device)
            stft_repr = torch.view_as_real(stft_repr)
            stft_repr = unpack_one(stft_repr, batch_audio_channel_packed_shape, '* f t c')
            stft_repr = rearrange(stft_repr, 'b s f t c -> b (f s) t c')
        x = rearrange(stft_repr, 'b f t c -> b t (f c)')
        if self.use_torch_checkpoint:
            x = checkpoint(self.band_split, x, use_reentrant=False)
        else:
            x = self.band_split(x)
        store = [None] * len(self.layers)
        for (i, transformer_block) in enumerate(self.layers):
            if len(transformer_block) == 3:
                (linear_transformer, time_transformer, freq_transformer) = transformer_block
                (x, ft_ps) = pack([x], 'b * d')
                if self.use_torch_checkpoint:
                    x = checkpoint(linear_transformer, x, use_reentrant=False)
                else:
                    x = linear_transformer(x)
                (x,) = unpack(x, ft_ps, 'b * d')
            else:
                (time_transformer, freq_transformer) = transformer_block
            if self.skip_connection:
                for j in range(i):
                    x = x + store[j]
            x = rearrange(x, 'b t f d -> b f t d')
            (x, ps) = pack([x], '* t d')
            if self.use_torch_checkpoint:
                x = checkpoint(time_transformer, x, use_reentrant=False)
            else:
                x = time_transformer(x)
            (x,) = unpack(x, ps, '* t d')
            x = rearrange(x, 'b f t d -> b t f d')
            (x, ps) = pack([x], '* f d')
            if self.use_torch_checkpoint:
                x = checkpoint(freq_transformer, x, use_reentrant=False)
            else:
                x = freq_transformer(x)
            (x,) = unpack(x, ps, '* f d')
            if self.skip_connection:
                store[i] = x
        x = self.final_norm(x)
        num_stems = len(self.mask_estimators)
        if self.use_torch_checkpoint:
            mask = torch.stack([checkpoint(fn, x, use_reentrant=False) for fn in self.mask_estimators], dim=1)
        else:
            mask = torch.stack([fn(x) for fn in self.mask_estimators], dim=1)
        mask = rearrange(mask, 'b n t (f c) -> b n f t c', c=2)
        if x_is_dml:
            # Complex masking and ISTFT stay on CPU; all learned real-valued
            # layers above remain on DirectML.
            stft_repr = rearrange(stft_repr_cpu, 'b f t c -> b 1 f t c')
            stft_repr = torch.view_as_complex(stft_repr.contiguous())
            mask = torch.view_as_complex(mask.float().cpu().contiguous())
            stft_repr = stft_repr * mask
            stft_repr = rearrange(
                stft_repr,
                'b n (f s) t -> (b n s) f t',
                s=self.audio_channels,
            )
            recon_audio = torch.istft(
                stft_repr,
                **self.stft_kwargs,
                window=stft_window,
                return_complex=False,
                length=raw_audio.shape[-1],
            ).to(device)
        else:
            stft_repr = rearrange(stft_repr, 'b f t c -> b 1 f t c')
            stft_repr = torch.view_as_complex(stft_repr)
            mask = torch.view_as_complex(mask)
            stft_repr = stft_repr * mask
            stft_repr = rearrange(stft_repr, 'b n (f s) t -> (b n s) f t', s=self.audio_channels)
            try:
                recon_audio = torch.istft(stft_repr, **self.stft_kwargs, window=stft_window, return_complex=False, length=raw_audio.shape[-1])
            except:
                recon_audio = torch.istft(stft_repr.cpu() if x_is_mps else stft_repr, **self.stft_kwargs, window=stft_window.cpu() if x_is_mps else stft_window, return_complex=False, length=raw_audio.shape[-1]).to(device)
        recon_audio = rearrange(recon_audio, '(b n s) t -> b n s t', s=self.audio_channels, n=num_stems)
        if num_stems == 1:
            recon_audio = rearrange(recon_audio, 'b 1 s t -> b s t')
        if not exists(target):
            return recon_audio
        if self.num_stems > 1:
            assert target.ndim == 4 and target.shape[1] == self.num_stems
        if target.ndim == 2:
            target = rearrange(target, '... t -> ... 1 t')
        target = target[..., :recon_audio.shape[-1]]
        loss_audio = recon_audio.cpu() if x_is_dml else recon_audio
        loss_target = target.cpu() if x_is_dml else target
        loss = F.l1_loss(loss_audio, loss_target)
        multi_stft_resolution_loss = 0.0
        for window_size in self.multi_stft_resolutions_window_sizes:
            spectral_device = 'cpu' if x_is_dml else device
            res_stft_kwargs = dict(n_fft=max(window_size, self.multi_stft_n_fft), win_length=window_size, return_complex=True, window=self.multi_stft_window_fn(window_size, device=spectral_device), **self.multi_stft_kwargs)
            recon_Y = torch.stft(rearrange(loss_audio, '... s t -> (... s) t'), **res_stft_kwargs)
            target_Y = torch.stft(rearrange(loss_target, '... s t -> (... s) t'), **res_stft_kwargs)
            multi_stft_resolution_loss = multi_stft_resolution_loss + F.l1_loss(recon_Y, target_Y)
        weighted_multi_resolution_loss = multi_stft_resolution_loss * self.multi_stft_resolution_loss_weight
        total_loss = loss + weighted_multi_resolution_loss
        if not return_loss_breakdown:
            return total_loss
        return (total_loss, (loss, multi_stft_resolution_loss))