"""Mastery Model training pipeline.

Trains a RandomForestClassifier for mastery label prediction.
Target: mastery_label (4 classes: 0=weak, 1=developing, 2=proficient, 3=mastered).
Features: attempt_count, accuracy, average_marks_ratio, average_time_seconds,
          hint_usage_rate, attendance_percentage, assignment_completion_rate,
          average_login_per_week, inactive_days_last_14.
Primary metric: macro F1.
"""

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

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 = [
    "attempt_count",
    "accuracy",
    "average_marks_ratio",
    "average_time_seconds",
    "hint_usage_rate",
    "attendance_percentage",
    "assignment_completion_rate",
    "average_login_per_week",
    "inactive_days_last_14",
]

TARGET_COLUMN = "mastery_label"

MASTERY_LABEL_MAP = {0: "weak", 1: "developing", 2: "proficient", 3: "mastered"}


class MasteryModelTrainer(BaseTrainer):
    """RandomForestClassifier for mastery label prediction.

    Target: mastery_label (4 classes: 0=weak, 1=developing, 2=proficient, 3=mastered)
    Features: attempt_count, accuracy, average_marks_ratio, average_time_seconds,
              hint_usage_rate, attendance_percentage, assignment_completion_rate,
              average_login_per_week, inactive_days_last_14
    Primary metric: macro F1
    """

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

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

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

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

        Algorithm:
        1. Extract feature columns (all numeric, no encoding needed)
        2. Target: mastery_label (integer 0-3)
        3. Fit RandomForestClassifier(n_estimators=100, 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,
            random_state=self._seed,
        )
        rf.fit(X_train, y_train)

        logger.info(
            "Mastery model trained — %d samples, %d features, %d classes",
            X_train.shape[0],
            X_train.shape[1],
            len(np.unique(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: macro F1, weighted F1, per-class precision/recall/f1,
        confusion matrix.
        """
        model = artifacts["model"]

        X = df[FEATURE_COLUMNS].values
        y_true = df[TARGET_COLUMN].values.astype(int)
        y_pred = model.predict(X)

        # F1 scores
        macro_f1 = f1_score(y_true, y_pred, average="macro", zero_division=0)
        weighted_f1 = f1_score(y_true, y_pred, average="weighted", zero_division=0)

        # Per-class metrics using label map for readable names
        target_names = [
            MASTERY_LABEL_MAP[i] for i in sorted(MASTERY_LABEL_MAP.keys())
        ]
        report = classification_report(
            y_true,
            y_pred,
            labels=sorted(MASTERY_LABEL_MAP.keys()),
            target_names=target_names,
            output_dict=True,
            zero_division=0,
        )

        per_class = {}
        for label_int, label_name in MASTERY_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(MASTERY_LABEL_MAP.keys())
        ).tolist()

        metrics = {
            "macro_f1": round(macro_f1, 4),
            "weighted_f1": round(weighted_f1, 4),
            "per_class": per_class,
            "confusion_matrix": cm,
        }

        logger.info(
            "%s metrics — macro_f1: %.4f, weighted_f1: %.4f",
            split_name, macro_f1, weighted_f1,
        )

        return metrics

    def _check_baseline(self, metrics: dict) -> None:
        """Verify macro F1 > 0.25 (random baseline for 4 classes).

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

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

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

        if macro_f1 <= 0.25:
            raise TrainingError(
                f"Macro F1 ({macro_f1:.4f}) does not exceed random baseline (0.25). "
                f"Model is not better than random for 4-class classification.",
                model_name=self.model_name,
            )

        logger.info("Baseline check passed — macro F1 %.4f > 0.25", macro_f1)

    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.",
                "4-class mastery labels derived from synthetic mastery_score thresholds.",
                "All features are numeric; no text or contextual features used.",
                "Class distribution may not reflect real-world mastery patterns.",
            ],
        }

    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,
                "random_state": self._seed,
                "algorithm": "RandomForestClassifier",
            },
            "feature_columns": FEATURE_COLUMNS,
            "target_column": TARGET_COLUMN,
            "label_map": MASTERY_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: Mastery Model

## Model Details

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

## Intended Use

Predict per-student per-LO mastery label (weak, developing, proficient, mastered)
based on behavioral and performance features. Used in the mastery prediction
endpoint to classify student mastery level for a given learning outcome.

## Training Data

- **Source:** training_mastery_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)
- **Target:** mastery_label (integer 0-3, mapped to weak/developing/proficient/mastered)

## Metrics

### Validation Set
- Macro F1: {val_metrics.get("macro_f1", "N/A")}
- Weighted F1: {val_metrics.get("weighted_f1", "N/A")}

### Test Set
- Macro F1: {test_metrics.get("macro_f1", "N/A")}
- Weighted F1: {test_metrics.get("weighted_f1", "N/A")}

## Per-Class Performance (Test Set)

| Class | Precision | Recall | F1 | Support |
|-------|-----------|--------|-----|---------|
"""
        test_per_class = test_metrics.get("per_class", {})
        for label_name in ["weak", "developing", "proficient", "mastered"]:
            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"
            )

        card += f"""
## Known Limitations

- Trained on synthetic data only — performance on real student data is unknown.
- 4-class mastery labels derived from synthetic mastery_score thresholds.
- All features are numeric; no text or contextual features used.
- Class distribution may not reflect real-world mastery patterns.
- No encoding needed since all features are already numeric.

## Fallback Behavior

When the model is not loaded or confidence is below the threshold (0.55),
the system falls back to rule-based mastery estimation using mastery_score
thresholds: <0.4 weak, <0.6 developing, <0.8 proficient, else mastered.
"""
        return card