<filename>src/beatmapsynth.py
from __future__ import print_function
import numpy as np
import pandas as pd
import librosa
import json
#import requests
import pickle
#import matplotlib.pyplot as plt
from io import BytesIO, TextIOWrapper, StringIO
from zipfile import ZipFile
import os
import soundfile as sf
import audioread
from pydub import AudioSegment
#from sklearn.multioutput import MultiOutputClassifier, ClassifierChain
#from sklearn.ensemble import RandomForestClassifier
import markovify
import sklearn.cluster
#import librosa.display
import scipy
import sys
import argparse
import shutil

#Main Function:
def beat_map_synthesizer(song_path, song_name, difficulty, model, k=5, version = 2):
    """
    Function to load a music file and generate a custom Beat Saber map based on the specified model and difficulty. Outputs a zipped folder of necessary files to play the custom map in the Beat Saber game.

    ***
    song_path = string file path to music file location

    song_name = string to name level as it will appear in the game

    difficulty = desired difficulty level, can be: 'easy', 'normal', 'hard', 'expert', or 'expertPlus'

    model = desired model to use for map generation, can be: 'random', 'HMM', 'segmented_HMM', or 'rate_modulated_segmented_HMM'

    k = number of song segments if using a segmented model. Default is 5, may want to increase or decrease based on song complexity

    version = for HMM models, can choose either 1 or 2. 1 was trained on a smaller, but potentially higher quality dataset (custom maps with over 90% rating on beatsaver.com), while 2 was trained on a larger dataset of custom maps with over 70% rating, so it may have a larger pool of "potential moves."
    ***
    """
    if model == 'random':
        random_mapper(song_path, song_name, difficulty)
    elif model == 'HMM':
        HMM_mapper(song_path, song_name, difficulty, version = version)
    elif model == 'segmented_HMM':
        segmented_HMM_mapper(song_path, song_name, difficulty, k = k, version = version)
    elif model == 'rate_modulated_segmented_HMM':
        rate_modulated_segmented_HMM_mapper(song_path, song_name, difficulty, version = version, k = k)
    else:
        print('Please specify model for mapping.')

#Basic File Writing Functions
def write_info(song_name, bpm, difficulty):
    """This function creates the 'info.dat' file that needs to be included in the custom folder."""

    difficulty_rank = None
    jump_movement = None
    if difficulty.casefold() == 'easy'.casefold():
        difficulty_rank = 1
        jump_movement = 8
        diff_name = 'Easy'
    elif difficulty.casefold() == 'normal'.casefold():
        difficulty_rank = 3
        jump_movement = 10
        diff_name = 'Normal'
    elif difficulty.casefold() == 'hard'.casefold():
        difficulty_rank = 5
        jump_movement = 12
        diff_name = 'Hard'
    elif difficulty.casefold() == 'expert'.casefold():
        difficulty_rank = 7
        jump_movement = 14
        diff_name = 'Expert'
    elif difficulty.casefold() == 'expertPlus'.casefold():
        difficulty_rank = 9
        jump_movement = 16
        diff_name = 'ExpertPlus'

    info = {'_version': '2.0.0',
            '_songName': f"{song_name}",
            '_songSubName': '',
            '_songAuthorName': '',
            '_levelAuthorName': 'BeatMapSynth',
            '_beatsPerMinute': round(bpm),
            '_songTimeOffset': 0,
            '_shuffle': 0,
            '_shufflePeriod': 0,
            '_previewStartTime': 10,
            '_previewDuration': 30,
            '_songFilename': 'song.egg',
            '_coverImageFilename': 'cover.jpg',
            '_environmentName': 'DefaultEnvironment',
            '_customData': {},
             '_difficultyBeatmapSets': [{'_beatmapCharacteristicName': 'Standard',
                                         '_difficultyBeatmaps': [{'_difficulty': diff_name,
                                                                  '_difficultyRank': difficulty_rank,
                                                                  '_beatmapFilename': f"{difficulty}.dat",
                                                                  '_noteJumpMovementSpeed': jump_movement,
                                                                  '_noteJumpStartBeatOffset': 0,
                                                                  '_customData': {}}]}]}
    with open('info.dat', 'w') as f:
        json.dump(info, f)

def write_level(difficulty, events_list, notes_list, obstacles_list):
    """This function creates the 'level.dat' file that contains all the data for that paticular difficulty level"""

    level = {'_version': '2.0.0',
             '_customData': {'_time': '', #not sure what time refers to
                             '_BPMChanges': [],
                             '_bookmarks': []},
             '_events': events_list,
             '_notes': notes_list,
             '_obstacles': obstacles_list}
    with open(f"{difficulty}.dat", 'w') as f:
        json.dump(level, f)

def music_file_converter(song_path):
    """This function makes sure the file type of the provided song will be converted to the music file type that
    Beat Saber accepts"""
    if song_path.endswith('.mp3'):
        AudioSegment.from_mp3(song_path).export('song.egg', format='ogg')
    elif song_path.endswith('.wav'):
        AudioSegment.from_wav(song_path).export('song.egg', format='ogg')
    elif song_path.endswith('.flv'):
        AudioSegment.from_flv(song_path).export('song.egg', format='ogg')
    elif song_path.endswith('.raw'):
        AudioSegment.from_raw(song_path).export('song.egg', format='ogg')
    elif song_path.endswith('.ogg') or song_path.endswith('.egg'):
        shutil.copy2(song_path, 'song.egg')
    else:
        print("Unsupported song file type. Choose a file of type .mp3, .wav, .flv, .raw, or .ogg.")

def events_writer(beat_times):
    """Placeholder function for writing a list of events to be incorporated into a beatmap file. May have future support."""
    events_list = []
    return events_list

def obstacles_writer(beat_times, difficulty):
    """Placeholder function for writing a list of obstacles to be incorporated into a beatmap file."""
    obstacles_list = []
    return obstacles_list

def zip_folder_exporter(song_name, difficulty):
    "This function exports the zip folder containing the info.dat, difficulty.dat, cover.jpg, and song.egg files."
    files = ['info.dat', f"{difficulty}.dat", 'cover.jpg', 'song.egg']
    with ZipFile(f"{song_name}.zip", 'w') as custom:
        for file in files:
            custom.write(file)
    for file in files:
        if file != 'cover.jpg':
            os.remove(file)

#Random Mapping Functions
def random_mapper(song_path, song_name, difficulty):
    """Function to output the automatically created completely random map (i.e. baseline model) for a provided song. Returns a zipped folder that can be unzipped and placed in the 'CustomMusic' folder in the Beat Saber game directory and played. CAUTION: This is completely random and is likely not enjoyable if even playable!"""
    #Load song and get beat features
    print("Loading Song...")
    bpm, beat_times, y, sr = beat_features(song_path)
    print("Song loaded successfully!")
    #Write lists for note placement, event placement, and obstacle placement
    print("Random mapping...")
    #notes_list = random_notes_writer(beat_times, difficulty)
    notes_list = random_notes_writer_v2(beat_times, difficulty, bpm) #fixes _time != beat time
    events_list = events_writer(beat_times)
    obstacles_list = obstacles_writer(beat_times, difficulty)
    print("Mapping done!")
    #Write and zip files
    print("Writing files to disk...")
    write_info(song_name, bpm, difficulty)
    write_level(difficulty, events_list, notes_list, obstacles_list)
    print("Converting music file...")
    music_file_converter(song_path)
    print("Zipping folder...")
    zip_folder_exporter(song_name, difficulty)
    print("Finished! Look for zipped folder in your current path, unzip the folder, and place in the 'CustomMusic' folder in the Beat Saber directory")

def beat_features(song_path):
    """This function takes in the song stored at 'song_path' and estimates the bpm and beat times."""
    #Load song and split into harmonic and percussive parts.
    y, sr = librosa.load(song_path)
    #y_harmonic, y_percussive = librosa.effects.hpss(y)
    #Isolate beats and beat times
    bpm, beat_frames = librosa.beat.beat_track(y=y, sr=sr, trim = False)
    beat_times = librosa.frames_to_time(beat_frames, sr=sr)
    return bpm, beat_times, y, sr

def random_notes_writer_v2(beat_times, difficulty, bpm):
    """This function randomly places blocks at approximately each beat or every other beat depending on the difficulty."""
    notes_list = []
    line_index = [0, 1, 2, 3]
    line_layer = [0, 1, 2]
    types = [0, 1, 2, 3]
    directions = list(range(0, 10))
    #beat_times = [float(x) for x in beat_times]
    beat_times = [x*(bpm/60) for x in beat_times] #list(range(len(beat_times)))

    if difficulty == 'Easy' or difficulty == 'Normal':
        for beat in beat_times:
            empty = np.random.choice([0,1])
            if empty == 1:
                note = {'_time': beat,
                        '_lineIndex': int(np.random.choice(line_index)),
                        '_lineLayer': int(np.random.choice(line_layer)),
                        '_type': int(np.random.choice(types)),
                        '_cutDirection': int(np.random.choice(directions))}
                notes_list.append(note)
            else:
                continue
    else:
        random_beats = np.random.choice(beat_times, np.random.choice(range(len(beat_times)))) #randomly choose beats to have more than one note placed
        randomly_duplicated_beat_times = np.concatenate([beat_times, random_beats])
        randomly_duplicated_beat_times.sort()
        randomly_duplicated_beat_times = [float(x) for x in randomly_duplicated_beat_times]
        for beat in randomly_duplicated_beat_times:
            note = {'_time': beat,
                    '_lineIndex': int(np.random.choice(line_index)),
                    '_lineLayer': int(np.random.choice(line_layer)),
                    '_type': int(np.random.choice(types)),
                    '_cutDirection': int(np.random.choice(directions))}
            notes_list.append(note)
    #Remove potential notes that come too early in the song:
    for i, x in enumerate(notes_list):
        if notes_list[i]['_time'] >= 0 and notes_list[i]['_time'] <= 1.5:
            del notes_list[i]
        elif notes_list[i]['_time'] > beat_times[-1]:
            del notes_list[i]

    return notes_list

