"""100 ms / 640-bit waveform channel, independent G, D and analyzer."""
import torch
from torch import nn
from torch.nn import functional as F
from torch.nn.utils.parametrizations import spectral_norm

class Residual(nn.Module):
    def __init__(self, channels, temporal=False, dilation=1):
        super().__init__()
        k = 3 if temporal else 7
        self.first = nn.Conv1d(channels, channels, k, padding=dilation*(k//2), dilation=dilation)
        self.second = nn.Conv1d(channels, channels, 1 if temporal else 7, padding=0 if temporal else 3)
    def forward(self,x):
        return F.leaky_relu(x+self.second(F.leaky_relu(self.first(x), .2)), .2)

class Generator(nn.Module):
    def __init__(self):
        super().__init__()
        self.net=nn.Sequential(nn.Conv1d(64,128,1),nn.LeakyReLU(.2),
            *[Residual(128,True,d) for d in (1,2,4)],
            nn.ConvTranspose1d(128,128,20,10,5),nn.LeakyReLU(.2),Residual(128),
            nn.ConvTranspose1d(128,64,8,4,2),nn.LeakyReLU(.2),Residual(64),
            nn.ConvTranspose1d(64,32,8,4,2),nn.LeakyReLU(.2),Residual(32),
            nn.Conv1d(32,1,7,padding=3),nn.Tanh())
    def forward(self,bits):
        return self.net((2*bits-1).transpose(1,2)).squeeze(1)

class Analyzer(nn.Module):
    def __init__(self):
        super().__init__()
        self.net=nn.Sequential(nn.Conv1d(1,32,7,padding=3),nn.LeakyReLU(.2),Residual(32),
            nn.Conv1d(32,64,8,4,2),nn.LeakyReLU(.2),Residual(64),
            nn.Conv1d(64,128,8,4,2),nn.LeakyReLU(.2),Residual(128),
            nn.Conv1d(128,128,20,10,5),nn.LeakyReLU(.2),
            *[Residual(128,True,d) for d in (1,2,4)],nn.Conv1d(128,64,1))
    def forward(self,audio):
        return self.net(audio.unsqueeze(1)).transpose(1,2)

class ScaleDiscriminator(nn.Module):
    def __init__(self):
        super().__init__()
        cs=(1,16,64,128,256,256,1)
        self.layers=nn.ModuleList([spectral_norm(nn.Conv1d(cs[i],cs[i+1],k,stride=s,padding=k//2,groups=g))
            for i,(k,s,g) in enumerate(zip((15,41,41,41,5,3),(1,4,4,4,1,1),(1,4,8,16,1,1)))])
    def forward(self,x):
        fs=[]
        for i,layer in enumerate(self.layers):
            x=layer(x)
            if i<len(self.layers)-1: x=F.leaky_relu(x,.2); fs.append(x)
        return x,fs

class PeriodDiscriminator(nn.Module):
    def __init__(self,period):
        super().__init__(); self.period=period
        cs=(1,16,64,128,256,1)
        self.layers=nn.ModuleList([spectral_norm(nn.Conv2d(cs[i],cs[i+1],(5 if i<4 else 3,1),
            stride=(3 if i<4 else 1,1),padding=(2 if i<4 else 1,0))) for i in range(5)])
    def forward(self,x):
        pad=(-x.shape[-1])%self.period
        if pad: x=F.pad(x,(0,pad),mode='reflect')
        x=x.reshape(x.shape[0],1,-1,self.period); fs=[]
        for i,layer in enumerate(self.layers):
            x=layer(x)
            if i<4: x=F.leaky_relu(x,.2); fs.append(x)
        return x,fs

class Discriminators(nn.Module):
    def __init__(self):
        super().__init__()
        self.scales=nn.ModuleList([ScaleDiscriminator() for _ in range(3)])
        self.periods=nn.ModuleList([PeriodDiscriminator(p) for p in (2,3,5,7,11)])
    def forward(self,wave):
        x=wave.unsqueeze(1); results=[]
        for i,d in enumerate(self.scales):
            if i: x=F.avg_pool1d(x,4,2,padding=1)
            results.append(d(x))
        results.extend(d(wave.unsqueeze(1)) for d in self.periods)
        return results

def pcm16(wave,ste=False):
    clamped=wave.clamp(-1,32767/32768)
    q=torch.floor(clamped*32768)/32768
    return clamped+(q-clamped).detach() if ste else q

def feature_moments(real,fake):
    # Unpaired distribution moments over batch/spatial dimensions; never paired L1.
    terms=[]
    for (_,rfs),(_,ffs) in zip(real,fake):
        layer_terms=[]
        for r,f in zip(rfs,ffs):
            dims=(0,)+tuple(range(2,r.ndim))
            layer_terms.append(F.l1_loss(f.mean(dims),r.detach().mean(dims)))
        terms.append(torch.stack(layer_terms).mean())
    return torch.stack(terms).mean()

def requires_grad(module,enabled):
    for p in module.parameters(): p.requires_grad_(enabled)