import numpy as np
import xarray as xr

from xmris.core.config import ATTRS, COORDS, DIMS
from xmris.core.utils import _check_dims, as_variable
from xmris.processing.fourier import fft, fftshift, ifft, ifftshift


def to_spectrum(
    da: xr.DataArray, dim: str = DIMS.time, out_dim: str = DIMS.frequency
) -> xr.DataArray:
    """
    Convert a time-domain Free Induction Decay (FID) to a frequency-domain spectrum.

    The stored, digital FID signal can be processed by a discrete Fourier
    transformation (DFT) to produce a digital MR spectrum.
    This function applies the FFT along the specified time dimension and shifts
    the zero-frequency component to the center of the spectrum.

    Parameters
    ----------
    da : xr.DataArray
        The input time-domain FID data.
    dim : str, optional
        The time dimension to transform, by default `DIMS.time`.
    out_dim : str, optional
        The name of the resulting frequency dimension, by default `DIMS.frequency`.

    Returns
    -------
    xr.DataArray
        The frequency-domain spectrum with centered zero-frequency coordinates.
    """
    _check_dims(da, dim, "to_spectrum")

    # 1. Standard FFT (handles transform and creates unshifted frequency coords)
    da_freq = fft(da, dim=dim, out_dim=out_dim)

    # 2. Shift the frequency domain to center the DC component
    da_spectrum = fftshift(da_freq, dim=out_dim)

    return da_spectrum


def to_fid(
    da: xr.DataArray, dim: str = DIMS.frequency, out_dim: str = DIMS.time
) -> xr.DataArray:
    """
    Convert a frequency-domain spectrum back to a time-domain FID.

    This is the mathematical inverse of `to_spectrum`. It inverse-shifts the
    data to position 0 Hz at the array boundary, computes the IFFT, and
    reconstructs strictly positive time coordinates.

    Parameters
    ----------
    da : xr.DataArray
        The input frequency-domain data.
    dim : str, optional
        The frequency dimension to transform, by default `DIMS.frequency`.
    out_dim : str, optional
        The name of the resulting time dimension, by default `DIMS.time`.

    Returns
    -------
    xr.DataArray
        The reconstructed time-domain FID.
    """
    _check_dims(da, dim, "to_fid")

    # 1. Inverse shift the frequency domain to put the DC component at index 0
    # This prepares the data for the standard IFFT algorithm
    da_unshifted = ifftshift(da, dim=dim)

    # 2. Apply IFFT
    # The output is naturally ordered [t=0, t=1, ... t=N-1]
    da_fid = ifft(da_unshifted, dim=dim, out_dim=out_dim)

    # 3. Reconstruct the strictly positive time coordinates [0, T_acq]
    if dim in da.coords:
        freqs = da.coords[dim].values
        n_points = len(freqs)
        if n_points > 1:
            # Calculate dwell time (dt) based on the sampling theorem:
            # Spectral Width (SW) = 1/dt
            # SW = n_points * df
            df = abs(freqs[1] - freqs[0])
            dt = 1.0 / (n_points * df)

            t_coords = np.arange(n_points) * dt

            # Re-inject metadata if mapping to standard DIMS.time
            term = COORDS.time if out_dim == DIMS.time else None

            if term:
                time_var = as_variable(term, out_dim, t_coords)
            else:
                time_var = xr.Variable(out_dim, t_coords)

            da_fid = da_fid.assign_coords({out_dim: time_var})

    return da_fid


def apodize_exp(da: xr.DataArray, dim: str = DIMS.time, lb: float = 1.0) -> xr.DataArray:
    """
    Apply an exponential weighting filter function for line broadening.

    During apodization, the time-domain FID signal $f(t)$ is multiplied with a filter
    function $f_{filter}(t) = e^{-t/T_L}$. This improves the Signal-to-Noise Ratio (SNR)
    because data points at the end of the FID, which primarily contain noise, are
    attenuated. The time constant $T_L$ is calculated from the desired line broadening
    in Hz.


    Parameters
    ----------
    da : xr.DataArray
        The input time-domain data.
    dim : str, optional
        The dimension corresponding to time, by default `DIMS.time`.
    lb : float, optional
        The desired line broadening factor in Hz, by default 1.0.

    Returns
    -------
    xr.DataArray
        A new apodized DataArray, preserving coordinates and attributes.
    """
    _check_dims(da, dim, "apodize_exp")

    t = da.coords[dim]

    # Calculate exponential filter: exp(-t / T_L) where T_L = 1 / (pi * lb)
    # This simplifies to: exp(-pi * lb * t)
    weight = np.exp(-np.pi * lb * t)

    # Functional application (transpose ensures broadcasting doesn't scramble axis order)
    da_apodized = (da * weight).transpose(*da.dims).assign_attrs(da.attrs)

    # Record lineage
    da_apodized.attrs[ATTRS.apodization_lb] = lb

    return da_apodized


