from glob import glob
import os
import numpy as np
from spire.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 cal_exph_for_parallel(args):
    working_leo = args[0]
    qq = args[1]
    rofile = args[2]
    tarsource = args[3]
    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)
        # TC: updated 2026/04/28 (1 line, return False as it failed)
        return False
    else:
        try:
            with np.load(rofile.replace('t01_RO', 't03_pair'), allow_pickle=True) as fid:
                reffiles = fid['reffile']
            calculate_exph(rofile, reffiles, tarsource)
            # TC: updated 2026/04/28 (1 line, return True as it succeeded)
            return True
        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)
            # TC: updated 2026/04/28 (1 line, return False as it failed)
            return False

def cal_exph(working_leo='', nProcess=0, tarsource=''):
    ropair = glob(f'./t01_RO_{working_leo}/*.npz')
    ropair.sort()
    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:
            # TC: updated 2026/04/28 (1 line, take the result status)
            results = pool.map(cal_exph_for_parallel, allargs)
            #pool.map(cal_exph_for_parallel, allargs)
    else:
        # TC: updated 2026/04/28 (4 lines, take the result status and return for further process)
        results = np.full((len(allargs),),False,dtype=bool)
        for nn, args in enumerate(allargs):
            results[nn] = cal_exph_for_parallel(args)
    return np.array(results)
        #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)