# ============================================================
# TEST / INFERENCE RIÊNG
# Load toàn bộ model đã train và dự đoán dữ liệu mới
# ============================================================

!pip -q install -U transformers accelerate scikit-learn scipy joblib sentencepiece

import os
import re
import gc
import math
import joblib

import numpy as np
import pandas as pd
import torch

from pathlib import Path
from scipy import sparse
from scipy.special import softmax

from sklearn.feature_extraction.text import TfidfVectorizer

from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
)

from torch.utils.data import Dataset, DataLoader


# ============================================================
# 1. PATH
# ============================================================

# Thư mục model đã train ở notebook trước
WORK_DIR = Path(
    "/content/drive/MyDrive/rvihsd_stacking_work"
)

# File dữ liệu MỚI cần dự đoán
TEST_CSV = "/content/drive/MyDrive/new_test.csv"

# File kết quả
OUTPUT_CSV = "/content/drive/MyDrive/new_test_prediction.csv"


# ============================================================
# 2. CONFIG
# ============================================================

MAX_LENGTH = 128
BATCH_SIZE = 64

LABELS = [
    "CLEAN",
    "OFFENSIVE",
    "HATE",
]

NOISE_LABELS = [
    "ORIGINAL",
    "NO_DIACRITICS",
    "TEENCODE",
    "CHAR_REPEAT",
    "PUNCT_NOISE",
    "OBFUSCATION",
    "MIXED",
]

id2label = {
    i: x
    for i, x in enumerate(LABELS)
}

id2noise = {
    i: x
    for i, x in enumerate(NOISE_LABELS)
}


# ============================================================
# 3. CLASS DualTfidf
#    PHẢI CÓ vì joblib cần class này khi load vectorizer
# ============================================================

class DualTfidf:

    def __init__(self):
        self.char = None
        self.word = None

    def transform(self, texts):

        texts = (
            pd.Series(texts)
            .fillna("")
            .astype(str)
            .tolist()
        )

        Xc = self.char.transform(texts)
        Xw = self.word.transform(texts)

        return sparse.hstack(
            [Xc, Xw],
            format="csr",
            dtype=np.float32,
        )


# ============================================================
# 4. SURFACE FEATURES
#    PHẢI GIỐNG HỆT LÚC TRAIN META-MODEL
# ============================================================

VI_DIACRITICS = set(
    "ăâđêôơưáàảãạắằẳẵặấầẩẫậ"
    "éèẻẽẹếềểễệ"
    "íìỉĩị"
    "óòỏõọốồổỗộớờởỡợ"
    "úùủũụứừửữự"
    "ýỳỷỹỵ"
    "ĂÂĐÊÔƠƯÁÀẢÃẠẮẰẲẴẶẤẦẨẪẬ"
    "ÉÈẺẼẸẾỀỂỄỆ"
    "ÍÌỈĨỊ"
    "ÓÒỎÕỌỐỒỔỖỘỚỜỞỠỢ"
    "ÚÙỦŨỤỨỪỬỮỰ"
    "ÝỲỶỸỴ"
)


def meta_surface_features(texts):

    rows = []

    for text in texts:

        s = str(text)

        n = max(len(s), 1)

        alpha = max(
            sum(ch.isalpha() for ch in s),
            1
        )

        punct = sum(
            (not ch.isalnum())
            and
            (not ch.isspace())
            for ch in s
        )

        diac = sum(
            ch in VI_DIACRITICS
            for ch in s
        )

        repeats = len(
            re.findall(
                r"(.)\1{2,}",
                s.lower()
            )
        )

        obfus = len(
            re.findall(
                r"(?<=[A-Za-zÀ-ỹ])[.*_~](?=[A-Za-zÀ-ỹ])",
                s
            )
        )

        rows.append([
            math.log1p(len(s)),
            punct / n,
            diac / alpha,
            min(repeats, 10) / 10.0,
            min(obfus, 10) / 10.0,
        ])

    return np.asarray(
        rows,
        dtype=np.float32
    )


# ============================================================
# 5. KIỂM TRA MODEL FILE
# ============================================================

VISO_DIR = WORK_DIR / "visobert_full"

VECTORIZER_PATH = (
    WORK_DIR / "full_vectorizer.joblib"
)

HATE_SVM_PATH = (
    WORK_DIR / "full_hate_svm.joblib"
)

NOISE_SVM_PATH = (
    WORK_DIR / "full_noise_svm.joblib"
)

META_PATH = (
    WORK_DIR / "meta_model.joblib"
)


required = [
    VISO_DIR / "config.json",
    VECTORIZER_PATH,
    HATE_SVM_PATH,
    NOISE_SVM_PATH,
    META_PATH,
]

for p in required:

    if not p.exists():
        raise FileNotFoundError(
            f"Không tìm thấy: {p}"
        )

print("Tất cả model files đều tồn tại.")


# ============================================================
# 6. LOAD TF-IDF + SVM + META MODEL
# ============================================================

print("\nLoading TF-IDF/SVM...")

full_vec = joblib.load(
    VECTORIZER_PATH
)

full_hate_svm = joblib.load(
    HATE_SVM_PATH
)

full_noise_svm = joblib.load(
    NOISE_SVM_PATH
)

meta_model = joblib.load(
    META_PATH
)

print("TF-IDF/SVM/meta-model loaded.")


# ============================================================
# 7. LOAD ViSoBERT
# ============================================================

