#!/usr/bin/env python3
# Copyright (c) 2026 DataRobot, Inc. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""
Generate a CSV template for prediction data.

Usage:
    python generate_prediction_data_template.py <deployment_id> [n_rows] [output_file]

Generates a CSV template with all required columns and sample values.
"""

import csv
import sys

import datarobot as dr


def generate_prediction_data_template(
    deployment_id: str, n_rows: int = 1, output_file: str | None = None
) -> str:
    """
    Generate a CSV template for prediction data.

    Args:
        deployment_id: The deployment ID
        n_rows: Number of template rows to generate (default: 1)
        output_file: Optional output file path (default: prints to stdout)

    Returns:
        CSV template content
    """
    # Initialize client
    dr.Client()

    # Get deployment features
    deployment = dr.Deployment.get(deployment_id)
    model = dr.Model.get(deployment.model["id"])
    features = model.get_features()

    # Filter out target feature
    prediction_features = [f for f in features if f.name != model.target_name]

    # Generate template rows with sample values
    import io

    output = io.StringIO()
    writer = csv.DictWriter(output, fieldnames=[f.name for f in prediction_features])
    writer.writeheader()

    # Generate sample rows based on feature types
    for i in range(n_rows):
        row = {}
        for feature in prediction_features:
            if feature.feature_type == "Numeric":
                row[feature.name] = 0.0
            elif feature.feature_type == "Categorical":
                row[feature.name] = "sample_category"
            elif feature.feature_type == "Text":
                row[feature.name] = "sample text"
            elif feature.feature_type == "Date":
                row[feature.name] = "2024-01-01"
            else:
                row[feature.name] = ""
        writer.writerow(row)

    csv_content = output.getvalue()

    # Add metadata comments
    metadata_comments = f"""# Prediction Data Template for Deployment: {deployment_id}
# Model: {model.project_name}
# Target: {model.target_name}
# Generated: {n_rows} template rows
# 
# Instructions:
# 1. Fill in the values for each feature
# 2. Ensure data types match feature types
# 3. Use validate_prediction_data.py to check before submitting
#
"""

    full_content = metadata_comments + csv_content

    if output_file:
        with open(output_file, "w") as f:
            f.write(full_content)
        return f"Template written to {output_file}"
    else:
        return full_content


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(
            "Usage: python generate_prediction_data_template.py <deployment_id> [n_rows] [output_file]",
            file=sys.stderr,
        )
        sys.exit(1)

    deployment_id = sys.argv[1]
    n_rows = int(sys.argv[2]) if len(sys.argv) > 2 else 1
    output_file = sys.argv[3] if len(sys.argv) > 3 else None

    result = generate_prediction_data_template(deployment_id, n_rows, output_file)
    print(result)

#!/usr/bin/env python3
# Copyright (c) 2026 DataRobot, Inc. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""
Get comprehensive information about features required by a deployment.

Usage:
    python get_deployment_features.py <deployment_id>

Outputs JSON with feature information, types, importance, and time series config.
"""

import json
import sys

import datarobot as dr


def get_deployment_features(deployment_id: str) -> dict:
    """
    Get comprehensive information about features required by a deployment.

    Args:
        deployment_id: The deployment ID

    Returns:
        Dictionary with feature information, types, importance, and time series config
    """
    # Initialize client
    dr.Client()

    deployment = dr.Deployment.get(deployment_id)
    model = dr.Model.get(deployment.model["id"])
    project = dr.Project.get(model.project_id)

    # Get feature information
    features = model.get_features()
    feature_importance = model.get_feature_impact()

    # Build feature list
    feature_list = []
    for feature in features:
        importance = 0.0
        for fi in feature_importance:
            if fi["featureName"] == feature.name:
                importance = fi.get("impactNormalized", 0.0)
                break

        feature_list.append(
            {
                "feature_name": feature.name,
                "feature_type": feature.feature_type,
                "importance": importance,
                "is_target": feature.name == model.target_name,
            }
        )

    # Get time series config if applicable
    time_series_config = None
    if project.use_time_series:
        try:
            time_series_info = project.get_time_series_info()
            time_series_config = {
                "datetime_column": time_series_info.datetime_partition_column,
                "forecast_window_start": time_series_info.forecast_window_start,
                "forecast_window_end": time_series_info.forecast_window_end,
                "series_id_columns": time_series_info.multiseries_id_columns or [],
            }
        except dr.errors.ClientError as e:
            print(
                f"Note: time series info unavailable: {e}",
                file=sys.stderr,
            )

    return {
        "deployment_id": deployment_id,
        "model_type": model.target_type,
        "target": model.target_name,
        "target_type": model.target_type,
        "features": feature_list,
        "time_series_config": time_series_config,
    }


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(
            "Usage: python get_deployment_features.py <deployment_id>", file=sys.stderr
        )
        sys.exit(1)

    deployment_id = sys.argv[1]
    result = get_deployment_features(deployment_id)
    print(json.dumps(result, indent=2))

