src/augment_data.py

"""
augment_data.py - Tăng cường dữ liệu cực kỳ phong phú
Tạo 15-25 biến thể câu hỏi từ mỗi câu gốc với nhiều phong cách:
- Học thuật / Sách giáo khoa
- Lịch sự / Trang trọng
- Thân mật / Bạn bè
- Gen Z / Tiếng lóng
- Viết tắt / Nhắn tin
- Thắc mắc / Ngạc nhiên
- Hỏi lại / Nhờ giải thích
Chạy: python src/augment_data.py
"""
import pandas as pd
import random
import os
import re
import sys

sys.stdout.reconfigure(encoding='utf-8')

INPUT_CSV  = 'data/raw/output.csv'
OUTPUT_CSV = 'data/augmented/augmented.csv'

# ============================================================
# BẢNG TỪ ĐỒNG NGHĨA CHO CÁC PREFIX / SUFFIX HỎI
# ============================================================

PREFIX_DINH_NGHIA = [
    "{base} là gì?",
    "Thế nào là {base}?",
    "Định nghĩa {base} là gì?",
    "Định nghĩa của {base} là gì?",
    "Khái niệm {base} là gì?",
    "Khái niệm {base} được hiểu như thế nào?",
    "{base} được định nghĩa như thế nào?",
    "Hãy định nghĩa {base}.",
    "Cho mình biết {base} là gì?",
    "Bạn hiểu {base} là gì?",
    "{base} nghĩa là gì?",
    "Ý nghĩa của {base} là gì?",
    # Học thuật
    "Phát biểu định nghĩa {base}.",
    "Trình bày khái niệm {base}.",
    "Nêu định nghĩa {base}.",
    # Thân mật / gen-Z
    "ê {base} là gì vậy?",
    "{base} là gì z?",
    "{base} là gì ạ?",
    "cho hỏi {base} là gì?",
    "xin hỏi {base} là gì ạ?",
    "mn ơi {base} là gì vậy?",
    "Hỏi ngu tí, {base} là gì?",
    "{base} là cái gì vậy?",
    "cứu bé {base} là gì?",
    "cíu với {base} là gì z?",
    "ai biết {base} là gì k?",
    "có ai biết {base} là gì khum?",
    "{base} là cái quái gì vậy?",
    "{base} là qq gì dị?",
    "giúp tui hiểu {base} là gì đi",
    "giải ngố tui phát, {base} là cái gì?",
    "trời ơi {base} là gì",
    # Viết tắt / nhắn tin nhanh
    "{base} la gi?",
    "kn {base} là gì?",
    "đn {base}?",
    "{base} la j?",
    "{base} la j v?",
    "{base} ntn?",
    "{base} nghia la s?",
    "{base} la sao v bro?",
    "{base} chốt lại là j?",
    "{base} là j dị",
    "{base} la gi v tr",
    "{base} là saoooo",
    "k/n {base} là j?",
    "đ/n {base} ntn?",
    # Ngạc nhiên / thắc mắc
    "Ủa {base} là gì vậy bạn?",
    "Tại sao lại gọi là {base}?",
    "Mình chưa hiểu {base} là gì, giải thích giúp mình với?",
    "Bạn ơi {base} có nghĩa là gì ạ?",
    "Ủa alo {base} rốt cuộc là sao?",
    "Không hiểu {base} là gì luôn á?",
]

PREFIX_GIAI_THICH = [
    "Giải thích {base}.",
    "Hãy giải thích {base}.",
    "Giải thích {base} cho mình hiểu với.",
    "Giải thích {base} là gì?",
    "Giải thích khái niệm {base}?",
    "Mô tả {base}.",
    "Mô tả {base} chi tiết hơn?",
    "Diễn giải {base}.",
    "Trình bày {base}.",
    "Cho mình hiểu về {base}?",
    # Thân mật
    "bạn giải thích {base} giúp mình với",
    "giải thích {base} đi bạn",
    "explain {base} cho mình với",
    "{base} hiểu kiểu gì vậy?",
    "mình cần hiểu về {base}, bạn giúp mình không?",
    # Gen Z / Viết tắt
    "alo bạn ơi giải thích {base} cho t vs",
    "glt {base} là gì v bạn?",
    "{base} là cái chi vậy tr?",
    "hú hú ai giải thích {base} cho mình với 😭",
    "cứu với, {base} là gì vậy??",
    "ok nhưng {base} là gì???",
    "gt {base} tui với",
    "gt {base} ntn v?",
    "ai gt {base} giúp t với",
    "giải thích nhanh {base} vs",
    "plz gt {base}",
    "bro gt {base} hộ t",
    "giảng lại {base} t nghe coi",
    "giải thích {base} bằng ngôn ngữ loài người đi",
    "{base} khó hiểu vãi, gt lại đi",
    "giảng hộ phát {base} là sao",
]

