import numpy as np
import json, os, glob, yaml
import argparse, pickle, h5py
from pathlib import Path
from sfdata import SFDataFile, SFDataFiles

from alvra_tools.channels import *
from alvra_tools.load_data import load_data_compact_pump_probe
#from alvra_tools.utils import *
#from alvra_tools.XAS_functions import *
#from alvra_tools.XAS_utils import *

def get_acq_from_pth(path):
    import re
    return re.search(r"acq(\d+)", path).group(1).zfill(4)

parser = argparse.ArgumentParser()
parser.add_argument("--input", "-i", required=True)
parser.add_argument("channels_file")
args = parser.parse_args()

filename = Path(args.input)
output_file = Path(str(filename).replace("/raw/", "/res/processed/").replace(".h5",".PROCESSED.h5") )
output_file.parent.mkdir(parents=True, exist_ok=True, mode=0o775)

acq = get_acq_from_pth(args.input)
acq_meta_pth = Path(str(filename).replace("/data/ac", "/meta/ac")).parent
acq_info_pth = acq_meta_pth.joinpath(f"acq{acq}.json")

#if not acq_info_pth.exists():
#    write_apo_meta(apo_meta_pth, apo_meta_dic)
#    msg = f"Acq info(.json)  file does not exist: {acq_info_pth}"
#    raise IOError(msg)

with open(acq_info_pth, "r") as f:
    acq_parameters = json.load(f)
    rbk_value = acq_parameters.get("scan_info", {}).get("scan_readbacks", 0)
    scan_value = acq_parameters.get("scan_info", {}).get("scan_values", 0)
    xlabel = acq_parameters.get("scan_info", {}).get("name", 'None')
    units = acq_parameters.get("scan_info", {}).get("units", 'None')

with open(args.channels_file, "r") as f:
    channels_aliases = yaml.safe_load(f)

channels_mapping = channels_aliases["channels"]
channels_pp = list(channels_mapping[ch] for ch in channels_aliases["channels_pump"])
channels_all = channels_pp

with SFDataFile(args.input) as step:
    resultsPP, results, _, _ = load_data_compact_pump_probe(channels_pp, channels_all, step)

res = {}
for name, chan in channels_mapping.items():
    ch = resultsPP[chan]
    if name in channels_aliases["channels_pump"]:
        res[name + "_pump"] = ch.pump
    if name in channels_aliases["channels_unpump"]:
        res[name + "_unpump"] = ch.unpump

with h5py.File(output_file, 'w') as f:
    gr1 = f.create_group("allShots")
    for key, value in results.items():
        gr1[key] = value
    gr2 = f.create_group("OnOff")
    for key, value in res.items():
        gr2[key] = value
    gr3 = f.create_group("meta")
    gr3["readback_value"] = np.atleast_1d(rbk_value)
    gr3["scan_value"] = np.atleast_1d(scan_value)
    gr3["xlabel"] = xlabel
    gr3["units"]  = units