"""Text generators for synthetic documents.

Word sources: /usr/share/dict/words (Webster's Second International, 1934, public domain) when present,
plus the small hand-written vocabularies below. Numbers, dates, codes and prices are generated.
"""

import re
import string
from functools import lru_cache
from pathlib import Path

import numpy as np

from ..charset import CHARS, scripted

COMMON = """the of and to in is for on that with as by this be are from at or an it was will not have all
can your more has our we new one their which about other time also any may out up only use these some
year first two over into after most such through each where much before must between well under those
date total amount account payment service customer order number price item quantity balance due please
note thank you report section page table figure summary results data project team meeting notes review
plan update policy contact address phone email office company invoice receipt store market product
annual financial statement period month week day today information details reference terms conditions""".split()

ITEMS = """coffee latte cappuccino espresso mocha tea chai juice water soda cola lemonade smoothie milk
bagel croissant muffin donut toast sandwich burger fries salad soup pizza pasta noodles rice curry wrap
taco burrito chicken beef pork fish shrimp tofu egg cheese bread butter jam honey yogurt cereal apple
banana orange grapes mango berries tomato onion potato carrot spinach lettuce garlic pepper flour sugar
salt oil vinegar sauce ketchup napkins batteries charger cable notebook pen pencil marker tape paper
envelope stamps shampoo soap toothpaste detergent towel tissue lotion vitamins aspirin bandage socks
shirt cap gloves umbrella bulb filter screws nails paint brush hammer drill glue""".split()
ITEM_MODS = """large small medium regular extra hot iced fresh organic whole sliced grilled spicy sweet
classic double single family pack bottle can box bag jar dozen kg lb""".split()
UNITS = ["", "", "", " kg", " lb", " pc", " pcs", " ea", " L", " ml", " g"]

NAME_A = """Blue Golden Green Silver Red Royal Sunny Happy Urban Corner Main Street Lucky Fresh Daily Prime
Star Ocean River Maple Oak Pine Cedar North South East West Central Metro City Village Grand Little Big
Bright Modern Classic United General Global Pacific Atlantic Summit Valley Harbor Lake""".split()
NAME_B = """Cafe Bakery Market Grocery Mart Store Shop Bistro Diner Kitchen Grill Pharmacy Hardware Books
Supplies Traders Foods Deli Pizzeria Restaurant Outlet Electronics Florist Bar Pub Coffee Tea House""".split()
COMPANY_SUFFIX = ["Ltd", "LLC", "Inc.", "Co.", "GmbH", "Pvt Ltd", "Corp.", "& Co", "Group", "Services", "Solutions", ""]
STREETS = """Main Oak Pine Maple Cedar Elm Washington Park Lake Hill Church High Market King Queen Station
Victoria Mill Bridge River Spring Garden Forest Sunset Lincoln Jackson Madison""".split()
STREET_T = ["St", "Street", "Ave", "Avenue", "Rd", "Road", "Blvd", "Lane", "Ln", "Dr", "Way", "Square"]
CITIES = """Springfield Riverside Fairview Georgetown Franklin Clinton Salem Madison Bristol Oxford Dover
Newport Ashland Burlington Manchester Kingston Arlington Lexington Milton Chester Hudson Auburn""".split()
STATES = ["CA", "NY", "TX", "WA", "IL", "MA", "FL", "GA", "OR", "NJ", "ON", "BC", "NSW", "VIC", "KA", "MH", "DL"]
FIRST = """James Mary John Patricia Robert Jennifer Michael Linda David Sarah Daniel Priya Rahul Anita Wei
Mei Carlos Maria Ahmed Fatima Olga Ivan Kenji Yuki Lucas Emma Noah Olivia Arjun Kavya Hari Meera""".split()
LAST = """Smith Johnson Brown Garcia Miller Davis Wilson Taylor Anderson Thomas Moore Martin Lee Clark
Lewis Walker Hall Young King Wright Lopez Hill Scott Green Adams Baker Nair Sharma Patel Kumar Chen""".split()
FOOTERS = ["Thank you!", "THANK YOU FOR SHOPPING", "Please come again", "Have a nice day!",
           "Thank you for your visit", "No refunds without receipt", "Returns within 30 days",
           "Customer copy", "Keep this receipt", "www.example.com", "Follow us online", "*** CUSTOMER COPY ***"]
