"""文本分类 baseline"""
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.svm import LinearSVC
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, classification_report


# 1. 数据
texts = [
    "我要申请退款",
    "这个订单可以退款吗",
    "钱什么时候退回来",
    "我要取消订单退款",
    "商品有问题申请退款",
    "怎么退货退款",

    "我的密码忘记了",
    "怎么修改登录密码",
    "账号登录不上去了",
    "如何重置密码",
    "账户无法登录",
    "密码在哪里修改",

    "快递什么时候到",
    "物流怎么还没更新",
    "我的包裹到哪里了",
    "查询一下配送进度",
    "快递一直没有收到",
    "帮我查一下物流"
]

labels = [
    "refund", "refund", "refund",
    "refund", "refund", "refund",

    "account", "account", "account",
    "account", "account", "account",

    "delivery", "delivery", "delivery",
    "delivery", "delivery", "delivery"
]


# 2. 切分训练集 / 验证集
X_train, X_val, y_train, y_val = train_test_split(
    texts,
    labels,
    test_size=0.33,
    random_state=42,
    stratify=labels
)


# 3. 文本 -> TF-IDF 稀疏向量
vectorizer = TfidfVectorizer(
    analyzer="char",
    ngram_range=(1, 3)
)

X_train_vec = vectorizer.fit_transform(X_train)
X_val_vec = vectorizer.transform(X_val)


# 4. 训练分类器
model = LinearSVC(C=1.0)
model.fit(X_train_vec, y_train)


# 5. 验证
pred = model.predict(X_val_vec)

print("accuracy:", accuracy_score(y_val, pred))
print(classification_report(y_val, pred))


# 6. 推理
test_texts = [
    "我的东西怎么还没送过来",
    "登录密码忘了怎么办",
    "我要把这个订单退掉"
]

X_test = vectorizer.transform(test_texts)
pred = model.predict(X_test)

for text, label in zip(test_texts, pred):
    print(text, "->", label)