import torch
import torch.nn.functional as F
import sys

ALPHABET = "0123456789+-*=()| "
N_DIM = len(ALPHABET)
CHAR_TO_INT = {c: i for i, c in enumerate(ALPHABET)}
INT_TO_CHAR = {i: c for i, c in enumerate(ALPHABET)}
SPACE_IDX = CHAR_TO_INT.get(' ', 0)
CONTEXT_SIZE = 64
MAX_GEN = 60

def generate(model, prompt_text):
    clean = prompt_text.replace(' ', '')
    stripped = ''.join(c for c in clean if c in CHAR_TO_INT)
    if '=' not in stripped:
        return prompt_text + '?'
    prompt = stripped[:stripped.index('=') + 1]
    generated = ''
    for _ in range(MAX_GEN):
        prefix = (prompt + generated).rjust(CONTEXT_SIZE, ' ')[:CONTEXT_SIZE]
        indices = torch.tensor([CHAR_TO_INT.get(c, SPACE_IDX) for c in prefix], dtype=torch.long)
        x = F.one_hot(indices, num_classes=N_DIM).float().view(1, -1)
        with torch.no_grad():
            logits = model(x)
        next_token = torch.argmax(logits, dim=-1, keepdim=True)
        char = INT_TO_CHAR.get(next_token.item(), '')
        if not char or char == ' ' or char == '|':
            break
        generated += char
    return prompt_text[:prompt_text.index('=') + 1] + generated

def main():
    if len(sys.argv) < 2:
        print("Usage: python infer.py <model.pt> [prompt]")
        print("       python infer.py <model.pt> (interactive mode)")
        sys.exit(1)

    model_path = sys.argv[1]
    model = torch.jit.load(model_path)
    model.eval()
    print(f"Loaded model from {model_path}", file=sys.stderr)

    if len(sys.argv) >= 3:
        prompt = sys.argv[2]
        result = generate(model, prompt)
        print(result)
    else:
        while True:
            try:
                prompt = input("Prompt > ")
                if prompt.lower() == 'q':
                    break
                print(generate(model, prompt))
            except (EOFError, KeyboardInterrupt):
                break
            except Exception as e:
                print(f"Error: {e}")

if __name__ == "__main__":
    main()