import os
import csv
import joblib
from models.narrative import NarrativeClassifier

BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

DATA_PATH = os.path.join(BASE_DIR, "data", "narrative.csv")
VEC_PATH = os.path.join(BASE_DIR, "artifacts", "vectorizer.joblib")
OUT_PATH = os.path.join(BASE_DIR, "artifacts", "narrative_model.joblib")


def load_dataset(path):
    texts, labels = [], []
    with open(path, "r", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for row in reader:
            texts.append(row["text"])
            labels.append(int(row["label"]))
    return texts, labels


def main():
    texts, labels = load_dataset(DATA_PATH)
    vectorizer = joblib.load(VEC_PATH)

    model = NarrativeClassifier(vectorizer)
    model.fit(texts, labels)

    joblib.dump(model, OUT_PATH)
    print(f"Narrative model trained and saved to {OUT_PATH}")


if __name__ == "__main__":
    main()
