#!/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
import shutil
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_MODEL_ARTIFACTS_DIR = "/opt/ml/processing/input/model_artifacts"

# Constants for file paths to ensure consistency between training and inference
# Match XGBoost training output format
RISK_TABLE_FILENAME = "risk_table_map.pkl"  # Matches XGBoost training output
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__)


# --- File I/O Helper Functions with Format Preservation ---


def _detect_file_format(split_dir: Path, split_name: str) -> tuple:
    """
    Detect the format of processed data file.

    Returns:
        Tuple of (file_path, format) where format is 'csv', 'tsv', or 'parquet'
    """
    # Try different formats in order of preference
    formats = [
        (f"{split_name}_processed_data.csv", "csv"),
        (f"{split_name}_processed_data.tsv", "tsv"),
        (f"{split_name}_processed_data.parquet", "parquet"),
    ]

    for filename, fmt in formats:
        file_path = split_dir / filename
        if file_path.exists():
            return file_path, fmt

    raise RuntimeError(
        f"No processed data file found in {split_dir}. "
        f"Looked for: {[f[0] for f in formats]}"
    )


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 with automatic format detection.

    For 'training': Loads data from train, test, and val subdirectories
    For others: Loads single job_type split

    Returns:
        Dictionary with DataFrames and detected format stored in '_format' key
    """
    input_path = Path(input_dir)
    result = {}

    if job_type == "training":
        # For training, we expect data in train/test/val subdirectories
        splits = ["train", "test", "val"]
        detected_format = None

        for split_name in splits:
            split_dir = input_path / split_name
            file_path, fmt = _detect_file_format(split_dir, split_name)

            # Store format from first split (they should all match)
            if detected_format is None:
                detected_format = fmt

            # Read based on format
            if fmt == "csv":
                df = pd.read_csv(file_path)
            elif fmt == "tsv":
                df = pd.read_csv(file_path, sep="\t")
            elif fmt == "parquet":
                df = pd.read_parquet(file_path)
            else:
                raise RuntimeError(f"Unsupported format: {fmt}")

            result[split_name] = df

        result["_format"] = detected_format  # Store detected format
        logger.info(
            f"Loaded training data splits (format={detected_format}): "
            f"train={result['train'].shape}, test={result['test'].shape}, val={result['val'].shape}"
        )
    else:
        # For other job types, we expect data in a single directory named after job_type
        split_dir = input_path / job_type
        file_path, detected_format = _detect_file_format(split_dir, job_type)

        # Read based on format
        if detected_format == "csv":
            df = pd.read_csv(file_path)
        elif detected_format == "tsv":
            df = pd.read_csv(file_path, sep="\t")
        elif detected_format == "parquet":
            df = pd.read_parquet(file_path)
        else:
            raise RuntimeError(f"Unsupported format: {detected_format}")

        result[job_type] = df
        result["_format"] = detected_format  # Store detected format
        logger.info(f"Loaded {job_type} data (format={detected_format}): {df.shape}")

    return result


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

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

    # Extract format from data_dict (stored during load)
    output_format = data_dict.get("_format", "csv")  # Default to CSV if not found

    for split_name, df in data_dict.items():
        # Skip the format metadata key
        if split_name == "_format":
            continue

        split_output_dir = output_path / split_name
        split_output_dir.mkdir(exist_ok=True, parents=True)

        # Save in detected format
        if output_format == "csv":
            output_file = split_output_dir / f"{split_name}_processed_data.csv"
            df.to_csv(output_file, index=False)
        elif output_format == "tsv":
            output_file = split_output_dir / f"{split_name}_processed_data.tsv"
            df.to_csv(output_file, sep="\t", index=False)
        elif output_format == "parquet":
            output_file = split_output_dir / f"{split_name}_processed_data.parquet"
            df.to_parquet(output_file, index=False)
        else:
            raise RuntimeError(f"Unsupported output format: {output_format}")

        logger.info(
            f"Saved {split_name} data to {output_file} (format={output_format}), 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.

    Output format matches XGBoost training's risk_table_map.pkl format:
    A dictionary mapping variable names to their risk table dictionaries.
    Each risk table dict contains 'bins' and 'default_bin' keys.

    Args:
        binner: OfflineBinning instance with fitted risk tables
        hyperparams: Hyperparameter dictionary
        output_path: Path to save artifacts to
    """
    # Save risk tables in XGBoost training format - using consistent filename
    # Format: {variable_name: {"bins": {...}, "default_bin": value, ...}}
    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 risk table mapping to {bin_output_path}")
    logger.info(f"Format: {list(binner.risk_tables.keys())} variables")
    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 copy_existing_artifacts(src_dir: str, dst_dir: str) -> None:
    """
    Copy all existing model artifacts from previous processing steps.

    This enables the parameter accumulator pattern where each step:
    1. Copies artifacts from previous steps
    2. Adds its own artifacts
    3. Passes all artifacts to the next step

    Args:
        src_dir: Source directory containing existing artifacts
        dst_dir: Destination directory to copy artifacts to
    """
    if not src_dir or not os.path.exists(src_dir):
        logger.info(f"No existing artifacts to copy from {src_dir}")
        return

    os.makedirs(dst_dir, exist_ok=True)
    copied_count = 0

    for filename in os.listdir(src_dir):
        src_file = os.path.join(src_dir, filename)
        dst_file = os.path.join(dst_dir, filename)

        if os.path.isfile(src_file):
            shutil.copy2(src_file, dst_file)
            copied_count += 1
            logger.info(f"  Copied existing artifact: {filename}")

    logger.info(f"✓ Copied {copied_count} existing artifact(s) to {dst_dir}")


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],
    model_artifacts_input_dir: Optional[str] = None,
    model_artifacts_output_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
        model_artifacts_input_dir: Directory containing model artifacts from previous steps
        model_artifacts_output_dir: Directory to save model artifacts for next steps
        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}")

    # Determine model artifacts output directory
    artifacts_output_dir = (
        Path(model_artifacts_output_dir)
        if model_artifacts_output_dir
        else output_path / "model_artifacts"
    )
    artifacts_output_dir.mkdir(parents=True, exist_ok=True)

    # Copy existing artifacts from previous steps (parameter accumulator pattern)
    if model_artifacts_input_dir:
        copy_existing_artifacts(model_artifacts_input_dir, str(artifacts_output_dir))

    # 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 model_artifacts_input_dir:
        # Use the consistent filename for loading risk tables
        risk_table_path = Path(model_artifacts_input_dir) / RISK_TABLE_FILENAME
        if risk_table_path.exists():
            risk_tables_dict = load_risk_tables(risk_table_path)
            logger.info(f"Loaded pre-trained risk tables from {risk_table_path}")
        else:
            logger.warning(f"Risk tables not found at {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, artifacts_output_dir)

    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
            - "data_input": Input data directory
            - "model_artifacts_input": Model artifacts from previous steps (standardized)
            - "hyperparameters_s3_uri": Path to hyperparameters (optional)
        output_paths: Dictionary of output paths with logical names
            - "processed_data": Output directory for processed data
            - "model_artifacts_output": Model artifacts output for next steps (standardized)
        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 "input_data" not in input_paths:
            raise ValueError("Missing required input path: input_data")
        if "processed_data" not in output_paths:
            raise ValueError("Missing required output path: processed_data")

        # 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["input_data"]
        output_dir = output_paths["processed_data"]

        # Get standardized model artifacts paths
        model_artifacts_input_dir = input_paths.get("model_artifacts_input")
        model_artifacts_output_dir = output_paths.get("model_artifacts_output")

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

        # 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,
            model_artifacts_input_dir=model_artifacts_input_dir,
            model_artifacts_output_dir=model_artifacts_output_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
        # Separate data and model artifacts into different subfolders
        input_paths = {
            "input_data": DEFAULT_INPUT_DIR,
            "hyperparameters_s3_uri": DEFAULT_CONFIG_DIR,
        }

        output_paths = {
            "processed_data": DEFAULT_OUTPUT_DIR + "/data",
            "model_artifacts_output": DEFAULT_OUTPUT_DIR + "/model_artifacts",
        }

        # For non-training jobs, add model artifacts input path
        if args.job_type != "training":
            input_paths["model_artifacts_input"] = DEFAULT_MODEL_ARTIFACTS_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)
