"""
Functions for metabolomics drift correction.
"""

import logging

import numpy as np
import pandas as pd
from scipy.interpolate import CubicSpline
from statsmodels.nonparametric.smoothers_lowess import lowess

logger = logging.getLogger(__name__)


def filter_features_by_qc(
    df: pd.DataFrame, qc_rows: list, threshold: float = 0.5
) -> pd.DataFrame:
    """
    Filter features in a DataFrame based on quality control (QC) completeness.

    This function removes columns (features) that do not meet a minimum number of valid
    (non-missing) QC values. The minimum number of required valid values is computed as
    `ceil(n_qc * (1 - threshold))`, where `n_qc` is the number of QC rows.

    :param pandas.DataFrame df:
        Input DataFrame with samples as rows and features as columns.
    :param list qc_rows:
        List of row indices corresponding to QC samples.
    :param float threshold:
        Fraction (between 0 and 1) indicating the maximum allowed proportion
        of missing QC values per feature. For example, `threshold=0.6` allows
        up to 60% of QC values to be missing; so a feature with 3 out of 5
        QC values present (40% present, 60% missing) would be retained.
        Defaults to 0.5.

    :returns pandas.DataFrame:
        Filtered DataFrame containing only columns with sufficient valid QC values.

    :raises ValueError:
        If `threshold` is not between 0 and 1.

    **Example**

        import pandas as pd
        import numpy as np

        df = pd.DataFrame({
            'A': [1.0, np.nan, 2.0],
            'B': [np.nan, 3.0, np.nan],
            'C': [4.0, np.nan, 6.0]
        }, index=['QC1', 'QC2', 'QC3'])

        filtered = filter_features_by_qc(df, qc_rows=['QC1', 'QC2', 'QC3'], threshold=0.5)
        print(filtered)
        # Output: columns with at least 2 valid QC values (since ceil(3 * (1 - 0.5)) = 2)

    """
    if not 0 <= threshold <= 1:  # check validity of threshold value
        raise ValueError(
            f"Threshold for QC filtering must be between 0 and 1, got {threshold}."
        )

    n_qc = len(qc_rows)
    min_valid = int(np.ceil(n_qc * (1 - threshold)))  # minimum valid QC points
    valid_counts = df.loc[qc_rows].notna().sum(axis=0)

    return df.loc[:, valid_counts >= min_valid]


def qc_rlsc_loess(
    x_qc,
    y_qc,
    x_all,
    always_use_default=False,
    default=0.75,
    alpha_candidates=np.arange(0.4, 1.01, 0.05),
):
    """
    Estimate a QC-based drift curve using LOESS smoothing with
    leave-one-out cross-validation (LOOCV) to select the optimal
    smoothing span (alpha).

    This function:
    1. Tests multiple LOESS spans (alpha values).
    2. Fits LOESS to QC points for each candidate span.
    3. Performs LOOCV to compute prediction error for each span.
    4. Selects the alpha producing the lowest LOOCV error.
    5. Fits LOESS once more using the best alpha.
    6. Interpolates the LOESS fit to all sample injection orders
       using a cubic spline, with clamping outside the QC range.

    Parameters
    ----------
    x_qc : array-like
        Injection order of QC samples (numeric).
    y_qc : array-like
        Intensity values of QC samples corresponding to `x_qc`.
    x_all : array-like
        Injection order of all samples (QCs + regular samples) in
        the same order as data rows.
    always_use_default: bool, optional
        If True, the alpha value 0.75 is used for all values.
        LOOCV is skipped. This option is less computationally heavy.
    alpha_candidates : list of float, optional
        List of LOESS smoothing parameters (fractions of data used
        in local regression) to evaluate during optimization.
        Default is [0.4, 0.6, 0.8, 1.0].

    Returns
    -------
    drift_curve : ndarray
        The estimated drift correction curve evaluated at each
        injection order in `x_all`. Values outside the QC range
        are clamped to the nearest in-range LOESS value.
    best_alpha : float
        The alpha value producing the lowest LOOCV error.

    Notes
    -----
    - LOESS fits enforce a minimum fraction of (λ + 1) / n, with λ=1
      for linear LOESS.
    - CubicSpline is used for interpolation without extrapolation.
      Out-of-range values are manually clamped.
    - Drift curve values are clipped to be strictly positive
      (minimum 1e-6) to prevent division instability.
    """

    best_alpha = None
    best_loocv_error = np.inf
    n = len(x_qc)

    if (
        not always_use_default
    ):  # do leave one out cross validation for determining best alpha
        for alpha in alpha_candidates:
            span = max(
                (1 + 1) / n, alpha
            )  # minimum alpha is (λ+1)/n, λ=1 for linear LOESS

            # Compute LOESS on full QC data
            loess_fit = lowess(y_qc, x_qc, frac=span, it=0, return_sorted=False)

            # Leave-one-out cross validation
            errors = []
            for i in range(n):
                x_cv = np.delete(x_qc, i)
                y_cv = np.delete(y_qc, i)
                loess_cv = lowess(
                    y_cv, x_cv, frac=span, it=0, return_sorted=False, xvals=[x_qc[i]]
                )
                errors.append((y_qc[i] - loess_cv[0]) ** 2)
            loocv_error = np.mean(errors)

            if loocv_error < best_loocv_error:
                best_loocv_error = loocv_error
                best_alpha = alpha

    # Fit LOESS again with best alpha for final curve
    if best_alpha is None or always_use_default:
        best_alpha = default  # fallback to reasonable default if optimization failed
        if not always_use_default:
            logger.warning(
                f"LOESS optimization failed for n={n}. Using default alpha={default}."
            )

    span = max((1 + 1) / n, best_alpha)
    loess_fit = lowess(y_qc, x_qc, frac=span, it=0, return_sorted=False)

    # Cubic spline interpolation for all points (samples + QCs)
    sort_idx = np.argsort(x_qc)
    x_qc_sorted = np.asarray(x_qc)[sort_idx]
    loess_fit_sorted = np.asarray(loess_fit)[sort_idx]
    cs = CubicSpline(
        x_qc_sorted,
        loess_fit_sorted,
        extrapolate=False,  # No extrapolation outside QC range (restrict to interpolation range only)
    )
    drift_curve = cs(x_all)

    # Clip drift_curve to the edge values to prevent NaNs or negatives
    x_min, x_max = x_qc_sorted[0], x_qc_sorted[-1]
    # For values outside QC range, hold the fitted value at the min/max QC positions (clamping)
    drift_curve[x_all < x_min] = loess_fit_sorted[0]
    drift_curve[x_all > x_max] = loess_fit_sorted[-1]
    drift_curve = np.clip(drift_curve, a_min=1e-6, a_max=None)  # Ensure no negatives

    return drift_curve, best_alpha


