#!/usr/bin/env python
"""
Risk Table Mapping Processing Script

This script creates and applies risk tables for categorical features based on 
target variable correlation, and handles missing value imputation for numeric features.
It supports both training mode (fit and transform) and inference mode (transform only).
"""

import argparse
import os
import sys
import pandas as pd
import numpy as np
import json
import pickle as pkl
import traceback
from pathlib import Path
from sklearn.impute import SimpleImputer
import logging
from typing import Dict, List, Tuple, Any, Optional, Callable

# Default paths (will be overridden by parameters in main function)
DEFAULT_INPUT_DIR = "/opt/ml/processing/input/data"
DEFAULT_CONFIG_DIR = "/opt/ml/code/hyperparams"  # Source directory path
DEFAULT_OUTPUT_DIR = "/opt/ml/processing/output"
DEFAULT_RISK_TABLE_DIR = "/opt/ml/processing/input/risk_tables"

# Constants for file paths to ensure consistency between training and inference
# These constants ensure the same filenames are used across all job types,
# facilitating proper connections between training and non-training steps
RISK_TABLE_FILENAME = "bin_mapping.pkl"  # Used by downstream steps as input dependency
HYPERPARAMS_FILENAME = (
    "hyperparameters.json"  # Expected by the script and generated by the builder
)

# Set up logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - [%(filename)s:%(lineno)d] - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger(__name__)


def load_json_config(config_path: str) -> Dict[str, Any]:
    """
    Loads a JSON configuration file.

    Args:
        config_path: Path to the JSON configuration file

    Returns:
        Dict containing the loaded JSON configuration

    Raises:
        FileNotFoundError: If the configuration file doesn't exist
        PermissionError: If the configuration file can't be accessed due to permissions
        json.JSONDecodeError: If the configuration file contains invalid JSON
        Exception: For other unexpected errors
    """
    try:
        with open(config_path, "r") as file:
            return json.load(file)
    except FileNotFoundError as e:
        logger.error(f"Configuration file not found at {config_path}: {str(e)}")
        raise
    except PermissionError as e:
        logger.error(
            f"Permission denied when accessing configuration at {config_path}: {str(e)}"
        )
        raise
    except json.JSONDecodeError as e:
        logger.error(f"Invalid JSON format in configuration at {config_path}: {str(e)}")
        raise
    except Exception as e:
        logger.error(
            f"Unexpected error loading configuration from {config_path}: {str(e)}"
        )
        raise


def validate_categorical_fields(
    df: pd.DataFrame, cat_field_list: List[str]
) -> List[str]:
    """
    Validate that fields in cat_field_list are suitable for risk mapping.

    Args:
        df: DataFrame containing the data
        cat_field_list: List of categorical field names

    Returns:
        List of valid fields for risk mapping
    """
    valid_fields = []

    for field in cat_field_list:
        if field not in df.columns:
            logger.warning(f"Field '{field}' not found in data, skipping")
            continue

        # Check if field is categorical or has few unique values
        unique_count = df[field].nunique()

        if pd.api.types.is_categorical_dtype(df[field]) or unique_count < 100:
            valid_fields.append(field)
            logger.info(
                f"Field '{field}' is valid for risk mapping ({unique_count} unique values)"
            )
        else:
            logger.warning(
                f"Field '{field}' may not be suitable for risk mapping ({unique_count} unique values)"
            )

    return valid_fields


