import logging
from typing import Optional

import h5py
import numpy as np
from ewokscore import Task

logger = logging.getLogger(__name__)


def _norm_path(path: str) -> str:
    return path if path.startswith("/") else "/" + path


def _load_arrays(
    filename: str,
    *,
    tth_path: str,
    det_path: str,
    mon_path: str,
    attn_path: Optional[str] = None,
    attn_factors: Optional[np.ndarray] = None,
):
    with h5py.File(filename, "r") as f:
        tth = np.asarray(f[_norm_path(tth_path)][()])
        det = np.asarray(f[_norm_path(det_path)][()])
        mon = np.asarray(f[_norm_path(mon_path)][()])
        if attn_path and attn_factors is not None:
            attn = np.asarray(f[_norm_path(attn_path)][()]).astype(int)
            det = det * np.take(attn_factors, attn)
    return tth, det, mon


class LoadRawData(
    Task,
    input_names=[
        "filename",  # HDF5 file path
        "tth_path",  # e.g. "1.1/measurement/tth"
        "det_path",  # e.g. "1.1/measurement/det"
        "mon_path",  # e.g. "1.1/measurement/mon4"
    ],
    optional_input_names=[
        "bg_filename",  # if provided, background is loaded from the SAME paths in this file
        "attn_path",  # e.g. "1.1/measurement/attn"
        "attn_factors",  # ndarray mapping attenuator index -> factor
        "itime",  # scalar seconds; forwarded to outputs
    ],
    output_names=[
        "tth",
        "det",
        "mon",
        "itime",
        "tth_bg",
        "det_bg",
        "mon_bg",
        "itime_bg",
    ],
):
    """
    Load raw arrays from an HDF5 file using direct dataset paths like '1.1/measurement/tth'.
    If bg_filename is provided, background is read from the SAME paths in that file.
    Otherwise, det_bg is zeros and mon_bg is ones (with tth_bg = tth).
    """

    def run(self):
        tth, det, mon = _load_arrays(
            self.inputs.filename,
            tth_path=self.inputs.tth_path,
            det_path=self.inputs.det_path,
            mon_path=self.inputs.mon_path,
            attn_path=getattr(self.inputs, "attn_path", None),
            attn_factors=getattr(self.inputs, "attn_factors", None),
        )
        self.outputs.tth = tth
        self.outputs.det = det
        self.outputs.mon = mon
        self.outputs.itime = getattr(self.inputs, "itime", None)

        if getattr(self.inputs, "bg_filename", None):
            tth_bg, det_bg, mon_bg = _load_arrays(
                self.inputs.bg_filename,
                tth_path=self.inputs.tth_path,
                det_path=self.inputs.det_path,
                mon_path=self.inputs.mon_path,
                attn_path=getattr(self.inputs, "attn_path", None),
                attn_factors=getattr(self.inputs, "attn_factors", None),
            )
        else:
            tth_bg = tth
            det_bg = np.zeros_like(det)
            mon_bg = np.ones_like(mon)

        self.outputs.tth_bg = tth_bg
        self.outputs.det_bg = det_bg
        self.outputs.mon_bg = mon_bg
        self.outputs.itime_bg = getattr(self.inputs, "itime", None)
