import re, json

CONTRACTIONS = {
    "I'M": "I AM", "I'LL": "I WILL", "I'VE": "I HAVE",
    "DON'T": "DO NOT", "DOESN'T": "DOES NOT", "CAN'T": "CANNOT",
    "WON'T": "WILL NOT", "IT'S": "IT IS", "HE'S": "HE IS",
    "SHE'S": "SHE IS", "THEY'RE": "THEY ARE", "WE'RE": "WE ARE",
    "YOU'RE": "YOU ARE", "THAT'S": "THAT IS", "LET'S": "LET US",
    "WASN'T": "WAS NOT", "WEREN'T": "WERE NOT",
    "HAVEN'T": "HAVE NOT", "HASN'T": "HAS NOT",
    "WOULDN'T": "WOULD NOT", "COULDN'T": "COULD NOT",
}

def load_ke_model(model_dir):
    """Load saved Orpheus Kenyan English model for inference."""
    from unsloth import FastLanguageModel
    import torch

    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name     = model_dir,
        max_seq_length = 2048,
        dtype          = None,
        load_in_4bit   = True,
    )
    FastLanguageModel.for_inference(model)

    with open(f"{model_dir}/ke_lexicon.json") as f:
        lexicon = json.load(f)

    return model, tokenizer, lexicon


def apply_g2p(text, lexicon):
    """Convert text to Kenyan English phonemes using saved lexicon."""
    words = text.split()
    expanded = []
    for w in words:
        clean = re.sub(r"^[^\w\']+|[^\w\']+$", "", w).upper()
        if clean in CONTRACTIONS:
            expanded.extend(CONTRACTIONS[clean].split())
        else:
            expanded.append(w)
    parts = []
    for token in expanded:
        word = re.sub(r"^[^\w\']+|[^\w\']+$", "", token).upper()
        if not word:
            continue
        parts.append(lexicon.get(word, word))
    return " || ".join(parts)


def synthesise(text, model, tokenizer, lexicon, snac_model,
               temperature=0.6, top_p=0.95):
    """Full TTS pipeline: text -> Kenyan phonemes -> audio."""
    import torch, numpy as np

    START_OF_HUMAN  = 128259
    END_OF_TEXT     = 128009
    END_OF_HUMAN    = 128260
    START_OF_SPEECH = 128257
    END_OF_SPEECH   = 128258

    phoneme_text = apply_g2p(text, lexicon)
    print(f"  G2P: {phoneme_text}")

    input_ids = tokenizer.encode(phoneme_text, return_tensors="pt")
    prompt = torch.cat([
        torch.tensor([[START_OF_HUMAN]]),
        input_ids,
        torch.tensor([[END_OF_TEXT, END_OF_HUMAN]]),
    ], dim=1).to("cuda")

    with torch.no_grad():
        gen_ids = model.generate(
            input_ids          = prompt,
            attention_mask     = torch.ones_like(prompt),
            max_new_tokens     = 1200,
            do_sample          = True,
            temperature        = temperature,
            top_p              = top_p,
            repetition_penalty = 1.1,
            eos_token_id       = END_OF_SPEECH,
            use_cache          = True,
        )

    speech_start = (gen_ids == START_OF_SPEECH).nonzero(as_tuple=True)
    if len(speech_start[1]) == 0:
        print("  WARNING: no speech tokens generated")
        return np.zeros(24000)

    cropped   = gen_ids[:, speech_start[1][-1].item() + 1:]
    row       = cropped[0][cropped[0] != END_OF_SPEECH]
    trimmed   = row[:(len(row) // 7) * 7]
    code_list = [t.item() - 128266 for t in trimmed]

    def redistribute(codes):
        l1, l2, l3 = [], [], []
        for i in range(len(codes) // 7):
            l1.append(max(0, min(4095, codes[7*i])))
            l2.append(max(0, min(4095, codes[7*i+1] - 4096)))
            l3.append(max(0, min(4095, codes[7*i+2] - 2*4096)))
            l3.append(max(0, min(4095, codes[7*i+3] - 3*4096)))
            l2.append(max(0, min(4095, codes[7*i+4] - 4*4096)))
            l3.append(max(0, min(4095, codes[7*i+5] - 5*4096)))
            l3.append(max(0, min(4095, codes[7*i+6] - 6*4096)))
        return (
            torch.tensor(l1).unsqueeze(0),
            torch.tensor(l2).unsqueeze(0),
            torch.tensor(l3).unsqueeze(0),
        )

    snac_model.to("cpu")
    t1, t2, t3 = redistribute(code_list)
    with torch.inference_mode():
        audio_hat = snac_model.decode([t1, t2, t3])

    return audio_hat.squeeze().cpu().numpy()


# Usage:
# from snac import SNAC
# snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz")
# model, tokenizer, lexicon = load_ke_model("./orpheus-ke-lora")
# audio = synthesise("Habari, how are you today?", model, tokenizer, lexicon, snac_model)
# import soundfile as sf
# sf.write("output.wav", audio, 24000)