import numpy as np
from pwnlib import tubes
import time
import matplotlib.pyplot as plt
from rmsd import calculate_rmsd
from scipy.spatial.transform import Rotation
import seaborn as sb
import itertools
from sklearn.neighbors import BallTree

def read_starfile(data):
    stars = []
    for line in data.strip().split('\n'):
        [x,y,z,m] = [float(s.strip()) for s in line.split(',')]
        stars.append({'v': np.array([x,y,z]), 'm':m})
    return stars

with open('./myspace-mike33686zulu/test.txt') as f:
    catalog = read_starfile(f.read())

def angle(x, y):
    x_ = x/np.linalg.norm(x)
    y_ = y/np.linalg.norm(y)
    return np.arccos(np.clip(np.dot(x_, y_), -1.0, 1.0))

def gen_fingerprints(X, bt):
    dist, ind = bt.query(X[:], k=3)
    fingerprints = []
    for i in range(X.shape[0]):
        a = X[i,:]
        [bi, ci] = ind[i,1:]
        b = X[bi,:]
        c = X[ci,:]
        fingerprints.append([
            *dist[i, 1:],
            angle(b-a, c-a),
        ])
    return fingerprints

TICKET = 'THE_TICKET'
r = tubes.remote.remote('myspace.satellitesabove.me', 5016)
r.send(TICKET+'\n')
time.sleep(0.5)
r.recvuntil('Ticket please:\n', drop=True)

X = np.vstack([s['v'] for s in catalog])
bt_catalog = BallTree(X, leaf_size=30)

fingerprints = gen_fingerprints(X, bt_catalog)

for _ in range(5):
    refstars = read_starfile(r.recvuntil('\n\n').decode())

    Y = np.vstack([s['v'] for s in refstars])
    bt_unknown = BallTree(Y, leaf_size=30)

    fingerprints_unknown = gen_fingerprints(Y, bt_unknown)

    matches = []
    for query in fingerprints_unknown:
        error = [np.linalg.norm(np.array(fp_known) - np.array(query))
                 for fp_known in fingerprints]
        match_idx = np.argmin(error)
        if error[match_idx] < 1e-4:
            print(match_idx, error[match_idx])
            matches.append(match_idx)

    r.send(','.join(map(str,matches)) + '\n')
    r.recvuntil('Left...\n')
print(r.clean())