"""
Metric catalog for Verity audit engine.

This is the "everything possible" list organized by mode.
"""
from typing import Dict, List
from audit.metric_contract import (
    MetricCategory, AuditMode, TaskType, MetricUnit, 
    MetricDirection, MetricAggregation
)


# Dataset Model Catalog
DATASET_MODEL_CATALOG = {
    # Performance Metrics
    "performance": {
        "regression": [
            {
                "metric_id": "nmae",
                "category": MetricCategory.PERFORMANCE,
                "unit": MetricUnit.NORMALIZED_ERROR,
                "direction": MetricDirection.LOWER_IS_BETTER,
                "aggregation": MetricAggregation.SCALAR
            },
            {
                "metric_id": "nrmse_by_std",
                "category": MetricCategory.PERFORMANCE,
                "unit": MetricUnit.NORMALIZED_ERROR,
                "direction": MetricDirection.LOWER_IS_BETTER,
                "aggregation": MetricAggregation.SCALAR
            },
            {
                "metric_id": "nrmse_by_mean",
                "category": MetricCategory.PERFORMANCE,
                "unit": MetricUnit.NORMALIZED_ERROR,
                "direction": MetricDirection.LOWER_IS_BETTER,
                "aggregation": MetricAggregation.SCALAR
            },
            {
                "metric_id": "mape",
                "category": MetricCategory.PERFORMANCE,
                "unit": MetricUnit.PERCENTAGE,
                "direction": MetricDirection.LOWER_IS_BETTER,
                "aggregation": MetricAggregation.SCALAR
            },
            {
                "metric_id": "smape",
                "category": MetricCategory.PERFORMANCE,
                "unit": MetricUnit.PERCENTAGE,
                "direction": MetricDirection.LOWER_IS_BETTER,
                "aggregation": MetricAggregation.SCALAR
            },
            {
                "metric_id": "r2",
                "category": MetricCategory.PERFORMANCE,
                "unit": MetricUnit.RATIO,
                "direction": MetricDirection.HIGHER_IS_BETTER,
                "aggregation": MetricAggregation.SCALAR
            },
            {
                "metric_id": "mae",
                "category": MetricCategory.PERFORMANCE,
                "unit": MetricUnit.ABSOLUTE_DIFFERENCE,
                "direction": MetricDirection.LOWER_IS_BETTER,
                "aggregation": MetricAggregation.SCALAR
            },
            {
                "metric_id": "rmse",
                "category": MetricCategory.PERFORMANCE,
                "unit": MetricUnit.ABSOLUTE_DIFFERENCE,
                "direction": MetricDirection.LOWER_IS_BETTER,
                "aggregation": MetricAggregation.SCALAR
            }
        ],
        "classification": [
            {
                "metric_id": "accuracy",
                "category": MetricCategory.PERFORMANCE,
                "unit": MetricUnit.RATIO,
                "direction": MetricDirection.HIGHER_IS_BETTER,
                "aggregation": MetricAggregation.SCALAR
            },
            {
                "metric_id": "auc",
                "category": MetricCategory.PERFORMANCE,
                "unit": MetricUnit.RATIO,
                "direction": MetricDirection.HIGHER_IS_BETTER,
                "aggregation": MetricAggregation.SCALAR
            },
            {
                "metric_id": "f1",
                "category": MetricCategory.PERFORMANCE,
                "unit": MetricUnit.RATIO,
                "direction": MetricDirection.HIGHER_IS_BETTER,
                "aggregation": MetricAggregation.SCALAR
            },
            {
                "metric_id": "precision",
                "category": MetricCategory.PERFORMANCE,
                "unit": MetricUnit.RATIO,
                "direction": MetricDirection.HIGHER_IS_BETTER,
                "aggregation": MetricAggregation.SCALAR
            },
            {
                "metric_id": "recall",
                "category": MetricCategory.PERFORMANCE,
                "unit": MetricUnit.RATIO,
                "direction": MetricDirection.HIGHER_IS_BETTER,
                "aggregation": MetricAggregation.SCALAR
            },
            {
                "metric_id": "calibration",
                "category": MetricCategory.PERFORMANCE,
                "unit": MetricUnit.STATISTICAL_DISTANCE,
                "direction": MetricDirection.LOWER_IS_BETTER,
                "aggregation": MetricAggregation.SCALAR
            },
            {
                "metric_id": "brier_score",
                "category": MetricCategory.PERFORMANCE,
                "unit": MetricUnit.RATIO,
                "direction": MetricDirection.LOWER_IS_BETTER,
                "aggregation": MetricAggregation.SCALAR
            }
        ]
    },
    # Group Fairness Metrics
    "fairness": {
        "classification": [
            {
                "metric_id": "demographic_parity_diff",
                "category": MetricCategory.FAIRNESS,
                "unit": MetricUnit.NORMALIZED,
                "direction": MetricDirection.LOWER_IS_BETTER,
                "aggregation": MetricAggregation.PER_GROUP
            },
            {
                "metric_id": "equal_opportunity_diff",
                "category": MetricCategory.FAIRNESS,
                "unit": MetricUnit.NORMALIZED,
                "direction": MetricDirection.LOWER_IS_BETTER,
                "aggregation": MetricAggregation.PER_GROUP
            },
            {
                "metric_id": "disparate_impact_ratio_min",
                "category": MetricCategory.FAIRNESS,
                "unit": MetricUnit.RATIO,
                "direction": MetricDirection.HIGHER_IS_BETTER,
                "aggregation": MetricAggregation.PER_GROUP
            },
            {
                "metric_id": "fpr_gap_by_group",
                "category": MetricCategory.FAIRNESS,
                "unit": MetricUnit.NORMALIZED,
                "direction": MetricDirection.LOWER_IS_BETTER,
                "aggregation": MetricAggregation.PER_GROUP
            },
            {
                "metric_id": "fnr_gap_by_group",
                "category": MetricCategory.FAIRNESS,
                "unit": MetricUnit.NORMALIZED,
                "direction": MetricDirection.LOWER_IS_BETTER,
                "aggregation": MetricAggregation.PER_GROUP
            },
            {
                "metric_id": "calibration_gap_by_group",
                "category": MetricCategory.FAIRNESS,
                "unit": MetricUnit.NORMALIZED,
                "direction": MetricDirection.LOWER_IS_BETTER,
                "aggregation": MetricAggregation.PER_GROUP
            },
            {
                "metric_id": "worst_group_performance",
                "category": MetricCategory.FAIRNESS,
                "unit": MetricUnit.RATIO,
                "direction": MetricDirection.HIGHER_IS_BETTER,
                "aggregation": MetricAggregation.SCALAR
            }
        ],
        "regression": [
            {
                "metric_id": "mean_prediction_error_by_group",
                "category": MetricCategory.FAIRNESS,
                "unit": MetricUnit.NORMALIZED,
                "direction": MetricDirection.LOWER_IS_BETTER,
                "aggregation": MetricAggregation.PER_GROUP
            },
            {
                "metric_id": "mae_rmse_by_protected_group",
                "category": MetricCategory.FAIRNESS,
                "unit": MetricUnit.RELATIVE_DIFFERENCE,
                "direction": MetricDirection.LOWER_IS_BETTER,
                "aggregation": MetricAggregation.PER_GROUP
            },
            {
                "metric_id": "over_under_estimation_bias",
                "category": MetricCategory.FAIRNESS,
                "unit": MetricUnit.NORMALIZED,
                "direction": MetricDirection.LOWER_IS_BETTER,
                "aggregation": MetricAggregation.PER_GROUP
            },
            {
                "metric_id": "residual_distribution_parity",
                "category": MetricCategory.FAIRNESS,
                "unit": MetricUnit.STATISTICAL_DISTANCE,
                "direction": MetricDirection.LOWER_IS_BETTER,
                "aggregation": MetricAggregation.PER_GROUP
            },
            {
                "metric_id": "worst_group_error_analysis",
                "category": MetricCategory.FAIRNESS,
                "unit": MetricUnit.RATIO,
                "direction": MetricDirection.LOWER_IS_BETTER,
                "aggregation": MetricAggregation.SCALAR
            }
        ]
    },
    # Data Quality Metrics
    "data_quality": [
        {
            "metric_id": "missing_value_rates",
            "category": MetricCategory.DATA_QUALITY,
            "unit": MetricUnit.PERCENTAGE,
            "direction": MetricDirection.LOWER_IS_BETTER,
            "aggregation": MetricAggregation.PER_FEATURE
        },
        {
            "metric_id": "schema_validation",
            "category": MetricCategory.DATA_QUALITY,
            "unit": MetricUnit.BOOLEAN,
            "direction": MetricDirection.BOOLEAN_TRUE_IS_BAD,
            "aggregation": MetricAggregation.SCALAR
        },
        {
            "metric_id": "label_distribution",
            "category": MetricCategory.DATA_QUALITY,
            "unit": MetricUnit.RATIO,
            "direction": MetricDirection.IN_RANGE,
            "aggregation": MetricAggregation.DISTRIBUTION
        },
        {
            "metric_id": "duplicates",
            "category": MetricCategory.DATA_QUALITY,
            "unit": MetricUnit.PERCENTAGE,
            "direction": MetricDirection.LOWER_IS_BETTER,
            "aggregation": MetricAggregation.SCALAR
        },
        {
            "metric_id": "train_test_leakage",
            "category": MetricCategory.DATA_QUALITY,
            "unit": MetricUnit.BOOLEAN,
            "direction": MetricDirection.BOOLEAN_TRUE_IS_BAD,
            "aggregation": MetricAggregation.SCALAR
        },
        {
            "metric_id": "distribution_shifts",
            "category": MetricCategory.DATA_QUALITY,
            "unit": MetricUnit.STATISTICAL_DISTANCE,
            "direction": MetricDirection.LOWER_IS_BETTER,
            "aggregation": MetricAggregation.PER_FEATURE
        }
    ],
    # Explainability Metrics
    "explainability": [
        {
            "metric_id": "feature_importance",
            "category": MetricCategory.EXPLAINABILITY,
            "unit": MetricUnit.RATIO,
            "direction": MetricDirection.IN_RANGE,
            "aggregation": MetricAggregation.PER_FEATURE
        },
        {
            "metric_id": "proxy_detection",
            "category": MetricCategory.EXPLAINABILITY,
            "unit": MetricUnit.CORRELATION_COEFFICIENT,
            "direction": MetricDirection.LOWER_IS_BETTER,
            "aggregation": MetricAggregation.PER_FEATURE
        },
        {
            "metric_id": "sensitive_leakage_detection",
            "category": MetricCategory.EXPLAINABILITY,
            "unit": MetricUnit.MUTUAL_INFORMATION,
            "direction": MetricDirection.LOWER_IS_BETTER,
            "aggregation": MetricAggregation.PER_FEATURE
        }
    ],
    # Robustness Metrics
    "robustness": [
        {
            "metric_id": "noise_sensitivity",
            "category": MetricCategory.ROBUSTNESS,
            "unit": MetricUnit.RELATIVE_DIFFERENCE,
            "direction": MetricDirection.LOWER_IS_BETTER,
            "aggregation": MetricAggregation.SCALAR
        },
        {
            "metric_id": "subgroup_stability",
            "category": MetricCategory.ROBUSTNESS,
            "unit": MetricUnit.RATIO,
            "direction": MetricDirection.HIGHER_IS_BETTER,
            "aggregation": MetricAggregation.SCALAR
        },
        {
            "metric_id": "perturbation_sensitivity",
            "category": MetricCategory.ROBUSTNESS,
            "unit": MetricUnit.RELATIVE_DIFFERENCE,
            "direction": MetricDirection.LOWER_IS_BETTER,
            "aggregation": MetricAggregation.SCALAR
        }
    ]
}

