# scripts/helpers/transforms.py
import numpy as np
import pandas as pd


# ---------------------------------------------
# Missing values
# ---------------------------------------------
def add_missing_values(df: pd.DataFrame, cols: list[str], ratio: float = 0.1) -> pd.DataFrame:
    """
    Randomly replace a percentage of values with NaN in the given columns.

    Parameters
    ----------
    df : pd.DataFrame
        Input dataset.
    cols : list[str]
        Columns where missing values should be introduced.
    ratio : float
        Fraction of values (0–1) to set as NaN.

    Returns
    -------
    pd.DataFrame
        Modified DataFrame.
    """
    df = df.copy()
    for col in cols:
        mask = np.random.rand(len(df)) < ratio
        df.loc[mask, col] = np.nan
    return df


# ---------------------------------------------
# Outliers
# ---------------------------------------------

def add_outliers(df, col,ratio: float = 0.01, magnitude: float = 3.0, jitter: float = 0.2):
    """
    Multiply a random subset of rows in `col` by a factor near `magnitude`.

    - If magnitude > 1: factor in [magnitude*(1-jitter), magnitude]
    - If magnitude < 1: factor in [magnitude, min(1-eps, magnitude*(1+jitter))]
    - If magnitude == 1: no-op (returns df copy)

    jitter: relative band width around magnitude (0.2 => ~20% range)
    """
    if not (0 <= ratio <= 1):
        raise ValueError("ratio must be between 0 and 1")
    if magnitude is None or magnitude <= 0:
        raise ValueError("magnitude must be > 0")
    if jitter < 0:
        raise ValueError("jitter must be >= 0")

    df = df.copy()
    n = len(df)
    k = int(n * ratio)
    if k <= 0 or magnitude == 1:
        return df

    idx = np.random.choice(df.index, size=k, replace=False)

    # ensure float dtype (so multiplying by floats is valid)
    df[col] = df[col].astype("float64")

    if magnitude > 1:
        low = magnitude * (1 - jitter)
        high = magnitude
        if low <= 0:
            low = 1e-12
        factors = np.random.uniform(low, high, size=k)

    else:  # 0 < magnitude < 1
        low = magnitude
        high = magnitude * (1 + jitter)
        # keep factors strictly < 1
        eps = 1e-12
        high = min(high, 1 - eps)
        # if magnitude is already extremely close to 1, clamp to a tiny band
        if high <= low:
            high = min(1 - eps, low + (1 - low) * 0.01)  # small nudge up but still < 1
        factors = np.random.uniform(low, high, size=k)

    df.loc[idx, col] *= factors
    return df





# ---------------------------------------------
# Correlated features
# ---------------------------------------------
def add_correlated_feature(
    df: pd.DataFrame,
    base_col: str,
    new_col: str,
    noise_std: float = 0.2
) -> pd.DataFrame:
    """
    Add a new numeric column correlated with another column.

    new_col = base_col * (1 + normal(0, noise_std))

    Parameters
    ----------
    df : pd.DataFrame
        Input dataset.
    base_col : str
        Reference column.
    new_col : str
        Name of new correlated feature.
    noise_std : float
        Standard deviation of multiplicative Gaussian noise.

    Returns
    -------
    pd.DataFrame
        Modified DataFrame with the new correlated feature.
    """
    df = df.copy()
    df[new_col] = df[base_col] * (1 + np.random.normal(0, noise_std, len(df)))
    return df
