from glob import glob
import os
import numpy as np
from planetiq.tools import calculate_exph  #Jun modified
from traceback import extract_tb
from datetime import datetime

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 ro_out_exist(rofile):
    from re import sub
    dt_rr = datetime.strptime(sub(r'.*\.(\d{4}\.\d{3})\..*', r'\1', rofile), '%Y.%j')
    ro_out = sub(r'.*_RO_(.*)/occTab_(.*?)\.\d{4}\.\d{3}\.(\d{14})\.(\d{14})\.(.*?)\.(.*?)\.npz',
                 f's04_atmPhs/{dt_rr:%Y-%m-%d}_\\1/RONA_L1A_\\2_\\3_\\4_\\6_\\5_c*.nc',rofile)
    return len(glob(ro_out))

def cal_exph_for_parallel(args):
    working_leo = args[0]
    qq = args[1]
    rofile = args[2]
    tarsource = args[3]
    if ro_out_exist(rofile):
        log_message(f'-- Existing skip for {qq}: {os.path.basename(rofile)}', working_leo)
    elif 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:
        #with np.load(rofile.replace('t01_RO', 't03_pair'), allow_pickle=True) as fid:
        #    reffiles = fid['reffile']
        #calculate_exph(rofile, reffiles, tarsource)

        try:
            with np.load(rofile.replace('t01_RO', 't03_pair'), allow_pickle=True) as fid:
                reffiles = fid['reffile']
            calculate_exph(rofile, reffiles, tarsource)
        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)

def cal_exph(working_leo, nProcess=0, tarsource=''):
    ropair = glob(f'./t01_RO_{working_leo}/*.npz')
    ropair.sort()
    # TC: updated 2026/10/01 (added 10 lines and modified 1 line, check existence of outfile)
    ropair = np.array(ropair)
    from re import sub
    idx = np.full((len(ropair),), True)
    for ii, rofile in enumerate(ropair):
        dt_rr = datetime.strptime(sub(r'.*\.(\d{4}\.\d{3})\..*', r'\1', rofile), '%Y.%j')
        out_dir = f's04_atmPhs/{dt_rr:%Y-%m-%d}'
        token = sub(r'.*_([A-Z][A-Z]\d\d)\.\d{4}\.\d{3}\.(\d{12})\d\d\.(\d{12})\d\d\.(A\d\d)\.([A-Z]\d\d)\.npz', r'\1_\2_\3_\5_\4', rofile)
        if any(glob(f'{out_dir}/*{token}*')):
            idx[ii] = False
    idx = np.where(idx)[0]
    allargs = list(zip(np.full((len(idx),), working_leo), idx, ropair[idx], np.full((len(idx),), tarsource)))
    # allargs = list(zip(np.full_like(ropair, working_leo), range(len(ropair)), ropair, np.full_like(ropair, tarsource)))
    if nProcess > 0:
        import multiprocessing as mp
        with mp.Pool(processes=nProcess) as pool:
            pool.map(cal_exph_for_parallel, allargs)
    else:
        for args in allargs:
            cal_exph_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)