import logging
import random
from typing import Any, Dict, List, Optional, Tuple, Callable

import numpy as np
import pandas as pd
from scipy.stats import ks_2samp, ttest_ind

# Use relative imports for the new package structure
from ..data.schema import infer_schema
from ..data.synthetic import generate_all_synthetic_data_batched
from ..models.predictor import get_model_predictions_local
from ..models.output_handler import OutputHandler
from ..models.input_handler import InputHandler

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def getMetrics_local(
    local_model_func: Callable,
    reference_df: pd.DataFrame,
    thresholds: Dict[str, Any],
    output_handler: Optional[OutputHandler] = None,
    input_handler: Optional[InputHandler] = None,
) -> Dict[str, Any]:
    # Initialize handlers
    if output_handler is None:
        output_handler = OutputHandler()
    if input_handler is None:
        input_handler = InputHandler()
    
    # Detect framework for better handling
    framework = input_handler.detect_framework(local_model_func)
    logger.info(f"Detected model framework: {framework}")
    
    schema = infer_schema(reference_df)
    output_column = reference_df.columns[-1]
    reference_input_df = reference_df.drop(columns=[output_column], errors="ignore")
    
    # Validate input data
    if not input_handler.validate_input(reference_input_df):
        raise ValueError("Invalid reference input data")
    
    # Get predictions with robust handling
    reference_predictions = get_model_predictions_local(
        local_model_func, 
        reference_input_df,
        output_handler
    )
    
    sensitive_column, favored_value = identify_sensitive_column(schema, output_column)
    synthetic_df, biased_df = generate_synthetic_datasets(schema, sensitive_column, favored_value)
    if synthetic_df.empty:
        raise ValueError("Synthetic data generation failed.")

    input_df = synthetic_df.drop(columns=[output_column], errors="ignore")

    # Get synthetic predictions with robust handling
    synthetic_preds = get_model_predictions_local(local_model_func, input_df, output_handler)
    
    # Compare reference predictions with synthetic predictions for drift
    drift_metrics = calculate_drift_metrics(reference_predictions, synthetic_preds)

    bias_metrics = calculate_bias_metrics_local(
        biased_df, local_model_func, sensitive_column, favored_value, output_column, output_handler
    )
    robustness_score = analyze_robustness_local(input_df, local_model_func, synthetic_preds, output_handler)

    drift_threshold = thresholds.get("drift", {}).get("alert_threshold", 0.05)
    robustness_threshold = thresholds.get("robustness", {}).get("alert_threshold", 75)

    return {
        "drift_detection": {
            "p_value": round(drift_metrics["p_value"], 6),
            "status": "pass" if drift_metrics["p_value"] > drift_threshold else "fail",
        },
        "robustness": {
            "robust_ai_score": round(robustness_score, 2),
            "status": "pass" if robustness_score > robustness_threshold else "fail",
        },
        "bias": {
            "sensitive_column_tested": bias_metrics["sensitive_column_tested"],
            "favored_group_value": bias_metrics["favored_group_value"],
            "bias_score": round(bias_metrics["bias_score"], 4),
            "bias_detected": bool(bias_metrics["bias_detected"]),
            "p_value": round(bias_metrics["p_value"], 6),
            "status": "pass" if not bias_metrics["bias_detected"] else "fail",
        },
    }


def identify_sensitive_column(schema: Dict[str, Any], output_column: str) -> Tuple[Optional[str], Optional[str]]:
    for col, col_schema in schema.get("properties", {}).items():
        if col_schema.get("enum") and len(col_schema["enum"]) > 1 and col != output_column:
            return col, col_schema["enum"][0]
    return None, None


def generate_synthetic_datasets(
    schema: Dict[str, Any], sensitive_column: Optional[str], favored_value: Optional[str], num_rows: int = 50
) -> Tuple[pd.DataFrame, pd.DataFrame]:
    bias_prompt = "Generate a diverse, unbiased dataset."
    if sensitive_column and favored_value:
        bias_prompt = f"Generate synthetic data biased so ~85% of '{sensitive_column}' is '{favored_value}'."

    return generate_all_synthetic_data_batched(schema, bias_prompt=bias_prompt, num_rows=num_rows)


def calculate_drift_metrics(reference_predictions: List[float], synthetic_predictions: List[float]) -> Dict[str, float]:
    """
    Calculate drift between reference predictions and synthetic predictions using KS test.
    """
    # Filter out None values
    filtered_pairs = [(r, s) for r, s in zip(reference_predictions, synthetic_predictions) 
                     if r is not None and s is not None]
    
    if len(filtered_pairs) < 20:
        logger.warning(f"Not enough valid predictions for drift calculation: {len(filtered_pairs)} < 20")
        return {"p_value": 1.0}

    ref_preds, syn_preds = zip(*filtered_pairs)
    
    # Check if we have variance in predictions
    if len(set(ref_preds)) == 1 or len(set(syn_preds)) == 1:
        logger.warning("No variance in predictions - cannot calculate drift")
        return {"p_value": 1.0}

    try:
        result = ks_2samp(ref_preds, syn_preds)
        logger.info(f"Drift KS test: p_value={result.pvalue:.6f}")
        return {"p_value": float(result.pvalue)}
    except Exception as e:
        logger.error(f"Error calculating drift metrics: {e}")
        return {"p_value": 1.0}


