"""
Adaptive scaling module - Feature-specific optimal scalers
Each feature uses the most appropriate scaling method based on its distribution
"""
import os, sys
import numpy as np
import pandas as pd
from sklearn.preprocessing import PowerTransformer, StandardScaler, RobustScaler as SklearnRobustScaler

# 프로젝트 루트 경로 추가
_current_dir = os.path.dirname(os.path.abspath(__file__))
_project_root = os.path.dirname(os.path.dirname(_current_dir))
if _project_root not in sys.path:
    sys.path.insert(0, _project_root)

from admet_ft._modules.utils import DIDB_FILTER_COLS

# 화학적 도메인 기반 feature 범위 정의
FEATURE_RANGES = {
    "logP": (-10, 10),                    # Octanol-water partition coefficient
    "pKa": (0, 14),                        # Acid dissociation constant (pH scale)
    "solubility": (0, 10000),              # mg/L
    "permeability": (0, 1000),             # 10^-6 cm/s (Caco-2)
    "plasma_protein_binding": (0, 100),    # percentage
    "fu_in_vitro": (0, 1),                 # fraction unbound
}

SCALER_FILE_NAME = {
    'zscore': 'scaler_config.csv',
    'power': 'scaler_power_config.csv',
    'minmax': 'scaler_minmax_config.csv',
    'adapt': 'scaler_adapt.csv'
}

# Feature-specific scaler configuration
ADAPTIVE_SCALER_CONFIG = {
    'logP': {
        'type': 'standard',
        'reason': 'Symmetric distribution (skew=-0.25), no transformation needed',
        'transform': None
    },
    'pKa': {
        'type': 'standard',
        'reason': 'Symmetric distribution (skew=0.29), no transformation needed',
        'transform': None
    },
    'solubility': {
        'type': 'log_standard',
        'reason': 'Extreme skewness (10.35), mean/median=142x, needs log transformation',
        'transform': 'log1p'
    },
    'permeability': {
        'type': 'log_standard',
        'reason': 'High skewness (4.79), positive values, log transformation optimal',
        'transform': 'log1p'
    },
    'plasma_protein_binding': {
        'type': 'minmax',
        'reason': 'Percentage [0-100], bounded range, physical meaning preserved',
        'transform': None,
        'range': [0, 100]
    },
    'fu_in_vitro': {
        'type': 'none',
        'reason': 'Fraction [0-1], already normalized',
        'transform': None
    }
}


def _build_power_transformer(lmbda: float) -> PowerTransformer:
    pt = PowerTransformer(method='yeo-johnson', standardize=False)
    pt.lambdas_ = np.array([lmbda], dtype=np.float64)
    pt.n_features_in_ = 1
    pt.n_samples_seen_ = None
    return pt


def save_scaler(df: pd.DataFrame, scaler_path: str) -> pd.DataFrame:
    """
    기존 표준 스케일러 (Z-score)를 유지하기 위해 이전 구현과 동일한 기능을 제공합니다.
    """
    scaled_df = df.copy()
    scaler_data = []

    for col in DIDB_FILTER_COLS:
        if col in df.columns:
            df[col] = pd.to_numeric(df[col], errors='coerce')
            mean = float(df[col].mean(skipna=True))
            std = float(df[col].std(skipna=True))
            if pd.isna(std) or std == 0:
                std = 1.0
            scaled_df[col] = (df[col] - mean) / std
            scaler_data.append({'feature': col, 'mean': mean, 'std': std})

    scaler_df = pd.DataFrame(scaler_data)
    os.makedirs(os.path.dirname(scaler_path), exist_ok=True)
    scaler_df.to_csv(scaler_path, index=False)

    return scaled_df


def reverse_scaling(df: pd.DataFrame, scaler_path: str) -> pd.DataFrame:
    scaler_df = pd.read_csv(scaler_path)
    for col in DIDB_FILTER_COLS:
        if col in df.columns and col in scaler_df['feature'].values:
            mean = float(scaler_df.loc[scaler_df['feature'] == col, 'mean'].values[0])
            std = float(scaler_df.loc[scaler_df['feature'] == col, 'std'].values[0])
            df[col] = (df[col] * std) + mean
    return df