#!/usr/bin/env python3
# Copyright (c) 2026 DataRobot, Inc. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""
Make a prediction from a deployment, optionally with prediction explanations.

Usage:
    python make_prediction.py <deployment_id> <data_json> [options]

Where <data_json> is either:
    - a JSON object with feature values (single row), or
    - a JSON array of objects (multiple rows).

Options:
    --max-explanations N        Number of top explanations per row (0 disables, default 0).
    --max-ngram-explanations N  Cap text-segment explanations per row (text models only).
    --threshold-high X          Only explain rows with prediction probability above X (0-1).
    --threshold-low X           Only explain rows with prediction probability below X (0-1).
    --explanation-algorithm A   'shap' or 'xemp' (omit to use deployment default).
    --passthrough-columns COLS  'all' or comma-separated input columns to copy through.

Examples:
    # Plain prediction
    python make_prediction.py abc123 '{"feature1": 10, "feature2": 20}'

    # Top-3 SHAP explanations per row
    python make_prediction.py abc123 '{"feature1": 10}' --max-explanations 3 \\
        --explanation-algorithm shap

When --max-explanations > 0, each row in the output includes an `explanations` list
of {feature, value, strength, qualitative_strength} entries describing why the model
produced that prediction.
"""

import argparse
import json
import sys

import datarobot as dr
import pandas as pd
from datarobot_predict.deployment import predict as dr_predict


def make_prediction(
    deployment_id: str,
    data,
    *,
    max_explanations: int = 0,
    max_ngram_explanations: int | None = None,
    threshold_high: float | None = None,
    threshold_low: float | None = None,
    explanation_algorithm: str | None = None,
    passthrough_columns: str | None = None,
) -> dict:
    """Score `data` against `deployment_id`, optionally returning prediction explanations."""
    dr.Client()

    deployment = dr.Deployment.get(deployment_id)

    rows = data if isinstance(data, list) else [data]
    df = pd.DataFrame(rows)

    predict_kwargs = {"deployment": deployment, "data_frame": df}
    if max_explanations and max_explanations > 0:
        predict_kwargs["max_explanations"] = max_explanations
    if max_ngram_explanations is not None:
        predict_kwargs["max_ngram_explanations"] = max_ngram_explanations
    if threshold_high is not None:
        predict_kwargs["threshold_high"] = threshold_high
    if threshold_low is not None:
        predict_kwargs["threshold_low"] = threshold_low
    if explanation_algorithm is not None:
        predict_kwargs["explanation_algorithm"] = explanation_algorithm
    if passthrough_columns is not None:
        predict_kwargs["passthrough_columns"] = (
            "all"
            if passthrough_columns == "all"
            else {c.strip() for c in passthrough_columns.split(",")}
        )

    result = dr_predict(**predict_kwargs)
    predictions_df = result.dataframe

    return {
        "deployment_id": deployment_id,
        "row_count": len(predictions_df),
        "predictions": predictions_df.to_dict(orient="records"),
    }


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawTextHelpFormatter
    )
    parser.add_argument("deployment_id")
    parser.add_argument(
        "data_json", help="JSON object (single row) or JSON array (multiple rows)"
    )
    parser.add_argument("--max-explanations", type=int, default=0)
    parser.add_argument("--max-ngram-explanations", type=int, default=None)
    parser.add_argument("--threshold-high", type=float, default=None)
    parser.add_argument("--threshold-low", type=float, default=None)
    parser.add_argument(
        "--explanation-algorithm", choices=["shap", "xemp"], default=None
    )
    parser.add_argument("--passthrough-columns", default=None)
    return parser.parse_args()


if __name__ == "__main__":
    args = _parse_args()

    try:
        data = json.loads(args.data_json)
    except json.JSONDecodeError as e:
        print(f"Error: Invalid JSON: {e}", file=sys.stderr)
        sys.exit(1)

    result = make_prediction(
        args.deployment_id,
        data,
        max_explanations=args.max_explanations,
        max_ngram_explanations=args.max_ngram_explanations,
        threshold_high=args.threshold_high,
        threshold_low=args.threshold_low,
        explanation_algorithm=args.explanation_algorithm,
        passthrough_columns=args.passthrough_columns,
    )
    print(json.dumps(result, indent=2, default=str))

#!/usr/bin/env python3
# Copyright (c) 2026 DataRobot, Inc. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""
Validate prediction data for a deployment.

Usage:
    python validate_prediction_data.py <deployment_id> <file_path>

Returns validation report with errors, warnings, and info messages.
"""

