import json, random, sys
from pathlib import Path
import torch
from kokoro import KModel

TEACHER, CORPUS, PACK, OUT, COUNT = sys.argv[1:6]
MAX_PHONEMES = 500
device = "cuda" if torch.cuda.is_available() else "cpu"
model = KModel(repo_id="hexgrad/Kokoro-82M", config=TEACHER + "/config.json", model=TEACHER + "/kokoro.pth").to(device).eval()
model.decoder = None
pack = torch.load(PACK, map_location=device, weights_only=True)
sentences = [json.loads(line)["phonemes"] for line in open(CORPUS, encoding="utf-8")]
rng = random.Random(7)


def paragraph():
    parts = [rng.choice(sentences)]
    while True:
        candidate = rng.choice(sentences)
        if len(" ".join(parts + [candidate])) > MAX_PHONEMES:
            return " ".join(parts)
        parts.append(candidate)


with open(OUT, "w") as sink, torch.no_grad():
    for _ in range(int(COUNT)):
        phonemes = paragraph()
        ids = [0, *[model.vocab[p] for p in phonemes if p in model.vocab], 0]
        tokens = torch.LongTensor([ids]).to(device)
        lengths = torch.LongTensor([len(ids)]).to(device)
        mask = torch.zeros(1, len(ids), dtype=torch.bool, device=device)
        d_en = model.bert_encoder(model.bert(tokens, attention_mask=(~mask).int())).transpose(-1, -2)
        d = model.predictor.text_encoder(d_en, pack[len(phonemes) - 1][:, 128:], lengths, mask)
        x, _ = model.predictor.lstm(d)
        dur = torch.round(torch.sigmoid(model.predictor.duration_proj(x)).sum(-1)).clamp(min=1).long().reshape(-1).tolist()
        sink.write(json.dumps({"ids": ids, "dur": dur, "lang": "pt"}) + "\n")
print("CONTEXT_DUMP_DONE", flush=True)