#!/usr/bin/env python
import numpy as np
from datetime import datetime, timedelta
from scipy.interpolate import CubicSpline
import matplotlib.pyplot as plt
import re
import os
#from datetime import datetime

__authors__ = ["Ming Chen"]
__license__ = "Public Domain"
__maintainer__ = "Ming Chen"
__email__ = "ming.chen@noaa.gov"
__date__  = "Nov 05, 2025"

'''
#####################################################################
# 
# gnsspod_standalone.py
#
# Purpose:
# This Python script provides utilities to generate high-rate GNSS orbit
# and clock products from standard IGS SP3 files. The input SP3 products
# are normally available at 5-minute intervals and are interpolated to a
# higher sampling rate (default: 30 seconds). Satellite velocities are
# derived from position records when unavailable in the input products.
##
# History:  
#       Created by: Ming Chen 2025/11/05
#                   ming.chen@noaa.gov
# 
# Usage: python gnsspod_standalone.py the-input-file -o output-file -r resolution
#
#
#####################################################################

'''
# ------------------------ Read SP3 safely ------------------------

def read_sp3_with_velocity_regex(file_path, clock_unit='us'):
    """
    Read SP3 including positions, velocities, and clock using regex.
    Returns data dictionary and header lines.
    """
    data = {}
    header_lines = []

    # Regular expressions
    epoch_pattern = re.compile(r'^\*\s+(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s+([\d.]+)')
    pos_pattern   = re.compile(r'^P(\w{3})\s*([-\d.Ee+]+)\s*([-\d.Ee+]+)\s*([-\d.Ee+]+)\s*([-\d.Ee+]+)')
    vel_pattern   = re.compile(r'^V(\w{3})\s*([-\d.Ee+]+)\s*([-\d.Ee+]+)\s*([-\d.Ee+]+)')

    current_time = None
    last_sat = None

    with open(file_path, 'r') as f:
        lines = f.readlines()

    # Extract header
    for line in lines:
        if line.startswith('*'):
            break
        header_lines.append(line.rstrip('\n'))

    # Parse data
    for line in lines:
        line = line.rstrip('\n')
        if m := epoch_pattern.match(line):
            year, month, day, hour, minute, second = m.groups()
            current_time = datetime(int(year), int(month), int(day),
                                    int(hour), int(minute), int(float(second)))
        elif m := pos_pattern.match(line):
            sat, x, y, z, clk = m.groups()
            x, y, z, clk = map(float, (x, y, z, clk))
            if clk > 99999: continue
            if sat not in data:
                data[sat] = {"times": [], "X": [], "Y": [], "Z": [], "VX": [], "VY": [], "VZ": [], "CLK": []}
           
            data[sat]["times"].append(current_time)
            data[sat]["X"].append(x)
            data[sat]["Y"].append(y)
            data[sat]["Z"].append(z)
            data[sat]["CLK"].append(clk)
            last_sat = sat
        elif m := vel_pattern.match(line):
            sat, vx, vy, vz = m.groups()
            vx, vy, vz = map(float, (vx, vy, vz))
            if sat in data:
                data[sat]["VX"].append(vx)
                data[sat]["VY"].append(vy)
                data[sat]["VZ"].append(vz)
    return data, header_lines


# ------------------------ Compute velocities if missing or at high resolution ------------------------
def compute_velocity_from_position(sat_data):
    """
    Compute VX, VY, VZ from positions and times using finite differences.
    """
    times_sec = np.array([(t - sat_data["times"][0]).total_seconds() for t in sat_data["times"]])
    scale = 1.0e4
    X = np.array(sat_data["X"])
    Y = np.array(sat_data["Y"])
    Z = np.array(sat_data["Z"])
    n = len(times_sec)
    VX = np.zeros(n)
    VY = np.zeros(n)
    VZ = np.zeros(n)

    # Central differences
    for i in range(1, n-1):
        dt = times_sec[i+1] - times_sec[i-1]
        VX[i] = (X[i+1] - X[i-1]) / dt * scale
        VY[i] = (Y[i+1] - Y[i-1]) / dt * scale
        VZ[i] = (Z[i+1] - Z[i-1]) / dt * scale

    # Forward difference for first point
    VX[0] = (X[1] - X[0]) / (times_sec[1] - times_sec[0]) * scale
    VY[0] = (Y[1] - Y[0]) / (times_sec[1] - times_sec[0]) * scale
    VZ[0] = (Z[1] - Z[0]) / (times_sec[1] - times_sec[0]) * scale

    # Backward difference for last point
    VX[-1] = (X[-1] - X[-2]) / (times_sec[-1] - times_sec[-2]) * scale
    VY[-1] = (Y[-1] - Y[-2]) / (times_sec[-1] - times_sec[-2]) * scale
    VZ[-1] = (Z[-1] - Z[-2]) / (times_sec[-1] - times_sec[-2]) * scale

    return VX.tolist(), VY.tolist(), VZ.tolist()