PREFIX_VI_DU = [
    "Cho ví dụ về {base}.",
    "Hãy cho ví dụ minh họa về {base}.",
    "Lấy ví dụ về {base}?",
    "Ví dụ cụ thể về {base} là gì?",
    "Đưa ra ví dụ về {base}.",
    "Minh họa {base} bằng ví dụ?",
    "cho vd về {base} đi",
    "ví dụ {base} là gì?",
    "có ví dụ nào về {base} không?",
    "vd {base}?",
    "cho 1 vd {base}",
    "lấy vd {base} t xem",
    "vd cụ thể {base} là ntn?",
    "cho tí vd {base} đi",
    "xin vd {base} nha",
]

PREFIX_PHAT_BIEU = [
    "Phát biểu {base}.",
    "Phát biểu nội dung của {base}.",
    "Nội dung của {base} là gì?",
    "Nêu {base}.",
    "Nêu nội dung {base}.",
    "Trình bày {base}.",
    "Phát biểu và giải thích {base}.",
    "{base} phát biểu như thế nào?",
    "{base} nói gì vậy?",
    "{base} nội dung ra sao?",
]

PREFIX_PHAN_BIET = [
    "Phân biệt {A} và {B}.",
    "Sự khác nhau giữa {A} và {B} là gì?",
    "Sự khác biệt giữa {A} và {B}?",
    "{A} khác {B} như thế nào?",
    "{A} và {B} giống và khác nhau điểm gì?",
    "So sánh {A} và {B}.",
    "Giải thích sự khác nhau giữa {A} và {B}?",
    "Nêu điểm giống và khác nhau giữa {A} và {B}.",
    "{A} vs {B} khác nhau chỗ nào?",
    "{A} với {B} cái nào tổng quát hơn?",
    "{A} vs {B} giống/khác j?",
    "phân biệt {A} vs {B} ntn?",
    "so sánh {A} và {B} đi",
    "{A} khác {B} chỗ nào v?",
    "để pb {A} vs {B} thì lm s?",
    "điểm khác biệt cốt lõi {A} vs {B} là j?",
]

SUFFIX_LICH_SU = [
    " Bạn giải thích giúp mình nhé.",
    " Cảm ơn bạn trước nha!",
    " Mình đang ôn thi, cần gấp.",
    " (Đây là câu hỏi ôn tập nhé)",
]

SUFFIX_GENZ = [
    " ạ 🙏",
    " vậy ạ?",
    " z bạn?",
    " vậy tr?",
    " dc ko bạn?",
    " z chèn?",
    " dọ?",
    " vại?",
    " á",
    " nhaa",
    " nha b",
    " zị",
    " trùi",
    " luôn á",
    " k v",
    " v tròi",
]


