import joblib
import csv

import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression, Ridge
from sklearn.metrics import (
    accuracy_score,
    classification_report,
    mean_absolute_error,
    r2_score,
)
from sklearn.model_selection import train_test_split


def load_data(path="data.csv"):
    texts, labels, scores = [], [], []
    with open(path, newline="", encoding="utf-8") as f:
        for row in csv.DictReader(f):
            texts.append(row["text"])
            labels.append(int(row["label"]))
            scores.append(float(row["spam_score"]))
    return texts, np.array(labels), np.array(scores)


def main():
    texts, labels, scores = load_data()

    X_train_txt, X_test_txt, y_train, y_test, s_train, s_test = train_test_split(
        texts, labels, scores, test_size=0.25, random_state=42, stratify=labels
    )

    vectorizer = TfidfVectorizer(lowercase=True, ngram_range=(1, 2), max_features=2000)
    X_train = vectorizer.fit_transform(X_train_txt)
    X_test = vectorizer.transform(X_test_txt)

    # Classification Model (Spam/Ham)
    clf = LogisticRegression(max_iter=1000)
    clf.fit(X_train, y_train)
    y_pred = clf.predict(X_test)

    print("Classification: Spam vs Ham")
    print(f"Accuracy: {accuracy_score(y_test, y_pred):.3f}")
    print(classification_report(y_test, y_pred, target_names=["ham", "spam"]))

    # Regression model
    reg = Ridge(alpha=1.0)
    reg.fit(X_train, s_train)
    s_pred = reg.predict(X_test)
    s_pred = np.clip(s_pred, 0, 1)

    print("Regression: Spam Score (0-1)")
    print(f"MAE: {mean_absolute_error(s_test, s_pred):.3f}")
    print(f"R^2: {r2_score(s_test, s_pred):.3f}")

    # Demo model
    print("\nTest Samples")
    samples = [
        "Congratulations! You've won a free vacation, click here now!",
        "Hey, can you send me the meeting notes from earlier?",
    ]
    X_demo = vectorizer.transform(samples)
    preds = clf.predict(X_demo)
    scores_demo = np.clip(reg.predict(X_demo), 0, 1)
    for text, label, score in zip(samples, preds, scores_demo):
        verdict = "SPAM" if label == 1 else "HAM"
        print(f"[{verdict}] score={score:.2f} :: {text}")


if __name__ == "__main__":
    # persist our model
    joblib.dump(vectorizer, "vectorizer.joblib")
    joblib.dump(clf, "classifier.joblib")
    print("\nSaved model artifacts.")
    main()