import random
import math
from functools import partial
import json

import pysndfx
import librosa
import numpy as np
import torch

from ops.audio import (
    read_audio, compute_stft, trim_audio, mix_audio_and_labels,
    shuffle_audio, cutout
)


SAMPLE_RATE = 44100


class Augmentation:
    """A base class for data augmentation transforms"""
    pass


class MapLabels:

    def __init__(self, class_map, drop_raw=True):

        self.class_map = class_map

    def __call__(self, dataset, **inputs):

        labels = np.zeros(len(self.class_map), dtype=np.float32)
        for c in inputs["raw_labels"]:
            labels[self.class_map[c]] = 1.0

        transformed = dict(inputs)
        transformed["labels"] = labels
        transformed.pop("raw_labels")

        return transformed


class MixUp(Augmentation):

    def __init__(self, p):

        self.p = p

    def __call__(self, dataset, **inputs):

        transformed = dict(inputs)

        if np.random.uniform() < self.p:
            first_audio, first_labels = inputs["audio"], inputs["labels"]
            random_sample = dataset.random_clean_sample()
            new_audio, new_labels = mix_audio_and_labels(
                first_audio, random_sample["audio"],
                first_labels, random_sample["labels"]
            )

            transformed["audio"] = new_audio
            transformed["labels"] = new_labels

        return transformed


class FlipAudio(Augmentation):

    def __init__(self, p):

        self.p = p

    def __call__(self, dataset, **inputs):

        transformed = dict(inputs)

        if np.random.uniform() < self.p:
            transformed["audio"] = np.flipud(inputs["audio"])

        return transformed


class AudioAugmentation(Augmentation):

    def __init__(self, p):

        self.p = p

    def __call__(self, dataset, **inputs):

        transformed = dict(inputs)

        if np.random.uniform() < self.p:
            effects_chain = (
                pysndfx.AudioEffectsChain()
                .reverb(
                    reverberance=random.randrange(50),
                    room_scale=random.randrange(50),
                    stereo_depth=random.randrange(50)
                )
                .pitch(shift=random.randrange(-300, 300))
                .overdrive(gain=random.randrange(2, 10))
                .speed(random.uniform(0.9, 1.1))
            )
            transformed["audio"] = effects_chain(inputs["audio"])

        return transformed


class LoadAudio:

    def __init__(self):

        pass

    def __call__(self, dataset, **inputs):

        audio, sr = read_audio(inputs["filename"])

        transformed = dict(inputs)
        transformed["audio"] = audio
        transformed["sr"] = sr

        return transformed


class STFT:

    eps = 1e-4

    def __init__(self, n_fft, hop_size):

        self.n_fft = n_fft
        self.hop_size = hop_size

    def __call__(self, dataset, **inputs):

        stft = compute_stft(
            inputs["audio"],
            window_size=self.n_fft, hop_size=self.hop_size,
            eps=self.eps)

        transformed = dict(inputs)
        transformed["stft"] = np.transpose(stft)

        return transformed


class AudioFeatures:

    eps = 1e-4

    def __init__(self, descriptor, verbose=True):

        name, *args = descriptor.split("_")

        self.feature_type = name

        if name == "stft":

            n_fft, hop_size = args
            self.n_fft = int(n_fft)
            self.hop_size = int(hop_size)

            self.n_features = self.n_fft // 2 + 1
            self.padding_value = 0.0

            if verbose:
                print(
                    "\nUsing STFT features with params:\n",
                    "n_fft: {}, hop_size: {}".format(
                        n_fft, hop_size
                    )
                )

        elif name == "mel":

            n_fft, hop_size, n_mel = args
            self.n_fft = int(n_fft)
            self.hop_size = int(hop_size)
            self.n_mel = int(n_mel)

            self.n_features = self.n_mel
            self.padding_value = 0.0

            if verbose:
                print(
                    "\nUsing mel features with params:\n",
                    "n_fft: {}, hop_size: {}, n_mel: {}".format(
                        n_fft, hop_size, n_mel
                    )
                )

        elif name == "raw":

            self.n_features = 1
            self.padding_value = 0.0

            if verbose:
                print(
                    "\nUsing raw waveform features."
                )


    def __call__(self, dataset, **inputs):

        transformed = dict(inputs)

        if self.feature_type == "stft":

            # stft = compute_stft(
            #     inputs["audio"],
            #     window_size=self.n_fft, hop_size=self.hop_size,
            #     eps=self.eps, log=True
            # )

            transformed["signal"] = np.expand_dims(inputs["audio"], -1)

        elif self.feature_type == "mel":

            stft = compute_stft(
                inputs["audio"],
                window_size=self.n_fft, hop_size=self.hop_size,
                eps=self.eps, log=False
            )

            transformed["signal"] = np.expand_dims(inputs["audio"], -1)

        elif self.feature_type == "raw":
            transformed["signal"] = np.expand_dims(inputs["audio"], -1)

        return transformed