#Hidden Markov Models Mapping Functions
def HMM_mapper(song_path, song_name, difficulty, version = 2):
    """This function generates a custom map based on a Hidden Markov Model."""
    #Load song and get beat features
    print("Loading Song...")
    bpm, beat_times, y, sr = beat_features(song_path)
    beat_times = [x*(bpm/60) for x in beat_times] #list(range(len(beat_times)))
    print("Song loaded successfully!")
    #Write lists for note placement, event placement, and obstacle placement
    print("Mapping with Hidden Markov Model...")
    notes_list = HMM_notes_writer(beat_times, difficulty, version)
    events_list = events_writer(beat_times)
    obstacles_list = obstacles_writer(beat_times, difficulty)
    print("Mapping done!")
    #Write and zip files
    print("Writing files to disk...")
    write_info(song_name, bpm, difficulty)
    write_level(difficulty, events_list, notes_list, obstacles_list)
    print("Converting music file...")
    music_file_converter(song_path)
    print("Zipping folder...")
    zip_folder_exporter(song_name, difficulty)
    print("Finished! Look for zipped folder in your current path, unzip the folder, and place in the 'CustomMusic' folder in the Beat Saber directory")

def walk_to_df(walk):
    """Function for turning a Markov walk sequence into a DataFrame of note placement predictions"""
    sequence = []
    for step in walk:
        sequence.append(step.split(","))
    constant = ['notes_type_0', 'notes_lineIndex_0', 'notes_lineLayer_0',
                    'notes_cutDirection_0', 'notes_type_1', 'notes_lineIndex_1', 'notes_lineLayer_1',
                    'notes_cutDirection_1', 'notes_type_3', 'notes_lineIndex_3',
                    'notes_lineLayer_3', 'notes_cutDirection_3']
    df = pd.DataFrame(sequence, columns = constant)
    return df

def HMM_notes_writer(beat_list, difficulty, version):
    """Writes a list of notes based on a Hidden Markov Model walk."""
    #Load model
    if version == 1:
        with open(f"../../models/HMM_{difficulty}.pkl", 'rb') as m:
            MC = pickle.load(m)
    elif version == 2:
        with open(f"../../models/HMM_{difficulty}_v2.pkl", 'rb') as m:
            MC = pickle.load(m)
    #Set note placement rate dependent on difficulty level
    counter = 2
    beats = []
    rate = None
    if difficulty == 'easy':
        rate = 3
    elif difficulty == 'normal':
        rate = 2
    else:
        rate = 1
    while counter <= len(beat_list):
        beats.append(counter)
        counter += rate
    #Get HMM walk long enough to cover number of beats
    random_walk = MC.walk()
    while len(random_walk) < len(beats):
        random_walk = MC.walk()
    df_walk = walk_to_df(random_walk)
    #Combine beat numbers with HMM walk steps
    df_preds = pd.concat([pd.DataFrame(beats, columns = ['_time']), df_walk], axis = 1, sort = True)
    df_preds.dropna(axis = 0, inplace = True)
    #Write notes dictionaries
    notes_list = []
    for index, row in df_preds.iterrows():
        for x in list(filter(lambda y: y.startswith('notes_type'), df_preds.columns)):
            if row[x] != '999':
                num = x[-1]
                note = {'_time': row['_time'],
                        '_lineIndex': int(row[f"notes_lineIndex_{num}"]),
                        '_lineLayer': int(row[f"notes_lineLayer_{num}"]),
                        '_type': num,
                        '_cutDirection': int(row[f"notes_cutDirection_{num}"])}
                notes_list.append(note)
   #Remove potential notes that come too early in the song:
    for i, x in enumerate(notes_list):
        if notes_list[i]['_time'] >= 0 and notes_list[i]['_time'] <= 1.5:
            del notes_list[i]
        elif notes_list[i]['_time'] > beats[-1]:
            del notes_list[i]

    return notes_list

#Segmented HMM Functions

def segmented_HMM_mapper(song_path, song_name, difficulty, k = 5, version = 2):
    """This function generates a custom map based on a HMM model that operates on song segments. First, Laplacian song segmentation is performed to identify similar portions of songs. Then, a HMM is used to generate a block sequence through the first of each of these identified song segments. If that segment is repeated later in the song, the block sequence will be repeated."""
    #Load song and get beat features
    print("Loading Song...")
    bpm, beat_times, y, sr = beat_features(song_path)
    beat_times = [x*bpm/60 for x in beat_times]
    print("Song loaded successfully!")
    #Write lists for note placement, event placement, and obstacle placement
    print("Mapping with segmented Hidden Markov Model...")
    notes_list = segmented_HMM_notes_writer(y, sr, k, difficulty, version)
    events_list = events_writer(beat_times)
    obstacles_list = obstacles_writer(beat_times, difficulty)
    print("Mapping done!")
    #Write and zip files
    print("Writing files to disk...")
    write_info(song_name, bpm, difficulty)
    write_level(difficulty, events_list, notes_list, obstacles_list)
    print("Converting music file...")
    music_file_converter(song_path)
    print("Zipping folder...")
    zip_folder_exporter(song_name, difficulty)
    print("Finished! Look for zipped folder in your current path, unzip the folder, and place in the 'CustomMusic' folder in the Beat Saber directory")

def laplacian_segmentation(y, sr, k = 5):
    """This function uses the Laplacian Segmentation method described in McFee and Ellis, 2014, and adapted from example code in the librosa documentation. It returns the segment boundaries (in frame number and time and segment ID's of isolated music file segments."""
    BINS_PER_OCTAVE = 12 * 3
    N_OCTAVES = 7
    C = librosa.amplitude_to_db(np.abs(librosa.cqt(y=y, sr=sr,
                                        bins_per_octave=BINS_PER_OCTAVE,
                                        n_bins=N_OCTAVES * BINS_PER_OCTAVE)),
                                        ref=np.max)
    tempo, beats = librosa.beat.beat_track(y=y, sr=sr, trim=False)
    Csync = librosa.util.sync(C, beats, aggregate=np.median)

    # For plotting purposes, we'll need the timing of the beats
    # we fix_frames to include non-beat frames 0 and C.shape[1] (final frame)
    beat_times = librosa.frames_to_time(librosa.util.fix_frames(beats,
                                                                x_min=0,
                                                                x_max=C.shape[1]),
                                                                sr=sr)

    R = librosa.segment.recurrence_matrix(Csync, width=3, mode='affinity',
                                          sym=True)
    # Enhance diagonals with a median filter (Equation 2)
    df = librosa.segment.timelag_filter(scipy.ndimage.median_filter)
    Rf = df(R, size=(1, 7))
    mfcc = librosa.feature.mfcc(y=y, sr=sr)
    Msync = librosa.util.sync(mfcc, beats)
    path_distance = np.sum(np.diff(Msync, axis=1)**2, axis=0)
    sigma = np.median(path_distance)
    path_sim = np.exp(-path_distance / sigma)
    R_path = np.diag(path_sim, k=1) + np.diag(path_sim, k=-1)
    deg_path = np.sum(R_path, axis=1)
    deg_rec = np.sum(Rf, axis=1)
    mu = deg_path.dot(deg_path + deg_rec) / np.sum((deg_path + deg_rec)**2)
    A = mu * Rf + (1 - mu) * R_path
    L = scipy.sparse.csgraph.laplacian(A, normed=True)
    # and its spectral decomposition
    evals, evecs = scipy.linalg.eigh(L)
    # We can clean this up further with a median filter.
    # This can help smooth over small discontinuities
    evecs = scipy.ndimage.median_filter(evecs, size=(9, 1))
    # cumulative normalization is needed for symmetric normalize laplacian eigenvectors
    Cnorm = np.cumsum(evecs**2, axis=1)**0.5
    # If we want k clusters, use the first k normalized eigenvectors.
    # Fun exercise: see how the segmentation changes as you vary k
    k = k
    X = evecs[:, :k] / Cnorm[:, k-1:k]
    KM = sklearn.cluster.KMeans(n_clusters=k)
    seg_ids = KM.fit_predict(X)
    bound_beats = 1 + np.flatnonzero(seg_ids[:-1] != seg_ids[1:])
    # Count beat 0 as a boundary
    bound_beats = librosa.util.fix_frames(bound_beats, x_min=0)
    # Compute the segment label for each boundary
    bound_segs = list(seg_ids[bound_beats])
    # Convert beat indices to frames
    bound_frames = beats[bound_beats]
    # Make sure we cover to the end of the track
    bound_frames = librosa.util.fix_frames(bound_frames,
                                           x_min=None,
                                           x_max=C.shape[1]-1)
    bound_times = librosa.frames_to_time(bound_frames)
    bound_times = [(x/60)*tempo for x in bound_times]
    beat_numbers = list(range(len(bound_frames)))
    bound_beats = np.append(bound_beats, list(range(len(beats)))[-1])
    segments = list(zip(zip(bound_times, bound_times[1:]), zip(bound_beats, bound_beats[1:]), bound_segs))

    return segments, beat_times, tempo

def segments_to_df(segments):
    """Helper function to translate a song semgmenation to a pandas DataFrame."""
    lengths = []
    for seg in segments:
        length = seg[1][1] - seg[1][0]
        lengths.append(length)
    df = pd.concat([pd.Series(lengths, name = 'length'), pd.Series([x[2] for x in segments], name = 'seg_no')], axis = 1)
    return df