class OfflineBinning:
    """
    A class to create risk tables for categorical features.

    Risk tables map categorical values to numerical risk scores based on
    their correlation with the target variable.
    """

    def __init__(self, cat_field_list: List[str], target_field: str):
        self.risk_tables = {}
        self.target = target_field
        self.variables = cat_field_list

    def fit(
        self, df: pd.DataFrame, smooth_factor: float = 0, count_threshold: int = 0
    ) -> None:
        """Fits the risk tables based on the provided dataframe."""
        # Drop any -1 or NaN target rows for fitting
        fit_df = df.loc[(df[self.target] != -1) & (~df[self.target].isnull())].copy()
        default_risk = float(fit_df[self.target].mean())
        smooth_samples = int(len(fit_df) * smooth_factor)

        for var in self.variables:
            if var not in fit_df.columns:
                continue

            # Infer data type from pandas
            if pd.api.types.is_numeric_dtype(fit_df[var]):
                data_mode = "numeric"
            else:
                data_mode = "categorical"

            self.risk_tables[var] = {
                "varName": var,
                "type": "categorical",
                "mode": data_mode,
                "default_bin": default_risk,
            }

            if fit_df[var].isnull().all():
                self.risk_tables[var]["bins"] = {}
                continue

            risk_table = self._create_risk_table(
                fit_df, var, default_risk, smooth_samples, count_threshold
            )
            self.risk_tables[var]["bins"] = risk_table

    def _create_risk_table(
        self,
        df: pd.DataFrame,
        variable: str,
        default_risk: float,
        samples: int,
        count_threshold: int,
    ) -> Dict:
        """Helper to calculate the risk table for a single variable."""
        cross_tab = pd.crosstab(
            df[variable],
            df[self.target].astype(object),
            margins=True,
            margins_name="_count_",
            dropna=False,
        )
        cross_tab["risk"] = cross_tab.apply(
            lambda x: x.get(1, 0.0) / (x.get(1, 0) + x.get(0, 0)), axis=1
        )
        cross_tab["smooth_risk"] = cross_tab.apply(
            lambda x: (
                (
                    (x["_count_"] * x["risk"] + samples * default_risk)
                    / (x["_count_"] + samples)
                )
                if x["_count_"] >= count_threshold
                else default_risk
            ),
            axis=1,
        )
        cross_tab = cross_tab.loc[cross_tab.index != "_count_"]
        return dict(zip(cross_tab.index, cross_tab["smooth_risk"]))

    def transform(self, df: pd.DataFrame) -> pd.DataFrame:
        """Transforms the dataframe using the fitted risk tables."""
        df_transformed = df.copy()
        for var, risk_table_info in self.risk_tables.items():
            if var in df_transformed.columns:
                bins = risk_table_info["bins"]
                default_bin = risk_table_info["default_bin"]
                df_transformed[var] = df_transformed[var].map(bins).fillna(default_bin)
        return df_transformed

    def load_risk_tables(self, risk_tables: Dict) -> None:
        """Loads pre-existing risk tables."""
        self.risk_tables = risk_tables
        logger.info(f"Loaded {len(risk_tables)} risk tables")


def load_split_data(job_type: str, input_dir: str) -> Dict[str, pd.DataFrame]:
    """
    Load data according to job_type.

    For 'training': Loads data from train, test, and val subdirectories
    For others: Loads single job_type split
    """
    input_path = Path(input_dir)

    if job_type == "training":
        # For training, we expect data in train/test/val subdirectories
        train_df = pd.read_csv(input_path / "train" / "train_processed_data.csv")
        test_df = pd.read_csv(input_path / "test" / "test_processed_data.csv")
        val_df = pd.read_csv(input_path / "val" / "val_processed_data.csv")
        logger.info(
            f"Loaded training data splits: train={train_df.shape}, test={test_df.shape}, val={val_df.shape}"
        )
        return {"train": train_df, "test": test_df, "val": val_df}
    else:
        # For other job types, we expect data in a single directory named after job_type
        df = pd.read_csv(input_path / job_type / f"{job_type}_processed_data.csv")
        logger.info(f"Loaded {job_type} data: {df.shape}")
        return {job_type: df}


def save_output_data(
    job_type: str, output_dir: str, data_dict: Dict[str, pd.DataFrame]
) -> None:
    """
    Save processed data according to job_type.

    For 'training': Saves data to train, test, and val subdirectories
    For others: Saves to single job_type directory
    """
    output_path = Path(output_dir)

    for split_name, df in data_dict.items():
        split_output_dir = output_path / split_name
        split_output_dir.mkdir(exist_ok=True, parents=True)

        output_file = split_output_dir / f"{split_name}_processed_data.csv"
        df.to_csv(output_file, index=False)
        logger.info(f"Saved {split_name} data to {output_file}, shape: {df.shape}")


