#!/usr/bin/env python3
"""
Comprehensive Basic Arithmetic & BODMAS Shard Generator
Generates:
1. Complete Multiplication Tables 1 to 100 (English & Hindi pahade)
2. Direct 1 to 4 digit Addition, Subtraction, Multiplication, Division (+ - * /)
3. Step-by-Step BODMAS Multi-operator Solutions
4. Percentage, Discount, GST, Fractions, Square roots, Cubes
"""

import random
import pandas as pd

random.seed(42)

def generate_tables():
    entries = []
    hindi_ordinals = ["ekam", "duni", "tiya", "chauke", "panche", "chhakke", "satte", "atthe", "namme", "daham"]
    for num in range(1, 101):
        lines = [f"### Multiplication Table of {num}:"]
        for m in range(1, 11):
            lines.append(f"{num} x {m} = {num * m}")
        full_text_en = "\n".join(lines)
        entries.append({"text": full_text_en, "category": "table_english", "num": num})

        # Hindi Pahada
        if num <= 30:
            lines_hi = [f"### {num} का पहाड़ा (Multiplication Table in Hindi):"]
            for m in range(1, 11):
                ans = num * m
                ord_name = hindi_ordinals[m - 1]
                lines_hi.append(f"{num} {ord_name} {ans} ({num} x {m} = {ans})")
            full_text_hi = "\n".join(lines_hi)
            entries.append({"text": full_text_hi, "category": "table_hindi", "num": num})
    return entries

def generate_direct_arithmetic(count=15000):
    entries = []
    for _ in range(count):
        op = random.choice(["+", "-", "*", "/"])
        if op == "+":
            digits = random.choice([1, 2, 3, 4])
            a = random.randint(10**(digits-1), 10**digits - 1)
            b = random.randint(10**(digits-1), 10**digits - 1)
            res = a + b
            text = (
                f"Question: Calculate {a} + {b}.\n"
                f"<soch>\nJod (Addition):\n  {a}\n+ {b}\n------\n  {res}\n</soch>\n"
                f"Answer: {a} + {b} = {res}."
            )
            entries.append({"text": text, "category": "addition"})
        elif op == "-":
            digits = random.choice([1, 2, 3, 4])
            a = random.randint(10**(digits-1), 10**digits - 1)
            b = random.randint(10**(digits-1), 10**digits - 1)
            res = a - b
            text = (
                f"Question: Calculate {a} - {b}.\n"
                f"<soch>\nGhatana (Subtraction):\n  {a}\n- {b}\n------\n  {res}\n</soch>\n"
                f"Answer: {a} - {b} = {res}."
            )
            entries.append({"text": text, "category": "subtraction"})
        elif op == "*":
            a = random.randint(2, 999)
            b = random.randint(2, 99)
            res = a * b
            text = (
                f"Question: Calculate {a} * {b} (Multiplication / Guna).\n"
                f"<soch>\nStep-by-step multiplication:\n{a} multiplied by {b} is calculated as:\n"
                f"{a} * {b} = {res}\n</soch>\n"
                f"Answer: {a} * {b} = {res}."
            )
            entries.append({"text": text, "category": "multiplication"})
        elif op == "/":
            b = random.randint(2, 50)
            q = random.randint(2, 200)
            r = random.choice([0, random.randint(1, b - 1)])
            a = b * q + r
            if r == 0:
                text = (
                    f"Question: Calculate {a} / {b} (Division / Bhag).\n"
                    f"<soch>\n{a} divided by {b}:\nQuotient = {q}, Remainder = 0.\nExact division: {a} / {b} = {q}\n</soch>\n"
                    f"Answer: {a} / {b} = {q}."
                )
            else:
                dec = round(a / b, 4)
                text = (
                    f"Question: Calculate {a} / {b} (Division / Bhag).\n"
                    f"<soch>\n{a} divided by {b}:\nQuotient (Bhagphal) = {q}, Remainder (Sheshphal) = {r}.\nIn decimals: {a} / {b} = {dec}\n</soch>\n"
                    f"Answer: {a} / {b} = {q} remainder {r} (approx {dec})."
                )
            entries.append({"text": text, "category": "division"})
    return entries