PAYMENTS = ["CASH", "Cash", "VISA", "Visa", "MASTERCARD", "Card", "CREDIT", "DEBIT", "UPI", "AMEX", "Mobile Pay"]
CURRENCIES = [("$", "USD"), ("$", "USD"), ("€", "EUR"), ("£", "GBP"), ("₹", "INR"), ("Rs.", "INR"),
              ("", "USD"), ("", "EUR"), ("¥", "JPY"), ("$", "CAD"), ("$", "AUD"), ("", "INR")]
MONTHS = ["January", "February", "March", "April", "May", "June", "July", "August", "September",
          "October", "November", "December"]

KEYS = {
    "invoice_number": ["Invoice No:", "Invoice #", "Invoice Number:", "INV #", "Bill No:", "Receipt #:",
                       "Receipt No.", "Order #", "Ref:", "Txn #", "Transaction:", "No."],
    "date": ["Date:", "DATE", "Invoice Date:", "Date", "Issued:", "Dated:"],
    "due_date": ["Due Date:", "Due:", "Payment Due:", "Pay By:"],
    "time": ["Time:", "TIME", "Time"],
    "subtotal": ["Subtotal", "SUBTOTAL", "Sub Total", "Sub-total:", "Net Amount", "Amount", "Total excl. tax"],
    "tax": ["Tax", "TAX", "VAT", "GST", "Sales Tax", "VAT {r}%", "Tax {r}%", "GST {r}%", "CGST+SGST"],
    "discount": ["Discount", "DISCOUNT", "Savings", "Promo", "Coupon"],
    "total": ["TOTAL", "Total", "Grand Total", "TOTAL DUE", "Amount Due", "Balance Due", "Total:", "TOTAL AMOUNT", "Net Payable"],
    "amount_paid": ["CASH", "Cash", "Paid", "Tendered", "VISA", "Card", "Amount Paid", "Payment"],
    "change": ["Change", "CHANGE", "Change Due"],
    "merchant_phone": ["Tel:", "Phone:", "Ph:", "TEL"],
    "merchant_tax_id": ["GSTIN:", "VAT No:", "Tax ID:", "TIN:", "ABN:", "EIN:"],
}
ITEM_HEADERS = {
    "item_desc": ["Description", "Item", "ITEM", "Product", "Details"],
    "item_qty": ["Qty", "QTY", "Quantity", "Units"],
    "item_unit_price": ["Unit Price", "Price", "Rate", "Unit Cost"],
    "item_amount": ["Amount", "Total", "AMOUNT", "Line Total"],
}


@lru_cache(maxsize=1)
def dictionary_words() -> list[str]:
    p = Path("/usr/share/dict/words")
    if p.exists():
        return [w.strip() for w in p.read_text().split() if 2 <= len(w.strip()) <= 14]
    return COMMON


def pick(rng, seq):
    return seq[rng.integers(len(seq))]


def word(rng):
    r = rng.random()
    if r < 0.6:
        return pick(rng, COMMON)
    w = pick(rng, dictionary_words())
    return w.capitalize() if r < 0.7 else w


def sentence(rng, n_min=4, n_max=14):
    n = int(rng.integers(n_min, n_max + 1))
    words = [word(rng) for _ in range(n)]
    for i in range(1, n - 1):
        if rng.random() < 0.08:
            words[i] += pick(rng, [",", ",", ";", ":"])
        if rng.random() < 0.05:
            words[i] = pick(rng, [number(rng), f"({word(rng)})", f"\"{word(rng)}\"", f"{rng.integers(1, 100)}%"])
    s = " ".join(words)
    return s[0].upper() + s[1:] + pick(rng, [".", ".", ".", "?", "!", ":", ""])


