"""Risk Model training pipeline.

Trains a RandomForestClassifier with class_weight='balanced' for risk prediction.
Target: risk_label (binary: 0=not at-risk, 1=at-risk).
Features: 19 numeric features covering mastery, performance, and engagement.
Primary metric: recall on positive class (at-risk).
"""

import logging
from datetime import datetime, timezone

import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import (
    classification_report,
    confusion_matrix,
    f1_score,
    precision_score,
    recall_score,
    roc_auc_score,
)

from app.core.config import settings
from app.core.exceptions import TrainingError
from training.base_trainer import BaseTrainer, TrainingResult

logger = logging.getLogger(__name__)

FEATURE_COLUMNS = [
    "avg_mastery_score",
    "weak_lo_count",
    "developing_lo_count",
    "mastered_lo_count",
    "avg_confidence",
    "avg_accuracy",
    "avg_marks_ratio",
    "avg_time_seconds",
    "hint_usage_rate",
    "total_attempts",
    "attendance_percentage",
    "assignment_completion_rate",
    "average_login_per_week",
    "inactive_days_last_14",
    "avg_active_minutes",
    "total_logins",
    "avg_video_watch_ratio",
    "total_content_completed",
    "total_quiz_attempts",
]

TARGET_COLUMN = "risk_label"

RISK_LABEL_MAP = {0: "not_at_risk", 1: "at_risk"}


