# -*- coding: utf-8 -*-
"""
Z-Born 4.5M — живой чат на CPU.

  pip install torch tokenizers
  python demo.py "Once upon a time"
"""
import os
import sys

os.environ['CUDA_VISIBLE_DEVICES'] = ''

if hasattr(sys.stdout, 'reconfigure'):
    sys.stdout.reconfigure(encoding='utf-8', errors='replace')

import torch
import torch.nn as nn

HERE = os.path.dirname(os.path.abspath(__file__))


class BornFactorRNN(nn.Module):
    """1L tanh-RNN, рождённый в Z-форме: Wx/Wh ранга k (U,S,V)."""

    def __init__(self, v, hid, k):
        super().__init__()
        self.v, self.hid, self.k = v, hid, k
        self.Ux = nn.Parameter(torch.zeros(hid, k))
        self.Vx = nn.Parameter(torch.zeros(v, k))
        self.Sx = nn.Parameter(torch.zeros(k))
        self.Uh = nn.Parameter(torch.zeros(hid, k))
        self.Vh = nn.Parameter(torch.zeros(hid, k))
        self.Sh = nn.Parameter(torch.zeros(k))
        self.b_ih = nn.Parameter(torch.zeros(hid))
        self.b_hh = nn.Parameter(torch.zeros(hid))
        self.head = nn.Linear(hid, v)

    def hidden(self, xt, h):
        xin = (self.Vx[xt] * self.Sx) @ self.Ux.T
        hh = ((h @ self.Vh) * self.Sh) @ self.Uh.T
        return torch.tanh(xin + hh + self.b_ih + self.b_hh)

    def step(self, xt, h):
        h = self.hidden(xt, h)
        return self.head(h), h

    @torch.no_grad()
    def generate(self, ids, n, temp=0.8, top_k=40, rep=1.12):
        self.eval()
        h = torch.zeros(1, self.hid)
        logits = None
        for tok in ids:
            logits, h = self.step(torch.tensor([tok]), h)
        out = list(ids)
        for _ in range(n):
            lg = logits[0]
            if rep and out:
                lg = lg.clone()
                for t in set(out[-48:]):
                    lg[t] = lg[t] / rep if lg[t] > 0 else lg[t] * rep
            lg = lg / max(temp, 1e-6)
            if top_k and top_k < lg.numel():
                thr = torch.topk(lg, top_k).values[-1]
                lg = lg.masked_fill(lg < thr, float('-inf'))
            nxt = torch.multinomial(torch.softmax(lg, -1), 1)
            out.append(int(nxt))
            logits, h = self.step(nxt.view(1), h)
        return out[len(ids):]


def load():
    ckpt = os.path.join(HERE, 'zborn.pt')
    if not os.path.isfile(ckpt):
        raise SystemExit('нет %s' % ckpt)
    pack = torch.load(ckpt, map_location='cpu', weights_only=False)
    meta = pack['meta']
    from tokenizers import Tokenizer
    tok = Tokenizer.from_file(os.path.join(HERE, 'tokenizer.json'))
    m = BornFactorRNN(meta['v'], meta['hid'], meta['k'])
    m.load_state_dict(pack['model'])
    m.eval()
    return m, tok, meta


def main():
    prompt = ' '.join(sys.argv[1:]).strip() or 'The meaning of life is'
    m, tok, meta = load()
    print('Z-Born %s  hid=%s k=%s n=%s  step=%s  CPU'
          % (meta.get('kind', '?'), meta.get('hid'), meta.get('k'),
             meta.get('n_born'), meta.get('step')))
    ids = tok.encode(prompt).ids
    out = m.generate(ids, n=120)
    print('PROMPT:', prompt)
    print(tok.decode(out))


if __name__ == '__main__':
    main()