from glob import glob
from matplotlib import pyplot as plt
from re import sub
import numpy as np
from datetime import datetime, timedelta
from tools import dt2float

clrs = np.array([[int(y['color'][x:x + 2], 16) / 255 for x in range(1, 7, 2)] for y in plt.rcParams['axes.prop_cycle']])

gfiles = glob('gnssFin/2025.310/*.npz')
for gfile in gfiles:
    gdt = datetime.strptime(sub(r'.*_(\d{11})_.*',r'\1',gfile),'%Y%j%H%M')
    gsrc = sub(r'.*/(...).*',r'\1',gfile)
    with np.load(gfile, allow_pickle=True) as fid:
        Gdata = fid['Gdata']
        Gdt = fid['Gdt']
        Gsats = fid['Gsats']
    sidx = np.argsort(Gsats)
    Gsats = Gsats[sidx]
    Gdata = Gdata[:,sidx,:]
    rmse = np.zeros((8,len(Gsats),2))
    ncum = np.zeros((8,len(Gsats),1))
    for tdif in range(11):
        dt0 = gdt - timedelta(hours=(tdif-3)*6)
        tfile = glob(f"gnssLrg/{dt0:%Y.%j}/{gsrc}_*_s{dt0:%Y%j%H%M}_*.npz")
        if len(tfile) == 0:
            continue
        with np.load(tfile[0], allow_pickle=True) as fid:
            GdataT = fid['Gdata']
            GdtT = fid['Gdt']
            GsatsT = fid['Gsats']
        sidx = np.argsort(GsatsT)
        GsatsT = GsatsT[sidx]
        GdataT = GdataT[:,sidx,:]
        tsidxs = np.isin(GsatsT,Gsats)
        stidxs = np.isin(Gsats,GsatsT)
        tsidxt = np.isin(GdtT, Gdt)
        stidxt = np.isin(Gdt, GdtT)
        err = GdataT[np.ix_(tsidxt,tsidxs,range(4))]-Gdata[np.ix_(stidxt,stidxs,range(4))]
        for tt in range(len(err)//72):
            rmse[max(tdif+tt-3,tt), stidxs, 0] += np.sum(err[tt * 72:(tt + 1) * 72, :, :3] ** 2, axis=(0,2))
            rmse[max(tdif+tt-3,tt), stidxs, 1] += np.sum(err[tt * 72:(tt + 1) * 72, :, 3] ** 2, axis=0)
            ncum[max(tdif+tt-3,tt), stidxs] += 72
    rmse /= ncum
    rmse[~np.isfinite(rmse)] = np.nan
    rmse **= 0.5

    xc = np.arange(len(Gsats))
    yc = rmse.transpose((2,0,1))

    plt.close('all')
    fig, axs = plt.subplots(2, 1, figsize=(12, 9))
    for ii in range(8):
        axs[0].plot(xc, yc[0,7-ii], '.-', color=clrs[ii], label=f"{8-ii:+d}")
        axs[1].plot(xc, yc[1,7-ii], '.-', color=clrs[ii])
    axs[0].set_xlim((-0.5, len(Gsats) - 0.5))
    axs[0].set_xticks(np.arange(len(Gsats)))
    axs[0].set_xticklabels(Gsats)
    axs[0].set_yscale('log')
    axs[0].tick_params(axis='x', labelrotation=90, labelsize=8)
    axs[0].set_ylabel('Position Distance (km)')
    axs[0].legend(ncol=8, loc='upper right')
    axs[0].set_title(sub(r".*/", r"Difference in ", gfile))
    axs[0].grid(True)
    axs[1].set_xlim((-0.5, len(Gsats) - 0.5))
    axs[1].set_xticks(np.arange(len(Gsats)))
    axs[1].set_xticklabels(Gsats)
    axs[1].set_yscale('log')
    axs[1].tick_params(axis='x', labelrotation=90, labelsize=8)
    axs[1].set_ylabel('Clock Bias Difference (us)')
    axs[1].grid(True)
    fig.savefig(sub(r".*/(.*)\.npz", f"\\1.png", gfile), dpi=300)