import time
import os
import pickle
from src.music.pipeline.music_pipeline import encode_music
from src.music2cocktailrep.pipeline.music2cocktailrep import music2cocktailrep, setup_translation_models, debug_translation
from src.cocktails.pipeline.cocktailrep2recipe import cocktailrep2recipe
from src.debugger import Debugger
from datetime import datetime
from shutil import copy

synestesia_path = '../data/synesthesia/'
debugger = Debugger()

def pianocktail(record=False, url=None, midi=None, audio=None, processed=None, crop=40, verbose=False, debug=False, level=0):
    assert url is not None or midi is not None or audio is not None or processed is not None
    if verbose: print('------\nNew synesthetic exploration!')
    init_time = time.time()
    music_ai_rep, music_handcoded_rep, all_paths, error = encode_music(record=record, url=url, audio_path=audio, midi_path=midi, nb_aug=0, noise_injection=False,
                                                                       augmentation=False, processed_path=processed, crop=crop, apply_filtering=False, verbose=verbose,
                                                                       level=level+2)
    if music_ai_rep is not None:
        cocktail_rep, affective_cluster_id, affect = music2cocktailrep(music_ai_rep, music_handcoded_rep, verbose=verbose, level=level+2)
        cocktail_recipes, scores = cocktailrep2recipe(cocktail_rep, target_affective_cluster=affective_cluster_id, verbose=verbose, full_verbose=verbose, level=level+2)
        cocktail_recipe = cocktail_recipes[0]
        recipe_score = scores[0]
        if debug:
            music_reconstruction = debug_translation(music_ai_rep)
            debugger.extract_info(all_paths, affective_cluster_id, affect, cocktail_rep, music_reconstruction, recipe_score, verbose=verbose, level=level+2)
            debug_info = debugger.debug_dict
        else:
            debug_info = None
        if verbose:
            print(cocktail_recipe.replace('Recipe', ' ' * (level + 2) + 'Generated recipe:').replace('None ()', ''))
            debugger.print_debug(level=level+2)
            print(f'\nEnd of synesthetic exploration ({int(time.time() - init_time)} secs).\n------')

    else:
        cocktail_recipe = None
        debug_info = None
    return cocktail_recipe, debug_info

def setup_and_run(url=None, midi=None, audio=None, verbose=False, debug=False, extra_code=None):
    assert url is not None or midi is not None or audio is not None
    now = datetime.now()
    folder_name = f'{now.year}-{now.month}-{now.day}_{now.hour}:{now.minute}:{now.second}'
    folder_path = synestesia_path + folder_name
    if extra_code is not None:
        folder_path += '_' + extra_code
    if os.path.exists(folder_path):
        folder_path += '_2'
    folder_path += '/'
    os.makedirs(folder_path, exist_ok=True)
    recipe, debug = pianocktail(url=url, verbose=verbose, debug=debug)
    with open(folder_path + 'debug.pk', 'wb') as f:
        pickle.dump(debug, f)
    with open(folder_path + 'recipe.txt', 'w') as f:
        f.write(recipe)
    paths = debug['all_paths']
    if paths['url'] is not None:
        with open(folder_path + 'url.txt', 'w') as f:
            f.write(paths['url'])
    for k in ['audio_path', 'midi_path']:
        origin = paths[k]
        copy(origin, folder_path + origin.split('/')[-1])


if __name__ == '__main__':
    urls = ["https://www.youtube.com/watch?v=PLFVGwGQcB0",
            "https://www.youtube.com/watch?v=VQmuAr93OlI",
            "https://www.youtube.com/watch?v=Nv2GgV34qIg&list=PLO9E3V4rGLD8_iWrCioJRWZXJJE3Fzu_J&index=4",
            "https://www.youtube.com/watch?v=qAEIjWYdoYc&list=PLO9E3V4rGLD8_iWrCioJRWZXJJE3Fzu_J&index=1",
            "https://www.youtube.com/watch?v=M73x3O7dhmg&list=PLO9E3V4rGLD8_iWrCioJRWZXJJE3Fzu_J&index=5"]
    setup_translation_models()
    setup_and_run(url=urls[0], verbose=True, debug=True)
    recipes = []
    for url in urls:
        recipe = pianocktail(url=url, verbose=True, debug=True)[0]
        recipes.append(recipe)
    stop = 1