tle_fitting.py

"""
Load and parse TLE from file (TLE API would be cute)

Load data from peak finding algorithm

Define a function that calculates the dopplershift curve based on TLE and time.

Make a function that calculates the residuals between peak data and TLE curve

Pass residual function to scipy least square method

Then profit, stonks

TODO: Compare to GNSS data and see if this fitting is rubbish or not
"""

import os
import numpy as np
from datetime import datetime, timezone
import matplotlib.pyplot as plt
from scipy.optimize import least_squares
from skyfield.api import load, wgs84, EarthSatellite
from unpack_data import load_doppler_csv


TRANSMIT_FREQ_MHZ = 437.0

TLE_SOURCE = "file"         # "file"      → load from TLE_FILE_PATH
TLE_FILE_PATH    = "disco-2.tle"   # path to file with 2-line or 3-line TLE blocks
fname = "raw_data/gqrx_20260428_134823_437075000_2000000_fc_doppler.csv" # you can make one of these with save_doppler_csv() form unpack_data.py

# GROUND_STATION = {
#     "name":      "Aarhus",
#     "latitude":  56.1629,   # degrees N
#     "longitude": 10.2039,   # degrees E
#     "elevation": 10,        # metres above sea level
# }
GROUND_STATION = {
    "name":      "CPH",
    "latitude":  55.6761,   # degrees N
    "longitude": 12.5683,   # degrees E
    "elevation": 10,        # metres above sea level
}

# THIS GETS OPTIMISED
param_names = ["epoch", "raan"] # list[str]

gs = wgs84.latlon(
    GROUND_STATION["latitude"],
    GROUND_STATION["longitude"],
    elevation_m=GROUND_STATION["elevation"],
)
ts = load.timescale()

# Load observed Doppler
parts = os.path.basename(fname).split("_")
dt = datetime.strptime(f"{parts[1]}_{parts[2]}", "%Y%m%d_%H%M%S")
dt = dt.replace(tzinfo=timezone.utc)
recording_start_unix = dt.timestamp()

timestamps_unix, doppler_obs_hz = load_doppler_csv(
    "raw_data/gqrx_20260428_134823_437075000_2000000_fc_doppler.csv",
    recording_start_unix)

print(f"Duration: {timestamps_unix[-1] - timestamps_unix[0]:.1f} seconds")
print(f"Observed Doppler range: {doppler_obs_hz.min()/1e3:.2f} to {doppler_obs_hz.max()/1e3:.2f} kHz")


# --- Here be dragons ---

def get_tle(stellite_id) -> list[tuple[str, str, str]]:
    raw = load_tles_from_file(TLE_FILE_PATH)

    filtered = _filter_by_id(raw, stellite_id)
    if stellite_id:
        found_ids = {l1[2:7].strip() for _, l1, _ in filtered}
        for sid in stellite_id:
            if str(sid).zfill(5) not in found_ids:
                print(f"  [WARN] ID {sid} not found in TLE source")
    return filtered

def load_tles_from_file(path: str) -> list[tuple[str, str, str]]:
    """Parse a plain-text file with 2-line or 3-line TLE blocks."""
    with open(path) as f:
        lines = [l.rstrip() for l in f if l.strip()]
    return _parse_tle_lines(lines)

def _parse_tle_lines(lines: list[str]) -> list[tuple[str, str, str]]:
    """
    Parse a list of stripped, non-empty text lines into (name, line1, line2) tuples.
    Handles both formats:
      • 3-line blocks: name / line-1 / line-2
      • 2-line blocks: line-1 / line-2  (name derived from catalog number)
    """
    tles = []
    i = 0
    while i < len(lines):
        # Peek ahead: is this line a TLE line-1?
        if lines[i].startswith("1 ") and i + 1 < len(lines) and lines[i + 1].startswith("2 "):
            # 2-line format – synthesise a name from the catalog number
            cat_id = lines[i][2:7].strip()
            tles.append((cat_id, lines[i], lines[i + 1]))
            i += 2
        elif (i + 2 < len(lines)
              and lines[i + 1].startswith("1 ")
              and lines[i + 2].startswith("2 ")):
            # 3-line format – first line is the name
            tles.append((lines[i], lines[i + 1], lines[i + 2]))
            i += 3
        else:
            i += 1
    return tles

