# Reference: # https://github.com/bytedance/Make-An-Audio-2

import torch
import torch.nn as nn
import torchaudio
from einops import rearrange
from librosa.filters import mel as librosa_mel_fn


def dynamic_range_compression_torch(x, C=1, clip_val=1e-5, norm_fn=torch.log10):
    return norm_fn(torch.clamp(x, min=clip_val) * C)


def spectral_normalize_torch(magnitudes, norm_fn):
    output = dynamic_range_compression_torch(magnitudes, norm_fn=norm_fn)
    return output


class STFTConverter(nn.Module):

    def __init__(
        self,
        *,
        sampling_rate: float = 16_000,
        n_fft: int = 1024,
        num_mels: int = 128,
        hop_size: int = 256,
        win_size: int = 1024,
        fmin: float = 0,
        fmax: float = 8_000,
        norm_fn=torch.log,
    ):
        super().__init__()
        self.sampling_rate = sampling_rate
        self.n_fft = n_fft
        self.num_mels = num_mels
        self.hop_size = hop_size
        self.win_size = win_size
        self.fmin = fmin
        self.fmax = fmax
        self.norm_fn = norm_fn

        mel = librosa_mel_fn(sr=self.sampling_rate,
                             n_fft=self.n_fft,
                             n_mels=self.num_mels,
                             fmin=self.fmin,
                             fmax=self.fmax)
        mel_basis = torch.from_numpy(mel).float()
        hann_window = torch.hann_window(self.win_size)

        self.register_buffer('mel_basis', mel_basis)
        self.register_buffer('hann_window', hann_window)

    @property
    def device(self):
        return self.hann_window.device

    def forward(self, waveform: torch.Tensor) -> torch.Tensor:
        # input: batch_size * length
        bs = waveform.shape[0]
        waveform = waveform.clamp(min=-1., max=1.)

        spec = torch.stft(waveform,
                          self.n_fft,
                          hop_length=self.hop_size,
                          win_length=self.win_size,
                          window=self.hann_window,
                          center=True,
                          pad_mode='reflect',
                          normalized=False,
                          onesided=True,
                          return_complex=True)

        spec = torch.view_as_real(spec)
        # print('After stft', spec.shape, spec.min(), spec.max(), spec.mean())

        power = (spec.pow(2).sum(-1))**(0.5)
        angle = torch.atan2(spec[..., 1], spec[..., 0])

        print('power 1', power.shape, power.min(), power.max(), power.mean())
        print('angle 1', angle.shape, angle.min(), angle.max(), angle.mean(), angle[:, :2, :2])

        # print('mel', self.mel_basis.shape, self.mel_basis.min(), self.mel_basis.max(),
        #       self.mel_basis.mean())

        # spec = self.mel_transform(spec)

        # power = torch.matmul(self.mel_basis, power)

        spec = rearrange(spec, 'b f t c -> (b c) f t')
        spec = self.mel_basis.unsqueeze(0) @ spec
        spec = rearrange(spec, '(b c) f t -> b f t c', b=bs)

        power = (spec.pow(2).sum(-1))**(0.5)
        angle = torch.atan2(spec[..., 1], spec[..., 0])

        print('power', power.shape, power.min(), power.max(), power.mean())
        print('angle', angle.shape, angle.min(), angle.max(), angle.mean(), angle[:, :2, :2])

        # print('After mel', spec.shape, spec.min(), spec.max(), spec.mean())

        # spec = spectral_normalize_torch(spec, self.norm_fn)

        # print('After norm', spec.shape, spec.min(), spec.max(), spec.mean())

        # compute magnitude
        # magnitude = torch.sqrt((spec**2).sum(-1))
        # normalize by magnitude
        # scaled_magnitude = torch.log10(magnitude.clamp(min=1e-5)) * 10
        # spec = spec / magnitude.unsqueeze(-1) * scaled_magnitude.unsqueeze(-1)

        # power = torch.log10(power.clamp(min=1e-5)) * 10
        power = torch.log10(power.clamp(min=1e-8))

        print('After scaling', power.shape, power.min(), power.max(), power.mean())

        # spec = torch.stack([power, angle], dim=-1)

        # spec = rearrange(spec, '(b c) f t -> b c f t', b=bs)
        # spec = rearrange(spec, 'b f t c -> b c f t', b=bs)

        # spec[:, :, 400:] = 0

        return power, angle
        # return spec[..., 0], spec[..., 1]

    def invert(self, spec: torch.Tensor, length: int) -> torch.Tensor:

        power, angle = spec

        bs = power.shape[0]

        # spec = rearrange(spec, 'b c f t -> (b c) f t')
        # print(spec.shape, self.mel_basis.shape)
        # spec = torch.linalg.lstsq(self.mel_basis.unsqueeze(0), spec).solution
        # spec = torch.linalg.pinv(self.mel_basis.unsqueeze(0)) @ spec

        # spec = self.invmel_transform(spec)

        # spec = rearrange(spec, 'b c f t -> b f t c', b=bs).contiguous()

        # spec[..., 0] = 10**(spec[..., 0] / 10)

        # power = spec[..., 0]
        power = 10**power

        # print('After unscaling', spec[..., 0].shape, spec[..., 0].min(), spec[..., 0].max(),
        #       spec[..., 0].mean())

        unit_vector = torch.stack([
            torch.cos(angle),
            torch.sin(angle),
        ], dim=-1)

        spec = power.unsqueeze(-1) * unit_vector

        # power = torch.linalg.lstsq(self.mel_basis.unsqueeze(0), power).solution
        spec = rearrange(spec, 'b f t c -> (b c) f t')
        spec = torch.linalg.pinv(self.mel_basis.unsqueeze(0)) @ spec
        # spec = torch.linalg.lstsq(self.mel_basis.unsqueeze(0), spec).solution
        spec = rearrange(spec, '(b c) f t -> b f t c', b=bs).contiguous()

        power = (spec.pow(2).sum(-1))**(0.5)
        angle = torch.atan2(spec[..., 1], spec[..., 0])

        print('power 2', power.shape, power.min(), power.max(), power.mean())
        print('angle 2', angle.shape, angle.min(), angle.max(), angle.mean(), angle[:, :2, :2])

        # spec = rearrange(spec, '(b c) f t -> b f t c', b=bs).contiguous()
        spec = torch.view_as_complex(spec)

        waveform = torch.istft(
            spec,
            self.n_fft,
            length=length,
            hop_length=self.hop_size,
            win_length=self.win_size,
            window=self.hann_window,
            center=True,
            normalized=False,
            onesided=True,
            return_complex=False,
        )

        return waveform