def apply_scaler(df: pd.DataFrame, scaler_path: str) -> pd.DataFrame:
    if not os.path.exists(scaler_path):
        raise FileNotFoundError(f"Scaler file not found: {scaler_path}")

    scaler_df = pd.read_csv(scaler_path)
    scaled_df = df.copy()

    for col in DIDB_FILTER_COLS:
        if col not in scaled_df.columns:
            continue

        if col not in scaler_df['feature'].values:
            continue

        mean = float(scaler_df.loc[scaler_df['feature'] == col, 'mean'].values[0])
        std = float(scaler_df.loc[scaler_df['feature'] == col, 'std'].values[0])

        scaled_df[col] = pd.to_numeric(scaled_df[col], errors='coerce')
        if pd.isna(std) or std == 0:
            scaled_df[col] = scaled_df[col] - mean
        else:
            scaled_df[col] = (scaled_df[col] - mean) / std

    return scaled_df


def save_scaler_power(df: pd.DataFrame, scaler_path: str) -> pd.DataFrame:
    """
    학습 데이터프레임을 scikit-learn PowerTransformer(Yeo-Johnson) 기반으로 스케일링하고,
    변환 파라미터(lambda, mean, std)를 CSV로 저장합니다.
    """
    scaled_df = df.copy()
    scaler_data = []

    for col in DIDB_FILTER_COLS:
        if col not in df.columns:
            continue

        series = pd.to_numeric(df[col], errors='coerce')
        valid_mask = series.notna()
        if valid_mask.sum() == 0:
            continue

        values = series.loc[valid_mask].to_numpy(dtype=np.float64).reshape(-1, 1)

        pt = PowerTransformer(method='yeo-johnson', standardize=False)
        transformed = pt.fit_transform(values).reshape(-1)
        lmbda = float(pt.lambdas_[0])

        mean = float(np.mean(transformed))
        std = float(np.std(transformed))
        if np.isnan(std) or std == 0.0:
            std = 1.0

        scaled_values = np.empty_like(series, dtype=np.float64)
        scaled_values[:] = np.nan
        scaled_values[valid_mask] = (transformed - mean) / std

        scaled_df[col] = scaled_values
        scaler_data.append({
            'feature': col,
            'lambda': lmbda,
            'mean': mean,
            'std': std
        })

    scaler_df = pd.DataFrame(scaler_data)
    os.makedirs(os.path.dirname(scaler_path), exist_ok=True)
    scaler_df.to_csv(scaler_path, index=False)

    return scaled_df


def reverse_scaling_power(df: pd.DataFrame, scaler_path: str) -> pd.DataFrame:
    scaler_df = pd.read_csv(scaler_path)
    result_df = df.copy()

    for col in DIDB_FILTER_COLS:
        if col not in result_df.columns:
            continue
        if col not in scaler_df['feature'].values:
            continue

        row = scaler_df.loc[scaler_df['feature'] == col].iloc[0]
        lmbda = float(row['lambda'])
        mean = float(row['mean'])
        std = float(row['std'])

        series = pd.to_numeric(result_df[col], errors='coerce')
        valid_mask = series.notna()
        if valid_mask.sum() == 0:
            continue

        values = series.to_numpy(dtype=np.float64)
        w = values[valid_mask] * std + mean

        pt = _build_power_transformer(lmbda)
        original = pt.inverse_transform(w.reshape(-1, 1)).reshape(-1)

        values[valid_mask] = original
        result_df[col] = values

    return result_df


