#!/usr/bin/env python
"""Step Specification for Model Calibration Step.

This file defines the step specification for the model calibration processing step,
including dependencies, outputs, and other metadata needed for pipeline integration.
"""

from ...core.base.specification_base import (
    StepSpecification,
    NodeType,
    DependencySpec,
    OutputSpec,
    DependencyType,
)
from ...registry.step_names import get_spec_step_type


def _get_model_calibration_contract():
    """Get the script contract for the ModelCalibration step.

    Returns:
        ScriptContract: The contract defining input/output paths and environment variables.
    """
    from ..contracts.model_calibration_contract import MODEL_CALIBRATION_CONTRACT

    return MODEL_CALIBRATION_CONTRACT


MODEL_CALIBRATION_SPEC = StepSpecification(
    step_type=get_spec_step_type("ModelCalibration"),
    node_type=NodeType.INTERNAL,
    script_contract=_get_model_calibration_contract(),
    dependencies={
        "evaluation_data": DependencySpec(
            logical_name="evaluation_data",
            dependency_type=DependencyType.PROCESSING_OUTPUT,
            required=True,
            compatible_sources=[
                "XGBoostTraining",
                "XGBoostModelEval",
                "XGBoostModelInference",
                "LightGBMTraining",
                "LightGBMModelEval",
                "LightGBMModelInference",
                "LightGBMMTTraining",
                "PyTorchTraining",
                "PyTorchModelEval",
                "PyTorchModelInference",
                "ModelEvaluation",
                "TrainingEvaluation",
                "CrossValidation",
            ],
            semantic_keywords=[
                "evaluation",
                "predictions",
                "scores",
                "results",
                "model_output",
                "performance",
                "inference",
                "output_data",
                "prediction_results",
                "training",
                "train",
                "training_evaluation",
                "train_predictions",
                "training_results",
                "validation",
                "val",
                "evaluate",
                "validation_evaluation",
                "val_predictions",
                "validation_results",
                "testing",
                "test",
                "testing_evaluation",
                "test_predictions",
                "testing_results",
                "calibration",
                "calib",
                "calibration_evaluation",
                "calib_predictions",
                "calibration_results",
            ],
            data_type="S3Uri",
            description="Evaluation dataset with ground truth labels and model predictions. Supports all job types: training, validation, testing, and calibration",
        )
    },
    outputs={
        "calibration_output": OutputSpec(
            logical_name="calibration_output",
            output_type=DependencyType.PROCESSING_OUTPUT,
            property_path="properties.ProcessingOutputConfig.Outputs['calibration_output'].S3Output.S3Uri",
            aliases=[
                "calibration_model",
                "calibration_artifacts",
                "probability_calibration",
                "calibrator",
                "score_transformer",
                "probability_adjustment",
                "training_calibration",
                "train_calibrator",
                "training_probability_calibration",
                "training_calibration_artifacts",
                "validation_calibration",
                "val_calibrator",
                "validation_probability_calibration",
                "validation_calibration_artifacts",
                "testing_calibration",
                "test_calibrator",
                "testing_probability_calibration",
                "testing_calibration_artifacts",
                "calibration_calibration",
                "calib_calibrator",
                "calibration_probability_calibration",
                "calibration_calibration_artifacts",
            ],
            data_type="S3Uri",
            description="Calibration mapping and artifacts. Compatible with all job types (training, validation, testing, calibration)",
        ),
        "metrics_output": OutputSpec(
            logical_name="metrics_output",
            output_type=DependencyType.PROCESSING_OUTPUT,
            property_path="properties.ProcessingOutputConfig.Outputs['metrics_output'].S3Output.S3Uri",
            aliases=[
                "calibration_metrics",
                "reliability_metrics",
                "probability_metrics",
                "calibration_performance",
                "calibration_evaluation",
                "training_calibration_metrics",
                "training_reliability_metrics",
                "training_probability_metrics",
                "training_calibration_performance",
                "training_calibration_evaluation",
                "validation_calibration_metrics",
                "validation_reliability_metrics",
                "validation_probability_metrics",
                "validation_calibration_performance",
                "validation_calibration_evaluation",
                "testing_calibration_metrics",
                "testing_reliability_metrics",
                "testing_probability_metrics",
                "testing_calibration_performance",
                "testing_calibration_evaluation",
                "calibration_calibration_metrics",
                "calibration_reliability_metrics",
                "calibration_probability_metrics",
                "calibration_calibration_performance",
                "calibration_calibration_evaluation",
            ],
            data_type="S3Uri",
            description="Calibration quality metrics and visualizations. Compatible with all job types (training, validation, testing, calibration)",
        ),
        "calibrated_data": OutputSpec(
            logical_name="calibrated_data",
            output_type=DependencyType.PROCESSING_OUTPUT,
            property_path="properties.ProcessingOutputConfig.Outputs['calibrated_data'].S3Output.S3Uri",
            aliases=[
                "calibrated_predictions",
                "calibrated_probabilities",
                "probability_scores",
                "adjusted_predictions",
                "calibrated_outputs",
                "training_calibrated_predictions",
                "training_calibrated_probabilities",
                "training_probability_scores",
                "training_adjusted_predictions",
                "training_calibrated_outputs",
                "validation_calibrated_predictions",
                "validation_calibrated_probabilities",
                "validation_probability_scores",
                "validation_adjusted_predictions",
                "validation_calibrated_outputs",
                "testing_calibrated_predictions",
                "testing_calibrated_probabilities",
                "testing_probability_scores",
                "testing_adjusted_predictions",
                "testing_calibrated_outputs",
                "calibration_calibrated_predictions",
                "calibration_calibrated_probabilities",
                "calibration_probability_scores",
                "calibration_adjusted_predictions",
                "calibration_calibrated_outputs",
            ],
            data_type="S3Uri",
            description="Dataset with calibrated probabilities. Compatible with all job types (training, validation, testing, calibration)",
        ),
    },
)
