import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn import Conv1d, ConvTranspose1d, AvgPool1d, Conv2d
from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm
from torch.utils.data import DataLoader, Dataset
from torch.optim import AdamW
from tqdm import tqdm
import numpy as np
import itertools
import glob

# Constants and Utilities
LRELU_SLOPE = 0.1

def init_weights(m, mean=0.0, std=0.01):
    classname = m.__class__.__name__
    if classname.find("Conv") != -1:
        m.weight.data.normal_(mean, std)

def get_padding(kernel_size, dilation=1):
    return int((kernel_size*dilation - dilation)/2)

# Model Components
class ResBlock1(nn.Module):
    def __init__(self, h, channels, kernel_size=3, dilation=(1, 3, 5)):
        super().__init__()
        self.h = h
        self.convs1 = nn.ModuleList([
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
                       padding=get_padding(kernel_size, dilation[0]))),
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
                       padding=get_padding(kernel_size, dilation[1]))),
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2],
                       padding=get_padding(kernel_size, dilation[2])))
        ])
        self.convs1.apply(init_weights)

        self.convs2 = nn.ModuleList([
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
                       padding=get_padding(kernel_size, 1))),
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
                       padding=get_padding(kernel_size, 1))),
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
                       padding=get_padding(kernel_size, 1)))
        ])
        self.convs2.apply(init_weights)

    def forward(self, x):
        for c1, c2 in zip(self.convs1, self.convs2):
            xt = F.leaky_relu(x, LRELU_SLOPE)
            xt = c1(xt)
            xt = F.leaky_relu(xt, LRELU_SLOPE)
            xt = c2(xt)
            x = xt + x
        return x

    def remove_weight_norm(self):
        for l in self.convs1:
            remove_weight_norm(l)
        for l in self.convs2:
            remove_weight_norm(l)

class EmbeddingGenerator(nn.Module):
    def __init__(self, h):
        super().__init__()
        self.h = h
        self.num_kernels = len(h.resblock_kernel_sizes)
        self.num_upsamples = len(h.upsample_rates)
        
        self.conv_pre = weight_norm(Conv1d(h.embedding_dim, h.upsample_initial_channel, 7, 1, padding=3))
        
        self.ups = nn.ModuleList()
        for i, (u, k) in enumerate(zip(h.upsample_rates, h.upsample_kernel_sizes)):
            self.ups.append(weight_norm(
                ConvTranspose1d(h.upsample_initial_channel//(2**i), 
                h.upsample_initial_channel//(2**(i+1)),
                k, u, padding=(k-u)//2)))

        self.resblocks = nn.ModuleList()
        for i in range(len(self.ups)):
            ch = h.upsample_initial_channel//(2**(i+1))
            for j, (k, d) in enumerate(zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes)):
                self.resblocks.append(ResBlock1(h, ch, k, d))

        self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, padding=3))
        self.ups.apply(init_weights)
        self.conv_post.apply(init_weights)

    def forward(self, x):
        x = x.transpose(1, 2)  # (B, 1024, T)
        x = self.conv_pre(x)
        for i in range(self.num_upsamples):
            x = F.leaky_relu(x, LRELU_SLOPE)
            x = self.ups[i](x)
            xs = None
            for j in range(self.num_kernels):
                if xs is None:
                    xs = self.resblocks[i*self.num_kernels+j](x)
                else:
                    xs += self.resblocks[i*self.num_kernels+j](x)
            x = xs / self.num_kernels
        x = F.leaky_relu(x)
        x = self.conv_post(x)
        return torch.tanh(x)

    def remove_weight_norm(self):
        for l in self.ups:
            remove_weight_norm(l)
        for l in self.resblocks:
            l.remove_weight_norm()
        remove_weight_norm(self.conv_pre)
        remove_weight_norm(self.conv_post)

