"""Isolated speed benchmark: teacher vs student decoder, 10s real audio."""
import torch
import time
import math
from torch.nn.utils import weight_norm
import torch.nn as nn


class Snake1d(nn.Module):
    def __init__(self, hidden_dim, logscale=True):
        super().__init__()
        self.alpha = nn.Parameter(torch.zeros(1, hidden_dim, 1))
        self.beta = nn.Parameter(torch.zeros(1, hidden_dim, 1))
        self.logscale = logscale

    def forward(self, x):
        shape = x.shape
        a = self.alpha if not self.logscale else torch.exp(self.alpha)
        b = self.beta if not self.logscale else torch.exp(self.beta)
        x = x.reshape(shape[0], shape[1], -1)
        x = x + (b + 1e-9).reciprocal() * torch.sin(a * x).pow(2)
        return x.reshape(shape)


class FastResidualUnit(nn.Module):
    def __init__(self, dim, dilation=1):
        super().__init__()
        pad = ((7 - 1) * dilation) // 2
        self.snake1 = Snake1d(dim)
        self.conv1 = weight_norm(nn.Conv1d(dim, dim, 7, dilation=dilation, padding=pad))
        self.snake2 = Snake1d(dim)
        self.conv2 = weight_norm(nn.Conv1d(dim, dim, 1))

    def forward(self, x):
        h = self.conv1(self.snake1(x))
        h = self.conv2(self.snake2(h))
        pad = (x.shape[-1] - h.shape[-1]) // 2
        if pad > 0:
            x = x[..., pad:-pad]
        return x + h


class FastDecoderBlock(nn.Module):
    def __init__(self, in_dim, out_dim, stride=1):
        super().__init__()
        self.snake1 = Snake1d(in_dim)
        self.conv_t = weight_norm(nn.ConvTranspose1d(in_dim, out_dim, 2 * stride, stride=stride, padding=math.ceil(stride / 2)))
        self.res1 = FastResidualUnit(out_dim, 1)
        self.res2 = FastResidualUnit(out_dim, 3)

    def forward(self, x):
        x = self.snake1(x)
        x = self.conv_t(x)
        x = self.res1(x)
        x = self.res2(x)
        return x


class FastOobleckDecoder(nn.Module):
    def __init__(self, channels=128, input_channels=64, audio_channels=2, upsampling_ratios=None, channel_multiples=None):
        super().__init__()
        upsampling_ratios = upsampling_ratios or [10, 6, 4, 4, 2]
        channel_multiples = channel_multiples or [1, 2, 4, 8, 8]
        cm = [1] + channel_multiples
        self.conv1 = weight_norm(nn.Conv1d(input_channels, channels * cm[-1], 7, padding=3))
        blocks = []
        for i, s in enumerate(upsampling_ratios):
            blocks.append(FastDecoderBlock(channels * cm[len(upsampling_ratios) - i], channels * cm[len(upsampling_ratios) - i - 1], s))
        self.blocks = nn.ModuleList(blocks)
        self.final_snake = Snake1d(channels)
        self.conv2 = weight_norm(nn.Conv1d(channels, audio_channels, 7, padding=3, bias=False))

    def forward(self, z):
        x = self.conv1(z)
        for b in self.blocks:
            x = b(x)
        return self.conv2(self.final_snake(x))


def main():
    device = "cuda"
    torch.cuda.empty_cache()

    # Load student
    print("Loading student...")
    ckpt = torch.load("checkpoints/fast_decoder_v3/student_step595000.pt", map_location=device, weights_only=False)
    student = FastOobleckDecoder().to(device).eval()
    student.load_state_dict(ckpt["student_state_dict"])
    del ckpt
    torch.cuda.empty_cache()

    # Load teacher
    print("Loading teacher...")
    from diffusers import AutoencoderOobleck
    vae = AutoencoderOobleck.from_pretrained("ACE-Step/Ace-Step1.5", subfolder="vae").to(device, dtype=torch.float32).eval()
    teacher = vae.decoder

    # Encode real audio to get latent
    print("Encoding real audio...")
    import soundfile as sf
    data, sr = sf.read("tests/fixtures/techno.wav", dtype="float32")
    wav = torch.tensor(data, dtype=torch.float32).T[:, :480000].unsqueeze(0).to(device)  # 10s
    with torch.no_grad():
        z = vae.encode(wav).latent_dist.sample()

    # Free everything except teacher decoder and student
    del vae, wav
    torch.cuda.empty_cache()

    print(f"Latent shape: {list(z.shape)} (10s of audio)")
    print(f"GPU memory: {torch.cuda.memory_allocated()/1e6:.0f}MB allocated\n")

    # Warmup
    print("Warming up...")
    for _ in range(10):
        with torch.no_grad():
            teacher(z)
            student(z)
    torch.cuda.synchronize()

    # Benchmark teacher
    print("Benchmarking teacher (20 trials)...")
    tt = []
    for _ in range(20):
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        with torch.no_grad():
            teacher(z)
        torch.cuda.synchronize()
        tt.append(time.perf_counter() - t0)

    # Benchmark student
    print("Benchmarking student (20 trials)...")
    st = []
    for _ in range(20):
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        with torch.no_grad():
            student(z)
        torch.cuda.synchronize()
        st.append(time.perf_counter() - t0)

    print(f"\n{'='*50}")
    print(f"Teacher: {sum(tt)/len(tt)*1000:.1f}ms avg  (min {min(tt)*1000:.1f}, max {max(tt)*1000:.1f})")
    print(f"Student: {sum(st)/len(st)*1000:.1f}ms avg  (min {min(st)*1000:.1f}, max {max(st)*1000:.1f})")
    print(f"Speedup: {(sum(tt)/len(tt))/(sum(st)/len(st)):.2f}x")
    print(f"{'='*50}")


if __name__ == "__main__":
    main()