def _filter_by_id(
    tles: list[tuple[str, str, str]],
    ids: list[int] | None,
) -> list[tuple[str, str, str]]:
    """Return only TLEs whose NORAD catalog ID (field in line-1) is in *ids*.
    If *ids* is None or empty, return all TLEs unchanged."""
    if not ids:
        return tles
    id_set = {str(i).zfill(5) for i in ids}
    return [
        (name, l1, l2)
        for name, l1, l2 in tles
        if l1[2:7].strip().zfill(5) in id_set
    ]

def parse_decimal_exponent(s):
    """
    Parse TLE decimal exponent format back to float.
    e.g. ' 00000-0' → 0.0, ' 30838-3' → 0.00030838
    """
    s = s.strip()
    if s == '00000-0' or s == '00000+0':
        return 0.0
    # Find the exponent sign
    for i in range(len(s)-1, 0, -1):
        if s[i] in '+-':
            mantissa = float("0." + s[:i].strip().lstrip())
            exp = int(s[i:])
            return mantissa * 10**exp
    return 0.0

def tle_to_dict(name, l1, l2):
    return {
        "satellite_number": int(l1[2:7]),
        "classification": l1[7],
        "launch_designator": l1[9:17].strip(),
        "epoch": float(l1[18:32]),
        "first_div_mean_motion": float(l1[33:43]),
        "second_div_mean_motion": parse_decimal_exponent(l1[44:52]),
        "drag_term": parse_decimal_exponent(l1[53:61]),
        "ephemeris_type": int(l1[62]),
        "element_set_number": int(l1[64:68]),
        "inclination": float(l2[8:16]),
        "raan": float(l2[17:25]),
        "eccentricity": float("0." + l2[26:33]),
        "arg_of_perigee": float(l2[34:42]),
        "mean_anomaly": float(l2[43:51]),
        "mean_motion": float(l2[52:63]),
        "revolution_number": int(l2[63:68]),
    }

def dict_to_tle(lines: dict):
    """
    Convert TLE parameter dict to properly formatted TLE strings.
    Checksums are set to 0 since skyfield ignores them.

    :param lines: dict of TLE parameters
    :return: l1, l2 — TLE line 1 and line 2 strings
    """

    def format_decimal_exponent(val):
        if val == 0.0:
            return "00000-0"
        import math
        exp = int(math.floor(math.log10(abs(val)))) + 1
        mantissa = int(val * 10 ** (-exp + 5))
        return f"{mantissa:+06d}{exp - 1:+d}".replace("+", " ")

    # TODO: remove debug code
    # print(type(lines['second_div_mean_motion']))
    # print(lines['second_div_mean_motion'])
    # print(lines['second_div_mean_motion'] == 0.0)

    first_div_str = f"{lines['first_div_mean_motion']:10.8f}".replace("0.", " .")
    ecc_str = f"{int(round(lines['eccentricity'] * 1e7)):07d}"
    l1 = (f"1 "
          f"{lines['satellite_number']:05d}{lines['classification']} "
          f"{lines['launch_designator']:<8s} "
          f"{lines['epoch']:014.8f} "
          f"{first_div_str}"
          f"  {format_decimal_exponent(lines['second_div_mean_motion'])} "
          f" {format_decimal_exponent(lines['drag_term'])} "
          f"{lines['ephemeris_type']:1d} "
          f"{lines['element_set_number']:4d}0")

    l2 = (f"2 "
          f"{lines['satellite_number']:05d} "
          f"{lines['inclination']:8.4f} "
          f"{lines['raan']:8.4f} "
          f"{ecc_str} "
          f"{lines['arg_of_perigee']:8.4f} "
          f"{lines['mean_anomaly']:8.4f} "
          f"{lines['mean_motion']:11.8f}"
          f"{lines['revolution_number']:5d}0")

    return l1, l2

