import os
from pathlib import Path
# import subprocess
# from glob import glob
from re import sub
import numpy as np
from datetime import datetime, timedelta
from tools import to_datetime, get_sat_freq, dt2float
from struct import unpack, calcsize

working_dir = Path(__file__).absolute().parent if '__file__' in dir() else Path.cwd()
# lsat_names = dict(C2E1=1, C2E2=2, C2E3=3, C2E4=4, C2E5=5, C2E6=6)
# for keys in list(lsat_names.keys()):
#     lsat_names[lsat_names[keys]] = keys

def log_message(msg, working_leo):
    with open(working_dir / f"log/progress_rt_{working_leo}_{datetime.now():%Y%m%d}.txt", 'a') as fid:
        fid.write(f"{datetime.now():%H:%M:%S} {msg}\n")

def gather_data(working_leo, tarfiles, filetypes, leo, cp_sep=False):
    log_message('- Start gathering REF profiles.', working_leo)
    (working_dir / 'p01_podRx3').mkdir(exist_ok=True)
    (working_dir / 'p01_podPair').mkdir(exist_ok=True)
    import pickle
    c0 = 299792458
    T0 = datetime(1980, 1, 6)
    if cp_sep:
        header = ['GPS_seconds', 'time_offset', 'L1C', 'L2X', 'S1C', 'S2X', 'C1C', 'C2X', 'P2X']
    else:
        header = ['GPS_seconds', 'time_offset', 'L1C', 'L2X', 'S1C', 'S2X', 'C1C', 'C2X']
    if 'ucar' in working_leo or 'tacc' in working_leo:
        crxfiles = np.array(list((working_dir / f'L0_extract_{working_leo}').glob('*/*/*.rnx')))
    else:
        crxfiles = tarfiles[filetypes==2]
    # dtstr = sub(r'.*_(\d{4}\.\d{3})\.\d{3}\.\d\d\.\d\d_.*', r'\1', crxfiles[-1])
    # dt = datetime.strptime(dtstr, '%Y.%j')
    # leo = lsat_names[int(sub(r'.*_\d{4}\.\d{3}\.(\d{3})\.\d\d\.\d\d_.*', r'\1', crxfiles[-1]))]
    # ant = int(sub(r'.*_\d{4}\.\d{3}\.\d{3}\.\d\d\.(\d\d)_.*', r'\1', crxfiles[-1]))

    dts = np.array([datetime.strptime(crxfile.parent.as_posix()[-8:], '%Y/%j') for crxfile in crxfiles])
    uniq_dts = np.unique(dts)
    # leos = np.unique([lsat_names[int(sub(r'.*_\d{4}\.\d{3}\.(\d{3})\.\d\d\.\d\d_.*', r'\1', crxfile.name))] for crxfile in crxfiles])
    if 'ucar' in working_leo or 'tacc' in working_leo:
        ants = np.array([int(sub(r'.*_\d{4}\.\d{3}\.\d{3}.\d\d\.(\d\d)\..*', r'\1', crxfile.name)) for crxfile in crxfiles])
    else:
        ants = np.array([int(sub(r'.*_\d{4}\.\d{3}\.\d\d\.\d\d\.\d{3}\.(\d\d)\..*', r'\1', crxfile.name)) for crxfile in crxfiles])
    # ants = np.array([int(sub(r'.*_\d{4}\.\d{3}\.\d{3}\.\d\d\.(\d\d)_.*', r'\1', crxfile.name)) for crxfile in crxfiles])
    uniq_ants = np.unique(ants)

    for dt in uniq_dts:
        (working_dir / f'p01_podPair/{dt:%Y-%m-%d}').mkdir(exist_ok=True)
        # for leo in leos:
        for ant in uniq_ants:
            crxfileout = working_dir / f'p01_podRx3/podRx3_{leo}.{dt:%Y.%j}.{ant:02d}.pkl'
            crxfileprev = working_dir / f'p01_podRx3/podRx3_{leo}.{dt-timedelta(days=1):%Y.%j}.{ant:02d}.pkl'
            if os.path.exists(crxfileout):
                log_message(f'-- Combining {crxfileout}', working_leo)
                with open(crxfileout, 'rb') as fid:
                    pkldata = pickle.load(fid)
                    dsat = pkldata['dsat']
                    dtime = pkldata['dtime']
                    dataall = pkldata['dataall']
            elif os.path.exists(crxfileprev):
                log_message(f'-- Retrieving {crxfileprev}', working_leo)
                with open(crxfileprev, 'rb') as fid:
                    pkldata = pickle.load(fid)
                    dsat = pkldata['dsat']
                    dtime = pkldata['dtime']
                    dataall = pkldata['dataall']
                for dnum in range(len(dsat)):
                    tsidx = np.where(np.diff(np.hstack([[T0], dtime[dnum]])) > timedelta(seconds=3))[0].astype('int')
                    teidx = np.hstack([tsidx[1:], tsidx[:1]]) - 1
                    tidx = tsidx[dtime[dnum][teidx] > dt - timedelta(hours=2)]
                    if tidx.size > 0:
                        dtime[dnum] = dtime[dnum][tidx[0]:]
                        dataall[dnum] = dataall[dnum][tidx[0]:, :]
                    else: # assign empty array with .size == 0 to be removed afterward
                        dtime[dnum] = dtime[dnum][:0]
                        dataall[dnum] = dataall[dnum][:0]
                didx = np.where([tt.size > 0 for tt in dtime])[0]
                dsat = [dsat[x] for x in didx]
                dtime = [np.array(dtime[x]) for x in didx]
                dataall = [np.array(dataall[x]) for x in didx]
            else:
                log_message(f'-- Creating {crxfileout}', working_leo)
                dsat = []
                dtime = []
                dataall = []
            crxfiled = crxfiles[np.logical_and(dts == dt, ants == ant)]
            crxfiled.sort()
            for crxfile in crxfiled:
                log_message(f'--- Processing {crxfile}', working_leo)
                is_head = True
                count = 0
                sat_append = False
                for line in open(crxfile, 'r'):
                    if is_head:
                        if 'END OF HEADER' in line:
                            is_head = False
                        continue
                    if sat_append:
                        sats = [line[32+3*x:35+3*x] for x in range(min(12, sat_num))][::-1] + sats
                        if sat_num>12:
                            sat_num -= 12
                        else:
                            sat_append = False
                    elif count == 0:
                        dt_now = to_datetime(('20'+line[1:]).split())
                        sat_num = int(line[29:32])
                        sats = [line[32+3*x:35+3*x] for x in range(min(12,sat_num))][::-1]
                        count = sat_num*2
                        if sat_num>12:
                            sat_num -= 12
                            sat_append = True
                    elif count%2==0:
                        count -= 1
                        gtmp = [line[ii:ii + 14] for ii in range(0, len(line) - 1, 16)]
                        if len(gtmp) == 4:
                            gtmp.append('')
                    else:
                        count -= 1
                        gtmp += [line[ii:ii + 14] for ii in range(0, len(line) - 1, 16)]
                        if cp_sep:
                            gtmp = [float(ss) if ss.strip() else np.nan for ss in gtmp]
                        else:
                            gtmp = [float(ss) for ss in gtmp if ss.strip()]
                        gsat = sats.pop()
                        if gsat.startswith('G') or gsat.startswith('R') or gsat.startswith('E'):
                            if gsat not in dsat:
                                dsat.append(gsat)
                                dtime.append(np.array([dt_now]))
                                dataall.append(np.array([gtmp]))
                            else:
                                gnum = dsat.index(gsat)
                                if dt_now > dtime[gnum][-1]:
                                    dtime[gnum]  = np.hstack([dtime[gnum], [dt_now]])
                                    dataall[gnum] = np.vstack([dataall[gnum],[gtmp]])
                                elif dt_now in dtime[gnum]:
                                    tidx = np.searchsorted(dtime[gnum], dt_now)
                                    dataall[gnum][tidx] = [pp if np.isfinite(pp) else qq for pp, qq in
                                                           zip(dataall[gnum][tidx], gtmp)]
                                else:
                                    tidx = np.searchsorted(dtime[gnum], dt_now)
                                    dtime[gnum] = np.hstack([dtime[gnum][:tidx], [dt_now], dtime[gnum][tidx:]])
                                    dataall[gnum] = np.vstack([dataall[gnum][:tidx], [gtmp], dataall[gnum][tidx:]])
            didx = np.argsort(dsat)
            dsat = [dsat[x] for x in didx]
            dtime = [np.array(dtime[x]) for x in didx]
            dataall = [np.array(dataall[x]) for x in didx]
            pkldata = {'dsat': dsat, 'dtime': dtime, 'dataall': dataall}
            log_message(f'--- Saving {crxfileout}', working_leo)
            with open(crxfileout, 'wb') as fid:
                pickle.dump(pkldata, fid)
            for dnum in range(len(dsat)):
                f1,f2,_ = get_sat_freq(dsat[dnum])
                idx = np.all(np.isfinite(dataall[dnum]), axis=1)
                dtime[dnum] = dtime[dnum][idx]
                dataall[dnum] = dataall[dnum][idx]
                # idx = [0] + list(np.where(np.diff(dtime[dnum]) > timedelta(seconds=3))[0] + 1) + [len(dtime[dnum])]
                idx = [0] + list(np.where(np.logical_or(np.diff(dtime[dnum]) > timedelta(seconds=3),
                                                        np.abs(np.diff(dataall[dnum][:, 1] * c0 / f2 - dataall[dnum][:, 0] * c0 / f1)) > 0.5))[0] + 1) + [len(dtime[dnum])]
                for ss in range(len(idx) - 1):
                    if idx[ss + 1] - idx[ss] < 60:
                        continue
                    gps_sec = (dtime[dnum][idx[ss]] - T0) / timedelta(seconds=1)
                    tdiff = (dtime[dnum][idx[ss]:idx[ss + 1]] - dtime[dnum][idx[ss]]) / timedelta(seconds=1)
                    if cp_sep:
                        data = np.column_stack((
                            gps_sec + np.zeros_like(tdiff),
                            tdiff,
                            dataall[dnum][idx[ss]:idx[ss + 1], [0, 1, 5, 6, 2, 3, 4]] * np.array(
                                [c0 / f1, c0 / f2, 1, 1, 1, 1, 1])
                        )).astype('float64')
                    else:
                        data = np.column_stack((
                            gps_sec + np.zeros_like(tdiff),
                            tdiff,
                            dataall[dnum][idx[ss]:idx[ss+1], [0, 1, 4, 5, 2, 3]] * np.array([c0 / f1, c0 / f2, 1, 1, 1, 1])
                        )).astype('float64')
                    tStart, tEnd = [T0 + timedelta(seconds=x[0] + x[1]) for x in data[[0, -1], :2]]
                    occFileName = working_dir / f'p01_podPair/{dt:%Y-%m-%d}/occTab_{leo}.{dt:%Y.%j}.{tStart:%Y%m%d%H%M%S}.{tEnd:%Y%m%d%H%M%S}.A{ant:02d}.{dsat[dnum]}.npz'
                    if occFileName.exists():
                        continue
                    for oldfile in (working_dir / f'p01_podPair/{dt:%Y-%m-%d}').glob(
                            f'occTab_{leo}.{dt:%Y.%j}.*.*.A{ant:02d}.{dsat[dnum]}.npz'):
                        ts_old = [datetime.strptime(x, '%Y%m%d%H%M%S') for x in sub(r'.*\.(\d{14})\.(\d{14})\..*\.npz', r'\1_\2', oldfile.name).split('_')]
                        if ts_old[0]<tStart or ts_old[1]>tEnd:
                            continue
                        token = sub(r'.*\.(\d{14})\.(\d{14})\.(A\d\d)\.([A-Z]\d\d)\.npz', r'\1_\2_\4_\3', oldfile.name)
                        for partfile in (working_dir / f'p05_podTec/{dt:%Y-%m-%d}').glob(f'*{leo}_{token}*'):
                            partfile.unlink()
                        oldfile.unlink()
                    np.savez(occFileName, data=data, header=header)

    log_message('- Start gathering leoOrb profiles.', working_leo)
    orbfiles = tarfiles[filetypes==0]
    orbfiles.sort()
    dts_orb = np.array([datetime.strptime(sub(r'(\d{4}).(\d{3})',r'\1/\2',orbfile.parent.as_posix()[-8:]), '%Y/%j') for orbfile in orbfiles])
    uniq_dts = np.unique(dts_orb)
    (working_dir / 'p01_leoOrb').mkdir(exist_ok=True)
    for dt in uniq_dts:
        log_message(f"-- Processing {dt:%Y-%m-%d}...", working_leo)
        orbfileout = working_dir / 'p01_leoOrb' / f'leoOrb_{leo}.{dt:%Y.%j}.npz'
        orbfileprev = working_dir / 'p01_leoOrb' / f'leoOrb_{leo}.{dt-timedelta(days=1):%Y.%j}.npz'
        orbfiled = orbfiles[dts_orb==dt]
        dts0 = []
        data0 = []
        tmp0 = np.full((8,),np.nan)
        for orbfile in orbfiled:
            isheader = True
            for line in open(orbfile,'r'):
                if line.startswith("*"):
                    if not isheader:
                        data0.append(tmp0.copy())
                    else:
                        isheader = False
                    tmp0[:] = np.nan
                    dts0.append(to_datetime(line[1:].split()))
                elif line.startswith("P"):
                    tmp0[:4] = np.array([line[x:x+14] for x in range(4,60,14)]).astype('float')
                elif line.startswith("V"):
                    tmp0[4:] = np.array([line[x:x+14] for x in range(4,60,14)]).astype('float')
            data0.append(tmp0.copy())
        dts0 = np.array(dts0)
        data0 = np.array(data0)
        idx = np.abs(data0[:, 3] - 1E6) > 1E-3
        dts0 = dts0[idx]
        data0 = data0[idx]
        # _, idx = np.unique(dts0, return_index=True)
        _, idx = np.unique(dts0[::-1], return_index=True)
        idx = len(dts0) - 1 - idx
        dts0 = dts0[idx]
        data0 = data0[idx]
        if orbfileout.is_file():
            with np.load(orbfileout, allow_pickle=True) as fid:
                data = fid['data']
                dts = fid['dts']
            idx = ~np.isin(dts, dts0)
            dts = np.hstack([dts[idx],dts0])
            data = np.vstack([data[idx], data0])
            idx = np.argsort(dts)
            dts = dts[idx]
            data = data[idx]
        elif orbfileprev.is_file():
            with np.load(orbfileprev, allow_pickle=True) as fid:
                data = fid['data']
                dts = fid['dts']
            idx = np.logical_and(~np.isin(dts, dts0),dts>dt-timedelta(hours=3))
            dts = np.hstack([dts[idx],dts0])
            data = np.vstack([data[idx], data0])
            idx = np.argsort(dts)
            dts = dts[idx]
            data = data[idx]
        else:
            dts = dts0
            data = data0
        np.savez(orbfileout, data=data, dts=dts)

    log_message('- Start gathering leoAtt profiles.', working_leo)
    attfiles = tarfiles[filetypes==1]
    vars_list = ['att_x', 'att_y', 'att_z', 'att_w', 'sca_x', 'sca_y', 'sca_z', 'sca_w', 'ang_x', 'ang_y', 'ang_z',
                 'pe_x', 'pe_y', 'pe_z', 've_x', 've_y', 've_z', 'pi_x', 'pi_y', 'pi_z', 'vi_x', 'vi_y', 'vi_z', 'tor', 'sad_x','sad_y']
    dts_att = np.array([datetime.strptime(attfile.parent.as_posix()[-8:], '%Y/%j') for attfile in attfiles])
    uniq_dts = np.unique(dts_att)
    (working_dir / 'p01_leoAtt').mkdir(exist_ok=True)
    for dt in uniq_dts:
        log_message(f"-- Processing {dt:%Y-%m-%d}...", working_leo)
        attfiled = attfiles[dts_att==dt]
        attfiled.sort()
        attfileout = working_dir / f'./p01_leoAtt/leoAtt_{leo}.{dt:%Y.%j}.npz'
        attfileprev = working_dir / f'./p01_leoAtt/leoAtt_{leo}.{dt-timedelta(days=1):%Y.%j}.npz'
        dts0 = []
        data0 = []
        tmp0 = np.full((len(vars_list),),np.nan)
        for attfile in attfiles:
            for line in open(attfile,'r'):
                if line.startswith("tim"):
                    if len(dts0)>len(data0):
                        data0.append(tmp0.copy())
                        tmp0[:] = np.nan
                    dts0.append(to_datetime(line[3:].split()))
                elif line.startswith("att"):
                    tmp0[:4] = np.fromstring(line[3:], sep=" ")[1:-1]
                elif line.startswith("sca"):
                    tmp0[4:8] = np.fromstring(line[3:], sep=" ")
                elif line.startswith("ang"):
                    tmp0[8:11] = np.fromstring(line[3:], sep=" ")
                elif line.startswith("pve"):
                    tmp0[11:17] = np.fromstring(line[3:], sep=" ")
                elif line.startswith("pvi"):
                    tmp0[17:23] = np.fromstring(line[3:], sep=" ")
                elif line.startswith("tor"):
                    tmp0[23:24] = np.fromstring(line[3:], sep=" ")
                elif line.startswith("sad"):
                    tmp0[24:] = np.fromstring(line[3:], sep=" ")
            data0.append(tmp0.copy())
        dts0 = np.array(dts0)
        data0 = np.array(data0)
        if os.path.isfile(attfileout):
            with np.load(attfileout, allow_pickle=True) as fid:
                data = fid['data']
                dts = fid['dts']
            idx = ~np.isin(dt2float(dts, dt), dt2float(dts0, dt))
            dts = np.hstack([dts[idx],dts0])
            data = np.vstack([data[idx], data0])
            idx = np.argsort(dts)
            dts = dts[idx]
            data = data[idx]
        elif os.path.isfile(attfileprev):
            with np.load(attfileprev, allow_pickle=True) as fid:
                data = fid['data']
                dts = fid['dts']
            idx = np.logical_and(~np.isin(dt2float(dts, dt), dt2float(dts0, dt)), dts > dt - timedelta(hours=3))
            dts = np.hstack([dts[idx],dts0])
            data = np.vstack([data[idx], data0])
            idx = np.argsort(dts)
            dts = dts[idx]
            data = data[idx]
        else:
            dts = dts0
            data = data0
        np.savez(attfileout, data=data, dts=dts, vars=vars_list)

    log_message('- Completed gathering profiles.', working_leo)
    return