class SampleSegment(Augmentation):

    def __init__(self, ratio=(0.3, 0.9), p=1.0):

        self.min, self.max = ratio
        self.p = p

    def __call__(self, dataset, **inputs):

        transformed = dict(inputs)

        if np.random.uniform() < self.p:
            original_size = inputs["audio"].size
            target_size = int(np.random.uniform(self.min, self.max) * original_size)
            start = np.random.randint(original_size - target_size - 1)
            transformed["audio"] = inputs["audio"][start:start+target_size]

        return transformed


class ShuffleAudio(Augmentation):

    def __init__(self, chunk_length=0.5, p=0.5):

        self.chunk_length = chunk_length
        self.p = p

    def __call__(self, dataset, **inputs):

        transformed = dict(inputs)

        if np.random.uniform() < self.p:
            transformed["audio"] = shuffle_audio(
                transformed["audio"], self.chunk_length, sr=transformed["sr"])

        return transformed


class CutOut(Augmentation):

    def __init__(self, area=0.25, p=0.5):

        self.area = area
        self.p = p

    def __call__(self, dataset, **inputs):

        transformed = dict(inputs)

        if np.random.uniform() < self.p:
            transformed["audio"] = cutout(
                transformed["audio"], self.area)

        return transformed


class SampleLongAudio:

    def __init__(self, max_length):

        self.max_length = max_length

    def __call__(self, dataset, **inputs):

        transformed = dict(inputs)

        if (inputs["audio"].size / inputs["sr"]) > self.max_length:

            max_length = self.max_length * inputs["sr"]

            start = np.random.randint(0, inputs["audio"].size - max_length)
            transformed["audio"] = inputs["audio"][start:start+max_length]

        return transformed


class OneOf:

    def __init__(self, transforms):

        self.transforms = transforms

    def __call__(self, dataset, **inputs):

        transform = random.choice(self.transforms)
        return transform(**inputs)


class DropFields:

    def __init__(self, fields):

        self.to_drop = fields

    def __call__(self, dataset, **inputs):

        transformed = dict()

        for name, input in inputs.items():
            if not name in self.to_drop:
                transformed[name] = input

        return transformed


class RenameFields:

    def __init__(self, mapping):

        self.mapping = mapping

    def __call__(self, dataset, **inputs):

        transformed = dict(inputs)

        for old, new in self.mapping.items():
            transformed[new] = transformed.pop(old)

        return transformed


class Compose:

    def __init__(self, transforms):
        self.transforms = transforms

    def switch_off_augmentations(self):
        for t in self.transforms:
            if isinstance(t, Augmentation):
                t.p = 0.0

    def __call__(self, dataset=None, **inputs):
        for t in self.transforms:
            inputs = t(dataset=dataset, **inputs)

        return inputs


class Identity:

    def __call__(self, dataset=None, **inputs):

        return inputs
