"""
Compute rules from midi files. 
Could be slightly different from online compute from piano roll because saved midi is a cleaner version of piano rolls.
"""
import os
import pretty_midi
import pandas as pd
import numpy as np
from tqdm import tqdm
import math
import glob
import torch as th
import multiprocessing
from argparse import ArgumentParser
from guided_diffusion import midi_util
import torch.nn.functional as F

def main():
    parser = ArgumentParser()
    parser.add_argument('--root_dir', type=str, default='loggings/cond_table/all/beam_50_1_2_cls_1',
                    help='Path to the folder that contains generated samples for rule guidance')
    parser.add_argument('--rule_name', type=str, default='chord_progression',
                    help='rule to compute from midis')
    args = parser.parse_args()

    gen_files = sorted(glob.glob(f'{args.root_dir}/*.midi') + glob.glob(f'{args.root_dir}/*.mid'))
    orig_dir = f'{args.root_dir}/gt'
    if args.rule_name is None:
        target_rules = ['pitch_hist', 'note_density', 'chord_progression']
    else:
        target_rules = [args.rule_name]

    all_results = pd.DataFrame()
    for file in tqdm(gen_files):
        gen_midi = pretty_midi.PrettyMIDI(file)
        gen_piano_roll = gen_midi.get_piano_roll(fs=100, pedal_threshold=None, onset=False)
        gen_piano_roll = th.from_numpy(gen_piano_roll)[None, None] / 63.5 - 1
        gen_piano_roll = F.pad(gen_piano_roll, (0, 1024 - gen_piano_roll.shape[3]), "constant", -1)
        basename = os.path.basename(file)
        orig_file = os.path.join(orig_dir, basename)
        try:
            orig_midi = pretty_midi.PrettyMIDI(orig_file)
        except:
            print(basename)
            continue
        orig_piano_roll = orig_midi.get_piano_roll(fs=100, pedal_threshold=None, onset=False)
        orig_piano_roll = th.from_numpy(orig_piano_roll)[None, None] / 63.5 - 1
        orig_piano_roll = F.pad(orig_piano_roll, (0, 1024 - orig_piano_roll.shape[3]), "constant", -1)

        results = midi_util.compute_rule(gen_piano_roll, orig_piano_roll, target_rules)
        all_results = pd.concat([all_results, results], ignore_index=True)

    all_results.to_csv(f'{args.root_dir}/results_computed.csv', index=False)


if __name__ == "__main__":
    multiprocessing.set_start_method('spawn', force=True)
    main()