class DiscriminatorP(nn.Module):
    def __init__(self, period, kernel_size=5, stride=3, use_spectral_norm=False):
        super().__init__()
        self.period = period
        norm_f = spectral_norm if use_spectral_norm else weight_norm
        self.convs = nn.ModuleList([
            norm_f(Conv2d(1, 32, (kernel_size, 1), (stride, 1), padding=(get_padding(5, 1), 0))),
            norm_f(Conv2d(32, 128, (kernel_size, 1), (stride, 1), padding=(get_padding(5, 1), 0))),
            norm_f(Conv2d(128, 512, (kernel_size, 1), (stride, 1), padding=(get_padding(5, 1), 0))),
            norm_f(Conv2d(512, 1024, (kernel_size, 1), (stride, 1), padding=(get_padding(5, 1), 0))),
            norm_f(Conv2d(1024, 1024, (kernel_size, 1), 1, padding=(2, 0))),
        ])
        self.conv_post = norm_f(Conv2d(1024, 1, (3, 1), 1, padding=(1, 0)))

    def forward(self, x):
        fmap = []
        b, c, t = x.shape
        if t % self.period != 0:
            n_pad = self.period - (t % self.period)
            x = F.pad(x, (0, n_pad), "reflect")
        x = x.view(b, c, t // self.period, self.period)
        for conv in self.convs:
            x = conv(x)
            x = F.leaky_relu(x, LRELU_SLOPE)
            fmap.append(x)
        x = self.conv_post(x)
        fmap.append(x)
        return torch.flatten(x, 1, -1), fmap

class MultiPeriodDiscriminator(nn.Module):
    def __init__(self):
        super().__init__()
        self.discriminators = nn.ModuleList([
            DiscriminatorP(2),
            DiscriminatorP(3),
            DiscriminatorP(5),
            DiscriminatorP(7),
            DiscriminatorP(11),
        ])

    def forward(self, y, y_hat):
        y_d_rs, y_d_gs, fmap_rs, fmap_gs = [], [], [], []
        for d in self.discriminators:
            y_d_r, fmap_r = d(y)
            y_d_g, fmap_g = d(y_hat)
            y_d_rs.append(y_d_r)
            y_d_gs.append(y_d_g)
            fmap_rs.append(fmap_r)
            fmap_gs.append(fmap_g)
        return y_d_rs, y_d_gs, fmap_rs, fmap_gs

class DiscriminatorS(nn.Module):
    def __init__(self, use_spectral_norm=False):
        super().__init__()
        norm_f = spectral_norm if use_spectral_norm else weight_norm
        self.convs = nn.ModuleList([
            norm_f(Conv1d(1, 128, 15, 1, padding=7)),
            norm_f(Conv1d(128, 128, 41, 2, groups=4, padding=20)),
            norm_f(Conv1d(128, 256, 41, 2, groups=16, padding=20)),
            norm_f(Conv1d(256, 512, 41, 4, groups=16, padding=20)),
            norm_f(Conv1d(512, 1024, 41, 4, groups=16, padding=20)),
            norm_f(Conv1d(1024, 1024, 41, 1, groups=16, padding=20)),
            norm_f(Conv1d(1024, 1024, 5, 1, padding=2)),
        ])
        self.conv_post = norm_f(Conv1d(1024, 1, 3, 1, padding=1))

    def forward(self, x):
        fmap = []
        for conv in self.convs:
            x = conv(x)
            x = F.leaky_relu(x, LRELU_SLOPE)
            fmap.append(x)
        x = self.conv_post(x)
        fmap.append(x)
        return torch.flatten(x, 1, -1), fmap

class MultiScaleDiscriminator(nn.Module):
    def __init__(self):
        super().__init__()
        self.discriminators = nn.ModuleList([
            DiscriminatorS(use_spectral_norm=True),
            DiscriminatorS(),
            DiscriminatorS(),
        ])
        self.meanpools = nn.ModuleList([
            AvgPool1d(4, 2, padding=2),
            AvgPool1d(4, 2, padding=2)
        ])

    def forward(self, y, y_hat):
        y_d_rs, y_d_gs, fmap_rs, fmap_gs = [], [], [], []
        for i, d in enumerate(self.discriminators):
            if i != 0:
                y = self.meanpools[i-1](y)
                y_hat = self.meanpools[i-1](y_hat)
            y_d_r, fmap_r = d(y)
            y_d_g, fmap_g = d(y_hat)
            y_d_rs.append(y_d_r)
            y_d_gs.append(y_d_g)
            fmap_rs.append(fmap_r)
            fmap_gs.append(fmap_g)
        return y_d_rs, y_d_gs, fmap_rs, fmap_gs

# Loss Functions
def feature_loss(fmap_r, fmap_g):
    loss = 0
    for dr, dg in zip(fmap_r, fmap_g):
        for rl, gl in zip(dr, dg):
            loss += torch.mean(torch.abs(rl - gl))
    return loss * 2

def discriminator_loss(disc_real_outputs, disc_generated_outputs):
    loss = 0
    for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
        loss += (torch.mean((1-dr)**2) + torch.mean(dg**2))
    return loss

def generator_loss(disc_outputs):
    loss = 0
    for dg in disc_outputs:
        loss += torch.mean((1-dg)**2)
    return loss

# Training Infrastructure
class Hparams:
    def __init__(self):
        self.embedding_dim = 1024
        self.upsample_rates = [4, 4, 4, 2]
        self.upsample_kernel_sizes = [8, 8, 4, 4]
        self.upsample_initial_channel = 512
        self.resblock_kernel_sizes = [3, 7, 11]
        self.resblock_dilation_sizes = [[1, 3, 5], [1, 3, 5], [1, 3, 5]]

class EmbeddingAudioDataset(Dataset):
    def __init__(self, embedding_files, audio_files):
        self.embedding_files = embedding_files
        self.audio_files = audio_files

    def __len__(self):
        return len(self.embedding_files)

    def __getitem__(self, idx):
        embedding = np.load(self.embedding_files[idx])
        audio = np.load(self.audio_files[idx])
        audio = audio / np.max(np.abs(audio))  # Normalize
        return torch.from_numpy(embedding).float(), torch.from_numpy(audio).float()

# Main Execution
if __name__ == "__main__":
    h = Hparams()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # Initialize models
    generator = EmbeddingGenerator(h).to(device)
    mpd = MultiPeriodDiscriminator().to(device)
    msd = MultiScaleDiscriminator().to(device)
    
    # Optimizers
    optim_g = AdamW(generator.parameters(), lr=0.0002, betas=(0.8, 0.99))
    optim_d = AdamW(itertools.chain(mpd.parameters(), msd.parameters()), 
                  lr=0.0002, betas=(0.8, 0.99))
    
    # Dataset
    embedding_files = sorted(glob.glob('/home/vikrant/Conversational-AI-Model/embedding_vocoder/embeddings/*.npy'))
    audio_files = sorted(glob.glob('/home/vikrant/Conversational-AI-Model/embedding_vocoder/non_empty_wavs/*.wav'))
    assert len(embedding_files) == len(audio_files), "Mismatched files"
    
    loader = DataLoader(
        EmbeddingAudioDataset(embedding_files, audio_files),
        batch_size=1, shuffle=True, num_workers=4, pin_memory=True
    )
    
    # Training loop
    num_epochs = 1
    for epoch in range(num_epochs):
        generator.train()
        mpd.train()
        msd.train()
        
        for batch_idx, (embeddings, audio) in enumerate(tqdm(loader)):
            embeddings, audio = embeddings.to(device), audio.to(device)
            
            if embeddings.shape[1] != h.embedding_dim:
                embeddings = embeddings.transpose(1, 2)
            
            # Generate audio
            fake_audio = generator(embeddings)
            
            # Pad/trim to match lengths
            if fake_audio.size(2) != audio.size(1):
                if fake_audio.size(2) > audio.size(1):
                    fake_audio = fake_audio[:, :, :audio.size(1)]
                else:
                    fake_audio = F.pad(fake_audio, (0, audio.size(1) - fake_audio.size(2)))
            
            # Discriminator training
            optim_d.zero_grad()
            _, disc_real_mpd, _, disc_real_msd = mpd(audio, fake_audio.detach())
            _, disc_fake_mpd, _, disc_fake_msd = msd(audio, fake_audio.detach())
            loss_disc = discriminator_loss(disc_real_mpd + disc_real_msd, 
                                          disc_fake_mpd + disc_fake_msd)
            loss_disc.backward()
            optim_d.step()
            
            # Generator training
            optim_g.zero_grad()
            disc_out_mpd, disc_out_msd, fmap_mpd, fmap_msd = mpd(audio, fake_audio)
            loss_gen = generator_loss(disc_out_mpd + disc_out_msd)
            loss_feat = feature_loss(fmap_mpd + fmap_msd)
            (loss_gen + loss_feat).backward()
            optim_g.step()
            
            if batch_idx % 100 == 0:
                print(f"Epoch {epoch} Batch {batch_idx} | G Loss: {loss_gen.item():.3f} | D Loss: {loss_disc.item():.3f}")
        
        # Save checkpoint
        torch.save({
            'generator': generator.state_dict(),
            'mpd': mpd.state_dict(),
            'msd': msd.state_dict(),
            'optim_g': optim_g.state_dict(),
            'optim_d': optim_d.state_dict(),
            'epoch': epoch,
        }, f"checkpoint_epoch_{epoch}.pt")
    
    # Finalize
    generator.eval()
    generator.remove_weight_norm()
    print("Training completed")