# API/GenAI Model Catalog
API_MODEL_CATALOG = {
    # Safety & Policy Metrics
    "safety": [
        {
            "metric_id": "toxicity_rate",
            "category": MetricCategory.SAFETY,
            "unit": MetricUnit.PERCENTAGE,
            "direction": MetricDirection.LOWER_IS_BETTER,
            "aggregation": MetricAggregation.SCALAR
        },
        {
            "metric_id": "hate_harassment_rate",
            "category": MetricCategory.SAFETY,
            "unit": MetricUnit.PERCENTAGE,
            "direction": MetricDirection.LOWER_IS_BETTER,
            "aggregation": MetricAggregation.SCALAR
        },
        {
            "metric_id": "sexual_content_rate",
            "category": MetricCategory.SAFETY,
            "unit": MetricUnit.PERCENTAGE,
            "direction": MetricDirection.LOWER_IS_BETTER,
            "aggregation": MetricAggregation.SCALAR
        },
        {
            "metric_id": "self_harm_policy_violation_rate",
            "category": MetricCategory.SAFETY,
            "unit": MetricUnit.PERCENTAGE,
            "direction": MetricDirection.LOWER_IS_BETTER,
            "aggregation": MetricAggregation.SCALAR
        },
        {
            "metric_id": "violence_policy_violation_rate",
            "category": MetricCategory.SAFETY,
            "unit": MetricUnit.PERCENTAGE,
            "direction": MetricDirection.LOWER_IS_BETTER,
            "aggregation": MetricAggregation.SCALAR
        },
        {
            "metric_id": "disallowed_content_violation_rate",
            "category": MetricCategory.SAFETY,
            "unit": MetricUnit.PERCENTAGE,
            "direction": MetricDirection.LOWER_IS_BETTER,
            "aggregation": MetricAggregation.SCALAR
        }
    ],
    # Privacy Metrics
    "privacy": [
        {
            "metric_id": "pii_leak_rate",
            "category": MetricCategory.PRIVACY,
            "unit": MetricUnit.PERCENTAGE,
            "direction": MetricDirection.LOWER_IS_BETTER,
            "aggregation": MetricAggregation.SCALAR
        },
        {
            "metric_id": "secret_leakage_rate",
            "category": MetricCategory.PRIVACY,
            "unit": MetricUnit.PERCENTAGE,
            "direction": MetricDirection.LOWER_IS_BETTER,
            "aggregation": MetricAggregation.SCALAR
        },
        {
            "metric_id": "data_exfil_jailbreak_success_rate",
            "category": MetricCategory.PRIVACY,
            "unit": MetricUnit.PERCENTAGE,
            "direction": MetricDirection.LOWER_IS_BETTER,
            "aggregation": MetricAggregation.SCALAR
        }
    ],
    # Reliability Metrics
    "reliability": [
        {
            "metric_id": "refusal_consistency_gap",
            "category": MetricCategory.RELIABILITY,
            "unit": MetricUnit.NORMALIZED,
            "direction": MetricDirection.LOWER_IS_BETTER,
            "aggregation": MetricAggregation.PER_GROUP
        },
        {
            "metric_id": "instruction_following_failure_rate",
            "category": MetricCategory.RELIABILITY,
            "unit": MetricUnit.PERCENTAGE,
            "direction": MetricDirection.LOWER_IS_BETTER,
            "aggregation": MetricAggregation.SCALAR
        },
        {
            "metric_id": "hallucination_rate",
            "category": MetricCategory.RELIABILITY,
            "unit": MetricUnit.PERCENTAGE,
            "direction": MetricDirection.LOWER_IS_BETTER,
            "aggregation": MetricAggregation.SCALAR
        }
    ],
    # Bias Metrics
    "bias": [
        {
            "metric_id": "stereotype_association_rate",
            "category": MetricCategory.FAIRNESS,
            "unit": MetricUnit.PERCENTAGE,
            "direction": MetricDirection.LOWER_IS_BETTER,
            "aggregation": MetricAggregation.SCALAR
        },
        {
            "metric_id": "sentiment_quality_gap_by_demographic",
            "category": MetricCategory.FAIRNESS,
            "unit": MetricUnit.NORMALIZED,
            "direction": MetricDirection.LOWER_IS_BETTER,
            "aggregation": MetricAggregation.PER_GROUP
        },
        {
            "metric_id": "refusal_inconsistency_gap",
            "category": MetricCategory.FAIRNESS,
            "unit": MetricUnit.NORMALIZED,
            "direction": MetricDirection.LOWER_IS_BETTER,
            "aggregation": MetricAggregation.PER_GROUP
        }
    ],
    # Security Metrics
    "security": [
        {
            "metric_id": "jailbreak_success_rate",
            "category": MetricCategory.SECURITY,
            "unit": MetricUnit.PERCENTAGE,
            "direction": MetricDirection.LOWER_IS_BETTER,
            "aggregation": MetricAggregation.SCALAR
        },
        {
            "metric_id": "prompt_injection_success_rate",
            "category": MetricCategory.SECURITY,
            "unit": MetricUnit.PERCENTAGE,
            "direction": MetricDirection.LOWER_IS_BETTER,
            "aggregation": MetricAggregation.SCALAR
        }
    ]
}

