#!/usr/bin/env python
import numpy as np
from datetime import datetime, timedelta
from scipy.interpolate import CubicSpline, interp1d
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__  = "Mar 24, 2023"

'''
#####################################################################
# 
# leopod_sp3.py
#
# Purpose:
# This python script has the functions to  combine the  orbit and clock 
# into a SP3 file of a specific range of epoches 
#
# History:  
#       Created by: Ming Chen 2026/02/05
#                   ming.chen@noaa.gov
# 
# Usage: python leopod_sp3.py orbit_file clock_file sp3_file 
# --start-time yyyymnddThhmmss --end-time yyyymnddThhmmss
#
#
#####################################################################

'''

# ------------------------------------------------------------
# Parse the clock file (RINEX CLK-like format)
# ------------------------------------------------------------

def parse_clock_file(clock_file):
    """
    Reads a clock file and returns a dictionary:
        clock[(sat, datetime)] = clock_value
    """
    clock_dict = {}

    with open(clock_file, "r") as f:
        for line in f:
            if not line.startswith("AR"):
                continue

            parts = line.split()

            # Satellite ID (e.g., G29)
            sat = parts[1].strip()
            if sat not in clock_dict:
               clock_dict[sat] = {"times": [], "CLK": []}

            # Date-time fields
            year  = int(parts[2])
            month = int(parts[3])
            day   = int(parts[4])
            hour  = int(parts[5])
            minute = int(parts[6])
            sec    = float(parts[7])

            # Build datetime object
            current_time = datetime(
                year, month, day, hour, minute,
                int(sec),
                int((sec % 1) * 1e6)
            )

            #current_time = datetime(int(year), int(month), int(day),
            #                        int(hour), int(minute), int(float(second)))
            # Clock offset (seconds)
            clk = float(parts[9])*1e6
            clock_dict[sat]["times"].append(current_time )
            clock_dict[sat]["CLK"].append(clk)

    return clock_dict



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