if __name__ == '__main__':

    converter = STFTConverter(sampling_rate=16000)

    signal = torchaudio.load('./output/ZZ6GRocWW38_000090.wav')[0]
    # resample signal at 44100 Hz
    # signal = torchaudio.transforms.Resample(16_000, 44_100)(signal)

    L = signal.shape[1]
    print('Input signal', signal.shape)
    spec = converter(signal)

    power, angle = spec

    # print(power.shape, angle.shape)
    # print(power, power.min(), power.max(), power.mean())
    # power = power.clamp(-1, 1)
    # angle = angle.clamp(-1, 1)

    import matplotlib.pyplot as plt

    # Visualize power
    plt.figure()
    plt.imshow(power[0].detach().numpy(), aspect='auto', origin='lower')
    plt.colorbar()
    plt.title('Power')
    plt.xlabel('Time')
    plt.ylabel('Frequency')
    plt.savefig('./output/power.png')

    # Visualize angle
    plt.figure()
    plt.imshow(angle[0].detach().numpy(), aspect='auto', origin='lower')
    plt.colorbar()
    plt.title('Angle')
    plt.xlabel('Time')
    plt.ylabel('Frequency')
    plt.savefig('./output/angle.png')

    # print('Final spec', spec.shape)

    signal_recon = converter.invert(spec, length=L)
    print('Output signal', signal_recon.shape, signal_recon.min(), signal_recon.max(),
          signal_recon.mean())

    print('MSE', torch.nn.functional.mse_loss(signal, signal_recon))
    torchaudio.save('./output/ZZ6GRocWW38_000090_recon.wav', signal_recon, 16000)