from tletools import TLE
import pandas as pd
from astropy import units as u


def get_tles(raw_tle_str):
    """Parses the raw TLE string and converts it to TLE objects.
    
    :param raw_tle_str: The raw string form of the TLEs
    :type raw_tle_str: str
    
    :return: The parsed object representations of the TLEs
    :rtype: [tletools.TLE]
    """
    cutnum = 2
    all_tle_lines = raw_tle_str.strip().splitlines()
    tles = []
    for i in range(len(all_tle_lines)//cutnum):
        # Calculate offset
        j = i*cutnum
        tle_lines = ["noaa-06"] + all_tle_lines[j:j+cutnum]
        print(tle_lines)
        # Strip line number from object name line
        # tle_lines[0] = tle_lines[0][2:]
        tle = TLE.from_lines(*tle_lines)
        tles.append(tle)
    # print(tles)
    return tles


def get_aso_data(tles):
    """Extracts the necessary data from the TLE objects for doing orbital
    prediction.

    :param tles: The list of TLE objects to extract orbit information from
    :type tles: [tletools.TLE]

    :return: A DataFrame of the extracted TLE data
    :rtype: pandas.DataFrame
    """
    tles_data = []
    for tle in tles:
        aso_data = {}
        aso_data['aso_name'] = tle.name
        aso_data['aso_id'] = tle.norad
        aso_data['epoch'] = tle.epoch.to_datetime()
        # Convert the TLE object to a poliastro.twobody.Orbit instance
        orbit = tle.to_orbit()
        # Calculate the position and velocity vectors
        r, v = orbit.rv()
        # Convert position vector from kilometers to meters
        r_m = r.to(u.m).to_value()
        # Convert the velocity vector from km/s to m/s
        v_ms = v.to(u.m/u.s).to_value()
        # Extract the components of the state vectiors
        aso_data['r_x'], aso_data['r_y'], aso_data['r_z'] = r_m
        aso_data['v_x'], aso_data['v_y'], aso_data['v_z'] = v_ms
        tles_data.append(aso_data)
    return pd.DataFrame(tles_data)


full_tle = []
with open("/home/lj020/Downloads/noaa.txt","r") as f:
    raw_tle = f.read()
    l = get_tles(raw_tle_str=raw_tle)
    full_tle += l

w = get_aso_data(full_tle)
w.to_parquet("/home/lj020/Downloads/data6.parquet")
with open("/home/lj020/Downloads/data6.txt","w") as f:
    f.write(w.to_string())
# print(w)