def apply_scaler_power(df: pd.DataFrame, scaler_path: str) -> pd.DataFrame:
    if not os.path.exists(scaler_path):
        raise FileNotFoundError(f"Scaler file not found: {scaler_path}")

    scaler_df = pd.read_csv(scaler_path)
    scaled_df = df.copy()

    for col in DIDB_FILTER_COLS:
        if col not in scaled_df.columns:
            continue
        if col not in scaler_df['feature'].values:
            continue

        row = scaler_df.loc[scaler_df['feature'] == col].iloc[0]
        lmbda = float(row['lambda'])
        mean = float(row['mean'])
        std = float(row['std'])

        series = pd.to_numeric(scaled_df[col], errors='coerce')
        valid_mask = series.notna()
        if valid_mask.sum() == 0:
            continue

        values = series.to_numpy(dtype=np.float64)
        pt = _build_power_transformer(lmbda)
        transformed = pt.transform(values[valid_mask].reshape(-1, 1)).reshape(-1)
        standardized = (transformed - mean) / std

        values[valid_mask] = standardized
        scaled_df[col] = values

    return scaled_df


def save_scaler_robust(df: pd.DataFrame, scaler_path: str) -> pd.DataFrame:
    """
    RobustScaler를 사용하여 데이터 스케일링 (median, IQR 기반 - outlier에 강건)

    Args:
        df: 입력 데이터프레임
        scaler_path: 스케일러 파라미터 저장 경로

    Returns:
        스케일링된 데이터프레임
    """
    scaled_df = df.copy()
    scaler_data = []

    for col in DIDB_FILTER_COLS:
        if col not in df.columns:
            continue

        series = pd.to_numeric(df[col], errors='coerce')
        valid_mask = series.notna()
        if valid_mask.sum() == 0:
            continue

        values = series[valid_mask].to_numpy(dtype=np.float64).reshape(-1, 1)

        # RobustScaler 적용 (median, IQR 사용)
        scaler = SklearnRobustScaler()
        scaled_values = scaler.fit_transform(values).reshape(-1)

        # 스케일링된 값을 데이터프레임에 저장
        scaled_array = np.empty_like(series, dtype=np.float64)
        scaled_array[:] = np.nan
        scaled_array[valid_mask] = scaled_values
        scaled_df[col] = scaled_array

        # 스케일러 파라미터 저장
        scaler_data.append({
            'feature': col,
            'median': float(scaler.center_[0]),
            'scale': float(scaler.scale_[0])  # IQR / 1.35
        })

    scaler_df = pd.DataFrame(scaler_data)
    os.makedirs(os.path.dirname(scaler_path), exist_ok=True)
    scaler_df.to_csv(scaler_path, index=False)

    return scaled_df


def reverse_scaling_robust(df: pd.DataFrame, scaler_path: str) -> pd.DataFrame:
    """
    RobustScaler로 스케일링된 데이터를 원본 스케일로 복원

    Args:
        df: 스케일링된 데이터프레임
        scaler_path: 스케일러 파라미터 경로

    Returns:
        원본 스케일로 복원된 데이터프레임
    """
    scaler_df = pd.read_csv(scaler_path)
    result_df = df.copy()

    for col in DIDB_FILTER_COLS:
        if col not in result_df.columns:
            continue
        if col not in scaler_df['feature'].values:
            continue

        row = scaler_df.loc[scaler_df['feature'] == col].iloc[0]
        median = float(row['median'])
        scale = float(row['scale'])

        series = pd.to_numeric(result_df[col], errors='coerce')
        valid_mask = series.notna()
        if valid_mask.sum() == 0:
            continue

        values = series.to_numpy(dtype=np.float64)
        # 역변환: (scaled_value * scale) + median
        values[valid_mask] = values[valid_mask] * scale + median
        result_df[col] = values

    return result_df


