"""Answer Scorer training pipeline.

Trains a TF-IDF + Ridge regression model for subjective answer scoring.
Target: teacher_marks (continuous). Features: TF-IDF of student_answer +
rubric_match_score + concept_coverage_score. Primary metric: MAE.
Predictions are clipped to [0, max_marks] range.
"""

import logging
from datetime import datetime, timezone

import numpy as np
import pandas as pd
import scipy.sparse
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import Ridge
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score

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

logger = logging.getLogger(__name__)


class AnswerScorerTrainer(BaseTrainer):
    """TF-IDF + Ridge regression for subjective answer scoring.

    Target: teacher_marks (continuous)
    Features: TF-IDF of student_answer + rubric_match_score + concept_coverage_score
    Primary metric: MAE
    """

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

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

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

    def train(self, train_df: pd.DataFrame, val_df: pd.DataFrame) -> dict:
        """Train TF-IDF + Ridge regression.

        Algorithm:
        1. Fit TF-IDF vectorizer on train student_answer text
        2. Build feature matrix: hstack([tfidf_features, rubric_match_score, concept_coverage_score])
        3. Target: teacher_marks
        4. Fit Ridge(alpha=1.0, random_state=seed)
        5. Return {"model": ridge, "vectorizer": tfidf}
        """
        tfidf = TfidfVectorizer(
            max_features=5000,
            ngram_range=(1, 2),
        )
        tfidf_features = tfidf.fit_transform(train_df["student_answer"])

        numeric_features = scipy.sparse.csr_matrix(
            train_df[["rubric_match_score", "concept_coverage_score"]].values
        )
        X_train = scipy.sparse.hstack([tfidf_features, numeric_features])

        y_train = train_df["teacher_marks"].values

        ridge = Ridge(alpha=1.0, random_state=self._seed)
        ridge.fit(X_train, y_train)

        logger.info(
            "Answer Scorer trained — %d TF-IDF features + 2 numeric features, %d samples",
            tfidf_features.shape[1],
            X_train.shape[0],
        )

        return {"model": ridge, "vectorizer": tfidf}

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

        Computes: MAE, RMSE, R-squared, distribution comparison (mean/std of
        predicted vs actual), percentage of predictions where
        |predicted - actual| > 1.0 (teacher review threshold).
        Predictions are clipped to [0, max_marks] range.
        """
        model = artifacts["model"]
        tfidf = artifacts["vectorizer"]

        tfidf_features = tfidf.transform(df["student_answer"])
        numeric_features = scipy.sparse.csr_matrix(
            df[["rubric_match_score", "concept_coverage_score"]].values
        )
        X = scipy.sparse.hstack([tfidf_features, numeric_features])

        y_true = df["teacher_marks"].values
        max_marks = df["max_marks"].values

        # Predict and clip to [0, max_marks] range
        y_pred_raw = model.predict(X)
        y_pred = np.clip(y_pred_raw, 0, max_marks)

        # Core metrics
        mae = mean_absolute_error(y_true, y_pred)
        rmse = float(np.sqrt(mean_squared_error(y_true, y_pred)))
        r_squared = r2_score(y_true, y_pred)

        # Distribution comparison
        pred_mean = float(np.mean(y_pred))
        pred_std = float(np.std(y_pred))
        actual_mean = float(np.mean(y_true))
        actual_std = float(np.std(y_true))

        # Teacher review threshold: percentage where |predicted - actual| > 1.0
        abs_errors = np.abs(y_pred - y_true)
        pct_above_threshold = float(np.mean(abs_errors > 1.0) * 100)

        metrics = {
            "mae": round(mae, 4),
            "rmse": round(rmse, 4),
            "r_squared": round(r_squared, 4),
            "distribution": {
                "predicted_mean": round(pred_mean, 4),
                "predicted_std": round(pred_std, 4),
                "actual_mean": round(actual_mean, 4),
                "actual_std": round(actual_std, 4),
            },
            "pct_above_review_threshold": round(pct_above_threshold, 2),
        }

        logger.info(
            "%s metrics — MAE: %.4f, RMSE: %.4f, R²: %.4f, %%>1.0: %.2f%%",
            split_name, mae, rmse, r_squared, pct_above_threshold,
        )

        return metrics

    def _check_baseline(self, metrics: dict) -> None:
        """Verify MAE < max_marks/2 (very lenient baseline for synthetic data).

        Uses the test metrics MAE. The max_marks average is estimated from the
        dataset; for a lenient check we use a generous threshold.
        """
        test_metrics = metrics.get("metrics", {}).get("test", {})
        mae = test_metrics.get("mae", float("inf"))

        # Fallback to validation if test not available
        if mae == float("inf"):
            val_metrics = metrics.get("metrics", {}).get("validation", {})
            mae = val_metrics.get("mae", float("inf"))

        # Use max_marks / 2 as baseline. Since max_marks varies per question,
        # we use a conservative estimate. Typical max_marks in the dataset is 5.
        # A lenient baseline: MAE < 2.5 (half of typical max_marks=5)
        max_marks_half = 2.5
        if mae >= max_marks_half:
            raise TrainingError(
                f"MAE ({mae:.4f}) does not meet baseline threshold "
                f"({max_marks_half:.1f} = max_marks/2). "
                f"Model is not better than naive prediction.",
                model_name=self.model_name,
            )

        logger.info(
            "Baseline check passed — MAE %.4f < baseline %.1f",
            mae, max_marks_half,
        )

    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.",
                "TF-IDF features do not capture deep semantic similarity.",
                "Predictions clipped to [0, max_marks] range.",
                "teacher_review_required is always True in V2 baseline.",
                "MAE is the primary metric; individual predictions may deviate significantly.",
            ],
        }

    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": {
                "tfidf_max_features": 5000,
                "ngram_range": [1, 2],
                "ridge_alpha": 1.0,
            },
            "feature_columns": [
                "student_answer (TF-IDF)",
                "rubric_match_score",
                "concept_coverage_score",
            ],
            "target_column": "teacher_marks",
            "algorithm": "TF-IDF + Ridge regression",
        }

    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: Answer Scorer

## Model Details

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

## Intended Use

Score subjective student answers against a rubric and model answer.
Produces a predicted marks value clipped to [0, max_marks]. Always sets
teacher_review_required=True in V2 baseline — predictions are advisory only.

## Training Data

- **Source:** training_answer_scoring.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:** student_answer (TF-IDF, max_features=5000, ngram_range=(1,2)) + \
rubric_match_score + concept_coverage_score
- **Target:** teacher_marks (continuous)

## Metrics

### Validation Set
- MAE: {val_metrics.get("mae", "N/A")}
- RMSE: {val_metrics.get("rmse", "N/A")}
- R-squared: {val_metrics.get("r_squared", "N/A")}
- % Above Review Threshold (>1.0): {val_metrics.get("pct_above_review_threshold", "N/A")}%

### Test Set
- MAE: {test_metrics.get("mae", "N/A")}
- RMSE: {test_metrics.get("rmse", "N/A")}
- R-squared: {test_metrics.get("r_squared", "N/A")}
- % Above Review Threshold (>1.0): {test_metrics.get("pct_above_review_threshold", "N/A")}%

## Known Limitations

- Trained on synthetic data only — performance on real student answers is unknown.
- TF-IDF features do not capture deep semantic similarity or paraphrasing.
- Predictions are clipped to [0, max_marks]; the model may predict outside this range before clipping.
- teacher_review_required is always True in V2 baseline.
- MAE is the primary metric; individual predictions may deviate significantly from teacher marks.

## Fallback Behavior

When the model is not loaded or confidence is below threshold,
the system falls back to rubric keyword coverage + length heuristic,
always setting teacher_review_required=True.
"""
        return card