def segment_predictions(segment_df, HMM_model):
    """This function predicts a Markov chain walk for each segment of a segmented music file. It will repeat a walk for segments that it has already mapped previously (truncating or extending as necessary)."""

    preds = pd.DataFrame([])
    completed_segments = {}
    for index, row in segment_df.iterrows():
        if row['seg_no'] not in completed_segments.keys():
            if index == 0:
                pred = HMM_model.walk()
                while len(pred) < row['length']:
                    pred = HMM_model.walk()
                completed_segments.update({row['seg_no']: {'start':0, 'end': len(pred)}})
                preds = pd.concat([preds, pd.Series(pred[0: row['length']])], axis = 0, ignore_index = True)

            else:
                try:
                    pred = HMM_model.walk(init_state = tuple(preds.iloc[-5:, 0]))
                    while len(pred) < row['length']:
                        pred = HMM_model.walk(init_state = tuple(preds.iloc[-5:, 0]))
                    completed_segments.update({row['seg_no']: {'start': len(preds)+1, 'end': len(preds)+len(pred)}})
                    preds = pd.concat([preds, pd.Series(pred[0: row['length']])], axis = 0, ignore_index = True)
                except:
                    pred = HMM_model.walk()
                    while len(pred) < row['length']:
                        pred = HMM_model.walk()
                    completed_segments.update({row['seg_no']: {'start': len(preds)+1, 'end': len(preds)+len(pred)}})
                    preds = pd.concat([preds, pd.Series(pred[0: row['length']])], axis = 0, ignore_index = True)

        else:
            if row['length'] <= (completed_segments[row['seg_no']]['end'] - completed_segments[row['seg_no']]['start']):
                pred = preds.iloc[completed_segments[row['seg_no']]['start']: completed_segments[row['seg_no']]['start'] + row['length'], 0]
                preds = pd.concat([preds, pred], axis = 0, ignore_index = True)
            else:
                try:
                    extend = HMM_model.walk(init_state = tuple(preds.iloc[completed_segments[row['seg_no']]['end'] - 5 : completed_segments[row['seg_no']]['end'], 0]))
                    pred = preds.iloc[completed_segments[row['seg_no']]['start']: completed_segments[row['seg_no']]['end'], 0]
                    diff = row['length'] - len(pred)
                    pred = pd.concat([pred, pd.Series(extend[0: diff+1])], axis = 0, ignore_index = True)
                    completed_segments.update({row['seg_no']: {'start': len(preds)+1, 'end': len(preds)+len(pred)}})
                    preds = pd.concat([preds, pred], axis = 0, ignore_index = True)
                except:
                    extend = HMM_model.walk()
                    pred = preds.iloc[completed_segments[row['seg_no']]['start']: completed_segments[row['seg_no']]['end'], 0]
                    diff = row['length'] - len(pred)
                    pred = pd.concat([pred, pd.Series(extend[0: diff+1])], axis = 0, ignore_index = True)
                    completed_segments.update({row['seg_no']: {'start': len(preds)+1, 'end': len(preds)+len(pred)}})
                    preds = pd.concat([preds, pred], axis = 0, ignore_index = True)

    preds_list = list(preds.iloc[:, 0])
    preds = walk_to_df(preds_list)
    return preds

def segmented_HMM_notes_writer(y, sr, k, difficulty, version = 2):
    """This function writes the list of notes based on the segmented HMM model."""
    #Load model:
    if version == 1:
        with open(f"../../models/HMM_{difficulty}.pkl", 'rb') as m:
            MC = pickle.load(m)
    elif version == 2:
        with open(f"../../models/HMM_{difficulty}_v2.pkl", 'rb') as m:
            MC = pickle.load(m)

    segments, beat_times, tempo = laplacian_segmentation(y, sr, k)
    segments_df = segments_to_df(segments)
    preds = segment_predictions(segments_df, MC)
    #Combine beat numbers with HMM walk steps
    beats = [(x/60)* tempo for x in beat_times]
    df_preds = pd.concat([pd.DataFrame(beats, columns = ['_time']), preds], axis = 1, sort = True)
    df_preds.dropna(axis = 0, inplace = True)
    #Write notes dictionaries
    notes_list = []
    for index, row in df_preds.iterrows():
        for x in list(filter(lambda y: y.startswith('notes_type'), df_preds.columns)):
            if row[x] != '999':
                num = x[-1]
                note = {'_time': row['_time'],
                        '_lineIndex': int(row[f"notes_lineIndex_{num}"]),
                        '_lineLayer': int(row[f"notes_lineLayer_{num}"]),
                        '_type': num,
                        '_cutDirection': int(row[f"notes_cutDirection_{num}"])}
                notes_list.append(note)
    #Remove potential notes that come too early in the song:
    for i, x in enumerate(notes_list):
        if notes_list[i]['_time'] >= 0 and notes_list[i]['_time'] <= 1.5:
            del notes_list[i]
        elif notes_list[i]['_time'] > beats[-1]:
            del notes_list[i]

    return notes_list

#Rate Modulated Segmented HMM mapping functions
def rate_modulated_segmented_HMM_mapper(song_path, song_name, difficulty, version = 2, k = 5):
    """This function generates the files for a custom map using a rate modulated segmented HMM model."""
    #Load song and get beat features
    print("Loading Song...")
    bpm, beat_times, y, sr = beat_features(song_path)
    print("Song loaded successfully!")
    #Write lists for note placement, event placement, and obstacle placement
    print("Mapping with rate modulated segmented Hidden Markov Model...")
    notes_list, modulated_beat_list = rate_modulated_segmented_HMM_notes_writer(y, sr, k, difficulty, version)
    events_list = events_writer(modulated_beat_list)
    obstacles_list = obstacles_writer(modulated_beat_list, difficulty)
    print("Mapping done!")
    #Write and zip files
    print("Writing files to disk...")
    write_info(song_name, bpm, difficulty)
    write_level(difficulty, events_list, notes_list, obstacles_list)
    print("Converting music file...")
    music_file_converter(song_path)
    print("Zipping folder...")
    zip_folder_exporter(song_name, difficulty)
    print("Finished! Look for zipped folder in your current path, unzip the folder, and place in the 'CustomMusic' folder in the Beat Saber directory")

def choose_rate(db, difficulty):
    """
    This function modulates the block placement rate by using the average amplitude (i.e., 'loudness') across beats to choose how many blocks per beat will be placed. Takes in the difficulty level and the amplitude and returns an integer in the set {0, 1, 2, 4, 8, 16}.

    If you are finding that your maps are too fast or too slow for you, you might want to play with the probabilities in this file.
    """
    db = np.abs(db)
    p = None
    if difficulty.casefold() == 'easy'.casefold():
        if db > 70:
            p = [0.95, 0.05, 0, 0, 0, 0]
        elif db <= 70 and db > 55:
            p = [0.90, 0.10, 0, 0, 0, 0]
        elif db <= 55 and db > 45:
            p = [0.80, 0.2, 0, 0, 0, 0]
        elif db <= 45 and db > 35:
            p = [0.4, 0.5, 0.1, 0, 0, 0]
        else:
            p = [0.3, 0.6, 0.1, 0, 0, 0]
    elif difficulty.casefold() == 'normal'.casefold():
        if db > 70:
            p = [0.95, 0.05, 0, 0, 0, 0]
        elif db <= 70 and db > 55:
            p = [0.5, 0.5, 0, 0, 0, 0]
        elif db <= 55 and db > 45:
            p = [0.3, 0.7, 0, 0, 0, 0]
        elif db <= 45 and db > 35:
            p = [0.2, 0.7, 0.1, 0, 0, 0]
        else:
            p = [0.05, 0.7, 0.25, 0, 0, 0]
    elif difficulty.casefold() == 'hard'.casefold():
        if db > 70:
            p = [0.95, 0.05, 0, 0, 0, 0]
        elif db <= 70 and db > 55:
            p = [0.5, 0.5, 0, 0, 0, 0]
        elif db <= 55 and db > 45:
            p = [0.2, 0.6, 0.2, 0, 0, 0]
        elif db <= 45 and db > 35:
            p = [0.1, 0.5, 0.4, 0, 0, 0]
        else:
            p = [0.05, 0.35, 0.6, 0, 0, 0]
    elif difficulty.casefold() == 'expert'.casefold():
        if db > 70:
            p = [0.8, 0.2, 0, 0, 0, 0]
        elif db <= 70 and db > 55:
            p = [0.2, 0.7, 0.1, 0, 0, 0]
        elif db <= 55 and db > 50:
            p = [0.1, 0.4, 0.3, 0.2, 0, 0]
        elif db <= 50 and db > 45:
            p = [0, 0.05, 0.6, 0.35, 0, 0]
        else:
            p = [0, 0, 0.35, 0.65, 0, 0]
    elif difficulty.casefold() == 'expertPlus'.casefold():
        if db > 70:
            p = [0, 0.5, 0.4, 0.1, 0, 0]
        elif db <= 70 and db > 55:
            p = [0, 0.3, 0.6, 0.1, 0, 0]
        elif db <= 55 and db > 50:
            p = [0, 0.1, 0.6, 0.3, 0, 0]
        elif db <= 50 and db > 45:
            p = [0, 0.05, 0.1, 0.6, 0.25, 0]
        else:
            p = [0, 0, 0, 0.5, 0.3, 0.2]
    return np.random.choice([0, 1, 2, 4, 8, 16], p = p)

def amplitude_rate_modulation(y, sr, difficulty):
    """This function uses the average amplitude (i.e., 'loudness') of a beat and the difficulty level to determine
    how many blocks will be placed within the beat. Returns a list of beat numbers."""
    #Make amplitude matrix
    D = np.abs(librosa.stft(y))
    db = librosa.amplitude_to_db(D, ref=np.max)
    #Get beat frames and sync with amplitudes
    tempo, beat_frames = librosa.beat.beat_track(y, sr, trim = False)
    beat_db = pd.DataFrame(librosa.util.sync(db, beat_frames, aggregate = np.mean))
    #Mean amplitude per beat
    avg_beat_db = beat_db.mean()
    #Choose rates and smooth rate transitions
    rates = [0]
    counter = 1
    while counter < len(avg_beat_db)-1:
        rate = choose_rate(np.mean([avg_beat_db.iloc[counter-1], avg_beat_db.iloc[counter], avg_beat_db.iloc[counter+1]]), difficulty)
        diff = np.abs(rate - rates[-1])
        if difficulty.casefold() == 'expert'.casefold():
            maxdiff = 4
        elif difficulty.casefold() == 'expertPlus'.casefold():
            maxdiff = 8
        else:
            maxdiff = 2
        while diff > maxdiff:
            rate = choose_rate(np.mean([avg_beat_db.iloc[counter-1], avg_beat_db.iloc[counter], avg_beat_db.iloc[counter+1]]), difficulty)
            diff = rates[-1] - rate
        if rate == 4 and rates[-1] == 4: #and rates[-2] == 4:
            rate = np.random.choice([0, 1, 2])
        rates.append(rate)
        counter +=1
    #Make list of beat numbers based on rates
    beat_num_list = []
    for ind, val in enumerate(rates):
        if val == 0:
            continue
        elif val == 1:
            beat_num_list.append(ind)
        else:
            num_list = [ind, ind+1]
            for x in range(1, val):
                num_list.append(ind+(x/val))
            for y in num_list:
                beat_num_list.append(y)
    beat_num_list = list(set(beat_num_list))
    beat_num_list.sort()
    return beat_num_list