def process_data(
    data_dict: Dict[str, pd.DataFrame],
    cat_field_list: List[str],
    label_name: str,
    job_type: str,
    risk_tables_dict: Optional[Dict] = None,
    smooth_factor: float = 0.01,
    count_threshold: int = 5,
) -> Tuple[Dict[str, pd.DataFrame], OfflineBinning]:
    """
    Core data processing logic for risk table mapping.

    Args:
        data_dict: Dictionary of dataframes keyed by split name
        cat_field_list: List of categorical field names
        label_name: Target column name
        job_type: Type of job (training, validation, testing, calibration)
        risk_tables_dict: Pre-existing risk tables (for non-training jobs)
        smooth_factor: Smoothing factor for risk tables (default: 0.01)
        count_threshold: Minimum count threshold (default: 5)

    Returns:
        Tuple containing:
        - Dictionary of transformed dataframes
        - OfflineBinning instance with fitted risk tables
    """
    # For training mode, we need to validate categorical fields
    if job_type == "training":
        # Validate categorical fields on training data
        valid_cat_fields = validate_categorical_fields(
            data_dict["train"], cat_field_list
        )

        if not valid_cat_fields:
            logger.warning(
                "No valid categorical fields found for risk mapping. Using original data."
            )
            transformed_data = data_dict

            # Create empty binner with no fields for artifact consistency
            binner = OfflineBinning([], label_name)
            logger.info(
                "Creating empty risk tables since no valid categorical fields found"
            )
        else:
            # Initialize transformers with valid fields
            binner = OfflineBinning(valid_cat_fields, label_name)

            logger.info(
                "Running in 'training' mode: fitting on train data, transforming all splits"
            )

            # Fit risk tables on training data only
            binner.fit(
                data_dict["train"],
                smooth_factor=smooth_factor,
                count_threshold=count_threshold,
            )

            # Transform all splits
            transformed_data = {}
            for split_name, df in data_dict.items():
                df_transformed = binner.transform(df)
                transformed_data[split_name] = df_transformed
                logger.info(
                    f"Transformed {split_name} data, shape: {df_transformed.shape}"
                )

    # --- Logic for Non-Training modes (validation, testing, calibration) ---
    else:
        if not risk_tables_dict:
            raise ValueError(
                "For non-training job types, risk_tables_dict must be provided"
            )

        # Create binner with categorical fields from loaded risk tables
        cat_fields = [
            rt_info.get("varName")
            for rt_info in risk_tables_dict.values()
            if "varName" in rt_info
        ]

        binner = OfflineBinning(cat_fields, label_name)
        binner.load_risk_tables(risk_tables_dict)

        logger.info(f"Using pre-trained risk tables with {len(cat_fields)} fields")

        # Transform the data
        transformed_data = {}
        for split_name, df in data_dict.items():
            df_transformed = binner.transform(df)
            transformed_data[split_name] = df_transformed
            logger.info(f"Transformed {split_name} data, shape: {df_transformed.shape}")

    return transformed_data, binner


def save_artifacts(
    binner: OfflineBinning, hyperparams: Dict[str, Any], output_path: Path
) -> None:
    """
    Save risk table artifacts to the specified output path.

    Args:
        binner: OfflineBinning instance with fitted risk tables
        hyperparams: Hyperparameter dictionary
        output_path: Path to save artifacts to
    """
    # Save risk tables - using consistent filename
    bin_output_path = output_path / RISK_TABLE_FILENAME
    with open(bin_output_path, "wb") as f:
        pkl.dump(binner.risk_tables, f)
    logger.info(f"Saved binning mapping to {bin_output_path}")
    logger.info(f"This file can be used as input for non-training jobs")

    # Save a copy of hyperparameters with the model
    hyperparams_output_path = output_path / HYPERPARAMS_FILENAME
    with open(hyperparams_output_path, "w") as f:
        json.dump(hyperparams, f, indent=2)
    logger.info(f"Saved hyperparameters to {hyperparams_output_path}")


