import sys
sys.path.append('.')
import argparse
import os
import time
from typing import Dict
import pathlib

import librosa
import numpy as np
import soundfile
import torch
import torch.nn as nn

from bytesep.models.lightning_modules import get_model_class
from bytesep.utils import read_yaml


class Separator:
    def __init__(
        self, model: nn.Module, segment_samples: int, batch_size: int, device: str
    ):
        r"""Separate to separate an audio clip into a target source.

        Args:
            model: nn.Module, trained model
            segment_samples: int, length of segments to be input to a model, e.g., 44100*30
            batch_size, int, e.g., 12
            device: str, e.g., 'cuda'
        """
        self.model = model
        self.segment_samples = segment_samples
        self.batch_size = batch_size
        self.device = device

    def separate(self, input_dict: Dict) -> np.array:
        r"""Separate an audio clip into a target source.

        Args:
            input_dict: dict, e.g., {
                waveform: (channels_num, audio_samples),
                ...,
            }

        Returns:
            sep_audio: (channels_num, audio_samples) | (target_sources_num, channels_num, audio_samples)
        """
        audio = input_dict['waveform']

        audio_samples = audio.shape[-1]

        # Pad the audio with zero in the end so that the length of audio can be
        # evenly divided by segment_samples.
        audio = self.pad_audio(audio)

        # Enframe long audio into segments.
        segments = self.enframe(audio, self.segment_samples)
        # (segments_num, channels_num, segment_samples)

        segments_input_dict = {'waveform': segments}

        if 'condition' in input_dict.keys():
            segments_num = len(segments)
            segments_input_dict['condition'] = np.tile(
                input_dict['condition'][None, :], (segments_num, 1)
            )
            # (batch_size, segments_num)

        # Separate in mini-batches.
        sep_segments = self._forward_in_mini_batches(
            self.model, segments_input_dict, self.batch_size
        )['waveform']
        # (segments_num, channels_num, segment_samples)

        # Deframe segments into long audio.
        sep_audio = self.deframe(sep_segments)
        # (channels_num, padded_audio_samples)

        sep_audio = sep_audio[:, 0:audio_samples]
        # (channels_num, audio_samples)

        return sep_audio

    def pad_audio(self, audio: np.array) -> np.array:
        r"""Pad the audio with zero in the end so that the length of audio can
        be evenly divided by segment_samples.

        Args:
            audio: (channels_num, audio_samples)

        Returns:
            padded_audio: (channels_num, audio_samples)
        """
        channels_num, audio_samples = audio.shape

        # Number of segments
        segments_num = int(np.ceil(audio_samples / self.segment_samples))

        pad_samples = segments_num * self.segment_samples - audio_samples

        padded_audio = np.concatenate(
            (audio, np.zeros((channels_num, pad_samples))), axis=1
        )
        # (channels_num, padded_audio_samples)

        return padded_audio

    def enframe(self, audio: np.array, segment_samples: int) -> np.array:
        r"""Enframe long audio into segments.

        Args:
            audio: (channels_num, audio_samples)
            segment_samples: int

        Returns:
            segments: (segments_num, channels_num, segment_samples)
        """
        audio_samples = audio.shape[1]
        assert audio_samples % segment_samples == 0

        hop_samples = segment_samples // 2
        segments = []

        pointer = 0
        while pointer + segment_samples <= audio_samples:
            segments.append(audio[:, pointer : pointer + segment_samples])
            pointer += hop_samples

        segments = np.array(segments)

        return segments

    def deframe(self, segments: np.array) -> np.array:
        r"""Deframe segments into long audio.

        Args:
            segments: (segments_num, channels_num, segment_samples)

        Returns:
            output: (channels_num, audio_samples)
        """
        (segments_num, _, segment_samples) = segments.shape

        if segments_num == 1:
            return segments[0]

        assert self._is_integer(segment_samples * 0.25)
        assert self._is_integer(segment_samples * 0.75)

        output = []

        output.append(segments[0, :, 0 : int(segment_samples * 0.75)])

        for i in range(1, segments_num - 1):
            output.append(
                segments[
                    i, :, int(segment_samples * 0.25) : int(segment_samples * 0.75)
                ]
            )

        output.append(segments[-1, :, int(segment_samples * 0.25) :])

        output = np.concatenate(output, axis=-1)

        return output

    def _is_integer(self, x: float) -> bool:
        if x - int(x) < 1e-10:
            return True
        else:
            return False

    def _forward_in_mini_batches(
        self, model: nn.Module, segments_input_dict: Dict, batch_size: int
    ) -> Dict:
        r"""Forward data to model in mini-batch.

        Args:
            model: nn.Module
            segments_input_dict: dict, e.g., {
                'waveform': (segments_num, channels_num, segment_samples),
                ...,
            }
            batch_size: int

        Returns:
            output_dict: dict, e.g. {
                'waveform': (segments_num, channels_num, segment_samples),
            }
        """
        output_dict = {}

        pointer = 0
        segments_num = len(segments_input_dict['waveform'])

        while True:
            if pointer >= segments_num:
                break

            batch_input_dict = {}

            for key in segments_input_dict.keys():
                batch_input_dict[key] = torch.Tensor(
                    segments_input_dict[key][pointer : pointer + batch_size]
                ).to(self.device)

            pointer += batch_size

            with torch.no_grad():
                model.eval()
                batch_output_dict = model(batch_input_dict)

            for key in batch_output_dict.keys():
                self._append_to_dict(
                    output_dict, key, batch_output_dict[key].data.cpu().numpy()
                )

        for key in output_dict.keys():
            output_dict[key] = np.concatenate(output_dict[key], axis=0)

        return output_dict

    def _append_to_dict(self, dict, key, value):
        if key in dict.keys():
            dict[key].append(value)
        else:
            dict[key] = [value]


