# %%
from datetime import datetime, timedelta
from os import stat
from attr import dataclass
from erfa.core import pas
from joblib.parallel import Parallel, delayed
from numpy import double
from tletools import TLE
from astropy import units as u
from spacetrack import SpaceTrackClient


def parse_tle(tle_lines: list[str]):
    assert (len(tle_lines) == 3)
    tle = TLE.from_lines(*tle_lines)
    orbit = tle.to_orbit()
    r, v = orbit.rv()
    r_m = r.to(u.m).to_value()
    v_ms = v.to(u.m / u.s).to_value()
    return r_m, v_ms


def to_TLE(tle_lines: list[str]):
    assert (len(tle_lines) == 3)
    return TLE.from_lines(*tle_lines)


def parse_tle_file(fd):
    lines = fd.readlines()
    res: list[TLE] = []
    for i in range(0, len(lines), 3):
        res.append(to_TLE(lines[i:i + 3]))
    return res


def get_tle(client: SpaceTrackClient, orbit_ids, days):
    return client.tle(**{
        'epoch': f'>now-{days}',
        'norad_cat_id': orbit_ids,
        'format': '3le'
    })


client = SpaceTrackClient("geulocrutusse-7937@yopmail.com", "ETionGeOmEnDiUMp")
# print(get_tle(client, [49138], 30))
# %%

with open("1.tle") as f:
    tles = parse_tle_file(f)

# %%
from poliastro.twobody.propagation import cowell
from tqdm import tqdm


@dataclass
class orbit_status:
    id: int
    epoch: datetime
    r_x: float
    r_y: float
    r_z: float
    v_x: float
    v_y: float
    v_z: float


stats = []
for i in tqdm(range(1, len(tles)), position=0):
    p1 = tles[i - 1]
    p2 = tles[i]
    t1 = p1.epoch.to_datetime()
    t2 = p2.epoch.to_datetime()
    o1 = p1.to_orbit()
    o2 = p2.to_orbit()
    dt: timedelta = t2 - t1
    if dt.seconds <= 60:
        continue

    for sec in tqdm(range(0, dt.seconds, 60), position=1, leave=False):
        ratio = sec / dt.seconds
        x1 = o1.propagate(sec * u.s, method=cowell)
        x2 = o2.propagate(-(dt.seconds - sec) * u.s, method=cowell)
        r1, v1 = x1.rv()
        r1_m = r1.to(u.m).to_value()
        v1_ms = v1.to(u.m / u.s).to_value()
        r2, v2 = x2.rv()
        r2_m = r2.to(u.m).to_value()
        v2_ms = v2.to(u.m / u.s).to_value()
        r_m = [x * (1 - ratio) + y * ratio for x, y in zip(r1_m, r2_m)]
        v_ms = [x * (1 - ratio) + y * ratio for x, y in zip(v1_ms, v2_ms)]
        stats.append(
            orbit_status(
                p1.norad,
                t1 + timedelta(seconds=sec),
                *r_m,
                *v_ms,
            ).__dict__)

import pandas as pd

df = pd.DataFrame(stats)
df.to_parquet("1.pq")

# %%
df
# %%