"""Baseline and classical models for cognitive-level classification.

Three models in increasing complexity: majority class, keyword heuristic,
and TF-IDF + logistic regression.
"""

from collections import Counter

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import GridSearchCV
from sklearn.pipeline import Pipeline


class MajorityClassModel:
    """Predicts the single most frequent class seen during fit."""

    def __init__(self):
        self.majority_label = None

    def fit(self, questions, labels):
        """Record the most common label."""
        self.majority_label = Counter(labels).most_common(1)[0][0]
        return self

    def predict(self, questions):
        """Return the majority label for every input."""
        return [self.majority_label] * len(questions)


class KeywordHeuristicModel:
    """Rule-based classifier using keyword cues."""

    CRITICAL_CUES = [
        "limitation", "trade-off", "tradeoff", "justified", "better than",
        "the best", "critique", "evaluate", "design a", "propose", "devise",
        "under what condition", "to what extent", "is it worth", "flaw",
        "weakness", "redesign", "invent", "how sound", "drawback", "should we",
    ]
    MECHANISTIC_CUES = [
        "how would you", "how do", "how does", "why does", "why is",
        "what causes", "calculate", "implement", "apply", "given ",
        "relationship between", "distinguish", "compare", "what process",
    ]
    SURFACE_CUES = [
        "what is", "define", "list", "name the", "state the", "identify",
        "when did", "where does", "who ", "what are the", "what does",
    ]

    def fit(self, questions, labels):
        """No training needed; present for interface symmetry."""
        return self

    def _classify_one(self, question):
        """Apply the cue rules to a single question."""
        text = question.lower()
        if any(cue in text for cue in self.CRITICAL_CUES):
            return "Critical"
        if any(cue in text for cue in self.MECHANISTIC_CUES):
            return "Mechanistic"
        if any(cue in text for cue in self.SURFACE_CUES):
            return "Surface"
        # fallback
        if "why" in text or "how" in text:
            return "Mechanistic"
        return "Surface"

    def predict(self, questions):
        """Classify each question by rules."""
        return [self._classify_one(q) for q in questions]


class TfidfLogRegModel:
    """TF-IDF features into logistic regression, tuned by cross-validation."""

    def __init__(self, training_config):
        """Store training config."""
        self.cfg = training_config
        self.search = None

    def fit(self, questions, labels):
        """Train with grid search over C and penalty."""
        pipeline = Pipeline([
            ("tfidf", TfidfVectorizer(
                ngram_range=(1, self.cfg.tfidf_ngram_max),
                max_features=self.cfg.tfidf_max_features,
                sublinear_tf=True,
            )),
            ("clf", LogisticRegression(
                solver="saga",  # supports l1 and l2
                max_iter=5000,
            )),
        ])
        param_grid = {
            "clf__C": self.cfg.logreg_C_grid,
            "clf__penalty": self.cfg.logreg_penalty_grid,
        }
        self.search = GridSearchCV(
            pipeline,
            param_grid,
            scoring="f1_macro",
            cv=self.cfg.cv_folds,
            n_jobs=-1,
        )
        self.search.fit(questions, labels)
        return self

    @property
    def best_params(self):
        """Best hyperparameters from grid search."""
        return self.search.best_params_

    @property
    def best_cv_score(self):
        """Best CV macro-F1 from grid search."""
        return self.search.best_score_

    def predict(self, questions):
        """Predict using best estimator."""
        return list(self.search.predict(questions))