from aoc.input import get_input
import copy
import itertools
import time
import collections
import re
from aoc.partselector import part_one, part_two
import functools

def pw(line):
    return int(line.strip())

def p1():
    inp = get_input(pw)
    inp.append(max(inp)+3)
    print(sorted(inp))
    last = 0
    result = { 0: 0, 1: 0, 2:0, 3:0 }
    for sample in sorted(inp):
        result[sample - last] +=1
        last = sample
    print(result)
    print(result[1] * result[3])
    segments = []
    for index, sample in enumerate(sorted(inp)):
        if (sample - last) == 3:
            segments.append(index)
        last = sample
    print(segments)

    return segments


def p2(segments):
    last = 0
    count = 0
    inp = get_input(pw)
    inp.append(max(inp)+3)
    inp = sorted(inp, reverse=True)
    last = 0
    result = { 0: 0, 1: 0, 2:0, 3:0 }
    lst = [(0, inp)]
    ok = True
    parts = []
    ls = 0
    for segment in segments:
        parts.append(inp[len(inp)-segment:len(inp)-ls])
        ls = segment

    pcount = []
    lasts = 0
    for part in parts:
        lst = [(lasts, part)]
        lasts = part[0]
        ok = True
        count = 0
        while ok:
            if len(lst) == 0:
                break
            sample = lst.pop()
            last = sample[0]
            newlst = sample[1][:]
            while sample[0] - newlst[-1] >= -3:
                if len(newlst) == 1:
                    count +=1
                    break
                lst.insert(0, (newlst[-1], newlst[:-1]))
                newlst.pop()
        pcount.append(count)
    print(pcount)
    print(functools.reduce(lambda a,b: a*b, pcount))
    return 0

result1 = None
if part_one():
    start = time.time()
    result1 = p1()
    print(round(1000*(time.time() - start), 2), 'ms')


if part_two():
    start = time.time()
    p2(result1)
    print(round(1000*(time.time() - start), 2), 'ms')