class SeparatorWrapper:
    def __init__(
        self, source_type='vocals', model=None, checkpoint_path=None, device='cuda'
    ):

        input_channels = 2
        target_sources_num = 1
        model_type = "ResUNet143_Subbandtime"
        segment_samples = 44100 * 10
        batch_size = 1

        self.checkpoint_path = self.download_checkpoints(checkpoint_path, source_type)

        if device == 'cuda' and torch.cuda.is_available():
            self.device = 'cuda'
        else:
            self.device = 'cpu'

        # Get model class.
        Model = get_model_class(model_type)

        # Create model.
        self.model = Model(
            input_channels=input_channels, target_sources_num=target_sources_num
        )

        # Load checkpoint.
        checkpoint = torch.load(self.checkpoint_path, map_location='cpu')
        self.model.load_state_dict(checkpoint["model"])

        # Move model to device.
        self.model.to(self.device)

        # Create separator.
        self.separator = Separator(
            model=self.model,
            segment_samples=segment_samples,
            batch_size=batch_size,
            device=self.device,
        )

    def download_checkpoints(self, checkpoint_path, source_type):

        if source_type == "vocals":
            checkpoint_bare_name = "resunet143_subbtandtime_vocals_8.8dB_350k_steps"

        elif source_type == "accompaniment":
            checkpoint_bare_name = (
                "resunet143_subbtandtime_accompaniment_16.4dB_350k_steps.pth"
            )

        else:
            raise NotImplementedError

        if not checkpoint_path:
            checkpoint_path = '{}/bytesep_data/{}.pth'.format(
                str(pathlib.Path.home()), checkpoint_bare_name
            )

        print('Checkpoint path: {}'.format(checkpoint_path))

        if (
            not os.path.exists(checkpoint_path)
            or os.path.getsize(checkpoint_path) < 4e8
        ):

            os.makedirs(os.path.dirname(checkpoint_path), exist_ok=True)

            zenodo_dir = "https://zenodo.org/record/5507029/files"
            zenodo_path = os.path.join(
                zenodo_dir, "{}?download=1".format(checkpoint_bare_name)
            )

            os.system('wget -O "{}" "{}"'.format(checkpoint_path, zenodo_path))

        return checkpoint_path

    def separate(self, audio):

        input_dict = {'waveform': audio}

        sep_wav = self.separator.separate(input_dict)

        return sep_wav


def inference(args):

    # Need to use torch.distributed if models contain inplace_abn.abn.InPlaceABNSync.
    import torch.distributed as dist

    dist.init_process_group(
        'gloo', init_method='file:///tmp/somefile', rank=0, world_size=1
    )

    # Arguments & parameters
    config_yaml = args.config_yaml
    checkpoint_path = args.checkpoint_path
    audio_path = args.audio_path
    output_path = args.output_path
    device = (
        torch.device('cuda')
        if args.cuda and torch.cuda.is_available()
        else torch.device('cpu')
    )

    configs = read_yaml(config_yaml)
    sample_rate = configs['train']['sample_rate']
    input_channels = configs['train']['channels']
    target_source_types = configs['train']['target_source_types']
    target_sources_num = len(target_source_types)
    model_type = configs['train']['model_type']

    segment_samples = int(30 * sample_rate)
    batch_size = 1

    print("Using {} for separating ..".format(device))

    # paths
    if os.path.dirname(output_path) != "":
        os.makedirs(os.path.dirname(output_path), exist_ok=True)

    # Get model class.
    Model = get_model_class(model_type)

    # Create model.
    model = Model(input_channels=input_channels, target_sources_num=target_sources_num)

    # Load checkpoint.
    checkpoint = torch.load(checkpoint_path, map_location='cpu')
    model.load_state_dict(checkpoint["model"])

    # Move model to device.
    model.to(device)

    # Create separator.
    separator = Separator(
        model=model,
        segment_samples=segment_samples,
        batch_size=batch_size,
        device=device,
    )

    # Load audio.
    audio, _ = librosa.load(audio_path, sr=sample_rate, mono=False)

    # audio = audio[None, :]

    input_dict = {'waveform': audio}

    # Separate
    separate_time = time.time()

    sep_wav = separator.separate(input_dict)
    # (channels_num, audio_samples)

    print('Separate time: {:.3f} s'.format(time.time() - separate_time))

    # Write out separated audio.
    soundfile.write(file='_zz.wav', data=sep_wav.T, samplerate=sample_rate)
    os.system("ffmpeg -y -loglevel panic -i _zz.wav {}".format(output_path))
    print('Write out to {}'.format(output_path))


if __name__ == "__main__":

    parser = argparse.ArgumentParser(description="")
    parser.add_argument("--config_yaml", type=str, required=True)
    parser.add_argument("--checkpoint_path", type=str, required=True)
    parser.add_argument("--audio_path", type=str, required=True)
    parser.add_argument("--output_path", type=str, required=True)
    parser.add_argument("--cuda", action='store_true', default=True)

    args = parser.parse_args()
    inference(args)