def number(rng):
    r = rng.random()
    if r < 0.3:
        return str(rng.integers(0, 100))
    if r < 0.6:
        return f"{rng.integers(0, 100000):,}"
    if r < 0.8:
        return f"{rng.random() * 1000:.{rng.integers(1, 4)}f}"
    return str(rng.integers(1900, 2031))


class MoneyFormat:
    """Consistent money formatting for one document."""

    def __init__(self, rng):
        self.symbol, self.code = pick(rng, CURRENCIES)
        self.euro_style = rng.random() < 0.18          # 1.234,56
        self.decimals = 0 if (self.code == "JPY" or rng.random() < 0.05) else 2
        self.symbol_mode = rng.choice(["none", "every", "totals"], p=[0.35, 0.3, 0.35]) if self.symbol else "none"
        self.space = rng.random() < 0.3
        self.thousands = rng.random() < 0.85

    def fmt(self, value, with_symbol=False):
        s = f"{abs(value):,.{self.decimals}f}" if self.thousands else f"{abs(value):.{self.decimals}f}"
        if self.euro_style:
            s = s.replace(",", "\0").replace(".", ",").replace("\0", ".")
        if value < 0:
            s = "-" + s
        if self.symbol and (self.symbol_mode == "every" or (self.symbol_mode == "totals" and with_symbol)):
            s = self.symbol + (" " if self.space else "") + s
        return s


def price_value(rng, scale=1.0):
    v = float(np.exp(rng.normal(1.8, 1.1))) * scale
    return round(v, 2)


def date_value(rng):
    y, m, d = int(rng.integers(2015, 2031)), int(rng.integers(1, 13)), int(rng.integers(1, 29))
    return y, m, d


def fmt_date(rng, y, m, d):
    style = int(rng.integers(0, 10))
    mon = MONTHS[m - 1]
    return [
        f"{m:02d}/{d:02d}/{y}", f"{d:02d}/{m:02d}/{y}", f"{y}-{m:02d}-{d:02d}", f"{d:02d}.{m:02d}.{y}",
        f"{d} {mon[:3]} {y}", f"{mon} {d}, {y}", f"{d:02d}-{mon[:3].upper()}-{y}", f"{m}/{d}/{str(y)[2:]}",
        f"{d:02d}/{m:02d}/{str(y)[2:]}", f"{mon[:3]} {d}, {y}",
    ][style]


def fmt_time(rng):
    h, mi, s = int(rng.integers(0, 24)), int(rng.integers(0, 60)), int(rng.integers(0, 60))
    r = rng.random()
    if r < 0.4:
        return f"{h:02d}:{mi:02d}"
    if r < 0.7:
        return f"{h:02d}:{mi:02d}:{s:02d}"
    return f"{(h % 12) or 12}:{mi:02d} {'AM' if h < 12 else 'PM'}"


def invoice_no(rng):
    r = rng.random()
    n = int(rng.integers(100, 99999999))
    if r < 0.3:
        return str(n)
    if r < 0.55:
        return f"INV-{rng.integers(2015, 2031)}-{rng.integers(1, 99999):05d}"
    if r < 0.75:
        return f"#{n}"
    letters = "".join(rng.choice(list(string.ascii_uppercase), int(rng.integers(1, 4))))
    return f"{letters}{rng.integers(1000, 999999)}"


def phone(rng):
    a, b, c = rng.integers(200, 999), rng.integers(200, 999), rng.integers(1000, 9999)
    return pick(rng, [f"({a}) {b}-{c}", f"{a}-{b}-{c}", f"+1 {a} {b} {c}", f"+91 {rng.integers(70000, 99999)} {rng.integers(10000, 99999)}",
                      f"0{rng.integers(10, 99)} {rng.integers(1000, 9999)} {rng.integers(1000, 9999)}", f"{a}.{b}.{c}"])