class RiskModelTrainer(BaseTrainer):
    """RandomForestClassifier with class_weight='balanced' for risk prediction.

    Target: risk_label (binary: 0=not at-risk, 1=at-risk)
    Features: avg_mastery_score, weak_lo_count, developing_lo_count, mastered_lo_count,
              avg_confidence, avg_accuracy, avg_marks_ratio, avg_time_seconds,
              hint_usage_rate, total_attempts, attendance_percentage,
              assignment_completion_rate, average_login_per_week, inactive_days_last_14,
              avg_active_minutes, total_logins, avg_video_watch_ratio,
              total_content_completed, total_quiz_attempts
    Primary metric: recall on positive class (at-risk)
    """

    @property
    def model_name(self) -> str:
        return "risk_model"

    @property
    def model_version(self) -> str:
        return "risk_model_v2_baseline_001"

    @property
    def table_name(self) -> str:
        return "training_risk_prediction"

    def train(self, train_df: pd.DataFrame, val_df: pd.DataFrame) -> dict:
        """Train RandomForestClassifier on numeric risk features.

        Algorithm:
        1. Extract feature columns (all numeric, no encoding needed)
        2. Target: risk_label (binary 0/1)
        3. Fit RandomForestClassifier(n_estimators=100, class_weight="balanced",
           random_state=seed)
        4. Return {"model": rf, "feature_columns.json": FEATURE_COLUMNS}
        """
        X_train = train_df[FEATURE_COLUMNS].values
        y_train = train_df[TARGET_COLUMN].values.astype(int)

        rf = RandomForestClassifier(
            n_estimators=100,
            class_weight="balanced",
            random_state=self._seed,
        )
        rf.fit(X_train, y_train)

        logger.info(
            "Risk model trained — %d samples, %d features, positive class ratio: %.2f%%",
            X_train.shape[0],
            X_train.shape[1],
            100.0 * y_train.sum() / len(y_train),
        )

        return {"model": rf, "feature_columns.json": FEATURE_COLUMNS}

    def evaluate(self, artifacts: dict, df: pd.DataFrame, split_name: str) -> dict:
        """Evaluate model on a split.

        Computes: recall (positive class), precision, F1, ROC-AUC via
        predict_proba[:, 1], per-class metrics, confusion matrix.
        Also computes recall for high/critical risk_level classes.
        """
        model = artifacts["model"]

        X = df[FEATURE_COLUMNS].values
        y_true = df[TARGET_COLUMN].values.astype(int)
        y_pred = model.predict(X)
        y_proba = model.predict_proba(X)[:, 1]

        # Binary metrics (positive class = 1 = at-risk)
        recall_pos = recall_score(y_true, y_pred, pos_label=1, zero_division=0)
        precision_pos = precision_score(y_true, y_pred, pos_label=1, zero_division=0)
        f1_pos = f1_score(y_true, y_pred, pos_label=1, zero_division=0)

        # ROC-AUC via predict_proba
        try:
            roc_auc = roc_auc_score(y_true, y_proba)
        except ValueError:
            # Only one class present in y_true
            roc_auc = 0.0

        # Per-class metrics
        target_names = [RISK_LABEL_MAP[i] for i in sorted(RISK_LABEL_MAP.keys())]
        report = classification_report(
            y_true,
            y_pred,
            labels=sorted(RISK_LABEL_MAP.keys()),
            target_names=target_names,
            output_dict=True,
            zero_division=0,
        )

        per_class = {}
        for label_int, label_name in RISK_LABEL_MAP.items():
            if label_name in report:
                per_class[label_name] = {
                    "precision": round(report[label_name]["precision"], 4),
                    "recall": round(report[label_name]["recall"], 4),
                    "f1": round(report[label_name]["f1-score"], 4),
                    "support": int(report[label_name]["support"]),
                }

        # Confusion matrix
        cm = confusion_matrix(
            y_true, y_pred, labels=sorted(RISK_LABEL_MAP.keys())
        ).tolist()

        # Recall for high/critical risk_level classes (from same table)
        risk_level_recall = self._compute_risk_level_recall(df, y_pred)

        metrics = {
            "recall_positive": round(recall_pos, 4),
            "precision_positive": round(precision_pos, 4),
            "f1_positive": round(f1_pos, 4),
            "roc_auc": round(roc_auc, 4),
            "per_class": per_class,
            "confusion_matrix": cm,
            "risk_level_recall": risk_level_recall,
        }

        logger.info(
            "%s metrics — recall: %.4f, precision: %.4f, F1: %.4f, ROC-AUC: %.4f",
            split_name, recall_pos, precision_pos, f1_pos, roc_auc,
        )

        return metrics

    def _compute_risk_level_recall(
        self, df: pd.DataFrame, y_pred: np.ndarray
    ) -> dict:
        """Compute recall for high/critical risk_level classes separately.

        These are subsets of the positive class (risk_label=1) where
        risk_level is 'high' or 'critical'. We check how many of those
        the model correctly predicted as positive (risk_label=1).
        """
        risk_level_recall = {}

        if "risk_level" not in df.columns:
            return risk_level_recall

        for level in ["high", "critical"]:
            mask = df["risk_level"].values == level
            if mask.sum() == 0:
                risk_level_recall[level] = {"recall": 0.0, "support": 0}
                continue

            y_true_subset = df[TARGET_COLUMN].values[mask].astype(int)
            y_pred_subset = y_pred[mask]

            # For high/critical, the true label should be 1 (at-risk)
            # Recall = how many of these did we correctly predict as 1
            recall = recall_score(
                y_true_subset, y_pred_subset, pos_label=1, zero_division=0
            )
            risk_level_recall[level] = {
                "recall": round(recall, 4),
                "support": int(mask.sum()),
            }

        return risk_level_recall

    def _check_baseline(self, metrics: dict) -> None:
        """Verify recall on positive class > 0.5 (lenient baseline).

        Raises TrainingError if not met.
        """
        test_metrics = metrics.get("metrics", {}).get("test", {})
        recall_pos = test_metrics.get("recall_positive")

        # Fallback to validation metrics if test not available
        if recall_pos is None:
            val_metrics = metrics.get("metrics", {}).get("validation", {})
            recall_pos = val_metrics.get("recall_positive")

        if recall_pos is None:
            raise TrainingError(
                "Cannot compute baseline: recall_positive not found in metrics.",
                model_name=self.model_name,
            )

        if recall_pos <= 0.5:
            raise TrainingError(
                f"Recall on positive class ({recall_pos:.4f}) does not exceed "
                f"baseline (0.5). Model fails to identify at-risk students adequately.",
                model_name=self.model_name,
            )

        logger.info("Baseline check passed — recall_positive %.4f > 0.5", recall_pos)

    def _build_metrics(
        self,
        val_metrics: dict,
        test_metrics: dict,
        train_df: pd.DataFrame,
        val_df: pd.DataFrame,
        test_df: pd.DataFrame,
    ) -> dict:
        """Assemble full metrics.json content."""
        return {
            "model_name": self.model_name,
            "model_version": self.model_version,
            "dataset_version": settings.ai_service_version,
            "trained_at": datetime.now(timezone.utc).isoformat(),
            "seed": self._seed,
            "split_counts": {
                "train": len(train_df),
                "validation": len(val_df),
                "test": len(test_df),
            },
            "metrics": {
                "validation": val_metrics,
                "test": test_metrics,
            },
            "limitations": [
                "Trained on synthetic data only.",
                "Binary risk_label derived from synthetic risk_score thresholds.",
                "All features are numeric; no text or contextual features used.",
                "Class imbalance (~16% positive) addressed via class_weight='balanced'.",
                "Critical class (~2%) recall should be monitored separately.",
            ],
        }

    def _build_training_config(
        self,
        train_df: pd.DataFrame,
        val_df: pd.DataFrame,
        test_df: pd.DataFrame,
    ) -> dict:
        """Build training_config.json with hyperparameters."""
        return {
            "model_name": self.model_name,
            "model_version": self.model_version,
            "dataset_version": settings.ai_service_version,
            "seed": self._seed,
            "split_counts": {
                "train": len(train_df),
                "validation": len(val_df),
                "test": len(test_df),
            },
            "hyperparameters": {
                "n_estimators": 100,
                "class_weight": "balanced",
                "random_state": self._seed,
                "algorithm": "RandomForestClassifier",
            },
            "feature_columns": FEATURE_COLUMNS,
            "target_column": TARGET_COLUMN,
            "label_map": RISK_LABEL_MAP,
            "algorithm": "RandomForestClassifier",
        }

    def _build_model_card(self, metrics: dict) -> str:
        """Generate model_card.md content."""
        val_metrics = metrics.get("metrics", {}).get("validation", {})
        test_metrics = metrics.get("metrics", {}).get("test", {})

        card = f"""# Model Card: Risk Model

## Model Details

- **Model Name:** {self.model_name}
- **Model Version:** {self.model_version}
- **Algorithm:** RandomForestClassifier (class_weight="balanced")
- **Framework:** scikit-learn
- **Trained At:** {metrics.get("trained_at", "N/A")}
- **Seed:** {self._seed}

## Intended Use

Predict whether a student is at-risk (binary: 0=not at-risk, 1=at-risk) based on
mastery, performance, and engagement features. Used in the risk prediction endpoint
to identify students who may need intervention. Primary optimization target is
recall on the positive class to minimize missed at-risk students.

## Training Data

- **Source:** training_risk_prediction.csv (synthetic dataset v2)
- **Split Counts:** train={metrics.get("split_counts", {}).get("train", "N/A")}, \
validation={metrics.get("split_counts", {}).get("validation", "N/A")}, \
test={metrics.get("split_counts", {}).get("test", "N/A")}
- **Features:** {", ".join(FEATURE_COLUMNS)} (all numeric, 19 features)
- **Target:** risk_label (binary 0/1)
- **Class Imbalance:** ~16% positive class, addressed via class_weight="balanced"

## Metrics

### Validation Set
- Recall (positive): {val_metrics.get("recall_positive", "N/A")}
- Precision (positive): {val_metrics.get("precision_positive", "N/A")}
- F1 (positive): {val_metrics.get("f1_positive", "N/A")}
- ROC-AUC: {val_metrics.get("roc_auc", "N/A")}

### Test Set
- Recall (positive): {test_metrics.get("recall_positive", "N/A")}
- Precision (positive): {test_metrics.get("precision_positive", "N/A")}
- F1 (positive): {test_metrics.get("f1_positive", "N/A")}
- ROC-AUC: {test_metrics.get("roc_auc", "N/A")}

## Per-Class Performance (Test Set)

| Class | Precision | Recall | F1 | Support |
|-------|-----------|--------|-----|---------|
"""
        test_per_class = test_metrics.get("per_class", {})
        for label_name in ["not_at_risk", "at_risk"]:
            cls_metrics = test_per_class.get(label_name, {})
            card += (
                f"| {label_name} | "
                f"{cls_metrics.get('precision', 'N/A')} | "
                f"{cls_metrics.get('recall', 'N/A')} | "
                f"{cls_metrics.get('f1', 'N/A')} | "
                f"{cls_metrics.get('support', 'N/A')} |\n"
            )

        # Risk level recall section
        test_risk_level = test_metrics.get("risk_level_recall", {})
        card += """
## Risk Level Recall (Test Set)

| Risk Level | Recall | Support |
|------------|--------|---------|
"""
        for level in ["high", "critical"]:
            level_metrics = test_risk_level.get(level, {})
            card += (
                f"| {level} | "
                f"{level_metrics.get('recall', 'N/A')} | "
                f"{level_metrics.get('support', 'N/A')} |\n"
            )

        card += f"""
## Known Limitations

- Trained on synthetic data only — performance on real student data is unknown.
- Binary risk_label derived from synthetic risk_score thresholds.
- All features are numeric; no text or contextual features used.
- Class imbalance (~16% positive) addressed via class_weight="balanced".
- Critical class (~2%) is very rare; recall on critical should be monitored.
- No temporal features (trend over time) included in this baseline.

## Fallback Behavior

When the model is not loaded or confidence is below the threshold (0.55),
the system falls back to rule-based risk estimation using:
- inactive_days_last_14 > 7 → high risk
- attendance_percentage < 60% → high risk
- avg_mastery_score < 0.4 → medium risk
- Otherwise → low risk
"""
        return card