import random
from data import Vocab, Control, Pitch, Duration, Polyphony, wavtool2midi
import pytest
import torch
import json
import math


@pytest.mark.parametrize(
    "input, expected",
    [
        (Pitch(0, 60, 28), "{C5 60 (28)}"),
        (Control(0, "pad"), "{pad}"),
        (Duration(0, 12, "1"), "{L1}"),
        (Polyphony(0, 2), "{P2}"),
    ],
)
def test_vocab_reprs(input, expected):
    assert str(input) == expected


def test_vocab():
    v = Vocab(24, 96, 24, 48, 20, 127, 2, 1, 12, 4, 4 * 4 * 12, 4, 12)
    assert str(v.pitches[60][80]) == "{C5 60 (80)}"
    assert str(v.durations[12]) == "{L1}"
    assert str(v.durations[4]) == "{L1/3}"
    assert set(s.index for s in v.symbols) == set(range(v.N))


vocab_test_inputs = [
    (
        '{"notes":[{"pitch":48,"start":0,"end":1,"velocity":0.7874015748031497,"lifted":true},{"pitch":49,"start":1,"end":2,"velocity":0.7874015748031497,"lifted":true},{"pitch":50,"start":2,"end":3,"velocity":0.7874015748031497,"lifted":true}]}',
        None,
    ),
    (
        '{"notes":[{"pitch":48,"start":1,"end":2,"velocity":0.7874015748031497,"lifted":true},{"pitch":49,"start":2,"end":3,"velocity":0.7874015748031497,"lifted":true},{"pitch":50,"start":3,"end":4,"velocity":0.7874015748031497,"lifted":true}]}',
        '{"notes":[{"pitch":40,"start":1,"end":1.125,"velocity":0.7874015748031497,"lifted":true},{"pitch":40,"start":1.75,"end":1.875,"velocity":0.7874015748031497,"lifted":true},{"pitch":41,"start":0.25,"end":0.375,"velocity":0.7874015748031497,"lifted":true},{"pitch":41,"start":0.75,"end":0.875,"velocity":0.7874015748031497,"lifted":true},{"pitch":41,"start":1.25,"end":1.375,"velocity":0.7874015748031497,"lifted":true},{"pitch":42,"start":1.5,"end":1.625,"velocity":0.7874015748031497,"lifted":true},{"pitch":43,"start":0,"end":0.125,"velocity":0.7874015748031497,"lifted":true},{"pitch":43,"start":0.5,"end":0.625,"velocity":0.7874015748031497,"lifted":true},{"pitch":48,"start":0,"end":1,"velocity":0.7874015748031497,"lifted":true},{"pitch":48,"start":1,"end":1.75,"velocity":0.7874015748031497,"lifted":true},{"pitch":50,"start":0.875,"end":1.875,"velocity":0.7874015748031497,"lifted":true},{"pitch":52,"start":0.125,"end":1,"velocity":0.7874015748031497,"lifted":true},{"pitch":52,"start":1.125,"end":2,"velocity":0.7874015748031497,"lifted":true},{"pitch":51,"start":0,"end":1,"velocity":0.7874015748031497,"lifted":true},{"pitch":47,"start":2,"end":2.75,"velocity":0.7874015748031497,"lifted":true},{"pitch":49,"start":2,"end":3,"velocity":0.7874015748031497,"lifted":true},{"pitch":54,"start":1.875,"end":2.75,"velocity":0.7874015748031497,"lifted":true},{"pitch":39,"start":2,"end":2.125,"velocity":0.7874015748031497,"lifted":true}]}',
    ),
    (
        '{"notes":[{"pitch":48,"start":0,"end":1,"velocity":0.7874015748031497,"lifted":true},{"pitch":49,"start":0,"end":1,"velocity":0.7874015748031497,"lifted":true},{"pitch":50,"start":0,"end":1,"velocity":0.7874015748031497,"lifted":true}]}',
        '{"notes":[{"pitch":46,"start":0,"end":0.125,"velocity":0.7874015748031497,"lifted":true},{"pitch":46,"start":0.125,"end":0.25,"velocity":0.7874015748031497,"lifted":true},{"pitch":46,"start":2.125,"end":2.5,"velocity":0.7874015748031497,"lifted":true},{"pitch":45,"start":0,"end":0.125,"velocity":0.7874015748031497,"lifted":true},{"pitch":45,"start":1.875,"end":2.25,"velocity":0.7874015748031497,"lifted":true},{"pitch":44,"start":0.25,"end":0.375,"velocity":0.7874015748031497,"lifted":true},{"pitch":44,"start":0.5,"end":0.625,"velocity":0.7874015748031497,"lifted":true},{"pitch":44,"start":1.375,"end":1.75,"velocity":0.7874015748031497,"lifted":true},{"pitch":43,"start":0.375,"end":0.5,"velocity":0.7874015748031497,"lifted":true},{"pitch":43,"start":1.125,"end":1.5,"velocity":0.7874015748031497,"lifted":true},{"pitch":35,"start":0,"end":0.25,"velocity":0.7874015748031497,"lifted":true},{"pitch":35,"start":1,"end":1.25,"velocity":0.7874015748031497,"lifted":true},{"pitch":36,"start":0.5,"end":0.75,"velocity":0.7874015748031497,"lifted":true},{"pitch":48,"start":0.75,"end":0.875,"velocity":0.7874015748031497,"lifted":true},{"pitch":50,"start":0.5,"end":0.625,"velocity":0.7874015748031497,"lifted":true},{"pitch":51,"start":0.375,"end":0.5,"velocity":0.7874015748031497,"lifted":true},{"pitch":51,"start":0.625,"end":0.75,"velocity":0.7874015748031497,"lifted":true},{"pitch":52,"start":0.125,"end":0.25,"velocity":0.7874015748031497,"lifted":true},{"pitch":53,"start":0,"end":0.125,"velocity":0.7874015748031497,"lifted":true},{"pitch":53,"start":0.25,"end":0.375,"velocity":0.7874015748031497,"lifted":true},{"pitch":41,"start":0.625,"end":0.75,"velocity":0.7874015748031497,"lifted":true},{"pitch":34,"start":1.5,"end":1.75,"velocity":0.7874015748031497,"lifted":true},{"pitch":47,"start":2.25,"end":2.625,"velocity":0.7874015748031497,"lifted":true}]}',
    ),
    (
        '{"notes":[{"pitch":51,"start":1.125,"end":2.125,"velocity":0.7874015748031497,"lifted":true},{"pitch":48,"start":0,"end":1,"velocity":0.7874015748031497,"lifted":true},{"pitch":49,"start":0.375,"end":1.375,"velocity":0.7874015748031497,"lifted":true},{"pitch":50,"start":0.75,"end":1.75,"velocity":0.7874015748031497,"lifted":true},{"pitch":52,"start":1.5,"end":2.5,"velocity":0.7874015748031497,"lifted":true}]}',
        None,
    ),
    (
        [
            '{"notes":[{"pitch":40,"start":1,"end":1.125,"velocity":0.7874015748031497,"lifted":true},{"pitch":40,"start":1.75,"end":1.875,"velocity":0.7874015748031497,"lifted":true},{"pitch":41,"start":0.25,"end":0.375,"velocity":0.7874015748031497,"lifted":true},{"pitch":41,"start":0.75,"end":0.875,"velocity":0.7874015748031497,"lifted":true},{"pitch":41,"start":1.25,"end":1.375,"velocity":0.7874015748031497,"lifted":true},{"pitch":42,"start":1.5,"end":1.625,"velocity":0.7874015748031497,"lifted":true},{"pitch":43,"start":0,"end":0.125,"velocity":0.7874015748031497,"lifted":true},{"pitch":43,"start":0.5,"end":0.625,"velocity":0.7874015748031497,"lifted":true},{"pitch":48,"start":0,"end":1,"velocity":0.7874015748031497,"lifted":true},{"pitch":48,"start":1,"end":1.75,"velocity":0.7874015748031497,"lifted":true},{"pitch":50,"start":0.875,"end":1.875,"velocity":0.7874015748031497,"lifted":true},{"pitch":52,"start":0.125,"end":1,"velocity":0.7874015748031497,"lifted":true},{"pitch":52,"start":1.125,"end":2,"velocity":0.7874015748031497,"lifted":true},{"pitch":51,"start":0,"end":1,"velocity":0.7874015748031497,"lifted":true},{"pitch":47,"start":2,"end":2.75,"velocity":0.7874015748031497,"lifted":true},{"pitch":49,"start":2,"end":3,"velocity":0.7874015748031497,"lifted":true},{"pitch":54,"start":1.875,"end":2.75,"velocity":0.7874015748031497,"lifted":true},{"pitch":39,"start":2,"end":2.125,"velocity":0.7874015748031497,"lifted":true}]}',
            '{"notes":[{"pitch":46,"start":0,"end":0.125,"velocity":0.7874015748031497,"lifted":true},{"pitch":46,"start":0.125,"end":0.25,"velocity":0.7874015748031497,"lifted":true},{"pitch":46,"start":2.125,"end":2.5,"velocity":0.7874015748031497,"lifted":true},{"pitch":45,"start":0,"end":0.125,"velocity":0.7874015748031497,"lifted":true},{"pitch":45,"start":1.875,"end":2.25,"velocity":0.7874015748031497,"lifted":true},{"pitch":44,"start":0.25,"end":0.375,"velocity":0.7874015748031497,"lifted":true},{"pitch":44,"start":0.5,"end":0.625,"velocity":0.7874015748031497,"lifted":true},{"pitch":44,"start":1.375,"end":1.75,"velocity":0.7874015748031497,"lifted":true},{"pitch":43,"start":0.375,"end":0.5,"velocity":0.7874015748031497,"lifted":true},{"pitch":43,"start":1.125,"end":1.5,"velocity":0.7874015748031497,"lifted":true},{"pitch":35,"start":0,"end":0.25,"velocity":0.7874015748031497,"lifted":true},{"pitch":35,"start":1,"end":1.25,"velocity":0.7874015748031497,"lifted":true},{"pitch":36,"start":0.5,"end":0.75,"velocity":0.7874015748031497,"lifted":true},{"pitch":48,"start":0.75,"end":0.875,"velocity":0.7874015748031497,"lifted":true},{"pitch":50,"start":0.5,"end":0.625,"velocity":0.7874015748031497,"lifted":true},{"pitch":51,"start":0.375,"end":0.5,"velocity":0.7874015748031497,"lifted":true},{"pitch":51,"start":0.625,"end":0.75,"velocity":0.7874015748031497,"lifted":true},{"pitch":52,"start":0.125,"end":0.25,"velocity":0.7874015748031497,"lifted":true},{"pitch":53,"start":0,"end":0.125,"velocity":0.7874015748031497,"lifted":true},{"pitch":53,"start":0.25,"end":0.375,"velocity":0.7874015748031497,"lifted":true},{"pitch":41,"start":0.625,"end":0.75,"velocity":0.7874015748031497,"lifted":true},{"pitch":34,"start":1.5,"end":1.75,"velocity":0.7874015748031497,"lifted":true},{"pitch":47,"start":2.25,"end":2.625,"velocity":0.7874015748031497,"lifted":true}]}',
        ],
        None,
    ),
    (
        '{"notes":[{"pitch":1036,"start":0,"end":1,"velocity":0.7874015748031497,"lifted":true},{"pitch":1037,"start":1,"end":2,"velocity":0.7874015748031497,"lifted":true}]}',
        None,
    ),
    (
        '{"notes":[]}',
        '{"notes":[{"pitch":40,"start":1,"end":1.125,"velocity":0.7874015748031497,"lifted":true},{"pitch":40,"start":1.75,"end":1.875,"velocity":0.7874015748031497,"lifted":true},{"pitch":41,"start":0.25,"end":0.375,"velocity":0.7874015748031497,"lifted":true},{"pitch":41,"start":0.75,"end":0.875,"velocity":0.7874015748031497,"lifted":true},{"pitch":41,"start":1.25,"end":1.375,"velocity":0.7874015748031497,"lifted":true},{"pitch":42,"start":1.5,"end":1.625,"velocity":0.7874015748031497,"lifted":true},{"pitch":43,"start":0,"end":0.125,"velocity":0.7874015748031497,"lifted":true},{"pitch":43,"start":0.5,"end":0.625,"velocity":0.7874015748031497,"lifted":true},{"pitch":48,"start":0,"end":1,"velocity":0.7874015748031497,"lifted":true},{"pitch":48,"start":1,"end":1.75,"velocity":0.7874015748031497,"lifted":true},{"pitch":50,"start":0.875,"end":1.875,"velocity":0.7874015748031497,"lifted":true},{"pitch":52,"start":0.125,"end":1,"velocity":0.7874015748031497,"lifted":true},{"pitch":52,"start":1.125,"end":2,"velocity":0.7874015748031497,"lifted":true},{"pitch":51,"start":0,"end":1,"velocity":0.7874015748031497,"lifted":true},{"pitch":47,"start":2,"end":2.75,"velocity":0.7874015748031497,"lifted":true},{"pitch":49,"start":2,"end":3,"velocity":0.7874015748031497,"lifted":true},{"pitch":54,"start":1.875,"end":2.75,"velocity":0.7874015748031497,"lifted":true},{"pitch":39,"start":2,"end":2.125,"velocity":0.7874015748031497,"lifted":true}]}',
    ),
    (
        '{"notes":[]}',
        '{"notes":[]}',
    ),
    (
        '{"notes":[]}',
        None,
    ),
]