def apply_scaler_robust(df: pd.DataFrame, scaler_path: str) -> pd.DataFrame:
    """
    저장된 RobustScaler 파라미터를 사용하여 새 데이터 스케일링

    Args:
        df: 입력 데이터프레임
        scaler_path: 스케일러 파라미터 경로

    Returns:
        스케일링된 데이터프레임
    """
    if not os.path.exists(scaler_path):
        raise FileNotFoundError(f"Scaler file not found: {scaler_path}")

    scaler_df = pd.read_csv(scaler_path)
    scaled_df = df.copy()

    for col in DIDB_FILTER_COLS:
        if col not in scaled_df.columns:
            continue
        if col not in scaler_df['feature'].values:
            continue

        row = scaler_df.loc[scaler_df['feature'] == col].iloc[0]
        median = float(row['median'])
        scale = float(row['scale'])

        series = pd.to_numeric(scaled_df[col], errors='coerce')
        valid_mask = series.notna()
        if valid_mask.sum() == 0:
            continue

        values = series.to_numpy(dtype=np.float64)
        # 변환: (value - median) / scale
        values[valid_mask] = (values[valid_mask] - median) / scale
        scaled_df[col] = values

    return scaled_df


def save_scaler_minmax(df: pd.DataFrame, scaler_path: str) -> pd.DataFrame:
    """
    도메인 지식 기반 MinMaxScaler (화학적 특성의 물리적 범위 고려)

    Args:
        df: 입력 데이터프레임
        scaler_path: 스케일러 파라미터 저장 경로

    Returns:
        스케일링된 데이터프레임 [0, 1] 범위
    """
    scaled_df = df.copy()
    scaler_data = []

    for col in DIDB_FILTER_COLS:
        if col not in df.columns:
            continue

        if col not in FEATURE_RANGES:
            print(f"Warning: {col} not in FEATURE_RANGES. Skipping.")
            continue

        min_val, max_val = FEATURE_RANGES[col]
        series = pd.to_numeric(df[col], errors='coerce')

        # MinMax 스케일링: (x - min) / (max - min)
        scaled_df[col] = (series - min_val) / (max_val - min_val)

        # 범위를 [0, 1]로 클리핑 (도메인 범위 밖의 값 처리)
        scaled_df[col] = scaled_df[col].clip(0, 1)

        scaler_data.append({
            'feature': col,
            'min': min_val,
            'max': max_val
        })

    scaler_df = pd.DataFrame(scaler_data)
    os.makedirs(os.path.dirname(scaler_path), exist_ok=True)
    scaler_df.to_csv(scaler_path, index=False)

    return scaled_df


def reverse_scaling_minmax(df: pd.DataFrame, scaler_path: str) -> pd.DataFrame:
    """
    MinMaxScaler로 스케일링된 데이터를 원본 스케일로 복원

    Args:
        df: 스케일링된 데이터프레임
        scaler_path: 스케일러 파라미터 경로

    Returns:
        원본 스케일로 복원된 데이터프레임
    """
    scaler_df = pd.read_csv(scaler_path)
    result_df = df.copy()

    for col in DIDB_FILTER_COLS:
        if col not in result_df.columns:
            continue
        if col not in scaler_df['feature'].values:
            continue

        row = scaler_df.loc[scaler_df['feature'] == col].iloc[0]
        min_val = float(row['min'])
        max_val = float(row['max'])

        series = pd.to_numeric(result_df[col], errors='coerce')
        # 역변환: scaled_value * (max - min) + min
        result_df[col] = series * (max_val - min_val) + min_val

    return result_df


def apply_scaler_minmax(df: pd.DataFrame, scaler_path: str) -> pd.DataFrame:
    """
    저장된 MinMaxScaler 파라미터를 사용하여 새 데이터 스케일링

    Args:
        df: 입력 데이터프레임
        scaler_path: 스케일러 파라미터 경로

    Returns:
        스케일링된 데이터프레임
    """
    if not os.path.exists(scaler_path):
        raise FileNotFoundError(f"Scaler file not found: {scaler_path}")

    scaler_df = pd.read_csv(scaler_path)
    scaled_df = df.copy()

    for col in DIDB_FILTER_COLS:
        if col not in scaled_df.columns:
            continue
        if col not in scaler_df['feature'].values:
            continue

        row = scaler_df.loc[scaler_df['feature'] == col].iloc[0]
        min_val = float(row['min'])
        max_val = float(row['max'])

        series = pd.to_numeric(scaled_df[col], errors='coerce')
        # MinMax 스케일링
        scaled_df[col] = (series - min_val) / (max_val - min_val)
        # 클리핑
        scaled_df[col] = scaled_df[col].clip(0, 1)

    return scaled_df


