import numpy as np
from ewokscore import Task
from .utils.signal import (
    smooth,
    find_minima_scipy,
    baseline,
    subtract_baseline,
    invert_signs,
)
from .utils.transforms import IFT_slow


class InverseFT(
    Task,
    input_names=[
        "q",
        "R",
        "zrange",
        "peak_distance",
        "peak_prominence",
        "baseline_type",
    ],
    output_names=[
        "z",
        "d_rho",
        "delta_rho",
        "delta_rho_bl",
        "delta_rho_nc",
        "baseline",
        "minima",
        "phase",
        "phase_bl",
        "Rq4_bl",
        "Rq4_sm",
    ],
):
    """
    Takes cleaned q & R, builds sqrt(R·q⁴), subtracts baseline & flips signs,
    then performs IFT and integrates to yield Δρ(z).
    """

    def _invert_baseline(self, y, q):
        """
        Given y = R·q⁴ and q, return:
          - sqrt_nc    : √y  (no baseline)
          - sqrt_bl    : √|y – baseline| with alternating sign
        Also returns the baseline & minima for diagnostics.
        """

        # plain sqrt
        sqrt_nc = np.sqrt(y)

        # smooth & find minima
        sm = smooth(y)
        minima = find_minima_scipy(
            sm,
            distance=self.inputs.peak_distance or 30,
            prominence=self.inputs.peak_prominence or 1e-8,
        )

        # build & subtract baseline
        b = baseline(sm, q, minima, method=self.inputs.baseline_type)
        y_bl = subtract_baseline(y, b)

        # sqrt + sign flips
        sqrt_bl = invert_signs(np.sqrt(y_bl), minima)
        phase = invert_signs(np.sqrt(y), minima)

        return sqrt_nc, sqrt_bl, b, minima, phase, y_bl, smooth(y_bl)

    def run(self):
        from scipy.integrate import cumulative_trapezoid

        # 1) Build y = R·q⁴
        y = self.inputs.R * self.inputs.q**4

        # 2) Invert baseline
        sqrt_nc, sqrt_bl, baseline_vals, minima, phase, Rq4_bl, Rq4_sm = (
            self._invert_baseline(y, self.inputs.q)
        )

        # 3) Inverse Fourier Transform on baseline‐subtracted
        d_rho, z = IFT_slow(phase, self.inputs.q, self.inputs.zrange)

        # 4) Integrate to get Δρ(z)
        delta_rho = cumulative_trapezoid(d_rho.imag, z, initial=0)

        # nc
        d_rho_nc, z_nc = IFT_slow(sqrt_nc, self.inputs.q, self.inputs.zrange)
        delta_rho_nc = cumulative_trapezoid(d_rho_nc.imag, z_nc, initial=0)

        d_rho_bl, z = IFT_slow(sqrt_bl, self.inputs.q, self.inputs.zrange)
        delta_rho_bl = cumulative_trapezoid(d_rho_bl.imag, z, initial=0)

        # 5) Assign outputs
        self.outputs.z = z
        self.outputs.d_rho = d_rho
        self.outputs.delta_rho = delta_rho
        self.outputs.delta_rho_bl = delta_rho_bl
        self.outputs.delta_rho_nc = delta_rho_nc
        self.outputs.baseline = baseline_vals
        self.outputs.minima = minima
        self.outputs.phase = phase
        self.outputs.phase_bl = sqrt_bl
        self.outputs.Rq4_bl = Rq4_bl
        self.outputs.Rq4_sm = Rq4_sm
