from glob import glob
import os
from pathlib import Path
import numpy as np
from tools import calculate_tec
from traceback import extract_tb
from datetime import datetime
from re import sub

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

def cal_tec_for_parallel(args):
    qq = args[0]
    reffile = args[1]
    try:
        sec_ref_0, tec, coszn, meps = calculate_tec(reffile, dumpnc=True)
    except Exception as e:
        log_message(f'-- Error detected for {qq}: {reffile.name}')
        log_message(f'---  {type(e).__name__} >> {str(e)}')
        for tb in extract_tb(e.__traceback__):
            log_message(f"---  line {tb.lineno} of {tb.filename}")

def calc_tec(leo, dt, nProcess):
    reffiles = list((working_dir / f'p01_podPair/{dt:%Y-%m-%d}').glob(f'occTab_{leo}.*.npz'))
    reffiles.sort()
    out_dir = working_dir / f'p05_podTec/{dt:%Y-%m-%d}'
    out_dir.mkdir(exist_ok=True)
    idx = np.full((len(reffiles),), True)
    for ii, reffile in enumerate(reffiles):
        # tStart = sub(r'.*\.(\d{14})\.\d{14}\.A\d\d\.[A-Z]\d\d\.npz',r'\1',reffile.name)
        # tEnd = sub(r'.*\.\d{14}\.(\d{14})\.A\d\d\.[A-Z]\d\d\.npz', r'\1', reffile.name)
        # gnss = sub(r'.*\.\d{14}\.\d{14}\.A\d\d\.([A-Z]\d\d)\.npz', r'\1', reffile.name)
        # ant = sub(r'.*\.\d{14}\.\d{14}\.A(\d\d)\.[A-Z]\d\d\.npz', r'\1', reffile.name)
        token = sub(r'.*\.(\d{14})\.(\d{14})\.(A\d\d)\.([A-Z]\d\d)\.npz', r'\1_\2_\4_\3', reffile.name)
        if any(out_dir.glob(f'*{leo}_{token}*')):
            idx[ii] = False
    idx = np.where(idx)[0]
    allargs = list(zip(idx,np.array(reffiles)[idx])) # list(zip(range(len(reffiles)), reffiles))
    if nProcess > 0:
        import multiprocessing as mp
        with mp.Pool(processes=nProcess) as pool:
            pool.map(cal_tec_for_parallel, allargs)
    else:
        for args in allargs:
            cal_tec_for_parallel(args)
        # for qq, rofile in enumerate(ropair):
        #     if not os.path.isfile(rofile.replace('t01_RO','t03_pair')):
        #         log_message(f'-- No-REF skip for {qq}: {os.path.basename(rofile)}', working_leo)
        #     else:
        #         try:
        #             with np.load(rofile.replace('t01_RO','t03_pair'), allow_pickle=True) as fid:
        #                 reffiles = fid['reffile']
        #             calculate_exph(rofile, reffiles)
        #         except Exception as e:
        #             log_message(f'-- Error detected for {qq}: {os.path.basename(rofile)}', working_leo)
        #             log_message(f'---  {type(e).__name__} >> {str(e)}', working_leo)
        #             for tb in extract_tb(e.__traceback__):
        #                 log_message(f"---  line {tb.lineno} of {tb.filename}", working_leo)