def save_scaler_adaptive(df: pd.DataFrame, scaler_path: str, feature_cols: list) -> pd.DataFrame:
    """
    Apply adaptive scaling - feature-specific optimal scalers

    Args:
        df: Input dataframe
        scaler_path: Path to save scaler config
        feature_cols: List of feature columns to scale

    Returns:
        Scaled dataframe
    """
    os.makedirs(os.path.dirname(scaler_path) or '.', exist_ok=True)

    df_scaled = df.copy()
    scaler_info = []

    for col in feature_cols:
        if col not in df.columns:
            continue

        config = ADAPTIVE_SCALER_CONFIG.get(col, {'type': 'standard', 'transform': None})
        scaler_type = config['type']
        transform = config.get('transform')

        # Get non-null values
        mask = df[col].notna()
        values = df.loc[mask, col].values.reshape(-1, 1)

        if len(values) == 0:
            continue

        # Step 1: Apply pre-transformation if needed
        transformed_values = values.copy()
        if transform == 'log1p':
            transformed_values = np.log1p(values)
            # Don't append to scaler_info yet - will be done in Step 2 with complete info

        # Step 2: Apply scaler
        if scaler_type == 'none':
            # No scaling
            scaled_values = transformed_values
            scaler_info.append({
                'feature': col,
                'transform': transform or 'none',
                'scaler_type': 'none',
                'mean': float(np.mean(values)),
                'std': float(np.std(values)),
                'reason': config.get('reason', '')
            })

        elif scaler_type == 'minmax' or scaler_type == 'log_minmax':
            # MinMaxScaler
            # Use predefined range (e.g., [0, 100] for PPB) instead of data min/max
            range_min, range_max = config.get('range', [values.min(), values.max()])

            # Scale using predefined range
            scaled_values = (transformed_values - range_min) / (range_max - range_min)

            scaler_info.append({
                'feature': col,
                'transform': transform or 'none',
                'scaler_type': 'minmax',
                'min': float(range_min),  # Use predefined range, not data min/max
                'max': float(range_max),  # Use predefined range, not data min/max
                'range_min': float(range_min),
                'range_max': float(range_max),
                'reason': config.get('reason', '')
            })

        elif scaler_type == 'standard' or scaler_type == 'log_standard':
            # StandardScaler (Z-score)
            scaler = StandardScaler()
            scaled_values = scaler.fit_transform(transformed_values)

            scaler_info.append({
                'feature': col,
                'transform': transform or 'none',
                'scaler_type': 'standard',
                'mean': float(scaler.mean_[0]),
                'std': float(np.sqrt(scaler.var_[0])),
                'reason': config.get('reason', '')
            })

        elif scaler_type == 'power':
            # PowerTransformer
            scaler = PowerTransformer(method='yeo-johnson', standardize=True)
            scaled_values = scaler.fit_transform(transformed_values)

            scaler_info.append({
                'feature': col,
                'transform': transform or 'none',
                'scaler_type': 'power',
                'lambda': float(scaler.lambdas_[0]),
                'mean': float(np.mean(scaled_values)),
                'std': float(np.std(scaled_values)),
                'reason': config.get('reason', '')
            })

        else:
            raise ValueError(f"Unknown scaler type: {scaler_type}")

        # Update scaled values
        df_scaled.loc[mask, col] = scaled_values.flatten()

    # Save scaler config
    scaler_df = pd.DataFrame(scaler_info)
    scaler_df.to_csv(scaler_path, index=False)
    print(f"📊 Adaptive scaler config saved to: {scaler_path}")

    return df_scaled