def augment_question(question):
    """Tạo nhiều biến thể câu hỏi từ 1 câu gốc."""
    variations = [question]
    q_lower    = question.lower().strip()

    # ——— Mẫu 1: "... là gì?" ———
    if q_lower.endswith("là gì?"):
        base = question[:-6].strip()
        b    = base[0].lower() + base[1:] if base else base
        for tmpl in PREFIX_DINH_NGHIA:
            v = tmpl.format(base=b)
            variations.append(v)
            if random.random() < 0.25:
                variations.append(v.rstrip("?!. ") + random.choice(SUFFIX_LICH_SU))
            if random.random() < 0.2:
                variations.append(v.rstrip("?!. ") + random.choice(SUFFIX_GENZ))

    # ——— Mẫu 2: "Thế nào là ...?" ———
    if q_lower.startswith("thế nào là"):
        base = question[10:].rstrip("?").strip()
        b    = base[0].lower() + base[1:]
        for tmpl in PREFIX_DINH_NGHIA[:12]:
            variations.append(tmpl.format(base=b))

    # ——— Mẫu 3: "Định nghĩa ...?" ———
    if q_lower.startswith("định nghĩa"):
        base = re.sub(r'^định nghĩa\s*', '', question, flags=re.IGNORECASE).rstrip("?").strip()
        b    = base[0].lower() + base[1:]
        if "là gì" not in b.lower():
            for tmpl in PREFIX_DINH_NGHIA[:10]:
                variations.append(tmpl.format(base=b))

    # ——— Mẫu 4: "Giải thích ...?" ———
    if q_lower.startswith("giải thích"):
        base = re.sub(r'^giải thích\s*', '', question, flags=re.IGNORECASE).rstrip("?").strip()
        b    = base[0].lower() + base[1:]
        for tmpl in PREFIX_GIAI_THICH:
            variations.append(tmpl.format(base=b))

    # ——— Mẫu 5: "Phát biểu ..." ———
    if q_lower.startswith("phát biểu"):
        base = re.sub(r'^phát biểu\s*', '', question, flags=re.IGNORECASE).rstrip("?").strip()
        b    = base[0].lower() + base[1:]
        for tmpl in PREFIX_PHAT_BIEU:
            v = tmpl.format(base=b)
            variations.append(v)

    # ——— Mẫu 6: "Nêu ..." ———
    if q_lower.startswith("nêu"):
        base = re.sub(r'^nêu\s*', '', question, flags=re.IGNORECASE).rstrip("?").strip()
        b    = base[0].lower() + base[1:]
        extras = [
            f"Hãy nêu {b}",
            f"Cho biết {b}",
            f"Liệt kê {b}",
            f"Trình bày {b}",
            f"nêu {b} đi bạn",
            f"nêu cho mình {b} với",
            f"cho mình xem {b}?",
        ]
        variations.extend(extras)

    # ——— Mẫu 7: "Như thế nào?" ———
    if "như thế nào?" in q_lower:
        for repl in ["ra sao?", "theo cách nào?", "thế nào?", "kiểu gì?"]:
            variations.append(question.replace("như thế nào?", repl))
            variations.append(question.replace("Như thế nào?", repl.capitalize()))

    # ——— Mẫu 8: "Tại sao ...?" ———
    if q_lower.startswith("tại sao"):
        base = question[7:].strip()
        for prefix in ["Vì sao", "Lý do nào", "Nguyên nhân nào khiến", "tại sao lại"]:
            variations.append(f"{prefix} {base}")

    # ——— Mẫu 9: "Có bao nhiêu ...?" ———
    if "có bao nhiêu" in q_lower:
        for repl in [("có bao nhiêu", "bao nhiêu"), ("có bao nhiêu", "số lượng"), ("Có bao nhiêu", "Bao nhiêu")]:
            variations.append(question.replace(repl[0], repl[1]))

    # ——— Mẫu 10: "Công thức ..." ———
    if "công thức" in q_lower:
        for repl in [("Công thức", "Biểu thức"), ("Công thức", "Cách tính"), ("công thức", "biểu thức"), ("công thức", "cách tính")]:
            variations.append(question.replace(repl[0], repl[1]))

    # ——— Mẫu 11: "Điều kiện ..." ———
    if "điều kiện" in q_lower:
        for repl in [("Điều kiện", "Yêu cầu"), ("điều kiện", "yêu cầu"), ("Điều kiện", "Tiêu chí"), ("điều kiện", "tiêu chí")]:
            variations.append(question.replace(repl[0], repl[1]))

    # ——— Mẫu 12: "Mối liên hệ / Mối quan hệ ..." ———
    if "mối liên hệ" in q_lower:
        variations.append(question.replace("mối liên hệ", "mối quan hệ").replace("Mối liên hệ", "Mối quan hệ"))
    if "mối quan hệ" in q_lower:
        variations.append(question.replace("mối quan hệ", "mối liên hệ").replace("Mối quan hệ", "Mối liên hệ"))

    # ——— Mẫu 13: Thêm prefix lịch sự cho câu ngắn ———
    if len(variations) < 6 and not q_lower.startswith(("hãy", "cho biết", "bạn", "ê", "alo", "mn")):
        q0 = question[0].lower() + question[1:]
        variations.extend([
            f"Cho em hỏi, {q0}",
            f"Bạn ơi, {q0}",
            f"Nhờ bạn giải thích: {q0}",
            f"Mình muốn hỏi: {q0}",
            f"Làm ơn giải thích {q0}",
        ])

    # ——— Mẫu 14: Gen Z / Nhắn tin phổ biến ———
    if len(variations) < 12:
        q0 = question[0].lower() + question[1:]
        variations.extend([
            f"bro ơi {q0}",
            f"alo ad ơi {q0}",
            f"hỏi nhanh: {q0}",
            f"cứu toi, {q0}",
            f"ét o ét {q0}",
            f"ai cứu tui khum, {q0}",
            f"ông ơi {q0}",
            f"hỏi ngu nhe, {q0}",
            f"xin in4: {q0}",
            f"help me {q0}",
            f"bấm quẻ xem {q0}",
            f"này nha, {q0}",
        ])

    # ——— Mẫu 15: Chuyển đổi chính tả / Viết tắt ngẫu nhiên ———
    for i in range(len(variations)):
        if random.random() < 0.2:
            # Thay thế ngẫu nhiên vài từ phổ biến thành viết tắt/teencode
            v = variations[i]
            v = v.replace("không", random.choice(["ko", "k", "kh", "khum", "hông"]))
            v = v.replace("như thế nào", random.choice(["ntn", "như nào", "như nèo"]))
            v = v.replace("bạn", random.choice(["b", "cậu", "bro", "ông"]))
            v = v.replace("được", "đc")
            v = v.replace("của", "cua")
            v = v.replace("gì", random.choice(["j", "zì", "z"]))
            v = v.replace("sao", "s")
            v = v.replace("những", "nhg")
            variations.append(v)

    # ——— Loại bỏ trùng lặp ———
    seen   = set()
    unique = []
    for v in variations:
        v_clean = v.strip()
        if v_clean and v_clean.lower() not in seen:
            seen.add(v_clean.lower())
            unique.append(v_clean)

    return unique