def segments_to_df_rate_modulated(segments, modulated_beat_list):
    """This function returns a DataFrame of the number of blocks needed for each song segment."""
    expanded_beat_list = []
    for x in segments:
        for y in modulated_beat_list:
            if y > x[1][0] and y <= x[1][1]:
                expanded_beat_list.append({'_time': y, 'segment': x[2]})

    df = pd.DataFrame([], columns = ['length', 'seg_no'])
    counter = 0
    first = None
    last = None
    while counter < len(expanded_beat_list):
        if counter == 0:
            first = counter
            counter += 1
        elif expanded_beat_list[counter]['segment'] != expanded_beat_list[counter-1]['segment']:
            first = counter
            counter += 1
        elif expanded_beat_list[counter] == expanded_beat_list[-1]:
            length = len(expanded_beat_list[first: -1])
            df = df.append(pd.DataFrame({'length': length, 'seg_no': expanded_beat_list[-1]['segment']}, index = [0]))
            break
        elif expanded_beat_list[counter]['segment'] == expanded_beat_list[counter+1]['segment']:
            counter += 1
        elif expanded_beat_list[counter]['segment'] != expanded_beat_list[counter+1]['segment']:
            last = counter
            length = len(expanded_beat_list[first: last+1])
            df = df.append(pd.DataFrame({'length': length, 'seg_no': expanded_beat_list[counter]['segment']}, index = [0]))
            counter += 1

    return df

def rate_modulated_segmented_HMM_notes_writer(y, sr, k, difficulty, version):
    """Function to write the notes to a list after predicting with the rate modulated segmented HMM model."""
    #Load model:
    if version == 1:
        with open(f"../../models/HMM_{difficulty}.pkl", 'rb') as m:
            MC = pickle.load(m)
    elif version == 2:
        with open(f"../../models/HMM_{difficulty}_v2.pkl", 'rb') as m:
            MC = pickle.load(m)

    segments, beat_times, bpm = laplacian_segmentation(y, sr, k)
    modulated_beat_list = amplitude_rate_modulation(y, sr, difficulty)
    segments_df = segments_to_df_rate_modulated(segments, modulated_beat_list)
    preds = segment_predictions(segments_df, MC)
    #Combine beat numbers with HMM walk steps
    beat_times = [(x/60)*bpm for x in beat_times]
    beat_count = list(range(len(beat_times)))
    beats = pd.concat([pd.Series(beat_times, name = '_time'), pd.Series(beat_count, name = 'beat_count')], axis = 1)
    for index, value in beats.iterrows():
        if value['beat_count'] not in modulated_beat_list:
            beats.drop(index = index, inplace=True)
    merged_beats = pd.merge(left = beats, right = pd.Series(modulated_beat_list, name = 'beat_count'), how='outer', on='beat_count', sort = True)
    merged_beats.interpolate(inplace=True)
    merged_beats.drop(columns = 'beat_count', inplace = True)

    df_preds = pd.concat([merged_beats, preds], axis = 1, sort = True)
    df_preds.dropna(axis = 0, inplace = True)
    #Write notes dictionaries
    notes_list = []
    for index, row in df_preds.iterrows():
        for x in list(filter(lambda y: y.startswith('notes_type'), df_preds.columns)):
            if row[x] != '999':
                num = x[-1]
                note = {'_time': row['_time'],
                        '_lineIndex': int(row[f"notes_lineIndex_{num}"]),
                        '_lineLayer': int(row[f"notes_lineLayer_{num}"]),
                        '_type': num,
                        '_cutDirection': int(row[f"notes_cutDirection_{num}"])}
                notes_list.append(note)
    #Remove potential notes that come too early in the song:
    for i, x in enumerate(notes_list):
        if notes_list[i]['_time'] >= 0 and notes_list[i]['_time'] <= 1.5:
            del notes_list[i]
        elif notes_list[i]['_time'] > beat_times[-1]:
            del notes_list[i]

    return notes_list, modulated_beat_list

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('song_path', metavar='path', type=str, help='File Path to song file')
    parser.add_argument('song_name', type=str, help='Name of song to be displayed in Beat Saber')
    parser.add_argument('difficulty', type=str, help="Desired difficulty level: 'easy', 'normal', 'hard', 'expert', or 'expertPlus'")
    parser.add_argument('model', type=str, help="Desired model for mapping: 'random', 'HMM', 'segmented_HMM', 'rate_modulated_segmented_HMM'")
    parser.add_argument('-k', type=int, help="Number of expected segments for segmented model. Default 5", default=5, required=False)
    parser.add_argument('--version', type=int, help="Version of HMM model to use: 1 (90% rating or greater) or 2 (70% rating or greater)", default=2, required=False)

    args = parser.parse_args()

    beat_map_synthesizer(args.song_path, args.song_name, args.difficulty, args.model, args.k, args.version)


#!/usr/bin/python3

import numpy as np
import pandas as pd

import random

from instance import Instance
from solver import Solver

import matplotlib
import matplotlib.pyplot as plt
from matplotlib.colors import colorConverter
from matplotlib.patches import Patch
from mpl_toolkits.axes_grid1 import make_axes_locatable
import matplotlib.patches as mpatches

from matplotlib import cm
from matplotlib import colors as mcolors

import visualizer
import histogram

import math

from numpy.random import default_rng

from scipy.stats import t
from scipy.stats import levy_stable
from scipy.stats import cauchy

from sklearn.preprocessing import StandardScaler
from sklearn.decomposition import PCA
import umap

import pylfit # perform a comparision with PRIDE implementation of LFIT