import csv
import json
import os
import sys

import datarobot as dr


def validate_prediction_data(deployment_id: str, file_path: str) -> dict:
    """
    Validate prediction data for a deployment.

    Args:
        deployment_id: The deployment ID
        file_path: Path to CSV file to validate

    Returns:
        Validation report with errors, warnings, and info
    """
    # Initialize client
    dr.Client()

    deployment = dr.Deployment.get(deployment_id)
    model = dr.Model.get(deployment.model["id"])
    features = model.get_features()

    # Get required features (exclude target)
    required_features = {
        f.name: f.feature_type for f in features if f.name != model.target_name
    }

    # Read CSV data
    if not os.path.exists(file_path):
        return {"valid": False, "errors": [f"File not found: {file_path}"]}

    with open(file_path, "r") as f:
        reader = csv.DictReader(f)
        rows = list(reader)

    if not rows:
        return {"valid": False, "errors": ["CSV file is empty"]}

    # Validate
    errors = []
    warnings = []
    info = []

    # Check for missing required features
    csv_columns = set(rows[0].keys())
    missing_features = set(required_features.keys()) - csv_columns

    if missing_features:
        errors.append(
            f"Missing required features: {', '.join(sorted(missing_features))}"
        )

    # Check for extra columns
    extra_columns = csv_columns - set(required_features.keys())
    if extra_columns:
        info.append(
            f"Extra columns (will be ignored): {', '.join(sorted(extra_columns))}"
        )

    # Check data types (simplified - actual validation would be more thorough)
    for row_num, row in enumerate(rows, start=2):  # Start at 2 (header is row 1)
        for feature_name, feature_type in required_features.items():
            if feature_name in row:
                value = row[feature_name]
                if value == "":
                    warnings.append(f"Row {row_num}, {feature_name}: Empty value")
                elif feature_type == "Numeric" and value:
                    try:
                        float(value)
                    except ValueError:
                        errors.append(
                            f"Row {row_num}, {feature_name}: Expected numeric, got '{value}'"
                        )

    # Get feature importance for warnings
    try:
        feature_importance = model.get_feature_impact()
        low_importance_features = [
            fi["featureName"]
            for fi in feature_importance
            if fi.get("impactNormalized", 0) < 0.05
        ]

        missing_low_importance = missing_features & set(low_importance_features)
        if missing_low_importance:
            warnings.append(
                f"Missing low-importance features: {', '.join(sorted(missing_low_importance))}"
            )
    except dr.errors.ClientError as e:
        print(
            f"Note: feature impact unavailable, skipping importance warnings: {e}",
            file=sys.stderr,
        )

    return {
        "valid": len(errors) == 0,
        "deployment_id": deployment_id,
        "row_count": len(rows),
        "errors": errors,
        "warnings": warnings,
        "info": info,
        "required_features": sorted(required_features.keys()),
        "provided_features": sorted(csv_columns),
    }


if __name__ == "__main__":
    if len(sys.argv) < 3:
        print(
            "Usage: python validate_prediction_data.py <deployment_id> <file_path>",
            file=sys.stderr,
        )
        sys.exit(1)

    deployment_id = sys.argv[1]
    file_path = sys.argv[2]

    result = validate_prediction_data(deployment_id, file_path)
    print(json.dumps(result, indent=2))