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

catalog = {}
with open('./attitude-papa21503yankee/test.txt') as f:
    i = 0
    for line in f:
        [x, y, z, m] = [float(s.strip()) for s in line.split(',')]
        catalog[i] = {'v': np.array([x,y,z]), 'm':m}
        i += 1

def parse_stars(stardata):
    stars = {}
    for line in stardata.strip().split('\n'):
        line = line.strip()
        star_id = int(line.split(':')[0].strip())
        direction = np.array([float(x) for x in line.split(':')[1].split(',\t')])
        stars[star_id] = direction
    return stars

def solve_orientation(stars, catalog):
    P = np.vstack(list(stars.values()))
    Q = np.vstack([catalog[i]['v'] for i in stars.keys()])
    print("rmsd: {}".format(calculate_rmsd.kabsch_rmsd(P,Q)))
    rotation_mtx = calculate_rmsd.kabsch(P, Q)
    rotation = Rotation.from_matrix(np.linalg.inv(rotation_mtx))
    return rotation

TICKET = 'THE_TICKET'
r = tubes.remote.remote('attitude.satellitesabove.me', 5012)
r.send(TICKET+'\n')
time.sleep(0.5)
for _ in range(20):
    r.recvuntil(b'--------------------------------------------------\n', drop=True)
    stars = parse_stars(r.recv().decode())
    rotation = solve_orientation(stars, catalog)
    r.send(','.join([str(x) for x in rotation.as_quat()]) + '\n')
    time.sleep(0.1)
print(r.clean())