def read_sp3_with_regex(file_path, clock_unit='us', t_start:datetime = None, t_end:datetime = None, check_clk=False):
    """
    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'))
    
    if t_start is None: t_start = datetime(1900, 1, 1) 
    if t_end is None:   t_end = datetime(2900, 1, 1)
 
    # 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)),  int((float(second) % 1) * 1e6))

        elif m := pos_pattern.match(line):
            if current_time < t_start or  current_time > t_end: continue
            sat, x, y, z, clk = m.groups()
            x, y, z, clk = map(float, (x, y, z, clk))
            if clk > 99999 and check_clk: 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):
            if current_time < t_start or  current_time > t_end: continue
            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_sat_data_cubic(sat_data, new_interval=30.0, new_times:list = [], ds_keys:list=[]):
    """
    Interpolates a single satellite:
    - X/Y/Z, VX/VY/VZ: cubic spline
    - CLK: linear
    """

    if len(ds_keys) < 1:
        ds_keys = sat_data.keys()
    if 'times' not in ds_keys:
        print(f"'times' not in sat_data ...")
        return
    
    tref =  datetime(2000,1,1)

    times_sec = np.array([(t - tref).total_seconds() for t in sat_data["times"]])
    
    if len(new_times)<1 :
       new_times_sec = np.arange(0, times_sec[-1] + new_interval, new_interval)
    else:
       new_times_sec = np.array([(t - tref).total_seconds() for t in new_times])

    ds_keys.remove('times')
    interp_data = {}

   # for comp in ["X","Y","Z","VX","VY","VZ"]:
    for comp in ds_keys:
        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["times"] = [tref + timedelta(seconds=t) for t in new_times_sec]
    #if 'CLK' in ds_keys:
    #  cs = CubicSpline(times_sec,sat_data['CLK'])
    #  csclk =  cs(new_times_sec)
    #  interp_data["CLK"] = csclk
    #  interp_data["times"] = [tref+ timedelta(seconds=t) for t in new_times_sec]
    return interp_data

# ------------------------ Interpolation ------------------------
def interpolate_sat_data_linear(sat_data, new_interval=30.0, new_times:list = [], ds_keys:list=[]):
    """
    Interpolates a single satellite:
    - X/Y/Z, VX/VY/VZ: cubic spline
    - CLK: linear
    """

    if len(ds_keys) < 1:
        ds_keys = sat_data.keys()
    if 'times' not in ds_keys:
        print(f"'times' not in sat_data ...")
        return
    
    tref =  datetime(2000,1,1)

    times_sec = np.array([(t - tref).total_seconds() for t in sat_data["times"]])
    
    if len(new_times)<1 :
       new_times_sec = np.arange(0, times_sec[-1] + new_interval, new_interval)
    else:
       new_times_sec = np.array([(t - tref).total_seconds() for t in new_times])

    ds_keys.remove('times')
    interp_data = {}

   # for comp in ["X","Y","Z","VX","VY","VZ"]:
    for comp in ds_keys:
        if comp in sat_data and len(sat_data[comp]) > 0:
            cs = interp1d(times_sec, sat_data[comp], kind="linear", fill_value="extrapolate")
            interp_data[comp] = cs(new_times_sec)
            interp_data["times"] = [tref + timedelta(seconds=t) for t in new_times_sec]


    return interp_data


def interpolate_all_satellites(data_dict, new_interval=30.0):
    return {sat: interpolate_satellite_data(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)


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
        if len(interpolated_data.keys()) > 1:
          all_sats = sorted(interpolated_data.keys(), key=sat_sort_key)
        else:
          all_sats =  list(interpolated_data.keys() )
        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
                    ))


# ------------------------ 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 normalize_to_datetime(t):
    """Ensure all times are Python datetime"""
    if isinstance(t, datetime):
        return t
    else:
        return t.astype('datetime64[ns]').astype(datetime)


def filter_dict_by_time(data, start_time, end_time):

    times = [normalize_to_datetime(t) for t in data["times"]]
 
    indices = [
        i for i, t in enumerate(times)
        if start_time <= t <= end_time
    ]

    filtered = {
        key: [values[i] for i in indices]
        for key, values in data.items()
    }

    return filtered   


def interpolate_leopod_finclk(fin_file, clk_file, fin_output, 
                            t_start:datetime, t_end:datetime, sliding=60, algo='linear',leostatus=''):

   if not os.path.exists(fin_file) or not os.path.exists(clk_file):
      exit
   
   print(f"FIN Input: {fin_file}")
   print(f"CLK Input: {clk_file}")
   print(f"FIN_Output: {fin_output}")
   #fname = os.path.basename(fin_file).split('_')
   #tend = fname[-2]
   #t_end = datetime.strptime(tend, "%Y%m%dT%H%M%S")
   #t_start = t_end - timedelta(minutes=sliding)

   # Load clock values
   clk_data = parse_clock_file(clk_file)
   sats = list(clk_data.keys()) ; sat = sats[0]
   clkpod = clk_data[sat]
   leosats = {'L23': 'C2E1', 'L24': 'C2E2','L25': 'C2E3',\
               'L26': 'C2E4', 'L27': 'C2E5','L28': 'C2E6'}

   leosats = {'C2E1':'L23', 'C2E2':'L24', 'C2E3':'L25',\
              'C2E4':'L26', 'C2E5':'L27', 'C2E6':'L28'}

   sat = leosats[sat]
   

   # Read file
   #t_start = kinpod['times'][0] ; t_end =  kinpod['times'][-1]
   #fin_data, fin_header_lines =  read_sp3_with_regex(fin_file, t_start = t_start, t_end=t_end)
   fin_data, fin_header_lines =  read_sp3_with_regex(fin_file)
   if leostatus:
      leostatus = f"/* LEOSTATUS {leostatus} \n"
      fin_header_lines[-1] = leostatus 

   finpod = fin_data[sat] 
   if len(finpod['VX']) < 1:
      finpod['VX'], finpod['VY'], finpod['VZ'] = compute_velocity_from_position(finpod)
  
   new_times = finpod['times']
   #interp_data = interpolate_sat_data(kinpod, new_interval=resolution)
   if algo == 'cubic':
      interp_data = interpolate_sat_data_cubic(clkpod, new_times=new_times, ds_keys=['times','CLK'] )
   else:
      interp_data = interpolate_sat_data_linear(clkpod, new_times=new_times, ds_keys=['times','CLK'] )

   finpod['CLK'] = interp_data['CLK']
   #print(finpod['times'])
   
   
   finpod = filter_dict_by_time(finpod, t_start, t_end)
   fin_data[sat] = finpod
   
   write_sp3_highres_with_velocity(fin_output, fin_data, fin_header_lines)
   


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

   parser = argparse.ArgumentParser(description='merge orbit and clock into SP3 files  ')
   parser.add_argument("-v","--verbose", action='store_true',  help="print info for debug")
   parser.add_argument("fin_input", type=str,  help="file with only position  ")
   parser.add_argument("clk_input", type=str,  help="file with leo clock")
   parser.add_argument("fin_output",type=str,  help="Output ORB SP3 file at higher resolution")
   parser.add_argument("-start", "--start-time", type=str,  help="%Y%m%dT%H%M%S")
   parser.add_argument("-end", "--end-time", type=str,  help="%Y%m%dT%H%M%S")
   parser.add_argument("-status", "--leo-status", type=str, help="LEO Bernese Solution Status")
   parser.add_argument("-sliding", "--sliding-minutes", type=int, default=30, help="sliding minutes")

   args = parser.parse_args()
    
   fin_input = args.fin_input  ; clk_input = args.clk_input
   fin_output = args.fin_output
   sliding = args.sliding_minutes
   t_start = args.start_time
   t_end = args.end_time
   tstart = datetime.strptime(t_start, "%Y%m%dT%H%M%S")
   tend =  datetime.strptime(t_end, "%Y%m%dT%H%M%S")
   leostatus = args.leo_status 
 
   interpolate_leopod_finclk(fin_input, clk_input, fin_output, tstart, tend, sliding=sliding, algo='linear',leostatus=leostatus)
