import gin
from scaper import Scaper, generate_from_jams
import copy
import logging
import p_tqdm
import nussl
import os
import numpy as np

def _reset_event_spec(sc):
    sc.reset_fg_event_spec()
    sc.reset_bg_event_spec()

def check_mixture(path_to_mix):
    mix_signal = nussl.AudioSignal(path_to_mix)
    if mix_signal.rms() < .01:
        return False
    return True

def make_one_mixture(sc, path_to_file, num_sources, 
                     event_parameters, allow_repeated_label):
    """
    Creates a single mixture, incoherent. Instantiates according to
    the event parameters for each source.
    """
    check = False
    while not check:
        for j in range(num_sources):
            sc.add_event(**event_parameters)
        
        sc.generate(
            path_to_file, 
            path_to_file.replace('.wav', '.jams'), 
            no_audio=False,
            allow_repeated_label=allow_repeated_label,
            save_isolated_events=True,
        )
        _reset_event_spec(sc)
        check = check_mixture(path_to_file)

def instantiate_and_get_event_spec(sc, master_label, event_parameters):
    _reset_event_spec(sc)

    _event_parameters = copy.deepcopy(event_parameters)
    _event_parameters['label'] = ('const', master_label)
    sc.add_event(**_event_parameters)
    event = sc._instantiate_event(sc.fg_spec[-1])

    _reset_event_spec(sc)
    return sc, event

def make_one_mixture_coherent(sc, path_to_file, labels, event_parameters, 
                              allow_repeated_label):
    check = False
    while not check:
        sc, event = instantiate_and_get_event_spec(
            sc, labels[0], event_parameters)
        for label in labels:
            try:
                sc.add_event(
                    label=('const', label),
                    source_file=('const', event.source_file.replace(labels[0], label)),
                    source_time=('const', event.source_time),
                    event_time=('const', 0),
                    event_duration=('const', sc.duration),
                    snr=event_parameters['snr'],
                    pitch_shift=('const', event.pitch_shift),
                    time_stretch=('const', event.time_stretch)
                )
            except:
                logging.exception(
                    f"Got an error for {label} @ {_source_file}. Moving on...")
        sc.generate(
            path_to_file, 
            path_to_file.replace('.wav', '.jams'), 
            no_audio=False, 
            allow_repeated_label=allow_repeated_label,
            save_isolated_events=True,
        )
        sc.fg_spec = []
        check = check_mixture(path_to_file)

@gin.configurable
def make_scaper_datasets(scopes=['train', 'val']):
    for scope in scopes:
        with gin.config_scope(scope):
            mix_with_scaper()

@gin.configurable
def mix_with_scaper(num_mixtures, foreground_path, background_path, 
                    scene_duration, sample_rate, target_folder, 
                    event_parameters, num_sources=None, labels=None, 
                    coherent=False, allow_repeated_label=False, 
                    ref_db=-40, bitdepth=16, seed=0, num_workers=1):
    nussl.utils.seed(seed)
    os.makedirs(target_folder, exist_ok=True)

    scaper_seed = np.random.randint(100)
    logging.info('Starting mixing.')

    if num_sources is None and labels is None:
        raise ValueError("One of labels or num_sources must be set!")

    if coherent and labels is None:
        raise ValueError("Coherent mixing requires explicit labels!")

    generators = []

    if background_path is None:
        background_path = foreground_path

    for i in range(num_mixtures):
        sc = Scaper(
            scene_duration, 
            fg_path=foreground_path,
            bg_path=background_path,
            random_state=scaper_seed,
        )
        sc.ref_db = ref_db
        sc.sr = sample_rate
        sc.bitdepth = bitdepth
        generators.append(sc)
        scaper_seed += 1

    mix_func = make_one_mixture_coherent if coherent else make_one_mixture

    def arg_tuple(i):
        _args = (
            generators[i],
            os.path.join(target_folder, f'{i:08d}.wav'),
            labels if coherent else num_sources,
            event_parameters,
            allow_repeated_label
        )
        return _args

    args = [arg_tuple(i) for i in range(num_mixtures)]
    # do one by itself for testing
    mix_func(*args[0])

    args = list(zip(*args[1:]))
    args = [list(a) for a in args]

    # now do the rest in parallel
    p_tqdm.p_map(mix_func, *args, num_cpus=num_workers)

[((15, 17, 15, 47), 'nussl.AudioSignal', 'nussl.AudioSignal', ({(15, 35, 15, 46): 'path_to_mix'}, {}), '(path_to_mix)', False, 'import nussl\n'), ((44, 24, 44, 55), 'copy.deepcopy', 'copy.deepcopy', ({(44, 38, 44, 54): 'event_parameters'}, {}), '(event_parameters)', False, 'import copy\n'), ((95, 4, 95, 26), 'nussl.utils.seed', 'nussl.utils.seed', ({(95, 21, 95, 25): 'seed'}, {}), '(seed)', False, 'import nussl\n'), ((96, 4, 96, 45), 'os.makedirs', 'os.makedirs', (), '', False, 'import os\n'), ((98, 18, 98, 40), 'numpy.random.randint', 'np.random.randint', ({(98, 36, 98, 39): '100'}, {}), '(100)', True, 'import numpy as np\n'), ((99, 4, 99, 36), 'logging.info', 'logging.info', ({(99, 17, 99, 35): '"""Starting mixing."""'}, {}), "('Starting mixing.')", False, 'import logging\n'), ((145, 4, 145, 55), 'p_tqdm.p_map', 'p_tqdm.p_map', (), '', False, 'import p_tqdm\n'), ((113, 13, 118, 9), 'scaper.Scaper', 'Scaper', (), '', False, 'from scaper import Scaper, generate_from_jams\n'), ((86, 13, 86, 36), 'gin.config_scope', 'gin.config_scope', ({(86, 30, 86, 35): 'scope'}, {}), '(scope)', False, 'import gin\n'), ((130, 12, 130, 55), 'os.path.join', 'os.path.join', ({(130, 25, 130, 38): 'target_folder', (130, 40, 130, 54): 'f"""{i:08d}.wav"""'}, {}), "(target_folder, f'{i:08d}.wav')", False, 'import os\n'), ((71, 16, 72, 79), 'logging.exception', 'logging.exception', ({(72, 20, 72, 78): 'f"""Got an error for {label} @ {_source_file}. Moving on..."""'}, {}), "(f'Got an error for {label} @ {_source_file}. Moving on...')", False, 'import logging\n')]