def apply_scaler_adaptive(df: pd.DataFrame, scaler_path: str, feature_cols: list) -> pd.DataFrame:
    """
    Apply saved adaptive scaling configuration

    Args:
        df: Input dataframe
        scaler_path: Path to scaler config
        feature_cols: List of feature columns to scale

    Returns:
        Scaled dataframe
    """
    if not os.path.exists(scaler_path):
        raise FileNotFoundError(f"Scaler config not found: {scaler_path}")

    scaler_df = pd.read_csv(scaler_path)
    df_scaled = df.copy()

    for _, row in scaler_df.iterrows():
        col = row['feature']

        if col not in df.columns or col not in feature_cols:
            continue

        # Get non-null values
        mask = df[col].notna()
        values = df.loc[mask, col].values.reshape(-1, 1)

        if len(values) == 0:
            continue

        # Step 1: Apply pre-transformation
        transformed_values = values.copy()
        if row['transform'] == 'log1p':
            transformed_values = np.log1p(values)

        # Step 2: Apply scaler
        if row['scaler_type'] == 'none':
            scaled_values = transformed_values

        elif row['scaler_type'] == 'minmax' or row['scaler_type'] == 'log_minmax':
            min_val = row['min']
            max_val = row['max']
            scaled_values = (transformed_values - min_val) / (max_val - min_val)

        elif row['scaler_type'] == 'standard' or row['scaler_type'] == 'log_standard':
            mean = row['mean']
            std = row['std']
            scaled_values = (transformed_values - mean) / std

        elif row['scaler_type'] == 'power':
            # For PowerTransformer, we need to use sklearn (complex transformation)
            # For now, approximate with standardization
            mean = row['mean']
            std = row['std']
            scaled_values = (transformed_values - mean) / std

        else:
            raise ValueError(f"Unknown scaler type: {row['scaler_type']}")

        df_scaled.loc[mask, col] = scaled_values.flatten()

    return df_scaled


def reverse_scaling_adaptive(df: pd.DataFrame, scaler_path: str) -> pd.DataFrame:
    """
    Reverse adaptive scaling for predictions (DataFrame version)

    Args:
        df: DataFrame with scaled predictions
        scaler_path: Path to scaler config

    Returns:
        DataFrame with original scale values
    """
    if not os.path.exists(scaler_path):
        raise FileNotFoundError(f"Scaler config not found: {scaler_path}")

    scaler_df = pd.read_csv(scaler_path)
    result_df = df.copy()

    for col in DIDB_FILTER_COLS:
        if col not in result_df.columns:
            continue

        scaler_row = scaler_df[scaler_df['feature'] == col]
        if len(scaler_row) == 0:
            # No scaling applied for this feature
            continue

        row = scaler_row.iloc[0]
        series = pd.to_numeric(result_df[col], errors='coerce')
        valid_mask = series.notna()
        if valid_mask.sum() == 0:
            continue

        values = series.to_numpy(dtype=np.float64)
        reversed_values = values[valid_mask].copy()

        # Step 1: Reverse scaler
        if row['scaler_type'] == 'none':
            pass

        elif row['scaler_type'] == 'minmax' or row['scaler_type'] == 'log_minmax':
            min_val = row['min']
            max_val = row['max']
            reversed_values = reversed_values * (max_val - min_val) + min_val

        elif row['scaler_type'] == 'standard' or row['scaler_type'] == 'log_standard':
            mean = row['mean']
            std = row['std']
            reversed_values = reversed_values * std + mean

        elif row['scaler_type'] == 'power':
            mean = row['mean']
            std = row['std']
            reversed_values = reversed_values * std + mean

        # Step 2: Reverse pre-transformation
        if row['transform'] == 'log1p':
            reversed_values = np.expm1(reversed_values)  # inverse of log1p

            # Safety: Handle invalid values from numerical issues
            reversed_values = np.nan_to_num(
                reversed_values,
                nan=0.01,       # NaN → 0.01 (minimum valid value)
                posinf=10000.0, # +inf → reasonable max
                neginf=0.01     # -inf → 0.01 (minimum valid value)
            )

            # Clip to ensure minimum valid value (0.01 for solubility/permeability)
            reversed_values = np.maximum(reversed_values, 0.01)

        values[valid_mask] = reversed_values
        result_df[col] = values

    return result_df