[((119, 20, 119, 50), 'ops.audio.read_audio', 'read_audio', ({(119, 31, 119, 49): "inputs['filename']"}, {}), "(inputs['filename'])", False, 'from ops.audio import read_audio, compute_stft, trim_audio, mix_audio_and_labels, shuffle_audio, cutout\n'), ((139, 15, 142, 25), 'ops.audio.compute_stft', 'compute_stft', (), '', False, 'from ops.audio import read_audio, compute_stft, trim_audio, mix_audio_and_labels, shuffle_audio, cutout\n'), ((145, 30, 145, 48), 'numpy.transpose', 'np.transpose', ({(145, 43, 145, 47): 'stft'}, {}), '(stft)', True, 'import numpy as np\n'), ((320, 20, 320, 50), 'random.choice', 'random.choice', ({(320, 34, 320, 49): 'self.transforms'}, {}), '(self.transforms)', False, 'import random\n'), ((54, 11, 54, 30), 'numpy.random.uniform', 'np.random.uniform', ({}, {}), '()', True, 'import numpy as np\n'), ((57, 36, 60, 13), 'ops.audio.mix_audio_and_labels', 'mix_audio_and_labels', ({(58, 16, 58, 27): 'first_audio', (58, 29, 58, 51): "random_sample['audio']", (59, 16, 59, 28): 'first_labels', (59, 30, 59, 53): "random_sample['labels']"}, {}), "(first_audio, random_sample['audio'], first_labels,\n    random_sample['labels'])", False, 'from ops.audio import read_audio, compute_stft, trim_audio, mix_audio_and_labels, shuffle_audio, cutout\n'), ((78, 11, 78, 30), 'numpy.random.uniform', 'np.random.uniform', ({}, {}), '()', True, 'import numpy as np\n'), ((79, 35, 79, 61), 'numpy.flipud', 'np.flipud', ({(79, 45, 79, 60): "inputs['audio']"}, {}), "(inputs['audio'])", True, 'import numpy as np\n'), ((94, 11, 94, 30), 'numpy.random.uniform', 'np.random.uniform', ({}, {}), '()', True, 'import numpy as np\n'), ((218, 36, 218, 71), 'numpy.expand_dims', 'np.expand_dims', ({(218, 51, 218, 66): "inputs['audio']", (218, 68, 218, 70): '-1'}, {}), "(inputs['audio'], -1)", True, 'import numpy as np\n'), ((247, 11, 247, 30), 'numpy.random.uniform', 'np.random.uniform', ({}, {}), '()', True, 'import numpy as np\n'), ((250, 20, 250, 70), 'numpy.random.randint', 'np.random.randint', ({(250, 38, 250, 69): 'original_size - target_size - 1'}, {}), '(original_size - target_size - 1)', True, 'import numpy as np\n'), ((267, 11, 267, 30), 'numpy.random.uniform', 'np.random.uniform', ({}, {}), '()', True, 'import numpy as np\n'), ((268, 35, 269, 78), 'ops.audio.shuffle_audio', 'shuffle_audio', (), '', False, 'from ops.audio import read_audio, compute_stft, trim_audio, mix_audio_and_labels, shuffle_audio, cutout\n'), ((285, 11, 285, 30), 'numpy.random.uniform', 'np.random.uniform', ({}, {}), '()', True, 'import numpy as np\n'), ((286, 35, 287, 48), 'ops.audio.cutout', 'cutout', ({(287, 16, 287, 36): "transformed['audio']", (287, 38, 287, 47): 'self.area'}, {}), "(transformed['audio'], self.area)", False, 'from ops.audio import read_audio, compute_stft, trim_audio, mix_audio_and_labels, shuffle_audio, cutout\n'), ((306, 20, 306, 75), 'numpy.random.randint', 'np.random.randint', ({(306, 38, 306, 39): '0', (306, 41, 306, 74): "inputs['audio'].size - max_length"}, {}), "(0, inputs['audio'].size - max_length)", True, 'import numpy as np\n'), ((104, 23, 104, 47), 'random.uniform', 'random.uniform', ({(104, 38, 104, 41): '0.9', (104, 43, 104, 46): '1.1'}, {}), '(0.9, 1.1)', False, 'import random\n'), ((222, 19, 226, 13), 'ops.audio.compute_stft', 'compute_stft', (), '', False, 'from ops.audio import read_audio, compute_stft, trim_audio, mix_audio_and_labels, shuffle_audio, cutout\n'), ((228, 36, 228, 71), 'numpy.expand_dims', 'np.expand_dims', ({(228, 51, 228, 66): "inputs['audio']", (228, 68, 228, 70): '-1'}, {}), "(inputs['audio'], -1)", True, 'import numpy as np\n'), ((231, 36, 231, 71), 'numpy.expand_dims', 'np.expand_dims', ({(231, 51, 231, 66): "inputs['audio']", (231, 68, 231, 70): '-1'}, {}), "(inputs['audio'], -1)", True, 'import numpy as np\n'), ((249, 30, 249, 67), 'numpy.random.uniform', 'np.random.uniform', ({(249, 48, 249, 56): 'self.min', (249, 58, 249, 66): 'self.max'}, {}), '(self.min, self.max)', True, 'import numpy as np\n'), ((103, 32, 103, 55), 'random.randrange', 'random.randrange', ({(103, 49, 103, 50): '2', (103, 52, 103, 54): '10'}, {}), '(2, 10)', False, 'import random\n'), ((102, 29, 102, 56), 'random.randrange', 'random.randrange', ({(102, 46, 102, 50): '-300', (102, 52, 102, 55): '300'}, {}), '(-300, 300)', False, 'import random\n'), ((96, 16, 96, 43), 'pysndfx.AudioEffectsChain', 'pysndfx.AudioEffectsChain', ({}, {}), '()', False, 'import pysndfx\n'), ((98, 33, 98, 53), 'random.randrange', 'random.randrange', ({(98, 50, 98, 52): '50'}, {}), '(50)', False, 'import random\n'), ((99, 31, 99, 51), 'random.randrange', 'random.randrange', ({(99, 48, 99, 50): '50'}, {}), '(50)', False, 'import random\n'), ((100, 33, 100, 53), 'random.randrange', 'random.randrange', ({(100, 50, 100, 52): '50'}, {}), '(50)', False, 'import random\n')]