@pytest.mark.parametrize("inputs,accompany", vocab_test_inputs)
@pytest.mark.parametrize("append_in_place", [True, False])
@pytest.mark.parametrize("put_desc_anchor", [True, False])
def test_midi_tensors(inputs, accompany, append_in_place, put_desc_anchor):
    if isinstance(inputs, str):
        inputs = [inputs]
    midis = [wavtool2midi(json.loads(input)) for input in inputs]
    if put_desc_anchor:
        for midi in midis:
            if len(midi) > 0:
                some_note = random.choice(midi)
                some_note["descriptionAnchor"] = True
    accompany_midi = (
        wavtool2midi(json.loads(accompany)) if accompany is not None else None
    )

    v = Vocab(24, 96, 24, 48, 20, 127, 2, 1, 24 * 4, 4, 4 * 4 * 24, 4, 24)

    recovered_midis_a = []
    recovered_keys_a = []
    for midi in midis:
        key = random.randint(0, 23)
        x = v.midi_to_tensor(midi, key, accompany=accompany_midi, strict=True)
        est_len = v.fast_estimate_length(midi, accompany=accompany_midi)
        assert abs(est_len - x.shape[0]) <= 2
        recovered_midi_a, recovered_key_a = v.tensor_to_midi(x)
        assert recovered_key_a == key
        recovered_midis_a.append(recovered_midi_a)
        recovered_keys_a.append(recovered_key_a)
        if not put_desc_anchor and len(recovered_midi_a) > 0:
            assert recovered_midi_a[0]["descriptionAnchor"] == True
            del recovered_midi_a[0]["descriptionAnchor"]
        assert recovered_midi_a == [
            {k: note[k] for k in note.keys() if k not in ["offVelocity", "sourceTrack"]}
            for note in midi
        ]
        if accompany is not None:
            rec_acc, rec_key = v.tensor_to_midi(x, return_accompany=True)
            assert rec_key == key
            assert rec_acc == [
                {
                    k: note[k]
                    for k in note.keys()
                    if k not in ["offVelocity", "sourceTrack"]
                }
                for note in accompany_midi
            ]

    for splices in [
        [1],
        [4],
        [1, 2, 3, 4],
        [3, 2, 4, 1],
        [8, 9, 2, 6, 7],
        [1] * 1000,
        [22, 23, 24],
    ]:
        symbolss = []
        for midi in midis:
            symbols = v.midi_to_symbols(midi, allow_anchor=True)
            symbolss.append(symbols)
        symbols_len = max(
            symbols.shape[0] if symbols is not None else 0 for symbols in symbolss
        )
        symbols = torch.zeros(len(midis), symbols_len, dtype=torch.long)
        for i, symbols_i in enumerate(symbolss):
            if symbols_i is not None:
                symbols[i, : symbols_i.shape[0]] = symbols_i

        ys = torch.tensor(
            [[[v.begin.index, 0, 0, 0, 0, 0, 0, 0, 0]]],
            dtype=torch.long,
        ).repeat(len(midis), 1 if not append_in_place else 1000, 1)

        seq_len = 1
        for splice in splices + [math.inf]:
            if symbols.shape[1] == 0:
                break
            msplice = min(splice, symbols.shape[1])
            symbols_pfx, symbols = symbols[:, :msplice], symbols[:, msplice:]
            if append_in_place:
                v.append_symbols(ys, symbols_pfx, seq_len=seq_len)
                seq_len += symbols_pfx.shape[1]
            else:
                ys = v.append_symbols(ys, symbols_pfx)

        for i, recovered_midi_a in enumerate(recovered_midis_a):
            recovered_midi_b = v.tensor_to_midi(
                ys[i, :seq_len] if append_in_place else ys[i]
            )[0]
            assert recovered_midi_a == recovered_midi_b