# Hybrid Model Catalog (union of both)
HYBRID_MODEL_CATALOG = {
    **DATASET_MODEL_CATALOG,
    **API_MODEL_CATALOG,
    "end_to_end": [
        {
            "metric_id": "end_to_end_leakage",
            "category": MetricCategory.PRIVACY,
            "unit": MetricUnit.BOOLEAN,
            "direction": MetricDirection.BOOLEAN_TRUE_IS_BAD,
            "aggregation": MetricAggregation.SCALAR
        },
        {
            "metric_id": "protected_attribute_exposure",
            "category": MetricCategory.PRIVACY,
            "unit": MetricUnit.BOOLEAN,
            "direction": MetricDirection.BOOLEAN_TRUE_IS_BAD,
            "aggregation": MetricAggregation.SCALAR
        },
        {
            "metric_id": "demographic_inconsistency",
            "category": MetricCategory.FAIRNESS,
            "unit": MetricUnit.BOOLEAN,
            "direction": MetricDirection.BOOLEAN_TRUE_IS_BAD,
            "aggregation": MetricAggregation.SCALAR
        }
    ]
}


def get_catalog_for_mode(mode: AuditMode) -> Dict:
    """Get metric catalog for a specific audit mode."""
    if mode == AuditMode.DATASET_MODEL:
        return DATASET_MODEL_CATALOG
    elif mode == AuditMode.API_MODEL:
        return API_MODEL_CATALOG
    elif mode == AuditMode.HYBRID:
        return HYBRID_MODEL_CATALOG
    return {}