def apodize_lg(
    da: xr.DataArray, dim: str = DIMS.time, lb: float = 1.0, gb: float = 1.0
) -> xr.DataArray:
    """
    Apply a Lorentzian-to-Gaussian transformation filter.

    This filter converts a Lorentzian line shape to a Gaussian line shape, which decays
    to the baseline in a narrower frequency range. The time-domain FID
    is multiplied by $e^{+t/T_L}e^{-t^2/T_G^2}$. The time constants $T_L$ and $T_G$
    are derived from the `lb` and `gb` frequency-domain parameters.

    Parameters
    ----------
    da : xr.DataArray
        The input time-domain data.
    dim : str, optional
        The dimension corresponding to time, by default `DIMS.time`.
    lb : float, optional
        The Lorentzian line broadening to cancel in Hz, by default 1.0.
    gb : float, optional
        The Gaussian line broadening to apply in Hz, by default 1.0.

    Returns
    -------
    xr.DataArray
        A new apodized DataArray, preserving coordinates and attributes.
    """
    _check_dims(da, dim, "apodize_lg")

    t = da.coords[dim]

    # Calculate Lorentzian cancellation: exp(+t / T_L)
    # T_L = 1 / (pi * lb)
    weight_lorentzian = np.exp(np.pi * lb * t)

    # Calculate Gaussian broadening: exp(-t^2 / T_G^2)
    # T_G = 2 * sqrt(ln(2)) / (pi * gb)
    if gb != 0:
        t_g = (2 * np.sqrt(np.log(2))) / (np.pi * gb)
        weight_gaussian = np.exp(-(t**2) / (t_g**2))
    else:
        weight_gaussian = 1.0

    weight = weight_lorentzian * weight_gaussian

    da_apodized = (da * weight).transpose(*da.dims).assign_attrs(da.attrs)

    # Record lineage
    da_apodized.attrs[ATTRS.apodization_lb] = lb
    da_apodized.attrs[ATTRS.apodization_gb] = gb

    return da_apodized


def zero_fill(
    da: xr.DataArray,
    dim: str = DIMS.time,
    target_points: int = 1024,
    position: str = "end",
) -> xr.DataArray:
    """
    Pad the specified dimension with zero amplitude points.

    Artificially extends the data with zeros to interpolate and increase the
    apparent digital resolution of the resulting spectrum.



    Parameters
    ----------
    da : xr.DataArray
        The input data.
    dim : str, optional
        The dimension along which to pad zeros, by default `DIMS.time`.
    target_points : int, optional
        The total number of points desired after padding, by default 1024.
    position : {"end", "symmetric"}, optional
        Where to apply the zeros. Use "end" for time-domain FIDs, and
        "symmetric" for spatial frequency domains like k-space. By default "end".

    Returns
    -------
    xr.DataArray
        A new DataArray padded with zeros to the target length, preserving metadata.
    """
    _check_dims(da, dim, "zero_fill")

    current_points = da.sizes[dim]
    if target_points <= current_points:
        return da.copy()

    pad_size = target_points - current_points

    # Determine padding distribution based on position
    if position == "end":
        pad_width = (0, pad_size)
    elif position == "symmetric":
        pad_left = pad_size // 2
        pad_right = pad_size - pad_left
        pad_width = (pad_left, pad_right)
    else:
        raise ValueError("`position` must be either 'end' or 'symmetric'.")

    # Pad with constant 0s
    da_padded = da.pad({dim: pad_width}, mode="constant", constant_values=0)

    # Extrapolate coordinates linearly if they exist
    if dim in da.coords:
        old_coords = da.coords[dim].values
        if len(old_coords) > 1:
            delta = old_coords[1] - old_coords[0]

            if position == "end":
                new_coords = old_coords[0] + np.arange(target_points) * delta
            else:  # symmetric
                start_coord = old_coords[0] - (pad_width[0] * delta)
                new_coords = start_coord + np.arange(target_points) * delta

            # Try to match dim against known config vocabulary to rebuild metadata
            term = None
            for candidate in [COORDS.time, COORDS.frequency, COORDS.chemical_shift]:
                if candidate == dim:
                    term = candidate
                    break

            if term:
                new_var = as_variable(term, dim, new_coords)
            else:
                # Retain old raw attributes if it's a custom dimension
                new_var = xr.Variable(dim, new_coords, attrs=da.coords[dim].attrs)

            da_padded = da_padded.assign_coords({dim: new_var})

    # Restore lost attrs and record lineage
    da_padded = da_padded.assign_attrs(da.attrs)
    da_padded.attrs[ATTRS.zero_fill_target] = target_points
    da_padded.attrs[ATTRS.zero_fill_position] = position

    return da_padded