def generate_bodmas(count=5000):
    entries = []
    for _ in range(count):
        a = random.randint(5, 50)
        b = random.randint(2, 12)
        c = random.randint(2, 8)
        d = random.randint(1, 10)
        e = random.randint(2, 6)

        pattern = random.choice([1, 2, 3])
        if pattern == 1:
            # a + b * (c + d)
            bracket = c + d
            mult = b * bracket
            ans = a + mult
            expr = f"{a} + {b} * ({c} + {d})"
            text = (
                f"Question: Solve the expression using BODMAS: {expr}\n"
                f"<soch>\nBODMAS Rule (Brackets, Orders, Division/Multiplication, Addition/Subtraction):\n"
                f"Step 1: Solve inside Bracket: ({c} + {d}) = {bracket}\n"
                f"Expression becomes: {a} + {b} * {bracket}\n"
                f"Step 2: Multiply: {b} * {bracket} = {mult}\n"
                f"Expression becomes: {a} + {mult}\n"
                f"Step 3: Add: {a} + {mult} = {ans}\n"
                f"</soch>\n"
                f"Answer: {expr} = {ans}."
            )
        elif pattern == 2:
            # (a * b) - (c * d)
            part1 = a * b
            part2 = c * d
            ans = part1 - part2
            expr = f"({a} * {b}) - ({c} * {d})"
            text = (
                f"Question: Solve the expression using BODMAS: {expr}\n"
                f"<soch>\nStep 1: Solve first bracket: {a} * {b} = {part1}\n"
                f"Step 2: Solve second bracket: {c} * {d} = {part2}\n"
                f"Step 3: Subtract: {part1} - {part2} = {ans}\n"
                f"</soch>\n"
                f"Answer: {expr} = {ans}."
            )
        else:
            # a * e + b - d
            mult = a * e
            ans = mult + b - d
            expr = f"{a} * {e} + {b} - {d}"
            text = (
                f"Question: Solve the expression: {expr}\n"
                f"<soch>\nBODMAS:\nStep 1: Multiply: {a} * {e} = {mult}\n"
                f"Step 2: Add: {mult} + {b} = {mult + b}\n"
                f"Step 3: Subtract: {mult + b} - {d} = {ans}\n"
                f"</soch>\n"
                f"Answer: {expr} = {ans}."
            )
        entries.append({"text": text, "category": "bodmas"})
    return entries

def generate_percentage_and_everyday(count=5000):
    entries = []
    gst_slabs = [5, 12, 18, 28]
    for _ in range(count):
        choice = random.choice(["gst", "discount", "simple_interest", "fraction", "powers"])
        if choice == "gst":
            amount = random.randint(10, 500) * 100
            rate = random.choice(gst_slabs)
            gst_val = int(amount * (rate / 100))
            total = amount + gst_val
            text = (
                f"Question: {amount} rupaye par {rate}% GST kitna hoga aur kul rashi (total amount) kitni hogi?\n"
                f"<soch>\nGST calculation:\n"
                f"GST Amount = Amount * (Rate / 100) = {amount} * ({rate} / 100) = {gst_val} rupaye.\n"
                f"Kul Rashi = Mool Rashi + GST = {amount} + {gst_val} = {total} rupaye.\n"
                f"</soch>\n"
                f"Answer: {amount} rupaye par {rate}% GST = {gst_val} rupaye hoga, aur kul rashi {total} rupaye hogi."
            )
            entries.append({"text": text, "category": "gst"})
        elif choice == "discount":
            mrp = random.randint(20, 200) * 50
            disc_rate = random.choice([10, 15, 20, 25, 30, 40, 50])
            disc_val = int(mrp * (disc_rate / 100))
            selling_price = mrp - disc_val
            text = (
                f"Question: MRP Rs. {mrp} par {disc_rate}% discount ke baad bikri muly (Selling Price) kya hoga?\n"
                f"<soch>\nDiscount = {mrp} * ({disc_rate} / 100) = {disc_val} rupaye.\n"
                f"Selling Price = MRP - Discount = {mrp} - {disc_val} = {selling_price} rupaye.\n"
                f"</soch>\n"
                f"Answer: {disc_rate}% discount ke baad selling price Rs. {selling_price} hoga (bachat: Rs. {disc_val})."
            )
            entries.append({"text": text, "category": "discount"})
        elif choice == "powers":
            base = random.randint(2, 25)
            exp = 2 if base > 15 else random.randint(2, 4)
            val = base ** exp
            text = (
                f"Question: What is {base}^{exp} (Square / Cube)?\n"
                f"<soch>\n{base} raised to the power {exp} = {val}\n</soch>\n"
                f"Answer: {base}^{exp} = {val}."
            )
            entries.append({"text": text, "category": "powers"})
        elif choice == "fraction":
            n1 = random.randint(1, 4)
            d1 = random.choice([2, 4, 8])
            n2 = random.randint(1, 4)
            d2 = d1
            ans_n = n1 + n2
            text = (
                f"Question: Solve fraction addition: {n1}/{d1} + {n2}/{d2}\n"
                f"<soch>\nDenominators are same ({d1}):\n({n1} + {n2}) / {d1} = {ans_n} / {d1}\n</soch>\n"
                f"Answer: {n1}/{d1} + {n2}/{d2} = {ans_n}/{d1}."
            )
            entries.append({"text": text, "category": "fraction"})
    return entries

def main():
    print("[*] Generating Basic Arithmetic, Tables, BODMAS, and Everyday Math...")
    all_data = []
    all_data.extend(generate_tables())
    all_data.extend(generate_direct_arithmetic(15000))
    all_data.extend(generate_bodmas(5000))
    all_data.extend(generate_percentage_and_everyday(5000))

    random.shuffle(all_data)
    df = pd.DataFrame(all_data)
    out_file = "basic_arithmetic_bodmas_tables.parquet"
    df.to_parquet(out_file, compression="zstd")
    print(f"[OK] Generated {out_file}: {len(df)} rows, {df.memory_usage().sum() / (1024**2):.2f} MB")

if __name__ == "__main__":
    main()