#---------------------------------------------------------------------
def singlecell_matrix_experiment():

    random.seed(42)
    np.random.seed(42)

    ####################################################################
    ###### read cell types
    ####################################################################
    file_name = '../dataset/IMAGINE/cell_types.csv'
    df_cell_types = pd.read_csv(file_name, index_col=0)
    df_cell_types.rename(columns={'cellType_final': 'Label'}, inplace=True)
    print(df_cell_types.head())

    ####################################################################
    ###### read macro cell types
    ####################################################################
    file_name = '../dataset/IMAGINE/cell_types_macro.csv'
    df_cell_types_macro = pd.read_csv(file_name, index_col=0)
    df_cell_types_macro.rename(columns={'cellType_macro': 'Label'}, inplace=True)
    print(df_cell_types_macro.head())

    ####################################################################
    ###### read the normalized dataset
    ####################################################################
    file_name = '../../../Learning/IMAGINE_dataset/dataset/ctls_normalised_counts.csv'

    df_normalised = pd.read_csv(file_name, index_col = 0).T
    print(df_normalised.head())

    ####################################################################
    ###### read the discretized dataset
    ####################################################################

    file_name = '../dataset/IMAGINE/IMAGINE_normalised_discrete_adaptive.csv'

    # Load the discrete dataset
    df = pd.read_csv(file_name, index_col=0)

    # select only the T cells
    #df = df.loc[df_cell_types_macro.loc[df_cell_types_macro['Label']=='T'].index]

    # remove axis name (name of the previously first column)
    df.rename_axis('Barcode', axis=0, inplace=True)

    # select a sub dataset
    # df = df.loc[cd4_negatives+cd4_positives]

    print(df.head())

    # classification of CD8+
    inst = Instance.create_cluster_instance(df, df_cell_types, 'NK')
    #inst = Instance.create_cluster_instance(df, df_cell_types, 'CD14')
    instance_name = 'NK'

    # inst = Instance.create_random_instance(df, 0.3)

    # extract positive/negative solutions for solution visualisation
    positive_cells = inst._pos_samples
    negative_cells = inst._neg_samples

    print('Original instance:')
    # print('n atoms: ', inst.n_atoms())
    print('n positive samples: ', inst.n_positives())
    print('n negative samples: ', inst.n_negatives())

    solver = Solver(inst)

    ############################################
    ### computation of supported front
    ############################################

    body_length = 20

    # create line equations from these points
    x1, y1 = 0.,0.
    x2, y2 = (inst.n_positives(), inst.n_negatives())
    a = 1.
    b = (x1-x2)/(y2-y1)
    c = -a*x1 - b*y1

    max_dist = a*0+b*y2+c

    # compute a line at distance 0.3
    p031 = [0,0]
    p031[1] = (0.3*max_dist-c)/b

    p032 = [0,y2]
    p032[0] = (0.3*max_dist-c-b*p032[1])/a

    # compute the list of distances to the diagonal
    atom_distance = [-1 for ind_atom in range(inst.n_atoms())]
    for ind_atom in range(inst.n_atoms()):
        score = inst.atom_score[ind_atom]
        atom_distance[ind_atom] = (a*score[0] + b*score[1] + c)/max_dist
    sorted_atoms = [ind for ind in range(inst.n_atoms())]
    sorted_atoms.sort(key=lambda ind:atom_distance[ind], reverse=True)

    stop_ind = 0
    for elt in sorted_atoms:
        if atom_distance[elt] < 0.3:
            break
        else:
            stop_ind += 1

    sorted_atoms = sorted_atoms[:stop_ind]

    #selected_atoms, scores = solver.select_k_best_atoms(50)
    selected_atoms, scores = solver.select_best_atoms_threshold(0.3)

    # print([(inst.get_atom(ind), atom_distance[ind]) for ind in sorted_atoms])

    print('n atoms selected: ', len(selected_atoms))
    #for elt in [(inst.get_atom(ind), atom_distance[ind]) for ind in selected_atoms]:
    #    print(elt)

    body_length = 20
    body = selected_atoms[:body_length]

    # plot the body
    print([inst.get_atom(ind) for ind in body])

    # create an histogram to visualize rule error over samples
    histo = histogram.Histogram(inst, body)

    # plot the histogram
    fig, ax = plt.subplots()
    ax.set_title('Histograms for ' + instance_name + ' classification')
    visualizer.plot_histograms(ax, histo, True)

    nPos = inst.n_positives()
    nNeg = inst.n_negatives()

    # plot each atom score
    fig,ax = plt.subplots()

    ax.set_xlim((0,nPos))
    ax.set_ylim((0,nNeg))

    #ax.set_aspect('equal')

    # ax.set_title('Atoms score')
    # ax.set_xlabel('Positive score')
    # ax.set_ylabel('Negative score')

    ax.set_title('Erreurs des atomes logiques (classification de NK)')
    ax.set_xlabel('Erreur positive')
    ax.set_ylabel('Erreur négative')

    # colors = ['royalblue' for _ in range(len(inst.atom_score))]
    ax.scatter([elt[0] for elt in inst.atom_score], [elt[1] for elt in inst.atom_score], marker='x')

    # y / x
    # ax.plot([0, inst.n_positives()], [0, inst.n_negatives()], color='red')

    # 0.3 threshold
    # ax.plot([p031[0], p032[0]], [p031[1], p032[1]], color='lightseagreen')

    # compute atom distance to the diagonal vs its lateral position on the diagonal
    nPos = inst.n_positives()
    nNeg = inst.n_negatives()

    a_diag = nNeg
    b_diag = -nPos
    c_diag = 0

    # diagonal
    ax.plot([0, nPos], [0, -a_diag*nPos/b_diag], color='red')

    ax.legend(['score = 0.0'], loc='lower right')

    # orthogonal projection
    x_mid = nPos/2
    y_mid = nNeg/2

    # compute the relative distances do diagonal and normal for each atoms
    atom_relative_distances = [(0,0) for _ in range(inst.n_atoms())]

    max_dist = -a_diag*0 -b_diag*nNeg

    k = math.sqrt(a_diag**2 + b_diag**2)
    for atom_index in range(inst.n_atoms()):
        pos = inst.atom_score[atom_index]
        c = -a_diag*pos[0] - b_diag*pos[1]
        dist_diag = c/max_dist
        if c > 0:
            x1 = 0
            y1 = (-a_diag*x1-c)/b_diag
            y2 = nNeg
            x2 = (-b_diag*y2-c)/a_diag
        else:
            x2 = nPos
            y2 = (-a_diag*x2-c)/b_diag
            y1 = 0
            x1 = (-b_diag*y1-c)/a_diag

        dist_left = math.sqrt((pos[0]-x1)**2 + (pos[1]-y1)**2)
        length = math.sqrt((x2-x1)**2 + (y2-y1)**2)
        lateral_dist = dist_left/length

        atom_relative_distances[atom_index] = (dist_diag, lateral_dist)

    colors = [max(inst.atom_score[ind_atom][0],inst.atom_score[ind_atom][1])/(inst.atom_score[ind_atom][0]+inst.atom_score[ind_atom][1]) for ind_atom in range(inst.n_atoms())]

    fig, ax = plt.subplots()
    ax.scatter([elt[0] for elt in atom_relative_distances], [elt[1] for elt in atom_relative_distances], marker='x', c=colors)
    ax.set_xlabel('distance to diagonal')
    ax.set_ylabel('lateral position')

    # add histograms

    # verif: atom with colors
    # plot each atom score
    fig,ax = plt.subplots()
    # ax.set_aspect('equal')
    ax.set_xlim((0,nPos))
    ax.set_ylim((0,nNeg))
    ax.set_title('Atoms score')
    ax.set_xlabel('Positive score')
    ax.set_ylabel('Negative score')

    ax.scatter([elt[0] for elt in inst.atom_score], [elt[1] for elt in inst.atom_score], marker='x', c=[elt[1] for elt in atom_relative_distances])

    ax.plot([0, nPos], [0, nNeg], color='red')

    selected_atoms = [ind for ind in range(inst.n_atoms())]
    selected_atoms.sort(key=lambda ind: atom_relative_distances[ind][0], reverse=False)

    fig, axs = plt.subplots(3)
    axs[0].hist([elt[0] for elt in atom_relative_distances], edgecolor='k', bins=100)
    axs[1].hist([elt[1] for elt in atom_relative_distances], edgecolor='k', bins=100)
    axs[2].hist([max(inst.atom_score[ind_atom][0],inst.atom_score[ind_atom][1])/(inst.atom_score[ind_atom][0]+inst.atom_score[ind_atom][1]) for atom_index in range(inst.n_atoms())], edgecolor='k', bins=100)

    plt.show()

###############################################################
def artificial_data_experiment():

    nPos = 1000
    nNeg = 1000

    nVar = 10000
    nAtoms = nVar*2

    # random_seed = 46
    random_seed = 52 # worst dataset

    random.seed(random_seed)
    np.random.seed(random_seed)

    # sampling over the normal distribution

    rng = default_rng(random_seed)

    # mu, sigma = 0, 0.1 # mean and standard deviation
    # sample_normal = rng.normal(mu, sigma, nAtoms)

    sample_uniform = rng.uniform(0, 1, nVar*2)

    df = 10
    sample_score = t.rvs(df, size=nVar)
    # sample_student = [elt for elt in sample_student if elt >= 0]

    # https://en.wikipedia.org/wiki/Stable_distribution
    # https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.levy_stable.html
    # stable distribution
    # alpha, beta = 1.995, 0.
    # sample_score = levy_stable.rvs(alpha, beta, size=nVar)

    # x0 = 0
    # gamma = 0.5
    # global sample_score
    # sample_score = cauchy.rvs(loc=x0, scale=gamma, size=nVar)

    # normal is normalized
    max_normal = np.max(sample_score)
    min_normal = np.min(sample_score)

    print(min_normal, max_normal)

    max_abs = max(max_normal, -min_normal) + 0.2
    for ind in range(len(sample_score)):
        sample_score[ind] = sample_score[ind]/max_abs
        sample_score[ind] *= 0.70
    # threshold = 50
    # for ind in range(len(sample_score)):
    #     sample_score[ind] = sample_score[ind]/threshold

    fig, axs = plt.subplots(2)
    axs[0].hist(sample_score, edgecolor='k', bins=100)
    axs[1].hist(sample_uniform, edgecolor='k', bins=100)

    #print(sample_normal)

    # fig,ax = plt.subplots()
    # ax.scatter(sample_normal, sample_uniform, marker='x')
    # ax.set_xlim([-1, 1])

    a = nNeg
    b = -nPos

    atoms = [(ind_var, 0) for ind_var in range(nVar)]
    atoms += [(ind_var, 1) for ind_var in range(nVar)]

    atom_scores = []

    for ind_atom in range(len(sample_score)):
        c = sample_score[ind_atom]*nPos*nNeg
        if c < 0:
            p1 = (-c/nPos, 0)
            p2 = (nPos, nPos+c/nNeg)
        else:
            p1 = (0, c/nPos)
            p2 = (nNeg-c/nPos, nNeg)
        v = (p2[0]-p1[0], p2[1]-p1[1])
        lateral = sample_uniform[ind_atom]
        p_star = (p1[0]+v[0]*lateral, p1[1]+v[1]*lateral)
        p_star = (math.floor(p_star[0]), math.floor(p_star[1]))
        atom_scores.append(p_star)

    # generating positiv and negative examples
    global positives
    global negatives
    positives = [ [0 for _ in range(nVar)] for _ in range(nPos) ]
    negatives = [ [0 for _ in range(nVar)] for _ in range(nNeg) ]

    for ind_atom in range(nVar):
        score = atom_scores[ind_atom]

        # positive examples
        pos_indexes = [ind for ind in range(nPos)]
        np.random.shuffle(pos_indexes)
        for ind_pos in range(nPos):
            if ind_pos < score[0]:
                positives[pos_indexes[ind_pos]][ind_atom] = 1
            else:
                positives[pos_indexes[ind_pos]][ind_atom] = 0

        # negative examples
        neg_indexes = [ind for ind in range(nNeg)]
        np.random.shuffle(neg_indexes)
        for ind_neg in range(nNeg):
            if ind_neg < score[1]:
                negatives[neg_indexes[ind_neg]][ind_atom] = 1
            else:
                negatives[neg_indexes[ind_neg]][ind_atom] = 0

    atom_scores_sym = []
    for score in atom_scores:
        atom_scores_sym.append( (nPos-score[0], nNeg-score[1]) )
    atom_scores += atom_scores_sym

    fig,ax = plt.subplots()
    ax.scatter([elt[0] for elt in atom_scores], [elt[1] for elt in atom_scores], marker='x')
    ax.plot([0, nPos], [0, nNeg], color='red')
    ax.set_xlim([0, nPos])
    ax.set_ylim([0, nNeg])

    ind_atoms_sorted = [ind for ind in range(len(atom_scores))]
    score = [-nNeg*error[0] +nPos*error[1] for error in atom_scores]
    ind_atoms_sorted.sort(key = lambda ind: score[ind], reverse=True)
    fig, ax = plt.subplots()
    ax.scatter([atom_scores[ind][0] for ind in ind_atoms_sorted], [atom_scores[ind][1] for ind in ind_atoms_sorted], c=[score[ind] for ind in ind_atoms_sorted], marker='x')
    ax.set_xlim([0, nPos])
    ax.set_ylim([0, nNeg])
    ax.plot([0, nPos], [0, nNeg], color='red')

    selected_atoms = ind_atoms_sorted[:20]

    pos_histogram = [0 for _ in range(len(selected_atoms)+1)]
    neg_histogram = [0 for _ in range(len(selected_atoms)+1)]

    for pos in positives:
        score = 0
        for ind_atom in selected_atoms:
            atom = atoms[ind_atom]
            if pos[atom[0]] != atom[1]:
                score += 1
        pos_histogram[score] += 1

    for neg in negatives:
        score = 0
        for ind_atom in selected_atoms:
            atom = atoms[ind_atom]
            if neg[atom[0]] != atom[1]:
                score += 1
        neg_histogram[score] += 1

    for score in range(len(pos_histogram)):
        pos_histogram[score] /= float(nPos)

    for score in range(len(neg_histogram)):
        neg_histogram[score] /= float(nNeg)

    print('positive histogram: ')
    print(pos_histogram)

    print('negative histogram:')
    print(neg_histogram)

    height = 10
    output = ''

    for y in range(height, 0, -1):
        for score in pos_histogram:
            if score*height*2 >= y:
                output += '*'
            else:
                output += ' '
        output += '\n'
    for score in pos_histogram:
        output += '.'
    print(output)

    output = ''
    for y in range(height, 0, -1):
        for score in neg_histogram:
            if score*height*2 >= y:
                output += '*'
            else:
                output += ' '
        output += '\n'
    for score in neg_histogram:
        output += '.'
    print(output)

    # create a csv file with the matrix (ordered positive and negative examples)
    matrix = positives + negatives
    variables = ['v_'+str(ind) for ind in range(nVar)]
    examples = ['s_' + str(ind) for ind in range(nPos+nNeg)]

    # exportation into a dataframe
    # global dataframe
    # dataframe = pd.DataFrame(matrix, columns = variables, index= examples)
    # dataframe.to_csv('../dataset/artificial_matrix_worst.csv')

    return