def trim_crx_files(podfile):
    import pickle
    c0 = 299792458
    T0 = datetime(1980, 1, 6)
    header = ['GPS_seconds', 'time_offset', 'L1C', 'L2X', 'S1C', 'S2X', 'C1C', 'C2X']
    print(f'{datetime.now():%H:%M:%S} - {podfile.stem}')
    dt = datetime.strptime(sub(r'.*(\d{4}\.\d{3})\..*',r'\1',podfile.stem),'%Y.%j')
    (working_dir / f'p01_podPair/{dt:%Y-%m-%d}').mkdir(exist_ok=True)
    leo = sub(r'.*(C2E\d).*',r'\1',podfile.stem)
    ant = int(podfile.stem[-1])
    with open(podfile, 'rb') as fid:
        pkldata = pickle.load(fid)
        dsat = pkldata['dsat']
        dtime = pkldata['dtime']
        dataall = pkldata['dataall']
    for dnum in range(len(dsat)):
        f1, f2, _ = get_sat_freq(dsat[dnum])
        idx = np.all(np.isfinite(dataall[dnum]), axis=1)
        dtime[dnum] = dtime[dnum][idx]
        dataall[dnum] = dataall[dnum][idx]
        # idx = [0] + list(np.where(np.diff(dtime[dnum]) > timedelta(seconds=3))[0] + 1) + [len(dtime[dnum])]
        idx = [0] + list(np.where(np.logical_or(np.diff(dtime[dnum]) > timedelta(seconds=3),
                                                np.abs(np.diff(dataall[dnum][:, 1] * c0 / f2 - dataall[dnum][
                                                    :, 0] * c0 / f1)) > 0.5))[0] + 1) + [len(dtime[dnum])]
        for ss in range(len(idx) - 1):
            if idx[ss + 1] - idx[ss] < 60:
                continue
            gps_sec = (dtime[dnum][idx[ss]] - T0) / timedelta(seconds=1)
            tdiff = (dtime[dnum][idx[ss]:idx[ss + 1]] - dtime[dnum][idx[ss]]) / timedelta(seconds=1)
            data = np.column_stack((
                gps_sec + np.zeros_like(tdiff),
                tdiff,
                dataall[dnum][idx[ss]:idx[ss + 1], [0, 1, 4, 5, 2, 3]] * np.array([c0 / f1, c0 / f2, 1, 1, 1, 1])
            )).astype('float64')
            tStart, tEnd = [T0 + timedelta(seconds=x[0] + x[1]) for x in data[[0, -1], :2]]
            occFileName = working_dir / f'p01_podPair/{dt:%Y-%m-%d}/occTab_{leo}.{dt:%Y.%j}.{tStart:%Y%m%d%H%M%S}.{tEnd:%Y%m%d%H%M%S}.A{ant:02d}.{dsat[dnum]}.npz'
            if occFileName.exists():
                continue
            print(f'{datetime.now():%H:%M:%S} - Patching {occFileName.stem}')
            for oldfile in (working_dir / f'p01_podPair/{dt:%Y-%m-%d}').glob(
                    f'occTab_{leo}.{dt:%Y.%j}.*.*.A{ant:02d}.{dsat[dnum]}.npz'):
                ts_old = [datetime.strptime(x, '%Y%m%d%H%M%S') for x in
                          sub(r'.*\.(\d{14})\.(\d{14})\..*\.npz', r'\1_\2', oldfile.name).split('_')]
                if ts_old[0] > tStart or ts_old[1] < tEnd:
                    continue
                print(f'{datetime.now():%H:%M:%S} - Removing {oldfile.stem}')
                token = sub(r'.*\.(\d{14})\.(\d{14})\.(A\d\d)\.([A-Z]\d\d)\.npz', r'\1_\2_\4_\3', oldfile.name)
                for partfile in (working_dir / f'p05_podTec/{dt:%Y-%m-%d}').glob(f'*{leo}_{token}*'):
                    partfile.unlink()
                oldfile.unlink()
            np.savez(occFileName, data=data, header=header)
    return

def process_trim(nProcess=20):
    podfiles = list(Path('p01_podRx3').glob('*.2025.365*.pkl'))
    import multiprocessing as mp
    with mp.Pool(processes=nProcess) as pool:
        pool.map(trim_crx_files, podfiles)
    return