v0 = Vocab(24, 96, 24, 48, 20, 127, 2, 1, 24 * 4, 4, 4 * 4 * 24, 4, 24)


@pytest.mark.parametrize(
    "input",
    [
        [[0, 0, 0, 0, 0, 0, 0, 0, 1]],
        [[0, 0, 0, 0, 1, 0, 0, 0, 1]],
        [[0, 0, 0, 0, 1, 0, 0, 0, 0]],
        [[0, 0, 0, 0, 0, 0, 0, 1, 0]],
        [[0, 0, 0, 1, 0, 0, 0, 0, 0]],
        [[0, 0, 0, 1, 0, 0, 0, 1, 0], [0, 0, 0, 2, 0, 0, 0, 2, 0]],
        [
            # wrong polyphony symbol
            [v0.polyphony[2].index, 0, 0, 0, 0, 0, 0, 0, 0],
            [v0.pitches[50][126].index, 0, 0, 0, 0, 0, 0, 0, 0],
            [v0.durations[1].index, 0, 0, 0, 0, 0, 0, 0, 1],
            [v0.rest.index, 0, 0, 0, 1, 0, 0, 0, 1],
            [v0.durations[1].index, 0, 0, 0, 1, 0, 0, 1, 0],
        ],
        [
            # wrong polyphony embed
            [v0.polyphony[2].index, 0, 0, 0, 0, 0, 0, 0, 0],
            [v0.pitches[50][126].index, 0, 0, 0, 0, 0, 0, 0, 0],
            [v0.durations[1].index, 0, 0, 0, 0, 0, 0, 0, 1],
            [v0.pitches[51][126].index, 0, 0, 0, 1, 0, 0, 0, 1],
            [v0.durations[1].index, 0, 0, 0, 1, 0, 0, 0, 3],
            [v0.rest.index, 0, 0, 0, 2, 0, 0, 0, 2],
            [v0.durations[1].index, 0, 0, 0, 2, 0, 0, 1, 0],
        ],
    ],
)
def test_tensor_to_midi_asserts(input):
    with pytest.raises(Exception):
        v0.tensor_to_midi(torch.tensor(input, dtype=torch.long))


