import numpy as np

def parse(data):
    low = [int(d[0].split('-')[0]) for d in data]
    high = [int(d[0].split('-')[1]) for d in data]
    seq = [d[2] for d in data]
    char = [d[1][0] for d in data]
    count = [count_obs(d[2], d[1][0]) for d in data]
    
    return low, high, seq, char, count

def valid1(low, high, count):
    return low <= count <= high

def valid1_alt(x):
    return x[0] <= x[2] <= x[1]

def get_valid_idx(X):
    return np.where(np.apply_along_axis(valid1_alt, 1, X))[0]

def valid2(seq, char, pos):
    x, y = pos[0], pos[1]
    valid = False

    cond1 = (seq[x-1] == char)
    cond2 = (seq[y-1] == char)
    if sum([cond1, cond2]) == 1:
        valid = True
    return valid

def test_1():
    cases = [1, 3, 'a', 'abcde', True], [1, 3, 'b', 'cdefg', False], [2, 9, 'c', 'ccccccccc', True]
    counts = [c[3].count(c[2]) for c in cases]
    for (i, case) in enumerate(cases):
        assert valid1(case[0], case[1], counts[i]) == case[-1]
        
def test_2():
    cases = [1, 3, 'a', 'abcde', True], [1, 3, 'b', 'cdefg', False], [2, 9, 'c', 'ccccccccc', False]

    for (i, case) in enumerate(cases):
        assert valid2(case[3], case[2], [case[0], case[1]]) == case[-1]

if __name__ == '__main__':
    # read data
    data = np.genfromtxt('input.txt', dtype = str)

    # part 1
    low, high, seq, char, count = parse(data)

    test_1()
    res = 0
    for i in range(len(seq)):
        res += valid1(low[i], high[i], count[i])
    
    print(f'Solution: {solution}')

    # part 2
    pos = np.array(list(zip(low, high))).astype('int')

    res = []
    for i in range(len(seq)):
        res.append(valid2(seq[i], char[i], pos[i]))

    test_2()
    print(f'Solution: {sum(res)}')

    # part 1 (alt)
    X = np.stack([low, high, count], axis = 1).astype('int')

    idx = get_valid_idx(X)
    solution = len(idx)

    print(f'Solution: {solution}')