def load_risk_tables(risk_table_path: Path) -> Dict:
    """
    Load risk tables from a pickle file.

    Args:
        risk_table_path: Path to the risk tables file

    Returns:
        Dictionary of risk tables
    """
    if not risk_table_path.exists():
        raise FileNotFoundError(f"Risk table file not found: {risk_table_path}")

    logger.info(f"Loading risk tables from {risk_table_path}")
    with open(risk_table_path, "rb") as f:
        risk_tables = pkl.load(f)

    logger.info(f"Successfully loaded risk tables with {len(risk_tables)} mappings")
    return risk_tables


def internal_main(
    job_type: str,
    input_dir: str,
    output_dir: str,
    hyperparams: Dict[str, Any],
    risk_table_input_dir: Optional[str] = None,
    load_data_func: Callable = load_split_data,
    save_data_func: Callable = save_output_data,
) -> Tuple[Dict[str, pd.DataFrame], OfflineBinning]:
    """
    Main logic for risk table mapping, modified to handle both training and inference modes.

    Args:
        job_type: Type of job (training, validation, testing, calibration)
        input_dir: Input directory for data
        output_dir: Output directory for processed data
        hyperparams: Hyperparameters dictionary loaded from hyperparameters.json
        risk_table_input_dir: Directory containing pre-trained risk tables (for non-training jobs)
        load_data_func: Function to load data (for dependency injection in tests)
        save_data_func: Function to save data (for dependency injection in tests)

    Returns:
        Tuple containing:
        - Dictionary of transformed dataframes
        - OfflineBinning instance with fitted risk tables
    """
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)

    # Extract required hyperparameters with defaults
    cat_field_list = hyperparams.get("cat_field_list", [])
    label_name = hyperparams.get("label_name", "target")
    smooth_factor = hyperparams.get("smooth_factor", 0.01)  # Default if not provided
    count_threshold = hyperparams.get("count_threshold", 5)  # Default if not provided

    logger.info(
        f"Using hyperparameters: label_name={label_name}, "
        + f"smooth_factor={smooth_factor}, count_threshold={count_threshold}"
    )
    logger.info(f"Categorical fields from hyperparameters: {cat_field_list}")

    # Load data according to job type
    data_dict = load_data_func(job_type, input_dir)

    # Load risk tables if needed (non-training modes)
    risk_tables_dict = None
    if job_type != "training" and risk_table_input_dir:
        # Use the consistent filename for loading risk tables
        risk_table_path = Path(risk_table_input_dir) / RISK_TABLE_FILENAME
        risk_tables_dict = load_risk_tables(risk_table_path)
        logger.info(f"Loaded pre-trained risk tables from {risk_table_path}")

    # Process the data
    transformed_data, binner = process_data(
        data_dict=data_dict,
        cat_field_list=cat_field_list,
        label_name=label_name,
        job_type=job_type,
        risk_tables_dict=risk_tables_dict,
        smooth_factor=smooth_factor,
        count_threshold=count_threshold,
    )

    # Save processed data
    save_data_func(job_type, output_dir, transformed_data)

    # Save fitted artifacts
    save_artifacts(binner, hyperparams, output_path)

    logger.info("Risk-table mapping complete.")
    return transformed_data, binner