def tax_id(rng):
    u = string.ascii_uppercase
    return pick(rng, [
        f"{rng.integers(10, 37)}{''.join(rng.choice(list(u), 5))}{rng.integers(1000, 9999)}{rng.choice(list(u))}1Z{rng.integers(0, 9)}",
        f"GB{rng.integers(100000000, 999999999)}", f"{rng.integers(10, 99)}-{rng.integers(1000000, 9999999)}",
        f"DE{rng.integers(100000000, 999999999)}", f"{rng.integers(10, 99)} {rng.integers(100, 999)} {rng.integers(100, 999)} {rng.integers(100, 999)}",
    ])


def merchant_name(rng):
    r = rng.random()
    if r < 0.5:
        n = f"{pick(rng, NAME_A)} {pick(rng, NAME_B)}"
    elif r < 0.75:
        n = f"{pick(rng, LAST)}'s {pick(rng, NAME_B)}"
    else:
        n = f"{pick(rng, NAME_A)} {pick(rng, LAST)} {pick(rng, COMPANY_SUFFIX)}".strip()
    return n.upper() if rng.random() < 0.4 else n


def person_name(rng):
    return f"{pick(rng, FIRST)} {pick(rng, LAST)}"


def address_lines(rng):
    l1 = f"{rng.integers(1, 9999)} {pick(rng, STREETS)} {pick(rng, STREET_T)}"
    if rng.random() < 0.25:
        l1 += pick(rng, [f", Suite {rng.integers(1, 999)}", f", Unit {rng.integers(1, 99)}", f", Floor {rng.integers(1, 30)}"])
    l2 = f"{pick(rng, CITIES)}, {pick(rng, STATES)} {rng.integers(10000, 99999)}"
    return [l1, l2]


def item_name(rng):
    n = pick(rng, ITEMS)
    if rng.random() < 0.45:
        n = f"{pick(rng, ITEM_MODS)} {n}"
    if rng.random() < 0.15:
        n += f" {rng.integers(1, 1000)}{pick(rng, ['g', 'ml', 'oz', ''])}"
    if rng.random() < 0.15:
        n = f"{rng.integers(100000, 999999)} {n}"  # SKU prefix
    r = rng.random()
    return n.upper() if r < 0.45 else (n.title() if r < 0.8 else n)


def random_chars(rng, n_min=3, n_max=20):
    """Uniform over the whole vocabulary so rare symbols are seen."""
    n = int(rng.integers(n_min, n_max + 1))
    return "".join(rng.choice(list(CHARS), n))


def line_text(rng):
    """Mixed text for recognizer training lines."""
    r = rng.random()
    mf = MoneyFormat(rng)
    if r < 0.18:
        return sentence(rng, 2, 10)
    if r < 0.30:
        return sentence(rng, 6, 14)   # long prose lines, as in book and newspaper columns
    if r < 0.42:
        return f"{item_name(rng)}"
    if r < 0.52:
        return mf.fmt(price_value(rng), with_symbol=rng.random() < 0.5)
    if r < 0.58:
        return fmt_date(rng, *date_value(rng)) + (f" {fmt_time(rng)}" if rng.random() < 0.3 else "")
    if r < 0.63:
        return pick(rng, sum(KEYS.values(), [])).replace("{r}", str(rng.integers(1, 28)))
    if r < 0.68:
        return pick(rng, [invoice_no(rng), phone(rng), tax_id(rng)])
    if r < 0.72:
        return merchant_name(rng)
    if r < 0.76:
        return pick(rng, address_lines(rng))
    if r < 0.82:
        return f"{rng.integers(1, 20)} x {item_name(rng)}"[:40]
    if r < 0.86:
        return " ".join(number(rng) for _ in range(int(rng.integers(1, 5))))
    if r < 0.90:
        return f"{word(rng)}@{word(rng).lower()}.{pick(rng, ['com', 'org', 'net', 'in', 'co.uk'])}"
    return random_chars(rng)


