"""Natural-language instruction templates for Flash-Instruct.
labels: dict(lang=<ibo|yor|hau|pcm|eng|cs:<l1>+<l2>>, rate=slow|normal|fast|None, pause=few|normal|many|None,
             energy=calm|neutral|lively|None, gender=male|female|None)
make(labels, rng) -> one instruction string (many surface forms per label combination).
DEFAULT = the original CosyVoice3 system prompt; ~18% of training items keep it.
"""
import random

DEFAULT = "You are a helpful assistant."
LN = {"ibo": "Igbo", "yor": "Yoruba", "hau": "Hausa", "pcm": "Nigerian Pidgin", "eng": "Nigerian English"}

RATE = {
    "slow": ["slowly", "at an unhurried pace", "at a slow, relaxed pace", "taking your time", "at a gentle, slow pace",
             "a little slower than usual", "slowly and deliberately"],
    "normal": ["at a natural pace", "at a normal speaking pace", "at a steady, everyday pace", "at a comfortable pace"],
    "fast": ["quickly", "at a brisk pace", "fast, like quick everyday conversation", "at a quick pace",
             "a little faster than usual", "briskly"],
}
PAUSE = {
    "few": ["with few pauses", "in one smooth flow with hardly any pauses", "without stopping much",
            "with minimal pausing", "running the phrases together"],
    "normal": ["with natural pauses", "with normal phrasing", "pausing where it feels natural"],
    "many": ["with plenty of pauses", "pausing often between phrases", "with frequent pauses",
             "with lots of breaks between phrases", "with deliberate pauses"],
}
ENERGY = {
    "calm": ["calm", "soft and calm", "gentle, even-toned", "relaxed and low-key", "quiet and steady"],
    "neutral": ["neutral", "plain, matter-of-fact", "even", "straightforward"],
    "lively": ["lively", "animated", "expressive and energetic", "bright and lively", "upbeat, expressive"],
}
GENDER = {"male": ["male", "man's"], "female": ["female", "woman's"]}


def art(w):
    return ("an " if w[:1].lower() in "aeiou" else "a ") + w


def lang_phrase(lang, rng):
    if lang.startswith("cs:"):
        a, b = lang[3:].split("+")
        A, B = LN.get(a, a), LN.get(b, b).replace("Nigerian English", "English")
        return rng.choice([f"an {A}-{B} mix" if A[0] in "AEIOU" else f"a {A}-{B} mix", f"mixed {A} and {B}",
                           f"{A} with some {B} mixed in", f"code-switched {A} and {B}"])
    return LN.get(lang, lang)


def make(lab, rng=random):
    parts = {}
    if lab.get("rate"):
        parts["rate"] = rng.choice(RATE[lab["rate"]])
    if lab.get("pause"):
        parts["pause"] = rng.choice(PAUSE[lab["pause"]])
    if lab.get("energy"):
        parts["energy"] = rng.choice(ENERGY[lab["energy"]])
    lang = lang_phrase(lab["lang"], rng) if lab.get("lang") else None
    g = rng.choice(GENDER[lab["gender"]]) if lab.get("gender") else None
    tone = parts.get("energy")
    adv = [parts[k] for k in ("rate", "pause") if k in parts]
    advs = ", ".join(adv)
    T = []
    # template family 1: imperative sentence
    s = "Speak"
    if lang:
        s += f" in {lang}" + (f", sounding {tone}," if tone else "")
    elif tone:
        s += f" in {art(tone)} tone"
    if advs:
        s += f" {advs}"
    T.append(s + ".")
    # family 2: read-this style
    s = "Read this"
    if advs:
        s += f" {advs}"
    if tone:
        s += f", in {art(tone)} voice"
    if lang:
        s += f". The language is {lang}"
    T.append(s + ".")
    # family 3: descriptive/caption style
    d = []
    if g:
        d.append(f"A {g} voice")
    else:
        d.append("A speaker")
    if lang:
        d.append(f"speaking {lang}")
    s = " ".join(d)
    if advs:
        s += f", {advs}"
    if tone:
        s += f"; the delivery is {tone}" if (lang or advs) else f" with {art(tone)} delivery"
    T.append(s + ".")
    # family 4: tag list
    tags = []
    if lang: tags.append(f"language: {lang}")
    if "rate" in parts: tags.append(f"pace: {lab['rate']}")
    if "pause" in parts: tags.append(f"pauses: {lab['pause']}")
    if tone: tags.append(f"tone: {lab['energy']}")
    if g: tags.append(f"voice: {lab['gender']}")
    if tags:
        T.append("Style - " + "; ".join(tags) + ".")
    # family 5: polite request
    s = "Please say this"
    if lang:
        s += f" in {lang}"
    if tone:
        s += f" with {art(tone)} delivery"
    if advs:
        s += f", {advs}"
    T.append(s + ".")
    out = rng.choice(T)
    return out[0].upper() + out[1:]


def sample_labels(full, rng, p=None):
    """drop each control independently so instructions mention subsets; keep at least one."""
    p = p or dict(lang=0.55, rate=0.75, pause=0.7, energy=0.6, gender=0.3)
    while True:
        lab = {k: (v if (v and rng.random() < p.get(k, 0.5)) else None) for k, v in full.items()}
        if any(lab.values()):
            return lab


if __name__ == "__main__":
    r = random.Random(0)
    full = dict(lang="cs:ibo+eng", rate="slow", pause="many", energy="calm", gender="female")
    for _ in range(12):
        print(make(sample_labels(full, r), r))
    for l in ["ibo", "yor", "hau", "pcm", "eng"]:
        print(make(dict(lang=l, rate="fast", pause="few", energy="lively"), r))