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 _run(model, prompt):
    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 + generated

def _count_operators(expr):
    count = 0
    for c in expr:
        if c in '+-*':
            count += 1
    return count

def _strip_outer_parens(expr):
    while expr and expr[0] == '(' and expr[-1] == ')':
        depth = 0
        balanced = True
        for i, c in enumerate(expr):
            if c == '(':
                depth += 1
            elif c == ')':
                depth -= 1
            if depth == 0 and i < len(expr) - 1:
                balanced = False
                break
        if balanced:
            expr = expr[1:-1]
        else:
            break
    return expr

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]
    expr = _strip_outer_parens(stripped[:stripped.index('=')])
    if _count_operators(expr) == 1:
        wrapped = f"0+({expr})="
        chain = _run(model, wrapped)
        parts = chain.split('=')
        answer = parts[-1].rstrip('| ')
        if answer and answer != expr and not answer.startswith('-'):
            return prompt_text[:prompt_text.index('=') + 1] + answer
    return _run(model, prompt)

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()