# ------------------------------------------------------------------------------------------ inline maths
# Written with ^{..} / _{..} markup; math_runs() splits it into (text, kind) runs so the renderer can typeset
# sub/superscripts small and raised/lowered, and the label spells them as charset.scripted() does for real data.
MATH_VARS = list("xyzabcnkmtpqrfgh")
GREEK_VARS = list("αβγδεθλμσφωπτρΔΣΩ")
GREEK_ALL = list("αβγδεζηθκλμνξπρστφχψωΓΔΘΛΞΠΣΦΨΩ")
CHEM = ["H_{2}O", "CO_{2}", "O_{2}", "N_{2}", "CH_{4}", "H_{2}SO_{4}", "NaHCO_{3}", "Fe_{2}O_{3}", "NO_{x}", "SO_{2}",
        "NH_{3}", "CaCO_{3}", "C_{6}H_{12}O_{6}", "H_{2}", "Al_{2}O_{3}", "SiO_{2}", "PO_{4}^{3-}", "Ca^{2+}", "Na^{+}",
        "Cl^{-}", "Fe^{3+}", "CO_{2}e", "PM_{2.5}", "PM_{10}"]
MATH_UNITS = ["m^{2}", "km^{2}", "cm^{2}", "m^{3}", "cm^{3}", "m/s^{2}", "kg m^{-3}", "s^{-1}", "mol^{-1}", "J mol^{-1}",
              "W m^{-2}", "g cm^{-3}", "mm^{2}", "ft^{2}", "μm", "μg", "μs", "Ω", "kΩ", "°C", "m s^{-1}", "h^{-1}"]
RELATIONS = ["=", "=", "=", "≤", "≥", "≠", "≈", "<", ">", "∝", "≡"]
SET_OPS = ["∈", "∉", "⊂", "⊆", "∪", "∩", "→", "←", "↔", "⇒", "⇔", "⊥", "∥"]
STAT_NAMES = ["R^{2}", "r", "p", "α", "χ^{2}", "σ", "μ", "β_{1}", "β_{0}", "n", "t", "F", "ρ", "κ", "η^{2}", "σ^{2}"]
_MARKUP = re.compile(r"([_^])\{([^{}]*)\}")


def _mvar(rng):
    return pick(rng, GREEK_VARS) if rng.random() < 0.3 else pick(rng, MATH_VARS)


def _mnum(rng):
    r = rng.random()
    if r < 0.5:
        return str(int(rng.integers(0, 100)))
    if r < 0.8:
        return f"{rng.uniform(0, 100):.{int(rng.integers(1, 4))}f}"
    return f"{rng.uniform(0, 1):.{int(rng.integers(2, 4))}f}"


def _sup(rng):
    return pick(rng, ["2", "2", "2", "3", "n", "-1", "-2", "4", "i", "n+1", "2n", "(n)", str(int(rng.integers(5, 30))),
                      f"-{rng.integers(3, 20)}", "T", "k", "*"])


def _sub(rng):
    return pick(rng, ["0", "1", "2", "3", "i", "j", "k", "n", "x", "t", "max", "min", "ij", "n-1", "i+1", "c", "B",
                      "eff", "L", "D", "s", "p", "m", "12", "ref", "0,1"])


