import random
from datetime import datetime

keys = {
    'C major': [1, 4, 7, 19, 50, 59, 73, 75, 98, 116, 117, 128, 136],
    'C minor': [24, 36, 83, 90, 111],
    'D major': [10, 39, 40, 54, 56, 57, 97, 105, 107, 121, 127, 143],
    'D minor': [14, 18, 31, 32, 37, 126, 130, 134],
    'E major': [45, 104, 120, 149],
    'E minor': [8, 26, 58, 67, 91, 95, 99, 109],
    'E-flat major': [70, 84, 124, 133, 140, 145],
    'F major': [20, 25, 28, 29, 64, 68, 76, 79, 82, 92, 118, 119, 122, 123, 129, 146, 148],
    'F minor': [53],
    'F-sharp minor': [44, 135],
    'G major': [11, 12, 21, 30, 35, 42, 43, 46, 52, 60, 101, 103, 106, 112, 114, 132],
    'G minor': [2, 6, 16, 27, 38, 41, 49, 55, 65, 77, 80, 81, 86, 89, 96, 102, 131, 150],
    'A major': [3, 34, 69, 78, 85, 88, 93],
    'A minor': [5, 9, 13, 15, 22, 63, 71, 72, 108, 110, 125, 138, 139, 141, 144],
    'B minor': [17, 66, 74, 100, 113, 137, 147],
    'B-flat major': [33, 47, 48, 51, 61, 62, 94, 115, 142],
    'B-flat minor': [23, 87],
}

key_reverse = {}
for key, lines in keys.items():
    for line in lines:
        key_reverse[line] = key

three_four_time = [4, 21, 27, 31, 33, 42, 44, 52, 54, 61, 62, 75, 76, 77, 79, 82, 89, 115, 117, 128, 131, 151, 159]

def extract(l, k):
    r = []

    for kk in k:
        r.extend(l.get(kk, []))

    return r

no_accidentals = extract(keys, ['C major', 'A minor'])
one_sharp = extract(keys, ['G major', 'E minor'])
one_flat = extract(keys, ['F major', 'D minor'])
two_sharps = extract(keys, ['D major', 'B minor'])
two_flats = extract(keys, ['B-flat major', 'G minor'])
three_four_or_five_sharps_or_flats = extract(keys, [
    'E-flat major', 'C minor',
    'A-flat major', 'F minor',
    'D-flat major', 'B-flat minor',
    
    'A major', 'F-sharp minor',
    'E major', 'C-sharp minor',
    'B major', 'G-sharp minor',
])

while True:
    selection = [
        random.choice(l)
        for l in [no_accidentals, one_sharp, one_flat, two_sharps, two_flats, three_four_or_five_sharps_or_flats]
    ]

    num_minor = 0
    num_major = 0

    for line in selection:
        if key_reverse[line].endswith('minor'):
            num_minor += 1
        else:
            num_major += 1

    if num_minor == 3 and num_major == 3:
        break

print('Example exam:', selection)

data = []

for line in open('data'):
    ts1, ts2, line, hands, result = line.split()
    line, diff = map(int, line.split('-'))
    result = 1 if result == 'pass' else 0
    key = key_reverse[line]
    ts = int(datetime.strptime(ts1 + ts2, '%Y-%m-%d%H:%M:%S').timestamp())
    data.append([ts, line, diff, key, hands, result])

print('Num entries:', len(data))