def run_loess_drift_correction(
    data,
    qc_rows,
    sample_rows,
    sample_order: pd.DataFrame,
    filter_percent: float = None,
    qc_min_threshold: int = 4,
    always_use_default=False,
    default=0.75,
):
    """
    Perform QC-based drift correction across multiple features using
    LOESS regression and spline interpolation.

    For each feature:
    1. Extract QC intensities and corresponding injection order.
    2. Optionally filter features based on QC completeness.
    3. Compute QC relative standard deviation (RSD).
    4. If sufficient QC points exist, estimate a drift curve using
       `qc_rlsc_loess`, finding the best alpha smoothing span
       with leave-one-out cross validation (LOOCV).
    5. Normalize all intensities by dividing by the drift curve and
       scaling to the QC median.
    6. Record drift parameters and correction metadata.

    Parameters
    ----------
    data : pandas.DataFrame
        Input intensity matrix with samples as rows and features as columns.
    qc_rows : list of str
        Row indices corresponding to QC injections.
    sample_rows : list of str
        Row indices corresponding to biological samples.
    sample_order : pandas.DataFrame
        Table mapping file names to injection order. Must contain
        columns "File Name" and "Sample ID". Sample ID must be
        numeric.
    filter_percent : float, optional
        Minimum proportion of QC values that must be non-missing for
        a feature to be retained (e.g., 0.6 means at least 60% of QCs
        must be present).
    qc_min_threshold : int, optional
        Minimum number of QC values required to perform drift
        correction. Features with fewer QCs are returned uncorrected.
    always_use_default: bool, optional
        If True, LOOCV is skipped and `default` is used for the smoothing span.
        This option is less computationally heavy.
    default : float, optional
        Default alpha to use when `always_use_default=True` or when LOOCV fails.
        Defaults to 0.75.

    Returns
    -------
    corrected_df : pandas.DataFrame
        Full input DataFrame with corrected values applied to
        sample_rows + qc_rows. Rows outside those arguments are
        returned unchanged. Features that fail QC requirements are
        returned unchanged.
    correction_info : dict
        Dictionary keyed by feature name, containing:
        - 'alpha': selected LOESS alpha (or None if skipped)
        - 'drift_curve': the evaluated drift correction vector
        - 'y_qc': QC intensities used
        - 'x_qc': QC injection orders
        - 'rsd_qc': QC relative standard deviation
        - 'median': QC median intensity (used for scaling)
        - 'y_all': original intensities
        - 'status': "corrected", "skipped_due_to_few_qcs", or error note

    Notes
    -----
    - Features with insufficient QC points are not corrected.
    - High QC RSD (>20%) is flagged but does not prevent correction.
    - Drift correction rescales intensities so that QC medians remain
      unchanged.
    - Sample names in `data` and `sample_order` must match exactly.
    """
    df_out = data.copy()
    df = data.copy()
    correction_info = {}  # Feature name -> dict with 'alpha' and 'drift_curve'
    all_rows = sample_rows + qc_rows

    try:
        sample_order["Sample ID"] = pd.to_numeric(
            sample_order["Sample ID"], errors="raise"
        ).astype(int)
    except ValueError as e:
        raise ValueError(
            "Sample ID column in sample order data contains non-integer values. All values in this column need to be numbers."
        ) from e

    # Injection order mapping (sample name -> injection order)
    injection_order_map = dict(
        zip(sample_order["File Name"], sample_order["Sample ID"])
    )
    x_all = np.array(
        [injection_order_map.get(sample, np.nan) for sample in all_rows]
    )  # Order of samples in order of file names

    if np.isnan(x_all).any():
        logger.warning(
            "Some of your samples don't have an associated sample order. They will be skipped.\nMake sure the names of your samples are identical in both your data frame and your sample order data."
        )

    if filter_percent is not None:  # Filter features by QC completeness
        df = filter_features_by_qc(
            df, qc_rows, threshold=(1 - filter_percent)
        )  # Only keeping features that have at least <percent> of all QCs not nan

    corrected_df = pd.DataFrame(
        index=all_rows, columns=df.columns, dtype=float
    )  # Initialises new df to be filled (for now has all nan), all rows including QC should be corrected

    # Loop through all features (columns) in filtered df
    for feature_name, col in df.loc[all_rows].items():

        # Y VALUES ARE INTENSITIES
        # X VALUES ARE RUN ORDER IDX

        # Intensities (y) for all rows in the samples (t, c and nan)
        y_all = col.values.astype(float)

        # Get intensities (y) and run order idx (x) for all QC data points
        y_qc = col.loc[qc_rows].values.astype(
            float
        )  # Gets all intensity values for feature for all qcs
        x_qc = np.array(
            [injection_order_map.get(s, np.nan) for s in qc_rows]
        )  # Gets the order

        # Remove NaNs from QC for fitting
        valid_mask = ~np.isnan(y_qc) & ~np.isnan(x_qc)
        y_qc_valid = y_qc[valid_mask]
        x_qc_valid = x_qc[valid_mask]

        # Calculate RSD
        eps = 1e-9
        rsd_qc = 100 * np.std(y_qc_valid) / (np.mean(y_qc_valid) + eps)
        if rsd_qc > 30:
            logger.info(
                f"Flagging feature {feature_name} due to too high QC RSD. (Feature not skipped). RSD: {rsd_qc}"
            )  # RSD>20% suggests qc intensities reflect instrumental drift, not biology

        if len(y_qc_valid) < qc_min_threshold:
            # Skip correction but preserve the column
            logger.info(f"Skipping correction due to few QCs: {feature_name}")
            corrected_df[feature_name] = y_all  # Insert the uncorrected values
            correction_info[feature_name] = {
                "alpha": None,
                "drift_curve": None,
                "y_qc": y_qc_valid.tolist(),
                "x_qc": x_qc_valid.tolist(),
                "rsd_qc": rsd_qc,
                "status": "skipped_due_to_few_qcs",
            }
            continue

        try:  # DRIFT CORRECTION
            # Calculate drift curve
            drift_curve, best_alpha = qc_rlsc_loess(
                x_qc_valid,
                y_qc_valid,
                x_all,
                always_use_default=always_use_default,
                default=default,
            )  # Calculate curve and alpha value

            # Remove NaNs from y_all for indexing (just keep for drift correction) -> only if they are not nan in either data or drift curve
            valid_mask = ~np.isnan(y_all) & ~np.isnan(drift_curve)

            # Normalize all intensities using drift curve
            median_qc = np.median(y_qc_valid)
            corrected = (
                y_all.copy()
            )  # Preserve original values where correction is not possible
            corrected[valid_mask] = (
                y_all[valid_mask] / drift_curve[valid_mask]
            ) * median_qc  # Only valid positions get corrected; skipped samples retain original intensity

            corrected_df[feature_name] = corrected

            correction_info[feature_name] = {
                "alpha": best_alpha,
                "drift_curve": drift_curve.tolist(),
                "y_qc": y_qc_valid.tolist(),  # Store QC values
                "x_qc": x_qc_valid.tolist(),  # Store QC injection orders
                "rsd_qc": rsd_qc,
                "median": median_qc,
                "y_all": y_all.tolist(),
                "new_values": corrected.tolist(),
                "status": "corrected",
            }

            logger.info(f"Corrected {feature_name} with alpha {best_alpha}.")

        except Exception as e:
            corrected_df[feature_name] = y_all  # Preserve original values
            correction_info[feature_name] = {
                "alpha": None,
                "drift_curve": None,
                "y_qc": y_qc_valid.tolist(),
                "x_qc": x_qc_valid.tolist(),
                "rsd_qc": rsd_qc,
                "status": f"error: {e}",
            }

            logger.error(f"Skipping feature {feature_name} due to error: {e}")
            continue

    df_out.loc[all_rows, corrected_df.columns] = corrected_df.values
    logger.info(
        "\nAll done. For further information on the corrected values, check the second returned object."
    )
    return df_out, correction_info