def math_expr(rng):
    v, w = _mvar(rng), _mvar(rng)
    forms = [
        lambda: f"{v}^{{{_sup(rng)}}}",
        lambda: f"{v}_{{{_sub(rng)}}}",
        lambda: f"{v}_{{{_sub(rng)}}}^{{{_sup(rng)}}}",
        lambda: f"{v}^{{2}} + {w}^{{2}} = {_mvar(rng)}^{{2}}",
        lambda: f"{_mnum(rng)} × 10^{{{pick(rng, ['-', ''])}{rng.integers(1, 24)}}}",
        lambda: f"10^{{{pick(rng, ['-', ''])}{rng.integers(1, 12)}}}",
        lambda: f"{_mnum(rng)} {pick(rng, MATH_UNITS)}",
        lambda: pick(rng, CHEM),
        lambda: f"{pick(rng, STAT_NAMES)} {pick(rng, RELATIONS)} {_mnum(rng)}",
        lambda: f"{v} {pick(rng, RELATIONS)} {_mnum(rng)}",
        lambda: f"{v} {pick(rng, SET_OPS)} {w}",
        lambda: f"{pick(rng, ['∀', '∃'])}{v} {pick(rng, ['∈', '>', '<', '≥'])} {w}",
        lambda: f"f({v}) = {rng.integers(2, 9)}{v}^{{{pick(rng, ['2', '3'])}}} {pick(rng, ['+', '-'])} "
                f"{rng.integers(1, 9)}{v} {pick(rng, ['+', '-'])} {rng.integers(1, 9)}",
        lambda: "√" + pick(rng, [str(rng.integers(2, 50)), v, f"({v}^{{2}} + {w}^{{2}})"]),
        lambda: f"Σ {v}_{{{pick(rng, ['i', 'k', 'n'])}}}",
        lambda: f"∫ f({v}) d{v}",
        lambda: f"∂{v}/∂{w}",
        lambda: f"∇{v}",
        lambda: f"Δ{v} = {_mnum(rng)}",
        lambda: f"{v} → ∞",
        lambda: f"∠{pick(rng, ['ABC', 'PQR', 'XYZ'])} = {rng.integers(10, 170)}°",
        lambda: f"{pick(rng, GREEK_ALL)} {pick(rng, RELATIONS)} {_mnum(rng)}"
                f"{pick(rng, ['', '', ' nm', ' s', ' Hz', '°', ' rad', ' K', ' ms'])}",
        lambda: "E = mc^{2}",
        lambda: f"{v}_{{1}}, {v}_{{2}}, …, {v}_{{n}}",
        lambda: f"{v}_{{{_sub(rng)}}} {pick(rng, ['+', '-', '×', '/'])} {w}_{{{_sub(rng)}}}",
        lambda: f"({v} + {w})^{{{pick(rng, ['2', '3', 'n'])}}}",
    ]
    return forms[int(rng.integers(len(forms)))]()


def script_mark(rng):
    """Footnote markers, ordinals and trademarks: small raised text inside ordinary prose."""
    r = rng.random()
    if r < 0.4:
        return f"{word(rng)}^{{{rng.integers(1, 40)}}}"
    if r < 0.55:
        return f"{word(rng)}^{{{pick(rng, ['*', 'a', 'b', 'c', '1,2', '3-5'])}}}"
    if r < 0.8:
        n = int(rng.integers(1, 32))
        suffix = "th" if 10 <= n % 100 <= 20 else {1: "st", 2: "nd", 3: "rd"}.get(n % 10, "th")
        return f"{n}^{{{suffix}}}"
    return f"{word(rng).capitalize()}^{{{pick(rng, ['TM', '®', '©'])}}}"


def math_line(rng):
    """Markup for one line of text with inline maths."""
    r = rng.random()
    if r < 0.3:
        return math_expr(rng)
    if r < 0.45:
        return f"{math_expr(rng)}{pick(rng, [', ', '; ', ' and '])}{math_expr(rng)}"
    if r < 0.8:
        return f"{sentence(rng, 1, 5).rstrip('.?!:')} {math_expr(rng)} {sentence(rng, 1, 4).lower()}"
    return f"{sentence(rng, 1, 5).rstrip('.?!:')} {script_mark(rng)} {sentence(rng, 1, 4).lower()}"


def math_runs(markup: str) -> list[tuple[str, str]]:
    """'H_{2}O' -> [('H', ''), ('2', '_'), ('O', '')]."""
    runs, pos = [], 0
    for m in _MARKUP.finditer(markup):
        if m.start() > pos:
            runs.append((markup[pos:m.start()], ""))
        runs.append((m.group(2), m.group(1)))
        pos = m.end()
    if pos < len(markup):
        runs.append((markup[pos:], ""))
    return runs


def runs_label(runs) -> str:
    return "".join(scripted(t, k) if k else t for t, k in runs)