@pytest.mark.parametrize(
    "allow_anchor, input, expected",
    [
        (
            False,
            [
                {
                    "note": 60,
                    "onBeat": 0,
                    "offBeat": 1,
                    "onVelocity": 20,
                    "descriptionAnchor": True,
                }
            ],
            [[60, 20, 0, 24, False]],
        ),
        (
            True,
            [
                {
                    "note": 60,
                    "onBeat": 0.01,
                    "offBeat": 0.99,
                    "onVelocity": 126,
                    "descriptionAnchor": True,
                }
            ],
            [[60, 126, 0, 24, True]],
        ),
        (
            True,
            [{"note": 60, "onBeat": 0, "offBeat": 1.01, "onVelocity": 125}],
            [[60, 124, 0, 24, False]],
        ),
        (
            True,
            [
                {"note": 60, "onBeat": 0.5, "offBeat": 0.6, "onVelocity": 20},
                {"note": 61, "onBeat": 0.625, "offBeat": 1.5, "onVelocity": 24},
                {
                    "note": 60,
                    "onBeat": 0.5,
                    "offBeat": 1.0,
                    "onVelocity": 30,
                    "descriptionAnchor": True,
                },
                {"note": 61, "onBeat": 0.625, "offBeat": 1.0, "onVelocity": 36},
            ],
            [[60, 20, 12, 14, True], [61, 36, 15, 24, False]],
        ),
        (
            True,
            [
                {"note": 60, "onBeat": 0, "offBeat": 0.5, "onVelocity": 40},
                {"note": 60, "onBeat": 0.5, "offBeat": 1.0, "onVelocity": 50},
            ],
            [[60, 40, 0, 12, False], [60, 50, 12, 24, False]],
        ),
        (
            True,
            [
                {"note": 60, "onBeat": 0, "offBeat": 1.0, "onVelocity": 80},
                {"note": 60, "onBeat": 0.125, "offBeat": 0.5, "onVelocity": 90},
                {"note": 60, "onBeat": 0.375, "offBeat": 0.75, "onVelocity": 100},
                {"note": 60, "onBeat": 0.8, "offBeat": 0.95, "onVelocity": 110},
            ],
            [
                [60, 80, 0, 3, False],
                [60, 90, 3, 9, False],
                [60, 100, 9, 18, False],
                [60, 110, 19, 23, False],
            ],
        ),
        (
            True,
            [
                {"note": 60, "onBeat": 0.01, "offBeat": 0.51, "onVelocity": 21},
                {"note": 60, "onBeat": 0.49, "offBeat": 0.99, "onVelocity": 127},
            ],
            [[60, 20, 0, 12, False], [60, 127, 12, 24, False]],
        ),
    ],
)
def test_quantize(allow_anchor, input, expected):
    v = Vocab(24, 96, 24, 48, 20, 127, 2, 1, 24 * 4, 4, 4 * 4 * 24, 4, 24)
    assert v.quantize_midi(input, allow_anchor=allow_anchor) == expected