def main():
    if not os.path.exists(INPUT_CSV):
        print(f"Không tìm thấy file: {INPUT_CSV}")
        return

    with open(INPUT_CSV, 'r', encoding='utf-8-sig') as f:
        first_line = f.readline()
    sep = ';' if first_line.count(';') >= 2 else ','

    df = pd.read_csv(INPUT_CSV, sep=sep, encoding='utf-8-sig')
    if len(df.columns) == 3:
        df.columns = ['topic', 'question', 'answer']
    else:
        print(f"Lỗi: File CSV có {len(df.columns)} cột, cần đúng 3 cột")
        return

    df = df.dropna(subset=['question', 'answer'])

    augmented_data = []
    for _, row in df.iterrows():
        for v in augment_question(str(row['question'])):
            augmented_data.append({
                'topic':    row['topic'],
                'question': v,
                'answer':   row['answer']
            })

    os.makedirs(os.path.dirname(OUTPUT_CSV), exist_ok=True)
    result_df = pd.DataFrame(augmented_data)
    result_df.to_csv(OUTPUT_CSV, sep=';', index=False, header=False, encoding='utf-8-sig')

    original_count  = len(df)
    augmented_count = len(result_df)
    print(f"\n{'='*55}")
    print(f"  AUGMENTATION HOÀN TẤT!")
    print(f"{'='*55}")
    print(f"  Câu hỏi gốc    : {original_count:,}")
    print(f"  Sau đa dạng hóa: {augmented_count:,}")
    print(f"  Tỷ lệ tăng     : x{augmented_count/original_count:.1f}")
    print(f"  File output     : {OUTPUT_CSV}")
    print(f"{'='*55}")

    # Thống kê theo chủ đề
    print("\n  Phân bổ theo chủ đề:")
    for topic, cnt in result_df.groupby('topic').size().sort_values(ascending=False).items():
        print(f"    {cnt:6,}  {topic}")


if __name__ == "__main__":
    main()