# ------------------------ Interpolation ------------------------
def interpolate_satellite(sat_data, new_interval=30.0):
    """
    Interpolates a single satellite:
    - X/Y/Z, VX/VY/VZ: cubic spline
    - CLK: linear
    """
    times_sec = np.array([(t - sat_data["times"][0]).total_seconds() for t in sat_data["times"]])
    new_times_sec = np.arange(0, times_sec[-1] + new_interval, new_interval)
    interp_data = {}

   # for comp in ["X","Y","Z","VX","VY","VZ"]:
    for comp in ["X","Y","Z" ]:
        if comp in sat_data and len(sat_data[comp]) > 0:
            cs = CubicSpline(times_sec, sat_data[comp])
            interp_data[comp] = cs(new_times_sec)

    interp_data["CLK"] = np.interp(new_times_sec, times_sec, sat_data["CLK"])
    interp_data["times"] = [sat_data["times"][0] + timedelta(seconds=t) for t in new_times_sec]
    return interp_data

def interpolate_all(data_dict, new_interval=30.0):
    return {sat: interpolate_satellite(sat_data, new_interval)
            for sat, sat_data in data_dict.items()}

# ------------------------ Satellite sorting ------------------------
def sat_sort_key(sat):
    """Sort satellites G E C then by number"""
    system_order = {'G':0,'E':1,'C':2}
    sys = sat[0]
    num = int(sat[1:])
    return (system_order.get(sys,99), num)


# ------------------------ Plotting ------------------------
def plot_satellite_data(original_data, interpolated_data, satellites_to_plot):
    for sat in satellites_to_plot:
        if sat not in original_data:
            continue
        orig = original_data[sat]
        interp = interpolated_data[sat]

        # Positions + Clock
        fig, axs = plt.subplots(4,1, figsize=(12,10), sharex=True)
        for i, comp in enumerate(["X","Y","Z","CLK"]):
            axs[i].plot(orig["times"], orig[comp], 'o-', label="Original")
            axs[i].plot(interp["times"], interp[comp], '-', label="Interpolated")
            axs[i].set_ylabel(comp)
            axs[i].grid(True)
            axs[i].legend()
        axs[3].set_xlabel("Time")
        plt.suptitle(f"Satellite {sat} Position + Clock")
        plt.tight_layout()
        plt.show()

        # Velocities
        if len(orig["VX"]) > 0:
            fig, axs = plt.subplots(3,1, figsize=(12,8), sharex=True)
            for i, comp in enumerate(["VX","VY","VZ"]):
                axs[i].plot(orig["times"], orig[comp], 'o-', label="Original")
                axs[i].plot(interp["times"], interp[comp], '-', label="Interpolated")
                axs[i].set_ylabel(comp)
                axs[i].grid(True)
                axs[i].legend()
            axs[2].set_xlabel("Time")
            plt.suptitle(f"Satellite {sat} Velocities")
            plt.tight_layout()
            plt.show()


def write_sp3_highres_with_velocity(output_file, interpolated_data, header_lines, resolution=30.0):
    """
    Write high-resolution SP3 including positions, velocities, clock
    - Position line is followed by velocity line
    - Velocity line gets default CLK value 999999.999999
    """
    default_clk_for_velocity = 999999.999999
    hres   = re.compile(r'(##\s*\d+\s+[-\d.Ee+]+)\s+([-\d.Ee+]+)(.+)')

    with open(output_file, 'w') as f:
        # Header
        for line in header_lines:
            resline = line.rstrip('\n')
            if m := hres.match(resline):
               h1,h2,h3 = m.groups()
               line = h1 + f"{resolution:15.8f}" + h3
            f.write(line + '\n')

        # Sort satellites
        all_sats = sorted(interpolated_data.keys(), key=sat_sort_key)
        num_epochs = len(interpolated_data[all_sats[0]]["times"])
        for i in range(num_epochs):
            t = interpolated_data[all_sats[0]]["times"][i]
            # Epoch line
            f.write("* {:4d} {:2d} {:2d} {:2d} {:2d} {:11.8f}\n".format(
                t.year, t.month, t.day, t.hour, t.minute, t.second + t.microsecond*1e-6
            ))
            for sat in all_sats:
                # Position line
                f.write("P{:<3s}{:14.6f}{:14.6f}{:14.6f}{:14.6f}\n".format(
                    sat,
                    interpolated_data[sat]["X"][i],
                    interpolated_data[sat]["Y"][i],
                    interpolated_data[sat]["Z"][i],
                    interpolated_data[sat]["CLK"][i]
                ))

                # Velocity line only if exists
                vx_array = interpolated_data[sat].get("VX")
                if vx_array is not None and len(vx_array) > 0:
                    f.write("V{:<3s}{:14.6f}{:14.6f}{:14.6f}{:14.6f}\n".format(
                        sat,
                        vx_array[i],
                        interpolated_data[sat]["VY"][i],
                        interpolated_data[sat]["VZ"][i],
                        default_clk_for_velocity
                    ))


