from ewokscore import Task
from lmfit.models import QuadraticModel


class CleanSignal(
    Task,
    input_names=["qraw", "Rraw", "qmax", "qvaluesRangesToRemove"],
    output_names=["q", "R"],
):
    """
    Ensure positivity of Rraw, crop to qmax if given, and
    remove specified q‐ranges via quadratic interpolation.
    """

    def run(self):
        # 1) Keep only positive Rraw
        qc = self.inputs.qraw
        Rc = self.inputs.Rraw
        pos = Rc > 0
        qc, Rc = qc[pos], Rc[pos]

        # 2) Crop to qmax if specified
        if self.inputs.qmax is not None:
            mask = qc < self.inputs.qmax
            qc, Rc = qc[mask], Rc[mask]

        # 3) Remove bad q‐ranges
        for qmin, qmax in self.inputs.qvaluesRangesToRemove or []:
            # keep points outside the bad range
            keep = (qc < qmin) | (qc > qmax)
            q2, R2 = qc[keep], Rc[keep]
            # select points just outside the range for interpolation
            sel = (q2 > qmin - 0.02) & (q2 < qmax + 0.02)
            x, y = q2[sel], q2[sel] ** 4 * R2[sel]
            # fit quadratic
            model = QuadraticModel()
            fit = model.fit(y, x=x)
            # replace inside-range values
            bad = (qc > qmin) & (qc < qmax)
            qc_bad = qc[bad]
            Rc[bad] = fit.eval(x=qc_bad) / qc_bad**4

        # 4) Assign outputs
        self.outputs.q = qc.copy()
        self.outputs.R = Rc.copy()
