"""Phase 4a: train and evaluate the three lightweight models.

Trains majority baseline, keyword heuristic, and TF-IDF + logistic regression,
then evaluates on both in-domain and OOD test sets.

Usage:
    python scripts/train_classical.py
"""

import json
import os
import sys

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import config  # noqa: E402
from scripts import metrics  # noqa: E402
from scripts.classical_models import (  # noqa: E402
    KeywordHeuristicModel,
    MajorityClassModel,
    TfidfLogRegModel,
)


def load_split(processed_dir, name):
    """Load one split file into parallel question and label lists."""
    path = os.path.join(processed_dir, f"{name}.jsonl")
    questions, labels = [], []
    with open(path, "r", encoding="utf-8") as handle:
        for line in handle:
            line = line.strip()
            if not line:
                continue
            record = json.loads(line)
            questions.append(record["question"])
            labels.append(record["bloom_class"])
    return questions, labels


class ClassicalExperiment:
    """Train and evaluate the three lightweight models."""

    def __init__(self, training_config):
        self.cfg = training_config
        self.results = {}

    def _evaluate_on(self, model, questions, labels):
        """Predict and compute metrics for one eval set."""
        predictions = model.predict(questions)
        return metrics.compute_metrics(labels, predictions)

    def run(self):
        """Load data, train each model, and evaluate on test and OOD."""
        train_q, train_y = load_split(self.cfg.processed_dir, "train")
        test_q, test_y = load_split(self.cfg.processed_dir, "test")
        ood_q, ood_y = load_split(self.cfg.processed_dir, "ood_test")
        print(f"train={len(train_q)}  test={len(test_q)}  ood_test={len(ood_q)}")

        model_specs = [
            ("majority_baseline", MajorityClassModel()),
            ("keyword_heuristic", KeywordHeuristicModel()),
            ("tfidf_logreg", TfidfLogRegModel(self.cfg)),
        ]

        for name, model in model_specs:
            print(f"\n=== {name} ===")
            model.fit(train_q, train_y)

            if name == "tfidf_logreg":
                print(f"  best params : {model.best_params}")
                print(f"  best CV F1  : {model.best_cv_score:.3f}")

            in_domain = self._evaluate_on(model, test_q, test_y)
            ood = self._evaluate_on(model, ood_q, ood_y)

            metrics.print_metrics(f"[{name}] in-domain test", in_domain)
            metrics.print_metrics(f"[{name}] OOD test", ood)

            metrics.plot_confusion_matrix(
                in_domain,
                f"{name} (in-domain)",
                os.path.join(self.cfg.output_dir, f"cm_{name}_indomain.png"),
            )
            metrics.plot_confusion_matrix(
                ood,
                f"{name} (OOD)",
                os.path.join(self.cfg.output_dir, f"cm_{name}_ood.png"),
            )

            entry = {"in_domain": in_domain, "ood": ood}
            if name == "tfidf_logreg":
                entry["best_params"] = model.best_params
                entry["best_cv_macro_f1"] = float(model.best_cv_score)
            self.results[name] = entry

        self._save_results()
        self._print_comparison()

    def _save_results(self):
        """Save all results to a JSON file."""
        os.makedirs(self.cfg.output_dir, exist_ok=True)
        path = os.path.join(self.cfg.output_dir, "classical_results.json")
        with open(path, "w", encoding="utf-8") as handle:
            json.dump(self.results, handle, indent=2)
        print(f"\nSaved results -> {path}")

    def _print_comparison(self):
        """Print in-domain and OOD comparison tables."""
        print("\n" + "=" * 60)
        print("In-domain test")
        print(f"  {'model':20s} {'acc':>6s} {'macroF1':>8s} {'QWK':>6s}")
        for name, entry in self.results.items():
            m = entry["in_domain"]
            print(f"  {name:20s} {m['accuracy']:6.3f} "
                  f"{m['macro_f1']:8.3f} {m['qwk']:6.3f}")

        print("\nCross-domain (OOD) test")
        print(f"  {'model':20s} {'inF1':>6s} {'oodF1':>6s} {'degrade':>8s}")
        for name, entry in self.results.items():
            in_f1 = entry["in_domain"]["macro_f1"]
            ood_f1 = entry["ood"]["macro_f1"]
            print(f"  {name:20s} {in_f1:6.3f} {ood_f1:6.3f} "
                  f"{in_f1 - ood_f1:8.3f}")
        print("=" * 60)


def main():
    """Run Phase 4a."""
    experiment = ClassicalExperiment(config.TrainingConfig())
    experiment.run()


if __name__ == "__main__":
    main()