"""آموزش ترنسفورمر کوچک نویسه‌ای از وزن‌های تصادفی، فقط روی data/train.txt و data/site_memory.txt (حافظهٔ سایت)."""
from pathlib import Path
import json, math, os, random, torch
from model_model import TinyGPT, encode
from text_utils import BLOCK, load_training_lines

BASE = Path(__file__).resolve().parent
MODEL = BASE / 'model'

# ---- تنظیمات (می‌توانید تغییر دهید) ----
# BLOCK (طول بافت، نویسه) در text_utils.py تنظیم می‌شود؛ بهتر است از طول بلندترین خط کمتر نباشد
EMBED, HEADS, LAYERS = 128, 4, 3
STEPS = 4000          # حداکثر تعداد گام
BATCH = 32
LR = 1e-3
TARGET_CHARS = 40000  # داده‌ی کم را تا این حجم تکرار می‌کنیم تا مدل زیاد ببیند
STOP_LOSS = 0.05      # اگر loss میانگین به این رسید، زودتر تمام می‌شود


def main():
    samples = [line + '\n' for line in load_training_lines()]
    if len(samples) < 5:
        raise SystemExit('حداقل ۵ خط در data/train.txt بنویسید (بهتر: چندصد خط).')
    longest = max(map(len, samples))
    if longest > BLOCK:
        print(f'هشدار: بلندترین خط {longest} نویسه است و از BLOCK={BLOCK} بزرگ‌تر است؛ BLOCK را بیشتر کنید.')

    random.seed(1337); torch.manual_seed(1337)
    reps = max(1, math.ceil(TARGET_CHARS / sum(map(len, samples))))
    lines = []
    for _ in range(reps):
        random.shuffle(samples); lines += samples
    text = ''.join(lines)

    chars = sorted(set(text)); vocab = {c: i for i, c in enumerate(chars)}
    ids = torch.tensor(encode(text, vocab), dtype=torch.long)
    nl = vocab['\n']
    # اندیس شروع خط‌ها: در زمان پرسش، متن همیشه از ابتدای یک خط شروع می‌شود
    starts = [0] + [i + 1 for i in (ids == nl).nonzero().flatten().tolist()]
    starts = torch.tensor([s for s in starts if s + BLOCK + 1 <= len(ids)])

    def get_batch():
        rand = torch.randint(0, len(ids) - BLOCK - 1, (BATCH,))
        line = starts[torch.randint(0, len(starts), (BATCH,))]
        use_line = torch.rand(BATCH) < 0.5
        ix = torch.where(use_line, line, rand)
        x = torch.stack([ids[i:i + BLOCK] for i in ix])
        y = torch.stack([ids[i + 1:i + BLOCK + 1] for i in ix])
        return x, y

    model = TinyGPT(len(vocab), BLOCK, EMBED, HEADS, LAYERS)
    opt = torch.optim.AdamW(model.parameters(), lr=LR)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=STEPS, eta_min=LR / 10)
    print(f'{len(samples)} خط | {len(text)} نویسه پس از تکرار | {len(vocab)} نویسهٔ یکتا')
    avg = None
    for step in range(1, STEPS + 1):
        x, y = get_batch()
        loss = torch.nn.functional.cross_entropy(model(x).reshape(-1, len(vocab)), y.reshape(-1))
        opt.zero_grad(set_to_none=True); loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0); opt.step(); sched.step()
        avg = loss.item() if avg is None else .98 * avg + .02 * loss.item()
        if step % 200 == 0:
            print(f'step {step}/{STEPS} | loss {avg:.3f}')
        if step >= 500 and avg < STOP_LOSS:
            print(f'loss به {avg:.3f} رسید؛ آموزش زودتر تمام شد.'); break

    MODEL.mkdir(exist_ok=True)
    # اول فایل موقت، بعد جایگزینی؛ تا سرور در حال اجرا هرگز فایل نیمه‌نوشته نخواند
    tmp_vocab, tmp_weights = MODEL / 'vocab.json.tmp', MODEL / 'weights.pt.tmp'
    tmp_vocab.write_text(json.dumps(vocab, ensure_ascii=False), encoding='utf-8')
    torch.save({'state_dict': model.state_dict(), 'block_size': BLOCK, 'n_embd': EMBED,
                'n_head': HEADS, 'n_layer': LAYERS}, tmp_weights)
    os.replace(tmp_vocab, MODEL / 'vocab.json')
    os.replace(tmp_weights, MODEL / 'weights.pt')
    print('مدل ذخیره شد. برای تست: python ask.py "سلام"   یا   python ask.py --test')


if __name__ == '__main__':
    main()