import numpy as np
from ewokscore import Task
from .utils.preprocessing import (
    correct_dark,
    norm_monitor,
    fit_direct_beam,
    footprint_correction,
)
from .utils.transforms import tth2q


class Preprocess(
    Task,
    input_names=[
        "tth",
        "det",
        "mon",
        "itime",
        "tth_bg",
        "det_bg",
        "mon_bg",
        "itime_bg",
        "dark_current",
        "direct_beam_amp",
        "beam_size",
        "sample_length",
        "tthmin",
        "energy",
    ],
    output_names=[
        "tthc",
        "detm",
        "detmn",
        "fpc",
        "detmnp",
        "tthcr",
        "Rraw",
        "qraw",
        "detm_bg",
        "detmn_bg",
        "tthc_bg",
        "detmnp_bg",
        "detmnpr",
        "tthcr_bg",
        "detmnpr_bg",
        "detmnpri_bg",
    ],
):
    """
    Applies dark-current correction, monitor normalization,
    direct-beam & footprint correction, BG subtraction → Rraw(qraw).
    """

    def _apply_pipeline(self, tth, det, mon, itime, amp_override=None):
        """
        Common pipeline: dark/monitor, direct beam, footprint correction.
        Returns (tthc, detm, detmn, fpc, detmnp, amp).
        """
        mondc = correct_dark(mon, self.inputs.dark_current, itime)
        detm = norm_monitor(det, mondc)

        # Use provided amplitude if not None; otherwise fit it
        amp = fit_direct_beam(tth, detm) if amp_override is None else amp_override

        detmn = detm / amp

        fpc = footprint_correction(
            tth / 2, self.inputs.beam_size, self.inputs.sample_length
        )
        # Safe division; fill 1.0 where fpc == 0
        detmnp = np.divide(detmn, fpc, out=np.full_like(detmn, 1.0), where=fpc != 0)
        return tth, detm, detmn, fpc, detmnp, amp

    def run(self):
        # Main XRR pipeline
        tthc, detm, detmn, fpc, detmnp, amp = self._apply_pipeline(
            self.inputs.tth,
            self.inputs.det,
            self.inputs.mon,
            self.inputs.itime,
            amp_override=self.inputs.direct_beam_amp,
        )

        # Background pipeline (same steps; reuse main amp)
        tthc_bg, detm_bg, detmn_bg, fpc_bg, detmnp_bg, _ = self._apply_pipeline(
            self.inputs.tth_bg,
            self.inputs.det_bg,
            self.inputs.mon_bg,
            self.inputs.itime_bg,
            amp_override=amp,
        )

        # Interpolate background onto XRR theta grid (sort BG grid first)
        if tthc_bg.size and detmnp_bg.size:
            order = np.argsort(tthc_bg)
            t_bg = tthc_bg[order]
            y_bg = detmnp_bg[order]
            detmnp_bg_i = np.interp(tthc, t_bg, y_bg, left=y_bg[0], right=y_bg[-1])
        else:
            detmnp_bg_i = np.zeros_like(tthc, dtype=detmnp.dtype)

        # Masks: main grid vs. BG's own grid
        mask_main = tthc > self.inputs.tthmin
        mask_bg = tthc_bg > self.inputs.tthmin

        tthcr = tthc[mask_main]
        detmnpr = detmnp[mask_main]
        detmnpri_bg = detmnp_bg_i[mask_main]  # BG on main grid (interpolated)

        tthcr_bg = tthc_bg[mask_bg]  # BG’s own grid after cut
        detmnpr_bg = detmnp_bg[mask_bg]  # BG corrected on its own grid

        Rraw = detmnpr - detmnpri_bg
        wavelength = 12398.42 / self.inputs.energy  # in Angstrom
        qraw = tth2q(tthcr, wavelength)

        # Outputs
        self.outputs.tthc = tthc
        self.outputs.tthc_bg = tthc_bg
        self.outputs.detm = detm
        self.outputs.detm_bg = detm_bg
        self.outputs.detmn = detmn
        self.outputs.detmn_bg = detmn_bg
        self.outputs.fpc = fpc
        self.outputs.detmnp = detmnp
        self.outputs.detmnpr = detmnpr
        self.outputs.detmnp_bg = detmnp_bg
        self.outputs.detmnpr_bg = detmnpr_bg
        self.outputs.detmnpri_bg = detmnpri_bg
        self.outputs.tthcr = tthcr
        self.outputs.tthcr_bg = tthcr_bg
        self.outputs.Rraw = Rraw
        self.outputs.qraw = qraw