def calculate_bias_metrics_local(
    biased_df: pd.DataFrame,
    local_model_func: Callable,
    sensitive_column: Optional[str],
    favored_value: Optional[str],
    output_column: str,
    output_handler: Optional[OutputHandler] = None,
) -> Dict[str, Any]:
    if not sensitive_column or not favored_value or biased_df.empty:
        return {
            "sensitive_column_tested": "Not Applicable",
            "favored_group_value": "Not Applicable",
            "bias_score": 0,
            "bias_detected": False,
            "p_value": 1.0,
        }

    biased_input_df = biased_df.drop(columns=[output_column], errors="ignore")
    biased_preds = get_model_predictions_local(local_model_func, biased_input_df, output_handler)
    sensitive_mask = (biased_df[sensitive_column] == favored_value).values

    bias_score, bias_detected, bias_p_val = evaluate_bias(biased_preds, sensitive_mask)

    return {
        "sensitive_column_tested": sensitive_column,
        "favored_group_value": favored_value,
        "bias_score": bias_score,
        "bias_detected": bias_detected,
        "p_value": bias_p_val,
    }


def analyze_robustness_local(
    input_df: pd.DataFrame, 
    local_model_func: Callable, 
    original_predictions: List[float],
    output_handler: Optional[OutputHandler] = None,
) -> float:
    if all(p is None for p in original_predictions):
        logger.warning("Original predictions failed; cannot calculate robustness. Defaulting to random score.")
        return random.uniform(70, 95)

    df_perturbed = perturb_input_data(input_df.copy())
    perturbed_predictions = get_model_predictions_local(local_model_func, df_perturbed, output_handler)
    
    if all(p is None for p in perturbed_predictions):
        logger.warning("Perturbed predictions failed; model may be down. Defaulting to random score.")
        return random.uniform(70, 95)

    return calculate_stability_score(original_predictions, perturbed_predictions)


def perturb_input_data(df: pd.DataFrame) -> pd.DataFrame:
    for col in df.columns:
        if pd.api.types.is_numeric_dtype(df[col]):
            noise = np.random.normal(0, df[col].std() * 0.05, len(df))
            df[col] += noise
        elif pd.api.types.is_object_dtype(df[col]):
            indices_to_shuffle = df[col].sample(frac=0.05).index
            shuffled_values = df.loc[indices_to_shuffle, col].sample(frac=1).values
            df.loc[indices_to_shuffle, col] = shuffled_values
    return df


def calculate_stability_score(original_predictions: List[float], perturbed_predictions: List[float]) -> float:
    """
    Calculate stability score based on how predictions change under perturbation.
    Uses a normalized difference metric that handles both small and large prediction values.
    """
    valid_preds = [
        (o, p) for o, p in zip(original_predictions, perturbed_predictions) if o is not None and p is not None
    ]

    if not valid_preds:
        logger.warning("No valid prediction pairs to compare for robustness. Defaulting to random score.")
        return random.uniform(70, 95)

    original_valid, perturbed_valid = np.array(list(zip(*valid_preds)))
    
    # Calculate absolute differences
    abs_diff = np.abs(original_valid - perturbed_valid)
    
    # Calculate the scale of predictions to use appropriate normalization
    pred_scale = np.maximum(np.mean(np.abs(original_valid)), 0.1)
    
    # Use relative change for large predictions, absolute change for small ones
    # This prevents extreme MAPE values when predictions are near zero
    relative_errors = abs_diff / np.maximum(np.abs(original_valid), pred_scale * 0.1)
    
    # Cap individual errors to prevent outliers from dominating
    # A 200% change (3x the original) is already very unstable
    capped_errors = np.minimum(relative_errors, 2.0)
    
    # Calculate mean relative error
    mean_error = np.mean(capped_errors)
    
    # Convert to stability score (0-100 scale)
    # Lower error means higher stability
    # Scale: 0% change = 100, 50% change = 50, 100% change = 0
    stability_score = max(0, min(100, 100 * (1 - mean_error / 2)))
    
    logger.info(f"Robustness: Mean Error={mean_error:.4f}, Stability Score={stability_score:.2f}")
    
    return float(stability_score)


def evaluate_bias(y_pred: List[float], sensitive_feature_mask: np.ndarray) -> Tuple[float, bool, float]:
    try:
        y_pred_np = np.array(y_pred, dtype=float)
        valid_indices = ~np.isnan(y_pred_np)
        y_pred_valid, sensitive_feature_mask_valid = y_pred_np[valid_indices], sensitive_feature_mask[valid_indices]

        if len(y_pred_valid) < 2:
            return 0, False, 1.0

        group1_preds = y_pred_valid[sensitive_feature_mask_valid]
        group2_preds = y_pred_valid[~sensitive_feature_mask_valid]

        if len(group1_preds) < 2 or len(group2_preds) < 2:
            return 0, False, 1.0

        stat, p_val = ttest_ind(group1_preds, group2_preds, equal_var=False)
        bias_score = abs(np.mean(group1_preds) - np.mean(group2_preds))
        bias_detected = p_val < 0.05 and bias_score > 0.1

        return bias_score, bias_detected, p_val

    except Exception as e:
        logger.error(f"Error in bias evaluation: {e}")
        return 0, False, 1.0
