#!/bin/python3
import sys,re,math, numpy as np
from collections import Counter, defaultdict
from pprint import pprint
sys.path.insert(0, '../../')
from fred import get_re,list2int

input_f = 'input'

part = 2
#########################################
#                                       #
#              Part 1                   #
#                                       #
#########################################

lines = []

if part == 1:
    with open(input_f) as file:
        for line in file:
            lines.append(get_re(r"^p=<(.*)>.*v=<(.*)>.*a=<(.*)>$",line.rstrip()))  

    min_acc = float('inf')
    acc_idx = 0

    for idx,i in enumerate(lines):
        tmp = math.sqrt(sum(component ** 2 for component in list2int(i.group(3).split(','))))
        if tmp < min_acc:
            min_acc = tmp
            acc_idx = idx
    print(acc_idx)

#########################################
#                                       #
#              Part 2                   #
#                                       #
#########################################

def calculate_position_xyz(initial_position, velocity, acceleration):

    velocity += acceleration

    new_position = initial_position + velocity # + (0.5 * acceleration * time**2)
    return np.array(new_position)

def find_duplicate_indices(data):
    value_to_indices = defaultdict(list)
    for index, value in data.items():
        value_tuple = tuple(value)  
        value_to_indices[value_tuple].append(index)
    
    duplicate_indices = [index for indices in value_to_indices.values() if len(indices) > 1 for index in indices]
    return duplicate_indices

if part == 2:
    with open(input_f) as file:
        for line in file:
            lines.append(get_re(r"^p=<(.*)>.*v=<(.*)>.*a=<(.*)>$",line.rstrip()))  

    for ldx,l in enumerate(lines):
        lines[ldx] = [np.array(list2int(l.group(1).split(','))), np.array(list2int(l.group(2).split(','))), np.array(list2int(l.group(3).split(',')))]

    particles = []

    iterations = 10000

    last_val = 0
    range_since_last = 0


    for iter in range(iterations):
        if len(lines) != last_val:
            last_val = len(lines)
        else:
            range_since_last += 1

        if range_since_last > 100: # we dont want to run forever
            break

        positions = {}
        length = len(lines)
        duplicates = []

        for ldx in range(length):
            pos = []
            p,v,a = lines[ldx]

            pos = calculate_position_xyz(p, v, a)
            positions[ldx] = pos
            lines[ldx][0] = pos

        duplicates = find_duplicate_indices(positions)

        if len(duplicates) > 0:
            lines = [item for i, item in enumerate(lines) if i not in duplicates]
    print(len(lines))