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
%matplotlib inline

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('./spacebook-golf56788echo/test.txt') as f:
    catalog = read_starfile(f.read())

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

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

    refstars_magsorted = sorted(refstars, key=lambda x:x['m'])[::-1]

    catalog_magnitudes = np.array([x['m'] for x in catalog])
    matches = []

    for idx, star in enumerate(refstars_magsorted):
        if star['m'] <= 500:
            break
        match = np.argmin(np.abs(catalog_magnitudes-star['m']))
        matches.append(match)
    print(matches)

    P = np.vstack([x['v'] for x in [catalog[i] for i in matches]])
    Q = np.vstack([x['v'] for x in refstars_magsorted[:4]])
    print("rmsd: {}".format(calculate_rmsd.kabsch_rmsd(P,Q)))
    rotation_mtx = calculate_rmsd.kabsch(P, Q)
    rotation = Rotation.from_matrix(rotation_mtx)

    rotated = [dict(v=rotation.apply(x['v']), m=x['m']) for x in refstars_magsorted]
    found_idxes = []
    for star in rotated:
        found_idxes.append(np.argmin([np.linalg.norm(star['v']-catalogstar['v']) for catalogstar in catalog]))

    r.send(','.join([str(x) for x in found_idxes]) + '\n')
    time.sleep(0.1)
    r.recvuntil('Left...\n')
print(r.clean())