import argparse
import json
from pathlib import Path

import numpy as np
from sfdata import SFDataFile, SFProcFile

from utils.helpers import get_acq_from_pth

parser = argparse.ArgumentParser()
parser.add_argument("--input", "-f", required=True)

parser.add_argument("--background", "-b", required=True)

args = parser.parse_args()

filename = Path(args.input)
out_pth = Path(
    str(filename).replace("/raw/", "/res/processed/").replace(".CAMERAS.h5", ".PROCESSED.h5")
    )
out_pth.parent.mkdir(parents=True, exist_ok=True)

bgs = {}
with SFDataFile(args.background) as fdata:
    for name, chan in fdata.items():
        if chan.ndim > 1:
            bgs[name] = np.mean(chan.data, axis=0)

sums = {}
with SFDataFile(args.input) as fdata, SFProcFile(out_pth) as fproc:
    for name, chan in fdata.items():
        if chan.ndim > 1:
            proj = np.sum(chan.data - bgs[name], axis=1)
            fproc[f"{name}.projection"] = (chan.pids, proj)
            sums[name] = np.sum(chan.data - bgs[name])

meta_pth = Path(str(filename).replace("/raw/", "/res/processed/").replace("/data/ac", "/meta/ac"))
acq = get_acq_from_pth(str(filename))
meta_pth = meta_pth.with_name(f"apo_acq{acq}.json")

meta_pth.parent.mkdir(parents=True, exist_ok=True)

meta_data = {
    "some key": "some value"
    }
meta_data.update(sums)

with open(meta_pth, "w") as f:
    json.dump(meta_data, f)