import logging
import numpy as np
import pandas as pd
import polars as pl
import pyarrow as pa
from sklearn.metrics import (
    mean_absolute_error,
    mean_squared_error,
    r2_score,
    mean_absolute_percentage_error,
    root_mean_squared_error,
)
from coded_flows.types import (
    Union,
    DataFrame,
    ArrowTable,
    DataSeries,
    List,
    NDArray,
    Float,
    Any,
)
from coded_flows.utils import CodedFlowsLogger

logger = CodedFlowsLogger(name="Regression Metrics", level=logging.INFO)


def _extract_vector(
    data: Any, column_name: str, verbose: bool = False, label: str = "Input"
) -> Union[List, np.ndarray]:
    """
    Extracts a 1D vector/list from various data structures (DF, Series, Arrow, List, Numpy).
    """
    result = None
    if isinstance(data, pd.DataFrame):
        if column_name not in data.columns:
            verbose and logger.error(
                f"Column '{column_name}' not found in Pandas DataFrame for {label}."
            )
            raise ValueError(
                f"Column '{column_name}' not found in Pandas DataFrame for {label}."
            )
        result = data[column_name].tolist()
    elif isinstance(data, pl.DataFrame):
        if column_name not in data.columns:
            verbose and logger.error(
                f"Column '{column_name}' not found in Polars DataFrame for {label}."
            )
            raise ValueError(
                f"Column '{column_name}' not found in Polars DataFrame for {label}."
            )
        result = data[column_name].to_list()
    elif isinstance(data, (pa.Table, pa.lib.Table)):
        if column_name not in data.column_names:
            verbose and logger.error(
                f"Column '{column_name}' not found in Arrow Table for {label}."
            )
            raise ValueError(
                f"Column '{column_name}' not found in Arrow Table for {label}."
            )
        result = data[column_name].to_pylist()
    elif isinstance(data, pd.Series):
        result = data.tolist()
    elif isinstance(data, pl.Series):
        result = data.to_list()
    elif isinstance(data, np.ndarray):
        result = data.flatten().tolist()
    elif isinstance(data, list):
        result = data
    if result is None:
        verbose and logger.error(
            f"Unsupported data type for {label}: {type(data)}. Expected DataFrame, Series, Array, or List."
        )
        raise ValueError(
            f"Unsupported data type for {label}: {type(data)}. Expected DataFrame, Series, Array, or List."
        )
    return result


def regression_metrics(
    y_true: Union[DataFrame, ArrowTable, DataSeries, NDArray, List],
    y_pred: Union[DataFrame, ArrowTable, DataSeries, NDArray, List],
    options: dict = None,
) -> Float:
    options = options or {}
    verbose = options.get("verbose", True)
    metric = options.get("metric", "Forecast Accuracy")
    y_true_col = options.get("y_true_column", "y_true")
    y_pred_col = options.get("y_pred_column", "y_pred")
    score = 0.0
    try:
        verbose and logger.info(f"Starting regression metric computation: {metric}")
        y_true_vec = _extract_vector(
            y_true,
            column_name=y_true_col,
            verbose=verbose,
            label="Real Target (y_true)",
        )
        y_pred_vec = _extract_vector(
            y_pred, column_name=y_pred_col, verbose=verbose, label="Prediction (y_pred)"
        )
        y_true_np = np.array(y_true_vec, dtype=float)
        y_pred_np = np.array(y_pred_vec, dtype=float)
        if len(y_true_np) != len(y_pred_np):
            raise ValueError(
                f"Length mismatch: y_true has {len(y_true_np)} rows, y_pred has {len(y_pred_np)} rows."
            )
        if metric == "Mean Absolute Error (MAE)":
            score = mean_absolute_error(y_true_np, y_pred_np)
        elif metric == "Mean Squared Error (MSE)":
            score = mean_squared_error(y_true_np, y_pred_np)
        elif metric == "Root Mean Squared Error (RMSE)":
            score = root_mean_squared_error(y_true_np, y_pred_np)
        elif metric == "R2 Score":
            score = r2_score(y_true_np, y_pred_np)
        elif metric == "Mean Absolute Percentage Error (MAPE)":
            score = mean_absolute_percentage_error(y_true_np, y_pred_np)
        elif metric in [
            "Forecast Accuracy",
            "Weighted Absolute Percentage Error (WAPE)",
        ]:
            eps = np.finfo(np.float64).eps
            sum_abs_error = np.sum(np.abs(y_true_np - y_pred_np))
            sum_abs_truth = np.maximum(np.sum(np.abs(y_true_np)), eps)
            if metric == "Forecast Accuracy":
                score = 1 - sum_abs_error / sum_abs_truth
            else:
                score = sum_abs_error / sum_abs_truth
        verbose and logger.info(f"Computation successful. {metric}: {score}")
    except Exception as e:
        verbose and logger.error(f"Error computing regression metrics: {str(e)}")
        raise e
    score = float(score)
    return score