from dataset_classes import DatasetFromPostgresLanes
from data import Vocab, VocabError

d = DatasetFromPostgresLanes("dbname=composer_new_dataset_v3 host=localhost", 1)
v = Vocab(1, 127, 1, 127, 1, 127, 1, 1, 96000, 8, 96000, 8, 24)

N = 1000


def percentiles(arr, ps):
    arr = list(sorted(arr))
    return [arr[int(len(arr) * p)] for p in ps]


def display_percentiles(arr, name, units):
    [low, mid, high] = percentiles(arr, [0.1, 0.5, 0.9])
    print(f"    {name}: {mid:.2f} {units} (10%: {low:.2f}, 90%: {high:.2f})")


for split in [0, 1]:
    print(f"Split {split}:")
    for permit_sets in [
        None,
        ["musescore"],
        ["bitmidi", "geocities", "lmd_full", "reddit_130k_midi", "tpb_50k", "vgmidi"],
    ]:
        print(f"  Datasets: {permit_sets if permit_sets is not None else 'all'}")
        d.permit_datasets = permit_sets
        total = 0
        succeeded = 0

        lengths = []
        max_pitch_durs = []
        max_rest_durs = []
        min_pitches = []
        max_pitches = []
        max_embed_times = []

        for i, e in zip(range(N), d.stream_examples(split, (0, 1), "random")):
            total += 1
            try:
                qm = v.quantize_midi(e.example)
                if len(qm) == 0:
                    continue

                # qm: [[pitch, velocity, qstart, qend, anchor], ...]
                max_pitch_dur = 0
                max_rest_dur = 0
                min_pitch = 127
                max_pitch = 0
                time = qm[0][2]
                for note in qm:
                    [pitch, _, start, end, _] = note
                    min_pitch = min(min_pitch, pitch)
                    max_pitch = max(max_pitch, pitch)
                    max_pitch_dur = max(max_pitch_dur, end - start)
                    max_rest_dur = max(max_rest_dur, start - time)
                    time = max(time, end)

                lengths.append(v.fast_estimate_length(e.example))
                max_pitch_durs.append(max_pitch_dur)
                max_rest_durs.append(max_rest_dur)
                min_pitches.append(min_pitch)
                max_pitches.append(max_pitch)
                max_embed_times.append(time)
                succeeded += 1
            except VocabError as ex:
                pass

        print(f"    Total: {total}, Succeeded: {succeeded} ({succeeded/total:.2f})")
        display_percentiles(lengths, "Length", "tokens")
        display_percentiles(max_pitch_durs, "Max pitch duration", "subbeats")
        display_percentiles(max_rest_durs, "Max rest duration", "subbeats")
        display_percentiles(min_pitches, "Min pitch", "MIDI pitch")
        display_percentiles(max_pitches, "Max pitch", "MIDI pitch")
        display_percentiles(max_embed_times, "Max embedding time", "subbeats")