import torch
from safetensors.torch import load_file
from torch import nn

WIDTH, LATENT, AUDIO_CHANNELS = 64, 64, 2
MULTS = (1, 1, 2, 4, 8, 16, 32)
STRIDES = (2, 2, 4, 4, 5, 6)
NORMED = (('weight_g', 'weight_v'), ('parametrizations.weight.original0', 'parametrizations.weight.original1'))


class Snake(nn.Module):
    def __init__(self, ch):
        super().__init__()
        self.alpha = nn.Parameter(torch.zeros(ch)); self.beta = nn.Parameter(torch.zeros(ch))

    def forward(self, x):
        freq, amp = self.alpha.exp().view(1, -1, 1), self.beta.exp().view(1, -1, 1)
        return x + torch.sin(freq * x).square() / (amp + 1e-9)


class Stack(nn.Module):
    def __init__(self, *layers):
        super().__init__()
        self.layers = nn.Sequential(*layers)

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


class Residual(Stack):
    def __init__(self, ch, dilation):
        super().__init__(Snake(ch), nn.Conv1d(ch, ch, 7, padding=3 * dilation, dilation=dilation), Snake(ch), nn.Conv1d(ch, ch, 1))

    def forward(self, x):
        return x + self.layers(x)


def down(cin, cout, s):
    return Stack(Residual(cin, 1), Residual(cin, 3), Residual(cin, 9), Snake(cin), nn.Conv1d(cin, cout, 2 * s, s, (s + 1) // 2))


def up(cin, cout, s):
    return Stack(Snake(cin), nn.ConvTranspose1d(cin, cout, 2 * s, s, (s + 1) // 2, output_padding=s % 2),
                 Residual(cout, 1), Residual(cout, 3), Residual(cout, 9))


class VAE(nn.Module):
    def __init__(self):
        super().__init__()
        ch = [WIDTH * m for m in MULTS]
        n = len(STRIDES)
        self.encoder = Stack(nn.Conv1d(AUDIO_CHANNELS, ch[0], 7, padding=3), *[down(ch[i], ch[i + 1], STRIDES[i]) for i in range(n)],
                             Snake(ch[-1]), nn.Conv1d(ch[-1], 2 * LATENT, 3, padding=1))
        self.decoder = Stack(nn.Conv1d(LATENT, ch[-1], 7, padding=3), *[up(ch[i + 1], ch[i], STRIDES[i]) for i in reversed(range(n))],
                             Snake(ch[0]), nn.Conv1d(ch[0], AUDIO_CHANNELS, 7, padding=3, bias=False))

    @torch.no_grad()
    def encode(self, audio):
        return self.encoder(audio)[:, :LATENT]

    @torch.no_grad()
    def decode(self, latents):
        return self.decoder(latents)


def plain_weights(state):
    state = {k.removeprefix('vae.'): v.float() for k, v in state.items()}
    out = {}
    for k, v in state.items():
        mag = next((m for m, d in NORMED if k.endswith(m)), None)
        if mag:
            stem = k[:-len(mag)]; direction = state[stem + dict(NORMED)[mag]]
            out[stem + 'weight'] = direction * (v / direction.flatten(1).norm(dim=1).view(-1, *[1] * (direction.dim() - 1)))
        elif not any(k.endswith(d) for _, d in NORMED):
            out[k] = v
    return out


def load_vae(path, device='cuda'):
    vae = VAE(); vae.load_state_dict(plain_weights(load_file(str(path))), strict=True)
    return vae.to(device, torch.float32).eval()