device = torch.device(
    "cuda"
    if torch.cuda.is_available()
    else "cpu"
)

print("\nDevice:", device)

print("Loading ViSoBERT...")

tokenizer = AutoTokenizer.from_pretrained(
    VISO_DIR
)

visobert = (
    AutoModelForSequenceClassification
    .from_pretrained(VISO_DIR)
)

visobert.to(device)
visobert.eval()

print("ViSoBERT loaded.")


# ============================================================
# 8. DATASET CHO ViSoBERT
# ============================================================

class TestDataset(Dataset):

    def __init__(self, texts):

        self.texts = (
            pd.Series(texts)
            .fillna("")
            .astype(str)
            .tolist()
        )

    def __len__(self):

        return len(self.texts)

    def __getitem__(self, idx):

        return self.texts[idx]


# ============================================================
# 9. ViSoBERT INFERENCE
# ============================================================

def predict_visobert(
    texts,
    batch_size=BATCH_SIZE
):

    dataset = TestDataset(texts)

    loader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=False,
    )

    all_probs = []

    with torch.no_grad():

        for batch_texts in loader:

            enc = tokenizer(
                list(batch_texts),
                padding=True,
                truncation=True,
                max_length=MAX_LENGTH,
                return_tensors="pt",
            )

            enc = {
                k: v.to(device)
                for k, v in enc.items()
            }

            outputs = visobert(**enc)

            probs = torch.softmax(
                outputs.logits,
                dim=1
            )

            all_probs.append(
                probs.cpu().numpy()
            )

    return np.concatenate(
        all_probs,
        axis=0
    ).astype(np.float32)


# ============================================================
# 10. LOAD TEST DATA
# ============================================================

test_df = pd.read_csv(
    TEST_CSV,
    dtype={"id": str}
)

required_cols = {
    "id",
    "text",
}

missing = required_cols - set(
    test_df.columns
)

if missing:
    raise ValueError(
        f"File test thiếu cột: {missing}"
    )

test_df = test_df[
    ["id", "text"]
].copy()

test_df["id"] = (
    test_df["id"]
    .astype(str)
)

test_df["text"] = (
    test_df["text"]
    .fillna("")
    .astype(str)
)

print(
    "\nSố mẫu test:",
    len(test_df)
)

display(
    test_df.head()
)


# ============================================================
# 11. TF-IDF
# ============================================================

print("\n[1/4] TF-IDF...")

X_test = full_vec.transform(
    test_df["text"]
)

print(
    "TF-IDF shape:",
    X_test.shape
)


# ============================================================
# 12. SVM PREDICTIONS
# ============================================================

print("\n[2/4] SVM inference...")

hate_svm_score = (
    full_hate_svm
    .decision_function(X_test)
    .astype(np.float32)
)

noise_score = (
    full_noise_svm
    .decision_function(X_test)
    .astype(np.float32)
)

noise_soft = softmax(
    noise_score,
    axis=1
).astype(np.float32)

print(
    "Hate SVM:",
    hate_svm_score.shape
)

print(
    "Noise SVM:",
    noise_score.shape
)


# ============================================================
# 13. ViSoBERT
# ============================================================

print("\n[3/4] ViSoBERT inference...")

viso_prob = predict_visobert(
    test_df["text"]
)

print(
    "ViSoBERT:",
    viso_prob.shape
)


# ============================================================
# 14. META SURFACE FEATURES
# ============================================================

surface = meta_surface_features(
    test_df["text"]
)

print(
    "Surface:",
    surface.shape
)


# ============================================================
# 15. STACKING
#
# Thứ tự PHẢI giống lúc train:
#
# ViSoBERT probability = 3
# SVM hate score       = 3
# noise softmax        = 7
# surface              = 5
#
# Tổng = 18 features
# ============================================================

print("\n[4/4] Stacking...")

meta_X = np.hstack([
    viso_prob,
    hate_svm_score,
    noise_soft,
    surface,
]).astype(np.float32)

print(
    "Meta feature shape:",
    meta_X.shape
)

assert meta_X.shape[1] == 18


# Hate classification
hate_pred_id = meta_model.predict(
    meta_X
).astype(int)

# Noise classification
noise_pred_id = noise_score.argmax(
    axis=1
)


# ============================================================
# 16. CHUYỂN ID -> LABEL
# ============================================================

pred_label = [
    id2label[int(x)]
    for x in hate_pred_id
]

pred_noise = [
    id2noise[int(x)]
    for x in noise_pred_id
]


# ============================================================
# 17. TẠO OUTPUT
# ============================================================

output = pd.DataFrame({

    "id":
        test_df["id"].values,

    "pred_label":
        pred_label,

    "pred_noise_type":
        pred_noise,
})


# ============================================================
# 18. KIỂM TRA
# ============================================================

assert len(output) == len(test_df)

assert output["id"].is_unique

assert set(
    output["pred_label"]
).issubset(LABELS)

assert set(
    output["pred_noise_type"]
).issubset(NOISE_LABELS)


print("\nPrediction distribution:")

print(
    "\nLABEL:"
)

print(
    output["pred_label"]
    .value_counts()
)

print(
    "\nNOISE:"
)

print(
    output["pred_noise_type"]
    .value_counts()
)


# ============================================================
# 19. SAVE
# ============================================================

output.to_csv(
    OUTPUT_CSV,
    index=False,
    encoding="utf-8",
)

print(
    "\nSaved:",
    OUTPUT_CSV
)

display(
    output.head(20)
)