def write_sp3_highres_with_velocity(output_file, interpolated_data, header_lines, resolution=30.0):
    """
    Write high-resolution SP3 including positions, velocities, clock
    - Position line is followed by velocity line
    - Velocity line gets default CLK value 999999.999999
    """
    default_clk_for_velocity = 999999.999999
    hres   = re.compile(r'(##\s*\d+\s+[-\d.Ee+]+)\s+([-\d.Ee+]+)(.+)')
    # Sort satellites
  
    all_sats = sorted(interpolated_data.keys(), key=sat_sort_key)
    sat_epochs = {sat: len(interpolated_data[sat]["times"]) for sat in all_sats}
    max_sat, max_sat_epochs = max(sat_epochs.items(), key=lambda item: item[1])
    
    satsub = [sat for sat in sat_epochs if sat_epochs[sat] < max_sat_epochs]
    
    with open(output_file, 'w') as f:
        # Header
        for line in header_lines:
            resline = line.rstrip('\n')
            if m := hres.match(resline):
               h1,h2,h3 = m.groups()
               line = h1 + f"{resolution:15.8f}" + h3
            f.write(line + '\n')

        # Sort satellites
        all_sats = sorted(interpolated_data.keys(), key=sat_sort_key)
        num_epochs = len(interpolated_data[max_sat]["times"])

        for i in range(num_epochs):
            t = interpolated_data[max_sat]["times"][i]
            # Epoch line
            f.write("* {:4d} {:2d} {:2d} {:2d} {:2d} {:11.8f}\n".format(
                t.year, t.month, t.day, t.hour, t.minute, t.second + t.microsecond*1e-6
            ))
            for sat in all_sats:
                # Position line
                if sat in satsub: continue
                f.write("P{:<3s}{:14.6f}{:14.6f}{:14.6f}{:14.6f}\n".format(
                    sat,
                    interpolated_data[sat]["X"][i],
                    interpolated_data[sat]["Y"][i],
                    interpolated_data[sat]["Z"][i],
                    interpolated_data[sat]["CLK"][i]
                ))

                # Velocity line only if exists
                vx_array = interpolated_data[sat].get("VX")
                if vx_array is not None and len(vx_array) > 0:
                    f.write("V{:<3s}{:14.6f}{:14.6f}{:14.6f}{:14.6f}\n".format(
                        sat,
                        vx_array[i],
                        interpolated_data[sat]["VY"][i],
                        interpolated_data[sat]["VZ"][i],
                        default_clk_for_velocity
                    ))
def interpolate_gnsspod(input_file, output_file, resolution=30 ):

   if not os.path.exists(input_file):
      exit
   
   print(f"Input: {input_file}")
   print(f"Output: {output_file}")
   print(f"Resolution: {resolution} seconds")
   new_interval = float(resolution)         # seconds

   # Read SP3
   data_all, header_lines =  read_sp3_with_velocity_regex(input_file)

   # Compute velocities if missing
   #for sat, sat_data in data_all.items():
   #    if len(sat_data["VX"]) == 0:  # No velocity data
   #        sat_data["VX"], sat_data["VY"], sat_data["VZ"] = compute_velocity_from_position(sat_data)

   # Interpolate all satellites
   interpolated_data = interpolate_all(data_all, new_interval=new_interval)
   # Compute velocities if missing or at the high resolution
   for sat, sat_data in interpolated_data.items():
        interpolated_data[sat]["VX"], interpolated_data[sat]["VY"], interpolated_data[sat]["VZ"] = compute_velocity_from_position(sat_data)

   # Write high-resolution SP3 including velocities
   write_sp3_highres_with_velocity(output_file, interpolated_data, header_lines, resolution = new_interval)
   print(f"High-resolution SP3 with velocities saved: {output_file}")

   # Plot positions, velocities, clock
   #plot_satellite_data(data_all, interpolated_data, satellites_to_plot)

# ------------------------ MAIN ------------------------
if __name__ == "__main__":
   import argparse

   parser = argparse.ArgumentParser(description='Standalone module to generate GNSSPOD from the orginal IGS SP3 files  ')
   parser.add_argument("-v","--verbose", action='store_true',  help="print info for debug")
   parser.add_argument("input", type=str,  help="IGS ORB SP3 file  ")
   parser.add_argument("-o", "--output",type=str, default='test.sp3', help="Output ORB SP3 file at higher resolution")
   parser.add_argument("-r", "--resolution", type=int, default=30, help="Output ORB file at higher resolution in seconds")
   args = parser.parse_args()
    
   input_file = args.input  ; output_file = args.output
   resolution = args.resolution

   extract_clk(input_file, output_file, resolution)
