# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
"""Command-line for audio compression."""
import argparse
import os
import sys
import typing as tp
from collections import OrderedDict
from pathlib import Path

import librosa
import soundfile as sf
import torch
from academicodec.models.encodec.net3 import SoundStream


def save_audio(wav: torch.Tensor,
               path: tp.Union[Path, str],
               sample_rate: int,
               rescale: bool=False):
    limit = 0.99
    mx = wav.abs().max()
    if rescale:
        wav = wav * min(limit / mx, 1)
    else:
        wav = wav.clamp(-limit, limit)
    wav = wav.squeeze().cpu().numpy()
    sf.write(path, wav, sample_rate)


def get_parser():
    parser = argparse.ArgumentParser(
        'encodec',
        description='High fidelity neural audio codec. '
        'If input is a .ecdc, decompresses it. '
        'If input is .wav, compresses it. If output is also wav, '
        'do a compression/decompression cycle.')
    parser.add_argument(
        '--input',
        type=Path,
        help='Input file, whatever is supported by torchaudio on your system.')
    parser.add_argument(
        '--output',
        type=Path,
        nargs='?',
        help='Output file, otherwise inferred from input file.')
    parser.add_argument(
        '--resume_path', type=str, default='resume_path', help='resume_path')
    parser.add_argument(
        '--sr', type=int, default=16000, help='sample rate of model')
    parser.add_argument(
        '-r',
        '--rescale',
        action='store_true',
        help='Automatically rescale the output to avoid clipping.')
    parser.add_argument(
        '--ratios',
        type=int,
        nargs='+',
        # probs(ratios) = hop_size
        default=[8, 5, 4, 2],
        help='ratios of SoundStream, shoud be set for different hop_size (32d, 320, 240d, ...)'
    )
    parser.add_argument(
        '--target_bandwidths',
        type=float,
        nargs='+',
        # default for 16k_320d
        default=[1, 1.5, 2, 4, 6, 12],
        help='target_bandwidths of net3.py')
    parser.add_argument(
        '--target_bw',
        type=float,
        # default for 16k_320d
        default=12,
        help='target_bw of net3.py')

    return parser


def fatal(*args):
    print(*args, file=sys.stderr)
    sys.exit(1)


# 这只是打印了但是没有真的 clip
def check_clipping(wav, rescale):
    if rescale:
        return
    mx = wav.abs().max()
    limit = 0.99
    if mx > limit:
        print(
            f"Clipping!! max scale {mx}, limit is {limit}. "
            "To avoid clipping, use the `-r` option to rescale the output.",
            file=sys.stderr)


def test_one(args, wav_root, store_root, rescale, soundstream):
    # torchaudio.load 的采样率为原始音频的采样率，不会自动下采样
    # wav, sr = torchaudio.load(wav_root)
    # # 取单声道, output shape [1, T]
    # wav = wav[0].unsqueeze(0)
    # # 重采样为模型的采样率
    # wav = torchaudio.transforms.Resample(orig_freq=sr, new_freq=args.sr)(wav)

    # load wav with librosa
    wav, sr = librosa.load(wav_root, sr=args.sr)
    wav = torch.tensor(wav).unsqueeze(0)

    # add batch axis
    wav = wav.unsqueeze(1).cuda()

    # compressing
    compressed = soundstream.encode(wav, target_bw=args.target_bw)
    print('finish compressing')
    out = soundstream.decode(compressed)
    out = out.detach().cpu().squeeze(0)
    check_clipping(out, rescale)
    save_audio(wav=out, path=store_root, sample_rate=args.sr, rescale=rescale)
    print('finish decompressing')


def remove_encodec_weight_norm(model):
    from academicodec.modules import SConv1d
    from academicodec.modules.seanet import SConvTranspose1d
    from academicodec.modules.seanet import SEANetResnetBlock
    from torch.nn.utils import remove_weight_norm

    encoder = model.encoder.model
    for key in encoder._modules:
        if isinstance(encoder._modules[key], SEANetResnetBlock):
            remove_weight_norm(encoder._modules[key].shortcut.conv.conv)
            block_modules = encoder._modules[key].block._modules
            for skey in block_modules:
                if isinstance(block_modules[skey], SConv1d):
                    remove_weight_norm(block_modules[skey].conv.conv)
        elif isinstance(encoder._modules[key], SConv1d):
            remove_weight_norm(encoder._modules[key].conv.conv)

    decoder = model.decoder.model
    for key in decoder._modules:
        if isinstance(decoder._modules[key], SEANetResnetBlock):
            remove_weight_norm(decoder._modules[key].shortcut.conv.conv)
            block_modules = decoder._modules[key].block._modules
            for skey in block_modules:
                if isinstance(block_modules[skey], SConv1d):
                    remove_weight_norm(block_modules[skey].conv.conv)
        elif isinstance(decoder._modules[key], SConvTranspose1d):
            remove_weight_norm(decoder._modules[key].convtr.convtr)
        elif isinstance(decoder._modules[key], SConv1d):
            remove_weight_norm(decoder._modules[key].conv.conv)


def test_batch():
    args = get_parser().parse_args()
    print("args.target_bandwidths:", args.target_bandwidths)
    if not args.input.exists():
        fatal(f"Input file {args.input} does not exist.")
    input_lists = os.listdir(args.input)
    input_lists.sort()
    soundstream = SoundStream(
        n_filters=32,
        D=512,
        ratios=args.ratios,
        sample_rate=args.sr,
        target_bandwidths=args.target_bandwidths)
    parameter_dict = torch.load(args.resume_path)
    new_state_dict = OrderedDict()
    # k 为 module.xxx.weight, v 为权重
    for k, v in parameter_dict.items():
        # 截取`module.`后面的xxx.weight
        name = k[7:]
        new_state_dict[name] = v
    soundstream.load_state_dict(new_state_dict)  # load model
    remove_encodec_weight_norm(soundstream)
    soundstream.cuda()
    soundstream.eval()
    os.makedirs(args.output, exist_ok=True)
    for audio in input_lists:
        test_one(
            args=args,
            wav_root=os.path.join(args.input, audio),
            store_root=os.path.join(args.output, audio),
            rescale=args.rescale,
            soundstream=soundstream)


if __name__ == '__main__':
    test_batch()