def main(
    input_paths: Dict[str, str],
    output_paths: Dict[str, str],
    environ_vars: Dict[str, str],
    job_args: Optional[argparse.Namespace] = None,
) -> Tuple[Dict[str, pd.DataFrame], OfflineBinning]:
    """
    Standardized main entry point for risk table mapping script.

    Args:
        input_paths: Dictionary of input paths with logical names
        output_paths: Dictionary of output paths with logical names
        environ_vars: Dictionary of environment variables
        job_args: Command line arguments (optional)

    Returns:
        Tuple containing:
        - Dictionary of transformed dataframes
        - OfflineBinning instance with fitted risk tables
    """
    try:
        # Extract paths from input parameters - required keys must be present
        if "data_input" not in input_paths:
            raise ValueError("Missing required input path: data_input")
        if "data_output" not in output_paths:
            raise ValueError("Missing required output path: data_output")

        # Extract job_type from args
        if job_args is None or not hasattr(job_args, "job_type"):
            raise ValueError("job_args must contain job_type parameter")

        job_type = job_args.job_type
        input_dir = input_paths["data_input"]
        output_dir = output_paths["data_output"]
        config_dir = input_paths.get("config_input", DEFAULT_CONFIG_DIR)

        # For non-training jobs, check if risk table input path is provided
        risk_table_input_dir = None
        if job_type != "training":
            risk_table_input_dir = input_paths.get("risk_table_input")
            if not risk_table_input_dir:
                logger.warning(
                    f"No risk_table_input path provided for non-training job {job_type}. "
                    + "Risk table mapping may fail."
                )

        # Log input/output paths for clarity
        logger.info(f"Input data directory: {input_dir}")
        logger.info(f"Output directory: {output_dir}")
        if risk_table_input_dir:
            logger.info(f"Risk table input directory: {risk_table_input_dir}")
            logger.info(
                f"Expected risk table path: {Path(risk_table_input_dir) / RISK_TABLE_FILENAME}"
            )

        # Load hyperparameters with source directory fallback (same pattern as XGBoost)
        # Use provided hyperparameters path, with source directory fallback
        if "hyperparameters_s3_uri" in input_paths:
            hparam_path = input_paths["hyperparameters_s3_uri"]
            if not hparam_path.endswith("hyperparameters.json"):
                hparam_path = os.path.join(hparam_path, "hyperparameters.json")
        else:
            # Fallback to source directory if not provided
            hparam_path = "/opt/ml/code/hyperparams/hyperparameters.json"

        logger.info(f"Loading hyperparameters from {hparam_path}")

        if os.path.exists(hparam_path):
            hyperparams = load_json_config(hparam_path)
        else:
            # FAIL with clear error instead of using arbitrary defaults
            raise FileNotFoundError(
                f"Hyperparameters file not found at {hparam_path}. "
                f"Risk table mapping requires hyperparameters to be provided either via "
                f"input channel or in source directory at /opt/ml/code/hyperparams/hyperparameters.json"
            )

        # Execute the internal main logic
        return internal_main(
            job_type=job_type,
            input_dir=input_dir,
            output_dir=output_dir,
            hyperparams=hyperparams,
            risk_table_input_dir=risk_table_input_dir,
        )

    except Exception as e:
        logger.error(f"Error in risk table mapping: {str(e)}")
        logger.error(traceback.format_exc())
        raise


if __name__ == "__main__":
    try:
        parser = argparse.ArgumentParser()
        parser.add_argument(
            "--job_type",
            type=str,
            required=True,
            choices=["training", "validation", "testing", "calibration"],
            help="Type of job to perform",
        )
        args = parser.parse_args()

        # Define standard SageMaker paths based on contract
        input_paths = {
            "data_input": DEFAULT_INPUT_DIR,
            "config_input": DEFAULT_CONFIG_DIR,
        }

        output_paths = {"data_output": DEFAULT_OUTPUT_DIR}

        # For non-training jobs, add risk table input path
        if args.job_type != "training":
            input_paths["risk_table_input"] = DEFAULT_RISK_TABLE_DIR

        # Environment variables dictionary (not used in this script)
        environ_vars = {}

        # Execute the main function with standardized inputs
        result, _ = main(input_paths, output_paths, environ_vars, args)

        logger.info(f"Risk table mapping completed successfully")
        sys.exit(0)
    except FileNotFoundError as e:
        logger.error(f"File not found error: {str(e)}")
        sys.exit(1)
    except ValueError as e:
        logger.error(f"Value error: {str(e)}")
        sys.exit(2)
    except Exception as e:
        logger.error(f"Error in risk table mapping script: {str(e)}")
        logger.error(traceback.format_exc())
        sys.exit(3)
