#!/usr/bin/env python
"""
Stage 12: Classical (non-LLM) baseline for comparison.

"Do not assume an LLM is automatically superior" -- this fits a TF-IDF +
Logistic Regression multi-label classifier on the REAL training split
(356 contracts) and evaluates on the REAL test split (77 contracts) using
the exact same (contract, category) presence protocol and evaluator as the
local-LLM baseline (src/evaluator.py::BaselineEvaluator), so the two
numbers are directly comparable.

Usage:
  python scripts/run_classical_baseline.py

No API, no local model server, no GPU required -- this is a CPU-only
scikit-learn pipeline, useful as a sanity floor: if the LLM baseline can't
beat this, that's a real, reportable finding, not something to hide.
"""

import os
import sys
import json
import logging
import subprocess
from pathlib import Path
from datetime import datetime

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.multiclass import OneVsRestClassifier
from sklearn.preprocessing import MultiLabelBinarizer

from src.schemas import CUAD_CATEGORIES
from src.evaluator import BaselineEvaluator

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

RANDOM_SEED = 42  # Same seed used throughout the project (data split, etc.)


def _git_commit_hash() -> str:
    try:
        return subprocess.check_output(
            ["git", "rev-parse", "--short", "HEAD"],
            cwd=Path(__file__).resolve().parent.parent,
            stderr=subprocess.DEVNULL,
        ).decode().strip()
    except Exception:
        return "unknown"


def load_split(name: str):
    path = Path(f"data/processed/cuad_{name}_contracts.json")
    if not path.exists():
        raise FileNotFoundError(f"{path} not found -- run scripts/prepare_datasets.py first")
    with open(path, encoding="utf-8") as f:
        return json.load(f)


def contract_labels(contract: dict) -> list:
    """The set of categories actually present in this contract (gold),
    matching the (contract, category) presence semantics used everywhere
    else in this project (ADR-018)."""
    return sorted({c["category"] for c in contract.get("clauses", [])})


def main():
    logger.info("=" * 70)
    logger.info("CLASSICAL BASELINE: TF-IDF + One-vs-Rest Logistic Regression")
    logger.info("=" * 70)

    train_contracts = load_split("train")
    test_contracts = load_split("test")
    logger.info(f"Train: {len(train_contracts)} contracts | Test: {len(test_contracts)} contracts")

    X_train_text = [c["text"] for c in train_contracts]
    X_test_text = [c["text"] for c in test_contracts]
    y_train_labels = [contract_labels(c) for c in train_contracts]
    y_test_labels = [contract_labels(c) for c in test_contracts]  # not used for fitting -- gold for eval only

    # Binarize labels over the FULL 41-category space (not just categories
    # seen in train) so a category absent from train but present in test
    # is handled correctly as "never predicted" rather than silently
    # dropped from the label space.
    mlb = MultiLabelBinarizer(classes=CUAD_CATEGORIES)
    Y_train = mlb.fit_transform(y_train_labels)
    logger.info(f"Label matrix: {Y_train.shape} (contracts x categories)")

    # TF-IDF fit ONLY on train text -- never on test, to avoid any leakage
    # of test vocabulary/IDF statistics into the model.
    vectorizer = TfidfVectorizer(
        max_features=20000,
        ngram_range=(1, 2),
        min_df=2,
        stop_words="english",
    )
    X_train = vectorizer.fit_transform(X_train_text)
    X_test = vectorizer.transform(X_test_text)
    logger.info(f"TF-IDF matrix: train {X_train.shape}, test {X_test.shape}")

    # class_weight='balanced' matters heavily here: CUAD's real category
    # frequencies range from 27 (Price Restrictions... actually 0 in test,
    # but present in train) to 2550 (Parties) instances -- see
    # data/DATA_SOURCES.md. Without balancing, rare categories would never
    # be predicted at all.
    clf = OneVsRestClassifier(
        LogisticRegression(
            max_iter=2000,
            class_weight="balanced",
            random_state=RANDOM_SEED,
        )
    )
    logger.info("Fitting classifier on train split only...")
    clf.fit(X_train, Y_train)

    logger.info("Predicting on test split...")
    Y_pred_proba = clf.predict_proba(X_test)
    Y_pred = clf.predict(X_test)

    # Convert to the same prediction-dict shape the evaluator expects,
    # tagged with contract_id exactly like the LLM pipeline's output, so
    # the identical BaselineEvaluator.compute_f1() call scores both fairly.
    predictions = []
    for i, contract in enumerate(test_contracts):
        for j, category in enumerate(CUAD_CATEGORIES):
            if Y_pred[i, j] == 1:
                predictions.append({
                    "contract_id": contract["id"],
                    "category": category,
                    "confidence": float(Y_pred_proba[i, j]),
                })

    gold = [
        {"contract_id": c["id"], "category": cat}
        for c in test_contracts
        for cat in contract_labels(c)
    ]

    evaluator = BaselineEvaluator()
    metrics = evaluator.compute_f1(predictions, gold)
    calibration = evaluator.compute_confidence_calibration(predictions, gold)
    evaluator.log_results(metrics)

    logger.info(f"\nClassical baseline Macro F1: {metrics['macro_f1']:.4f}")
    logger.info(f"Classical baseline Micro F1: {metrics['micro_f1']:.4f}")

    output = {
        "method": "TF-IDF + One-vs-Rest Logistic Regression",
        "sklearn_version": __import__("sklearn").__version__,
        "random_seed": RANDOM_SEED,
        "git_commit": _git_commit_hash(),
        "timestamp": datetime.utcnow().isoformat(),
        "train_contracts": len(train_contracts),
        "test_contracts": len(test_contracts),
        "tfidf_params": {"max_features": 20000, "ngram_range": [1, 2], "min_df": 2},
        "classifier_params": {"class_weight": "balanced", "max_iter": 2000},
        "metrics": metrics,
        "calibration": calibration,
    }

    out_path = Path("data/processed/classical_baseline_results.json")
    with open(out_path, "w", encoding="utf-8") as f:
        json.dump(output, f, indent=2)
    logger.info(f"Results saved to {out_path}")

    return output


if __name__ == "__main__":
    main()