import numpy as np
from pwnlib import tubes
import time
import matplotlib.pyplot as plt
from scipy.spatial.transform import Rotation
import seaborn as sb

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())

def match_brightest_stars(unknown, catalog):
    brightest = []
    catalog_matches = []
    for star in unknown:
        if star['m'] < 500:
            continue

        # find closest magnitude in catalog
        match = np.argmin(np.abs(np.array([s['m'] for s in catalog]) - star['m']))
        print(f'[+] matched {star} to {catalog[match]}')
        brightest.append(star)
        catalog_matches.append(catalog[match])

    return brightest, catalog_matches

def orient(ref, catalog):
    P = [s['v'] for s in catalog]
    Q = [s['v'] for s in ref]

    rot, rmsd = Rotation.align_vectors(P, Q)
    print(f'[+] aligned; rmsd = {rmsd}')
    return rot

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):
    unknown_stars = read_starfile(r.recvuntil('\n\n').decode())

    # find pairs
    unknown_ref, catalog_ref = match_brightest_stars(unknown_stars, catalog)
    # get attitude
    attitude = orient(unknown_ref, catalog_ref)
    # rotate each star to catalog-referenced coordinates and match by L2 norm
    index_guesses = []
    for star in unknown_stars:
        v = attitude.apply(star['v'])
        index_guesses.append(np.argmin([np.linalg.norm(catalog_star['v'] - v)
                                        for catalog_star in catalog]))

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