from __future__ import annotations

import math
import warnings

import numpy as np
from scipy.fft import irfft, next_fast_len, rfft
from scipy.signal import fftconvolve

from dp_accounting.pld.common import compute_self_convolve_bounds

from PLD_accounting.core_utils import (
    convolve_infinite_masses,
    enforce_mass_conservation,
    self_convolve_infinite_mass,
    stable_isclose,
)
from PLD_accounting.discrete_dist import (
    LinearDiscreteDist,
)
from PLD_accounting.types import BoundType
from PLD_accounting.utils import binary_self_convolve


def calc_FFT_window_size(pmf: np.ndarray, T: int, tail_truncation: float) -> tuple[int, int]:
    """Calculate FFT window bounds for T-fold self-convolution with fallback."""
    # `compute_self_convolve_bounds` gives a Chernoff-style window [L, U] that
    # should contain all but `tail_truncation` mass of the T-fold convolution.
    L, U = compute_self_convolve_bounds(pmf, T, tail_truncation)
    window_size = U - L + 1

    if not (0 < window_size < float("inf")):
        L = 0
        n = len(pmf)
        # Fallback to the exact full-support FFT length when the bound becomes
        # numerically unusable for extreme truncation parameters.
        window_size = T * (n - 1) + 1
        warnings.warn(
            "calc_FFT_window_size: Chernoff bounds failed "
            f"(tail_truncation={tail_truncation:.3e}, T={T}). "
            f"Using fallback L=0, window_size={window_size:,} (n={n})."
        )

    return int(L), int(window_size)


def FFT_convolve(
    dist_1: LinearDiscreteDist,
    dist_2: LinearDiscreteDist,
    tail_truncation: float,
    bound_type: BoundType,
) -> LinearDiscreteDist:
    """Convolve two linear-grid distributions via FFT."""
    if not isinstance(dist_1, LinearDiscreteDist) or not isinstance(dist_2, LinearDiscreteDist):
        raise TypeError("FFT_convolve requires LinearDiscreteDist inputs")

    if not np.any(dist_1.PMF_array) or not np.any(dist_2.PMF_array):
        raise ValueError("FFT convolution requires nonzero finite mass in both inputs")
    if not stable_isclose(dist_1.x_gap, dist_2.x_gap):
        raise ValueError(f"Grid spacing must match: w1={dist_1.x_gap:.12g} vs w2={dist_2.x_gap:.12g}")

    width = dist_1.x_gap
    conv_x_min = dist_1.x_min + dist_2.x_min
    L = dist_1.PMF_array.size + dist_2.PMF_array.size - 1

    conv_PMF = fftconvolve(dist_1.PMF_array, dist_2.PMF_array, mode="full")
    conv_PMF[conv_PMF < 0] = 0.0

    # Zero any mass outside the mathematically reachable support. FFT roundoff
    # can otherwise leave tiny "ghost" mass in bins that should be impossible.
    nz1 = np.nonzero(dist_1.PMF_array)[0]
    nz2 = np.nonzero(dist_2.PMF_array)[0]
    min_idx = int(nz1[0] + nz2[0])
    max_idx = int(nz1[-1] + nz2[-1])
    conv_PMF[:min_idx] = 0.0
    if max_idx + 1 < conv_PMF.size:
        conv_PMF[max_idx + 1 :] = 0.0

    finite_prob_1 = math.fsum(map(float, dist_1.PMF_array))
    finite_prob_2 = math.fsum(map(float, dist_2.PMF_array))
    current_finite_mass = math.fsum(map(float, conv_PMF))
    if current_finite_mass <= 0.0:
        raise ValueError("FFT convolution produced zero finite mass")
    # Renormalize finite mass before reattaching the analytically computed
    # infinity masses. This corrects small drift from FFT arithmetic/clipping.
    conv_PMF *= finite_prob_1 * finite_prob_2 / current_finite_mass

    expected_neg_inf, expected_pos_inf = convolve_infinite_masses(
        dist_1.p_neg_inf,
        dist_1.p_pos_inf,
        dist_2.p_neg_inf,
        dist_2.p_pos_inf,
    )
    conv_PMF, p_neg_inf, p_pos_inf = enforce_mass_conservation(
        PMF_array=conv_PMF,
        expected_neg_inf=expected_neg_inf,
        expected_pos_inf=expected_pos_inf,
        bound_type=bound_type,
    )

    return LinearDiscreteDist(
        x_min=conv_x_min,
        x_gap=width,
        PMF_array=conv_PMF,
        p_neg_inf=p_neg_inf,
        p_pos_inf=p_pos_inf,
    ).truncate_edges(tail_truncation, bound_type)


