import torch, sys, time
sys.path.insert(0, 'src')
from timm.models.convnext import ConvNeXtBlock
from gigatorch.ops import ConvBlockND, SAPETransformerBlockND

device = 'cuda:0'
ps = 16
decoder_dim = 768
n_ch = 6
input_channels = 3
H, W = 480, 480

class MergedAutoencoder(torch.nn.Module):
    def __init__(self, decoder_blocks):
        super().__init__()
        self.encoder = torch.nn.Sequential(
            torch.nn.Conv2d(input_channels, n_ch, kernel_size=ps, stride=ps),
        )
        self.decoder = torch.nn.Sequential(
            torch.nn.Conv2d(n_ch, decoder_dim, kernel_size=5, padding=2, padding_mode='reflect'),
            *decoder_blocks,
            torch.nn.Conv2d(decoder_dim, input_channels * ps * ps, kernel_size=1),
            torch.nn.ConvTranspose2d(input_channels * ps * ps, input_channels, kernel_size=ps, stride=ps),
            torch.nn.Hardtanh(),
        )

    def forward(self, x):
        z = self.encoder(x)
        if self.training:
            z = z + torch.rand_like(z) - 0.5
        return self.decoder(z)


def measure(model, x, n_warmup=200, n_measure=1000):
    model.train()
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

    for _ in range(n_warmup):
        optimizer.zero_grad()
        model(x).sum().backward()
        optimizer.step()

    torch.cuda.synchronize()
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)

    start.record()
    for _ in range(n_measure):
        optimizer.zero_grad()
        model(x).sum().backward()
        optimizer.step()
    end.record()
    torch.cuda.synchronize()

    ms = start.elapsed_time(end)
    return n_measure / (ms / 1000)


x = torch.randn(1, input_channels, H, W, device=device)

C_timm = lambda: ConvNeXtBlock(decoder_dim, decoder_dim, kernel_size=5, mlp_ratio=4, conv_mlp=True)
C = lambda: ConvBlockND(2, decoder_dim)
T = lambda: SAPETransformerBlockND(2, decoder_dim, num_heads=8)

configs = {
    'ConvNeXt baseline (CCCCCC)': [C_timm() for _ in range(6)],
    'LayerNorm (CCCCCC)':         [C() for _ in range(6)],
    'LayerNorm (CCCT)':           [C(), C(), C(), T()],
    'LayerNorm (CCCCT)':          [C(), C(), C(), C(), T()],
    'LayerNorm (CCTT)':           [C(), C(), T(), T()],
    'LayerNorm (CCTCT)':          [C(), C(), T(), C(), T()],
    'LayerNorm (CCCTCT)':         [C(), C(), C(), T(), C(), T()],
    'LayerNorm (CCCCTT)':         [C(), C(), C(), C(), T(), T()],
    'LayerNorm (CCTCTT)':         [C(), C(), T(), C(), T(), T()],
}

print(f'Device: {torch.cuda.get_device_name(0)}')
print(f'Input: {input_channels}x{H}x{W}, ps={ps}, decoder_dim={decoder_dim}, n_ch={n_ch}')
print(f'Latent: {n_ch}x{H//ps}x{W//ps} = {n_ch}x{H//ps}x{W//ps}')
print()

for name, blocks in configs.items():
    model = MergedAutoencoder(blocks).to(device)
    n_params = sum(p.numel() for p in model.parameters()) / 1e6
    throughput = measure(model, x)
    print(f'{name}: {throughput:.2f} it/s  |  {n_params:.2f}M params')
    del model
    torch.cuda.empty_cache()
