from glob import glob
import os
from pathlib import Path
import numpy as np
from cosmic2.tools import calculate_exph
from traceback import extract_tb
from datetime import datetime, timedelta
from shutil import copyfile

working_dir = Path(__file__).absolute().parent if '__file__' in dir() else Path.cwd()
def log_message(msg, working_leo):
    with open(working_dir / 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 (rofile.parent.with_name(f't03_pair_{working_leo}')/rofile.name).is_file():
        if rofile.parent.name.startswith('t01'):
            log_message(f'-- No-REF skip for {qq}: {rofile.name} (newly onhold)', working_leo)
            copyfile(rofile,rofile.parent.with_name(f't03_ROhold_{working_leo}')/rofile.name)
        elif datetime.now() - datetime.fromtimestamp(rofile.stat().st_mtime) > timedelta(minutes=90):
            log_message(f'-- No-REF skip for {qq}: {rofile.name} (hold release)', working_leo)
            rofile.unlink()
        else:
            log_message(f'-- No-REF skip for {qq}: {rofile.name} (keep on list)', working_leo)
    else:
        try:
            with np.load(rofile.parent.with_name(f't03_pair_{working_leo}')/rofile.name, allow_pickle=True) as fid:
                reffiles = fid['reffile']
            calculate_exph(rofile, reffiles, tarsource)
            if rofile.parent.name.startswith('t03'):
                rofile.unlink()
        except Exception as e:
            if rofile.parent.name.startswith('t01'):
                log_message(f'-- Error detected for {qq}: {rofile.name} (newly onhold)', working_leo)
                copyfile(rofile, rofile.parent.with_name(f't03_ROhold_{working_leo}') / rofile.name)
            elif datetime.now() - datetime.fromtimestamp(rofile.stat().st_mtime) > timedelta(minutes=90):
                log_message(f'-- Error detected for {qq}: {rofile.name} (hold release)', working_leo)
                rofile.unlink()
            else:
                log_message(f'-- Error detected for {qq}: {rofile.name} (keep on list)', 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(leo, working_leo, nProcess=0, tarsource=''):
    ropair = np.hstack([list((working_dir / f't01_RO_{working_leo}').glob(f'*{leo}*.npz')),
                        list((working_dir / f't03_ROhold_{working_leo}').glob(f'*{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:
            pool.map(cal_exph_for_parallel, allargs)
    else:
        for args in allargs:
            cal_exph_for_parallel(args)