def residuals(params, param_names, doppler_obs, gs, ts, unix_timestamps, carrier_freq_hz):
    name, l1, l2 = load_tles_from_file(TLE_FILE_PATH)[0]
    lines = tle_to_dict(name, l1, l2)

    for pname, value in zip(param_names, params):
        lines[pname] = value

    l1, l2 = dict_to_tle(lines)
    doppler_pred = calculate_doppler(l1, l2, gs, ts, unix_timestamps, carrier_freq_hz)
    return doppler_obs - doppler_pred


def calculate_doppler(line1, line2, gs, ts, unix_timestamps, carrier_freq_hz):

    sat = EarthSatellite(line1, line2, name = "DISCO-2", ts=ts)

    times = ts.ut1(jd=unix_timestamps / 86400.0 + 2440587.5)
    # times = ts.from_datetimes([datetime.fromtimestamp(t, tz=timezone.utc) for t in unix_timestamps])

    diff        = sat - gs
    topocentric = diff.at(times)
    alt, az, distance = topocentric.altaz()

    C_KM_S = 299792.458
    pos_km = topocentric.position.km
    vel_km_s = topocentric.velocity.km_per_s
    range_km = distance.km
    pos_unit = pos_km / range_km
    range_rate = (pos_unit * vel_km_s).sum(axis=0)
    doppler_hz = -carrier_freq_hz * range_rate / C_KM_S

    return doppler_hz

def optimize_tle(params, param_names, doppler_obs, gs, ts, unix_timestamps, carrier_freq_hz):
    res_lsq = least_squares(residuals, x0=params, args=(param_names, doppler_obs, gs, ts, unix_timestamps, carrier_freq_hz))
    return res_lsq


# Load TLE from file
name, real_l1, real_l2 = load_tles_from_file(TLE_FILE_PATH)[0]
lines = tle_to_dict(name, real_l1, real_l2)

print(f"Loaded TLE for {name}")
print(f"Epoch: {lines['epoch']}")

# Predict with initial TLE
doppler_initial = calculate_doppler(real_l1, real_l2, gs, ts, timestamps_unix, 437.075e6)
print(f"Predicted Doppler range: {doppler_initial.min()/1e3:.2f} to {doppler_initial.max()/1e3:.2f} kHz")

# Optimize
initial_params = np.array([lines[p] for p in param_names])
print(f"Initial params: {initial_params}")

result = optimize_tle(initial_params, param_names, doppler_obs_hz, gs, ts,
                      timestamps_unix, 437.075e6)

print(f"Success: {result.success}")
print(f"Message: {result.message}")
print(f"Optimized params: {result.x}")
for pname, init, opt in zip(param_names, initial_params, result.x):
    print(f"  {pname}: {init:.6f} → {opt:.6f} (Δ = {opt-init:.6f})")
print(f"Initial cost: {np.sum(residuals(initial_params, param_names, doppler_obs_hz, gs, ts, timestamps_unix, 437.075e6)**2):.2f}")
print(f"Final cost: {np.sum(result.fun**2):.2f}")

# Build optimized TLE
opt_lines = tle_to_dict(name, real_l1, real_l2)
for pname, value in zip(param_names, result.x):
    opt_lines[pname] = value
opt_l1, opt_l2 = dict_to_tle(opt_lines)
doppler_optimized = calculate_doppler(opt_l1, opt_l2, gs, ts, timestamps_unix, 437.075e6)

print(f"\nOptimized TLE:")
print(opt_l1)
print(opt_l2)

# Plot
t_rel = timestamps_unix - timestamps_unix[0]
plt.figure(figsize=(12, 4))
plt.plot(t_rel, doppler_obs_hz / 1e3, 'k.', markersize=2, label='observed')
plt.plot(t_rel, doppler_initial / 1e3, 'r-', linewidth=1, label='initial TLE')
plt.plot(t_rel, doppler_optimized / 1e3, 'b-', linewidth=1, label='optimized TLE')
plt.xlabel("Time (s)")
plt.ylabel("Doppler shift (kHz)")
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.show()