#------------------------------------------------------------------------------
def compute_pride_bodies(instance):

    # computation of the bodies of all rules returned from pride algorithm on the instance

    variable = -1
    value = -1 # not important, just for rules'head

    nb_features = instance.dataset.shape[0] # nb columns in the dataframe

    positives = instance.dataset.loc[instance._pos_samples].values.tolist()
    negatives = instance.dataset.loc[instance._neg_samples].values.tolist()

    verbose = 0
    output = pylfit.algorithms.pride.PRIDE.fit_var_val(variable, value, nb_features, positives, negatives, verbose)

    # compute the bodies of PRIDE rules: indexes of the selected atoms from the instance of NK cells
    pride_bodies = []
    for rule in output:

        pride_body = []

        for var_index in rule._body_variables:
            gene = instance.dataset.columns[var_index]
            value = rule._body_values[var_index]
            atom_index = instance.get_atom_index((gene, value))
            pride_body.append(atom_index)

        pride_bodies.append(pride_body)

    return pride_bodies

#------------------------------------------------------------------------------
def LOLH_artificial_data():

    random.seed(42)
    np.random.seed(42)

    ####################################################################
    ###### read the discretized dataset
    ####################################################################

    file_name = '../dataset/artificial_matrix.csv'
    # file_name = '../dataset/artificial_matrix_worst.csv'

    # Load the discrete dataset
    df = pd.read_csv(file_name, index_col=0)

    print(df.head())

    # compute a UMAP from the data
    X = df.values.copy()
    # standard scaling before PCA
    X=StandardScaler().fit_transform(X)
    # PCA with ten principal components
    X_pca = PCA(n_components=10).fit_transform(X)

    reducer = umap.UMAP(min_dist=0.3,n_neighbors=50,spread=1.0)
    embedding = reducer.fit_transform(X_pca)

    fig, axs = plt.subplots(1,2)

    col = ['forestgreen' for _ in range(int(df.shape[0]/2))] + ['darkred' for _ in range(int(df.shape[0]/2))]

    axs[0].scatter(embedding[:, 0], embedding[:, 1], s=2, c=col)
    axs[0].set_aspect('equal', 'datalim')
    axs[0].set_title('projection UMAP')
    axs[0].set_xlabel('UMAP 1')
    axs[0].set_ylabel('UMAP 2')
    axs[0].legend(handles=[mpatches.Patch(color='forestgreen', label='exemples positifs'), mpatches.Patch(color='darkred', label='exemples négatifs')], loc='lower right')

    # PCA projection
    X = df.values.copy()
    X=StandardScaler().fit_transform(X)
    X_pca = PCA(n_components=2).fit_transform(X)
    axs[1].scatter(X_pca[:,0], X_pca[:,1], c=col, s=2)
    axs[1].set_title('projection PCA')
    axs[1].set_xlabel('PCA 1')
    axs[1].set_ylabel('PCA 2')
    axs[1].legend(handles=[mpatches.Patch(color='forestgreen', label='exemples positifs'), mpatches.Patch(color='darkred', label='exemples négatifs')])

    # classification from positive vs negative
    pos_samples = ['s_'+str(ind) for ind in range(0, int(df.shape[0]/2.))]
    neg_samples = ['s_'+str(ind) for ind in range(int(df.shape[0]/2.)+1, df.shape[0])]

    instance_name = 'artificial'

    inst = Instance.create_instance_explicit(df, pos_samples, neg_samples)

    # extract positive/negative solutions for solution visualisation
    positive_cells = inst._pos_samples
    negative_cells = inst._neg_samples

    print('Original instance:')
    # print('n atoms: ', inst.n_atoms())
    print('n positive samples: ', inst.n_positives())
    print('n negative samples: ', inst.n_negatives())

    solver = Solver(inst)

    ############################################
    ### computation of supported front
    ############################################

    body_length = 20

    # create line equations from these points
    x1, y1 = 0.,0.
    x2, y2 = (inst.n_positives(), inst.n_negatives())
    a = 1.
    b = (x1-x2)/(y2-y1)
    c = -a*x1 - b*y1

    max_dist = a*0+b*y2+c

    # compute a line at distance 0.3
    p031 = [0,0]
    p031[1] = (0.3*max_dist-c)/b

    p032 = [0,y2]
    p032[0] = (0.3*max_dist-c-b*p032[1])/a

    # compute the list of distances to the diagonal
    atom_distance = [-1 for ind_atom in range(inst.n_atoms())]
    for ind_atom in range(inst.n_atoms()):
        score = inst.atom_score[ind_atom]
        atom_distance[ind_atom] = (a*score[0] + b*score[1] + c)/max_dist
    sorted_atoms = [ind for ind in range(inst.n_atoms())]
    sorted_atoms.sort(key=lambda ind:atom_distance[ind], reverse=True)

    stop_ind = 0
    for elt in sorted_atoms:
        if atom_distance[elt] < 0.3:
            break
        else:
            stop_ind += 1

    sorted_atoms = sorted_atoms[:stop_ind]

    # selection score = 0.75
    selection_score = 0.75

    selected_atoms, scores = solver.select_best_atoms_threshold(selection_score)

    print('n atoms selected: ', len(selected_atoms))

    body = selected_atoms

    # plot the body
    print([inst.get_atom(ind) for ind in body])

    nPos = inst.n_positives()
    nNeg = inst.n_negatives()

    ########################################################
    # compute pride rules
    pride_bodies = compute_pride_bodies(inst)
    pride_atoms = list(np.unique([ind_atom for pride_body in pride_bodies for ind_atom in pride_body]))
    print('n rules pride: ', len(pride_bodies))
    print('rule 0: ', pride_bodies[0])

    # inspection des indices des atomes des règles pride

    ###

    # number of atoms for pride rules
    # n_atoms = [len(elt) for elt in pride_bodies]
    # print('number of atoms pride: ', np.bincount(n_atoms))
    # print(sorted(n_atoms))

    # plot each atom score
    fig,ax = plt.subplots()

    ax.set_xlim((0,nPos))
    ax.set_ylim((0,nNeg))

    #ax.set_aspect('equal')

    ax.set_title('Atoms score')
    ax.set_xlabel('erreur positive')
    ax.set_ylabel('erreur négative')

    other_atoms = [ind_atom for ind_atom in range(inst.n_atoms()) if not ind_atom in body and not ind_atom in pride_atoms]

    # plot LOLH atoms
    ax.scatter([inst.atom_score[ind_atom][0] for ind_atom in body], [inst.atom_score[ind_atom][1] for ind_atom in body], marker='o', facecolor='darkorange', edgecolor='black', linewidth=0.5, zorder=2, label='atomes LOLH')

    # plot pride atoms
    ax.scatter([inst.atom_score[ind_atom][0] for ind_atom in pride_atoms], [inst.atom_score[ind_atom][1] for ind_atom in pride_atoms], marker='s', facecolor='dodgerblue', edgecolor='black', linewidth=0.5, zorder=2, label='atomes PRIDE')

    # other (not used) atoms
    ax.scatter([inst.atom_score[ind_atom][0] for ind_atom in range(inst.n_atoms())], [inst.atom_score[ind_atom][1] for ind_atom in range(inst.n_atoms())], alpha=0.5, marker='x', zorder=0, label='autres atomes')

    # score = 0.0
    ax.plot([0, nPos], [0, nNeg], '--', color='black', label='score = 0.0', lw=1.5)

    # compute the line corresponding to the selection score
    nPos = inst.n_positives()
    nNeg = inst.n_negatives()

    # line for score = 0.75
    A = (0,selection_score*nNeg)
    B = (nPos*(1.-selection_score), nNeg)
    ax.plot([A[0], B[0]], [A[1], B[1]], '--', color='forestgreen', label='score = 0.75', lw=1.5)

    ax.legend(loc='lower right')

    # compute a global score for each rule
    pride_rule_scores = []
    for pride_body in pride_bodies:
        score = (0, 0)
        for atom_index in pride_body:
            atom_score = inst.atom_score[atom_index]
            score = (score[0] + atom_score[0]/nPos, score[1] + atom_score[1]/nNeg)
        score = (score[0] / len(pride_body), score[1] / len(pride_body))
        pride_rule_scores.append(score)

    # compute the rule score for the LOLH rule
    lolh_rule_score = (0,0)
    for index in body:
        atom_score = inst.atom_score[index]
        lolh_rule_score = (lolh_rule_score[0] + atom_score[0]/nPos, lolh_rule_score[1] + atom_score[1]/nNeg)
    lolh_rule_score = (lolh_rule_score[0]/len(body), lolh_rule_score[1]/len(body))

    ############################################
    # plot the rule scores
    # display the normalized scores for all rules (between (0,0) and (n_positives, n_negatives))
    fig, ax = plt.subplots()

    ax.set_xlim((0, 1.))
    ax.set_ylim((0, 1.))
    ax.set_xlabel('score positif')
    ax.set_ylabel('score négatif')
    # ax.set_title('Comparaison des règles de PRIDE et LOLH')

    # diagonal indicating independant values
    ax.plot([0, 1], [0, 1], color='black', zorder=0)

    # plot the pride rules
    pride_indexes = [ind for ind in range(len(pride_bodies))]
    pride_indexes.sort(key=lambda ind: len(pride_bodies[ind]))
    pride_rule_scores_sorted = [pride_rule_scores[ind] for ind in pride_indexes] # sort all the points
    longest_rule = np.max([len(body) for body in pride_bodies])
    colors = [len(pride_bodies[ind]) for ind in pride_indexes]
    cnorm = mcolors.Normalize(vmin=0, vmax=longest_rule)

    ax.scatter([score[0] for score in pride_rule_scores_sorted], [score[1] for score in pride_rule_scores_sorted], c=colors, norm=cnorm, cmap=plt.get_cmap('viridis'), marker='x',  zorder=1, label='règles PRIDE', facecolor='darkblue')
    cbar = fig.colorbar(cm.ScalarMappable(norm=cnorm, cmap=plt.get_cmap('viridis')), ax=ax)
    cbar.set_label('longueur des règles PRIDE')
    ax.scatter(lolh_rule_score[0], lolh_rule_score[1], marker='*', s=70, color='chocolate', zorder=1, label='règle LOLH') # plot the optimized rule
    # plot the domination cone of the optimized rule
    ax.plot([lolh_rule_score[0], lolh_rule_score[0]], [lolh_rule_score[1], 0], linestyle='--', color='darkgrey', zorder=0)
    ax.plot([lolh_rule_score[0], 1.], [lolh_rule_score[1], lolh_rule_score[1]], linestyle='--', color='darkgrey', zorder=0)
    legend = ax.legend(loc='lower right')
    legend.legendHandles[0].set_color('darkcyan')

    ######################################################
    # create an histogram to visualize rule error over samples
    histo = histogram.Histogram(inst, body)
    # plot the histogram
    fig, ax = plt.subplots()
    #ax.set_title('Histograms for ' + instance_name + ' classification')
    visualizer.plot_histograms(ax, histo, True)

    ######################################################
    # histogramme de PRIDE à partir de l'apprentissage sur un sous jeu de données
    # histo = histogram.Histogram(inst, body)
    # # plot the histogram
    # fig, ax = plt.subplots()
    # #ax.set_title('Histograms for ' + instance_name + ' classification')
    # visualizer.plot_histograms(ax, histo, True)

    # histogramme de la meilleure règle PRIDE
    scores_pride = []
    ind_pride_bodies_sorted =  list(range(len(pride_bodies)))
    ind_pride_bodies_sorted.sort(key=lambda ind: pride_rule_scores[ind][1] - pride_rule_scores[ind][0], reverse=True)
    # print('pride scores sorted: ', [pride_rule_scores[ind][1] - pride_rule_scores[ind][0] for ind in ind_pride_bodies_sorted])

    ind_best_pride_body = ind_pride_bodies_sorted[0]
    print('best pride body: ', pride_bodies[ind_best_pride_body])

    # study atom indexes in pride bodies
    atom_indexes_pride = [inst.get_atom(ind) for body in pride_bodies for ind in body]
    print('pride body indexes (unique): ', np.unique(atom_indexes_pride, axis=0))

    # occ = np.bincount(atom_indexes_pride)
    # print('occurences: ')
    # for ind_atom in range(len(occ)):
    #     print(ind_atom, ', ', inst.atom_score[ind_atom][1]/nNeg - inst.atom_score[ind_atom][0]/nPos, ', ', occ[ind_atom])

    # fig, ax_ = plt.subplots()
    # ax_.scatter([inst.atom_score[ind_atom][0]/nPos for ind_atom in range(len(occ))], [inst.atom_score[ind_atom][1]/nPos for ind_atom in range(len(occ))], c=occ)
    # ax_.set_title('indexes des atomes de PRIDE')
    # ax_.set_xlabel('positive score')
    # ax_.set_ylabel('negative score')
    # cbar = fig.colorbar(cm.ScalarMappable(cmap=plt.get_cmap('viridis')), ax=ax_)
    # cbar.set_label('nb occurences')

    # histo = histogram.Histogram(inst, pride_bodies[ind_best_pride_body])
    # fig, ax = plt.subplots()
    # ax.set_title('Meilleure règle PRIDE')
    # visualizer.plot_histograms(ax, histo, True)

    ################################################################
    # histogramme de la plus longue règle PRIDE et de la meilleure règle PRIDE
    longest_pride = pride_bodies[0]
    for pride_body in pride_bodies:
        if len(pride_body) > len(longest_pride):
            longest_pride = pride_body
    histo1 = histogram.Histogram(inst, longest_pride)
    histo2 = histogram.Histogram(inst, pride_bodies[ind_best_pride_body])

    # plot the histogram
    fig, axs = plt.subplots(1, 2, gridspec_kw={'width_ratios': [1,0.9]})

    fig.tight_layout(w_pad=2)

    # ax.set_title('Histogrammes pour PRIDE')
    visualizer.plot_histograms(axs[0], histo1, True)
    axs[0].set_title(r'Plus longue règle PRIDE')

    visualizer.plot_histograms(axs[1], histo2, True)
    axs[1].set_title(r'Meilleure règle PRIDE')

    ################################################################
    # visualisation des règles en multi-objectif
    body_length = 10
    biobj_score, biobj_bodies, biobj_weight = solver.compute_supported(body_length, 1, 1)

    biobj_score = [ (score[0]/float(inst.n_positives()), score[1]/float(inst.n_negatives())) for score in biobj_score]

    # sélection de deux règles "extrêmes"
    ind_left = 0
    for ind in range(len(biobj_score)):
        score = biobj_score[ind]
        if score[0] >= -0.2 and score[0] <= 0.2:
            if score[1] >= 6.4 and score[1] <= 6.6:
                print('left: ', ind)
                ind_left = ind
    print('left body: ', biobj_bodies[ind_left])
    left_score = biobj_score[ind_left]

    ind_right = 0
    for ind in range(len(biobj_score)):
        score = biobj_score[ind]
        if score[0] >= 2.9 and score[0] <= 3.0:
            if score[1] >= 9.8 and score[1] <= 10.0:
                print('right: ', ind)
                ind_right = ind
    print('right body: ', biobj_bodies[ind_right])
    right_score = biobj_score[ind_right]

    fig, ax = plt.subplots()
    # ax.set_title('Score des règles bi-objectives')
    ax.set_xlabel('score positif')
    ax.set_ylabel('score négatif')
    ind_reduced = [ind for ind in range(len(biobj_score)) if ind != ind_left and ind != ind_right]
    ax.scatter([biobj_score[ind][0] for ind in ind_reduced], [biobj_score[ind][1] for ind in ind_reduced], marker='x', s=25)
    ax.scatter([biobj_score[ind_left][0], biobj_score[ind_right][0]], [biobj_score[ind_left][1], biobj_score[ind_right][1]], color='red', marker='x', s=30)
    ax.text(5-0.3, 5+0.3, 'score négatif - score positif = 0', horizontalalignment='center', verticalalignment='center', rotation=45)
    ax.plot([0, body_length], [0, body_length], '--', color='k')

    ax.text(left_score[0] + 0.5, left_score[1]-0.2, r'$p_1$')
    ax.text(right_score[0] + 0.0, right_score[1]-0.6, r'$p_2$')

    delta = 0.5
    ax.set_xlim((0-delta, body_length+delta))
    ax.set_ylim((0-delta, body_length+delta))
    ax.set_aspect('equal')

    histo_left = histogram.Histogram(inst, biobj_bodies[ind_left])
    histo_right = histogram.Histogram(inst, biobj_bodies[ind_right])

    fig,axs = plt.subplots(1,2)
    fig.tight_layout(w_pad=3)
    axs[0].set_title(r'Histogramme de la règle $p_1$')
    visualizer.plot_histograms(axs[0], histo_left, True)
    axs[1].set_title(r'Histogramme de la règle $p_2$')
    visualizer.plot_histograms(axs[1], histo_right, True)

    ################################################################
    # display several rules with histograms

    fig, ax = plt.subplots()
    ax.set_xlim([0, nPos])
    ax.set_ylim([0, nNeg])
    ax.scatter([inst.atom_score[ind][0] for ind in range(inst.n_atoms())], [inst.atom_score[ind][1] for ind in range(inst.n_atoms())], marker='x', zorder=1, alpha=0.5)
    # ax.set_title('Visualisation de plusieurs règles logiques')
    # ax.plot([0, nPos], [0, nNeg], '--', color='grey', zorder=3, alpha=0.7)
    ax.set_xlabel('erreur positive')
    ax.set_ylabel('erreur négative')

    fig, axsh = plt.subplots(2,2)
    fig.tight_layout(h_pad=4)

    ##########################################
    def compute_A_B(score, nPos, nNeg):
        return (0, score*nNeg), ((1.-score)*nPos, nNeg)

    score1 = 0.3
    score2 = 0.4
    A1, B1 = compute_A_B(score1, nPos, nNeg)
    A2, B2 = compute_A_B(score2, nPos, nNeg)
    selected_atoms, scores = solver.select_best_atoms_threshold(0.0)
    atoms_sandwich = [selected_atoms[ind] for ind in range(len(selected_atoms)) if scores[ind] >= score1 and scores[ind] <= score2]
    np.random.shuffle(atoms_sandwich)
    nAtomsSel = 10
    body_worst = atoms_sandwich[:nAtomsSel]
    print('body: ', body_worst)
    ax.scatter([inst.atom_score[ind_atom][0] for ind_atom in body_worst], [inst.atom_score[ind_atom][1] for ind_atom in body_worst], marker='x', color='orangered', zorder=4)
    # histogramme
    histo = histogram.Histogram(inst, body_worst)
    axsh[0,0].set_title('Règle sous-optimale')
    visualizer.plot_histograms(axsh[0,0], histo, True)

    ax.fill([A1[0], B1[0], B2[0], A2[0]], [A1[1], B1[1], B2[1], A2[1]], zorder=3, color='green', alpha=0.3, label='règle sous-optimale')
    ax.plot([A1[0], B1[0]], [A1[1], B1[1]], '--', color='k', zorder=3)
    ax.plot([A2[0], B2[0]], [A2[1], B2[1]], '--', color='k', zorder=3)

    score1 = -0.8
    score2 = -0.7
    A1, B1 = compute_A_B(score1, nPos, nNeg)

    A2, B2 = compute_A_B(score2, nPos, nNeg)

    selected_atoms, scores = solver.select_best_atoms_threshold(-1.)
    atoms_sandwich = [selected_atoms[ind] for ind in range(len(selected_atoms)) if scores[ind] >= score1 and scores[ind] <= score2]
    np.random.shuffle(atoms_sandwich)
    nAtomsSel = 10
    body_worst = atoms_sandwich[:nAtomsSel]
    print('body: ', body_worst)
    ax.scatter([inst.atom_score[ind_atom][0] for ind_atom in body_worst], [inst.atom_score[ind_atom][1] for ind_atom in body_worst], marker='x', color='orangered', zorder=4)
    # histogramme
    histo = histogram.Histogram(inst, body_worst)
    axsh[0,1].set_title('Règle inverse')
    visualizer.plot_histograms(axsh[0,1], histo, True)

    ax.fill([A1[0], B1[0], B2[0], A2[0]], [A1[1], B1[1], B2[1], A2[1]], zorder=3, color='orange', alpha=0.3, label='règle inverse')
    ax.plot([A1[0], B1[0]], [A1[1], B1[1]], '--', color='k', zorder=3)
    ax.plot([A2[0], B2[0]], [A2[1], B2[1]], '--', color='k', zorder=3)

    score1 = -0.05
    score2 = 0.05
    A1, B1 = compute_A_B(score1, nPos, nNeg)
    A2, B2 = compute_A_B(score2, nPos, nNeg)
    selected_atoms, scores = solver.select_best_atoms_threshold(-1.)
    atoms_sandwich = [selected_atoms[ind] for ind in range(len(selected_atoms)) if scores[ind] >= score1 and scores[ind] <= score2]
    np.random.shuffle(atoms_sandwich)
    nAtomsSel = 10
    body_worst = atoms_sandwich[:nAtomsSel]
    print('body: ', body_worst)
    ax.scatter([inst.atom_score[ind_atom][0] for ind_atom in body_worst], [inst.atom_score[ind_atom][1] for ind_atom in body_worst], marker='x', color='orangered', zorder=4, label='atomes sélectionnés')
    # histogramme
    histo = histogram.Histogram(inst, body_worst)
    axsh[1,0].set_title('Règle inappropriée')
    visualizer.plot_histograms(axsh[1,0], histo, True)

    ax.fill([A1[0], B1[0], B2[0], A2[0]], [A1[1], B1[1], B2[1], A2[1]], zorder=3, color='red', alpha=0.3, label='règle inappropriée')
    ax.plot([A1[0], B1[0]], [A1[1], B1[1]], '--', color='k', zorder=3)
    ax.plot([A2[0], B2[0]], [A2[1], B2[1]], '--', color='k', zorder=3)

    atomes = [ind for ind in range(inst.n_atoms())]
    np.random.shuffle(atomes)
    body_random = atomes[:10]
    histo = histogram.Histogram(inst, body_random)
    axsh[1,1].set_title('Règle aléatoire')
    visualizer.plot_histograms(axsh[1,1], histo, True)

    legend = ax.legend(loc='lower left')

    handles_old, labels_old = ax.get_legend_handles_labels()

    new_handles = []
    for ind_handle in range(len(handles_old)-1):
        patch = handles_old[ind_handle]
        new_handles.append(Patch(edgecolor='k', facecolor=colorConverter.to_rgba(patch.get_facecolor(), alpha=0.3), linewidth=0.5, label=labels_old[ind_handle]))
    new_handles.append(handles_old[-1])
    ax.legend(loc='lower left', handles=new_handles, labels=labels_old)

    ################################################################
    # K-fold validation on PRIDE and LOLH
    k = 10

    # create nsamples/k subset of the data
    pos_indexes = inst._pos_samples
    pos_sub_indexes = [[] for _ in range(k)]
    for ind in range(len(pos_indexes)):
        k_prime = int((ind/nPos)*k)
        pos_sub_indexes[k_prime].append(pos_indexes[ind])

    neg_indexes = inst._neg_samples
    neg_sub_indexes = [[] for _ in range(k)]
    for ind in range(len(neg_indexes)):
        k_prime = int((ind/nNeg)*k)
        neg_sub_indexes[k_prime].append(neg_indexes[ind])

    LOLH_perf = [[],[]]
    PRIDE_perf = [[],[]]

    for ind_k in range(1):

        # create the training and test set
        #ind_K = 0
        train_pos = pos_sub_indexes[ind_k]
        train_neg = neg_sub_indexes[ind_k]
        test_pos = []
        test_neg = []
        for ind in range(k):
            if ind != ind_k:
                test_pos += pos_sub_indexes[ind]
                test_neg += neg_sub_indexes[ind]

        # test_pos = train_pos
        # test_neg = train_neg

        df_train = df.loc[train_pos+train_neg]
        df_test = df.loc[test_pos+test_neg]

        # training set
        inst_train = Instance.create_instance_explicit(df_train, train_pos, train_neg)

         # test set
        inst_test = Instance.create_instance_explicit(df_test, test_pos, test_neg)

        # compute the model on the training set
        solver = Solver(inst_train)

        # computation of the solution with LOLH and PRIDE
        # the atom indexes are converted from the train instance to the test instance

        # compute LOLH body
        LOLH_body, scores = solver.select_best_atoms_threshold(0.75)
        LOLH_body = [inst_train.get_atom(atom_index) for atom_index in LOLH_body]
        LOLH_body = [inst_test.get_atom_index(atom) for atom in LOLH_body]

        # compute PRIDE bodies
        PRIDE_bodies = compute_pride_bodies(inst_train)
        PRIDE_bodies = [[inst_train.get_atom(atom_index) for atom_index in pride_body] for pride_body in PRIDE_bodies]
        PRIDE_bodies = [[inst_test.get_atom_index(atom) for atom in PRIDE_body] for PRIDE_body in PRIDE_bodies]

        # atoms in the train and test instance
        atoms_train_test = [inst_train.get_atom(atom_index) for atom_index in range(inst_train.n_atoms()) if inst_test.has_atom(inst_train.get_atom(atom_index))]
        col = []
        for atom in atoms_train_test:
            ind_atom_train = inst_train.get_atom_index(atom)
            score = -inst_train.atom_score[ind_atom_train][0]*len(train_neg)+inst_train.atom_score[ind_atom_train][1]*len(train_pos)
            col.append(score)

        # fig, ax = plt.subplots()
        # ax.scatter([inst_test.atom_score[inst_test.get_atom_index(atom)][0] for atom in atoms_train_test], [inst_test.atom_score[inst_test.get_atom_index(atom)][1] for atom in atoms_train_test], marker='x', c=col)
        # ax.set_xlim([0, len(test_pos)])
        # ax.set_ylim([0, len(test_neg)])
        # ax.plot([0, len(test_pos)], [0, len(test_neg)], '--', color='k')

        # histogramme de LOLH
        # histo = histogram.Histogram(inst_test, LOLH_body)
        # fig, ax = plt.subplots()
        # ax.set_title('histo LOLH données test')
        # visualizer.plot_histograms(ax, histo, True)

        # LOLH performance
        pos_match = 0
        neg_match = 0
        for pos in test_pos:
            matching_error = 0
            for ind_atom in LOLH_body:
                atom = inst_test.get_atom(ind_atom)
                if df_test[atom[0]][pos] != atom[1]:
                    matching_error += 1
            if matching_error < len(LOLH_body)/2.:
                pos_match += 1
        for neg in test_neg:
            matching_error = 0
            for ind_atom in LOLH_body:
                atom = inst_test.get_atom(ind_atom)
                if df_test[atom[0]][neg] != atom[1]:
                    matching_error += 1
            if matching_error < len(LOLH_body)/2.:
                neg_match += 1
        # print('pos match: ', pos_match)
        # print('neg match: ', neg_match)
        LOLH_perf[0].append(float(pos_match)/len(test_pos))
        LOLH_perf[1].append(float(neg_match)/len(test_neg))

        # PRIDE performance
        pos_match = 0
        neg_match = 0

        for pos in test_pos:
            match = False
            for pride_body in PRIDE_bodies:
                matching_error = 0
                for ind_atom in pride_body:
                    atom = inst_test.get_atom(ind_atom)
                    if df_test[atom[0]][pos] != atom[1]:
                        matching_error += 1
                if matching_error == 0:
                    match = True
                    pos_match += 1
                    break
        for neg in test_neg:
            match = False
            for pride_body in PRIDE_bodies:
                matching_error = 0
                for ind_atom in pride_body:
                    atom = inst_test.get_atom(ind_atom)
                    if df_test[atom[0]][neg] != atom[1]:
                        matching_error += 1
                if matching_error == 0:
                    match = True
                    neg_match += 1
                    break
        PRIDE_perf[0].append(pos_match/len(test_pos))
        PRIDE_perf[1].append(neg_match/len(test_neg))

    print(LOLH_perf)
    print(PRIDE_perf)

    return

# singlecell_matrix_experiment()

# artificial_data_experiment()

LOLH_artificial_data()