import json
from pathlib import Path

import joblib
import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import classification_report, confusion_matrix
from sklearn.model_selection import StratifiedKFold, cross_val_predict
from sklearn.pipeline import Pipeline

DATA_PATH = Path("data/train.jsonl")
MODEL_PATH = Path("model.joblib")


def load_data(path: Path) -> pd.DataFrame:
    records = []

    with path.open("r", encoding="utf-8") as file:
        for line in file:
            if line.strip():
                records.append(json.loads(line))

    return pd.DataFrame(records)


def build_pipeline() -> Pipeline:
    return Pipeline(
        [
            (
                "tfidf",
                TfidfVectorizer(
                    ngram_range=(1, 2),
                    lowercase=True,
                    min_df=2,
                    max_df=0.95,
                    sublinear_tf=True,
                ),
            ),
            (
                "classifier",
                LogisticRegression(
                    max_iter=2000,
                    class_weight="balanced",
                    random_state=42,
                ),
            ),
        ]
    )


def main():
    df = load_data(DATA_PATH)

    X = (
        "PROMPT: "
        + df["prompt"].astype(str)
        + "\nRESPONSE: "
        + df["response"].astype(str)
    )

    y = df["quality_label"].astype(str)

    print(f"Examples: {len(df)}")
    print(y.value_counts().sort_index())
    print()

    cv = StratifiedKFold(
        n_splits=5,
        shuffle=True,
        random_state=42,
    )

    pipeline = build_pipeline()

    predictions = cross_val_predict(
        pipeline,
        X,
        y,
        cv=cv,
    )

    print("Stratified 5-Fold Cross-Validation")
    print("=" * 44)

    print(
        classification_report(
            y,
            predictions,
            digits=3,
            zero_division=0,
        )
    )

    labels = sorted(y.unique())

    matrix = confusion_matrix(
        y,
        predictions,
        labels=labels,
    )

    print("Confusion matrix")
    print(f"Labels: {labels}")
    print(matrix)
    print()

    pipeline.fit(X, y)

    joblib.dump(
        pipeline,
        MODEL_PATH,
    )

    print(f"Saved model to: {MODEL_PATH.resolve()}")


if __name__ == "__main__":
    main()