aoc_2021_19

# Day 19
# Finding overlapped scanners and beacons
# Try minimizing the cost function (parameters = x, y, z and theta phi?)

with open('input_day19.txt') as f:
    report = f.readlines()
    beacons = []
    for i, l in enumerate(report):
        report[i] = l.strip().split(',')
        if '---' in report[i][0]:
            beacons.append([])
        elif len(report[i])==3:
            beacons[-1].append(tuple([int(v) for v in report[i]]))

'''
Distances between points are translational and rotational invariant,
thus a good clue to identify overlapping - what if there are points
with equal distance? --> checked, no biggie

After getting the 12 overlapping beacons, we can work out the
alignment and positions of sensors...by the wonderful work of
minization

Then rotate the coordinates of the beacons..
'''

def dist(arr1, arr2):
    total = 0.
    for i in range(3):
        total += (arr1[i]-arr2[i])**2
    return total**0.5


# for each sensor, calculate the distance between each
# pair of beacons, save as set and dict
b_dist_set = [[] for _ in beacons]
b_dist_dict = [{} for _ in beacons]
for k in range(len(beacons)):
    for i in range(len(beacons[k])-1):
        for j in range(i+1, len(beacons[k])):
            d = dist(beacons[k][i], beacons[k][j])
            b_dist_set[k].append(d)
            if d not in b_dist_dict[k]:
                b_dist_dict[k][d] = [beacons[k][i], beacons[k][j]]
            else:
                print('Warning', k, i, j)
                print(b_dist_dict[k][d])
                print(beacons[k][i], beacons[k][j])
                b_dist_dict[k][d] += [beacons[k][i], beacons[k][j]]

# turn it to a set for easier comparison
for k in range(len(beacons)):
    b_dist_set[k] = set(b_dist_set[k])

from scipy.spatial.transform import Rotation as R
from scipy.optimize import minimize

def transform(coor, alpha, beta, gamma, center):
    # align the axis rot()*coor
    # Please specify degrees
    # then perform translation
    r_a = R.from_euler('z', [alpha], degrees=True)
    r_b = R.from_euler('x', [beta], degrees=True)
    r_g = R.from_euler('z', [gamma], degrees=True)
    coor = r_a.apply(coor)
    coor = r_b.apply(coor)
    coor = r_g.apply(coor)

    return tuple([coor[0][i] + center[i] for i in range(3)])

# create a dict to store the transformed sensors
checked = {0:True}
# a list to store the location of the sensors
sensor = [[0,0,0] for _ in range(len(beacons))]

# check until all sensors are checked
while len(checked) != 25:
    for k in range(len(beacons)-1):
        for l in range(1, len(beacons)):
            if k == l:
                continue
            if k in checked and l in checked:
                continue
            if k not in checked and l not in checked:
                continue
            # compare the distance set
            overlap = b_dist_set[k].intersection(b_dist_set[l])
            # if 12 overlap, there will be exactly 66 distances
            # present in both sets
            if len(overlap) >= 66:
                # find the 1 to 1 correspondence for each point
                corr_map = {}
                if k in checked:
                    base = k
                    target = l
                else:
                    base = l
                    target = k
                print(base, target)
                for d in overlap:
                    for pt in b_dist_dict[base][d]:
                        if pt not in corr_map:
                            corr_map[pt] = set(b_dist_dict[target][d])
                        else:
                            corr_map[pt] = corr_map[pt].intersection(b_dist_dict[target][d])

                for pt in corr_map:
                    corr_map[pt] = list(corr_map[pt])[0]

                #print(corr_map)

                def cost_dist(par):
                    cost = 0.
                    for k in corr_map:
                        cost += dist(k, transform(corr_map[k], par[0], par[1], par[2], par[3:]))
                    return cost

                # find the transformation that minimize the cost
                result = minimize(cost_dist, [0.,0.,0.,0.,0.,0.])
                #print(result)
                par = result['x']
                sensor[target] = [int(np.rint(v)) for v in par[3:]]

                # align the target coordinates
                aligned = []
                for b in beacons[target]:
                    coor = [int(np.rint(v)) for v in transform(b, par[0], par[1], par[2], par[3:])]
                    aligned.append(tuple(coor))
                beacons[target] = aligned

                # update the key of the b_dist_dict
                for d in b_dist_dict[target]:
                    new_coor = []
                    for coor in b_dist_dict[target][d]:
                        nc = [int(np.rint(v)) for v in transform(coor, par[0], par[1], par[2], par[3:])]
                        new_coor.append(tuple(nc))
                    b_dist_dict[target][d] = new_coor

                checked[target] = True

# part 1, how many unique beacons
unique_beacon = set(beacons[0])
for bs in beacons:
    unique_beacon = unique_beacon.union(bs)

print(len(unique_beacon))

# part 2, what is the largest distance apart
max_dist = 0
for i in range(len(sensor)-1):
    for j in range(1, len(sensor)):
        max_dist = max(max_dist, sum([abs(sensor[i][k]-sensor[j][k]) for k in range(3)]))
print(max_dist)