import logging
import pandas as pd
import polars as pl
import pyarrow as pa
from coded_flows.types import Union, DataFrame, ArrowTable, Any, List, Str
from coded_flows.utils import CodedFlowsLogger

logger = CodedFlowsLogger(name="Apply Scaler", level=logging.INFO)


def apply_scaler(
    data: Union[DataFrame, ArrowTable], Scaler: Any, options: dict = None
) -> DataFrame:
    options = options or {}
    verbose = options.get("verbose", True)
    safe_mode = options.get("safe_mode", True)
    scaled_data = None
    try:
        verbose and logger.info("Starting Scaler application process.")
        if isinstance(data, pl.DataFrame):
            verbose and logger.info("Detected Polars DataFrame. Converting to Pandas.")
            scaled_data = data.to_pandas()
        elif isinstance(data, (pa.Table, pa.lib.Table)):
            verbose and logger.info("Detected PyArrow Table. Converting to Pandas.")
            scaled_data = data.to_pandas()
        elif isinstance(data, pd.DataFrame):
            verbose and logger.info("Detected Pandas DataFrame.")
            scaled_data = data.copy()
        else:
            raise ValueError(
                "Input data must be a Pandas DataFrame, Polars DataFrame, or Arrow Table."
            )
        target_columns = scaled_data.columns.tolist()
        verbose and logger.info("No specific columns selected. Targeting all columns.")
        if safe_mode:
            numeric_cols = scaled_data.select_dtypes(
                include=["number", "bool"]
            ).columns.tolist()
            valid_cols = [col for col in target_columns if col in numeric_cols]
            dropped_cols = set(target_columns) - set(valid_cols)
            if dropped_cols and verbose:
                logger.warning(
                    f"Safe Mode: Ignored non-numerical columns: {dropped_cols}"
                )
            target_columns = valid_cols
            verbose and logger.info(
                f"Safe Mode enabled. Final target columns: {target_columns}"
            )
        if not target_columns:
            verbose and logger.warning(
                "No valid columns remaining to scale. Returning original data."
            )
        else:
            missing_cols = [
                col for col in target_columns if col not in scaled_data.columns
            ]
            if missing_cols:
                raise ValueError(f"Columns not found in dataset: {missing_cols}")
            verbose and logger.info(
                f"Applying scaler ({type(Scaler).__name__}) transform."
            )
            data_subset = scaled_data[target_columns]
            scaled_values = Scaler.transform(data_subset)
            scaled_data[target_columns] = scaled_values
            verbose and logger.info("Scaler transform applied successfully.")
    except Exception as e:
        verbose and logger.error(f"Error during scaler application.")
        raise e
    return scaled_data