def FFT_self_convolve(
    dist: LinearDiscreteDist,
    T: int,
    tail_truncation: float,
    bound_type: BoundType,
    use_direct: bool,
) -> LinearDiscreteDist:
    """T-fold self-convolution via FFT with optional direct exponentiation path."""
    if not isinstance(dist, LinearDiscreteDist):
        raise TypeError("FFT_self_convolve requires LinearDiscreteDist input")

    if use_direct:
        return _fft_self_convolve_direct(
            dist=dist,
            T=T,
            tail_truncation=tail_truncation,
            bound_type=bound_type,
        )

    self_conv = binary_self_convolve(
        dist=dist,
        T=T,
        tail_truncation=tail_truncation,
        bound_type=bound_type,
        convolve=FFT_convolve,
    )
    if not isinstance(self_conv, LinearDiscreteDist):
        raise TypeError(f"Expected LinearDiscreteDist from FFT self-convolution, got {type(self_conv)}")
    return self_conv


def _fft_self_convolve_direct(
    dist: LinearDiscreteDist,
    T: int,
    tail_truncation: float,
    bound_type: BoundType,
) -> LinearDiscreteDist:
    # Half the tail budget is spent inside the FFT windowing logic, and the
    # remainder is consumed by the final truncation step below.
    tail_truncation /= 2

    finite_mass = math.fsum(map(float, dist.PMF_array))
    # The Chernoff window calculation expects a normalized finite PMF, so the
    # tail target must be rescaled when some mass already sits at infinity.
    normalized_PMF = dist.PMF_array / finite_mass
    tail_truncation_rescaled = tail_truncation / finite_mass

    L, window_size = calc_FFT_window_size(normalized_PMF, T, tail_truncation_rescaled)

    fft_size = next_fast_len(max(window_size, dist.PMF_array.size))
    raw_conv = irfft(rfft(dist.PMF_array, n=fft_size) ** T, n=fft_size)
    raw_conv[raw_conv < 0] = 0.0
    # `L` is the left edge of the retained convolution window. Rolling aligns
    # that window to index 0 so truncation logic can work in-place.
    rolled_conv = np.roll(raw_conv, -L)

    conv_neg_inf, conv_pos_inf = self_convolve_infinite_mass(dist.p_neg_inf, dist.p_pos_inf, T)
    if bound_type == BoundType.DOMINATES:
        # For an upper bound, any dropped left-tail mass is pushed to +inf.
        cumsum = np.cumsum(rolled_conv)
        left_tail_ind = int(np.searchsorted(cumsum, tail_truncation, side="right"))
        shifted_mass = math.fsum(map(float, rolled_conv[:left_tail_ind]))
        rolled_conv[:left_tail_ind] = 0.0
        right_tail_mass = math.fsum(map(float, rolled_conv[window_size:]))
        conv_pos_inf += shifted_mass + right_tail_mass
    elif bound_type == BoundType.IS_DOMINATED:
        # For a lower bound, dropped right-tail mass moves to -inf, while any
        # overflow beyond the retained FFT window is folded onto the last kept
        # finite bin to preserve domination direction.
        cumsum = np.cumsum(rolled_conv[::-1])
        right_tail_ind = rolled_conv.size - 1 - int(np.searchsorted(cumsum, tail_truncation, side="right"))
        shifted_mass = math.fsum(map(float, rolled_conv[right_tail_ind + 1 :]))
        rolled_conv[right_tail_ind + 1 :] = 0.0
        conv_neg_inf += shifted_mass

        right_tail_mass = math.fsum(map(float, rolled_conv[window_size:]))
        rolled_conv[min(window_size, right_tail_ind) - 1] += right_tail_mass
    else:
        raise ValueError(f"Unknown BoundType: {bound_type}")

    x_min = dist.x_min * T + L * dist.x_gap
    PMF_conv = rolled_conv[:window_size]
    PMF_conv, p_neg_inf_final, p_pos_inf_final = enforce_mass_conservation(
        PMF_array=PMF_conv,
        expected_neg_inf=conv_neg_inf,
        expected_pos_inf=conv_pos_inf,
        bound_type=bound_type,
    )

    return LinearDiscreteDist(
        x_min=x_min,
        x_gap=dist.x_gap,
        PMF_array=PMF_conv,
        p_neg_inf=p_neg_inf_final,
        p_pos_inf=p_pos_inf_final,
    ).truncate_edges(0.0, bound_type)
