import torch
import os
from model import GPTLanguageModel

device = 'cuda' if torch.cuda.is_available() else 'cpu'

checkpoint_path = 'model_checkpoint.pth'
if not os.path.exists(checkpoint_path):
    print("Checkpoint not found. Please run train.py first.")
    exit(1)

# Load the checkpoint to get chars and model state
checkpoint = torch.load(checkpoint_path, map_location=device)
chars = checkpoint['chars']
vocab_size = len(chars)

# Recreate the exact same mapping
stoi = { ch:i for i,ch in enumerate(chars) }
itos = { i:ch for i,ch in enumerate(chars) }
encode = lambda s: [stoi[c] for c in s]
decode = lambda l: ''.join([itos[i] for i in l])

# Initialize the model with the same hyperparameters used in training
n_embd = 128
n_head = 4
n_layer = 8
block_size = 64
dropout = 0.1

model = GPTLanguageModel(vocab_size, n_embd, n_head, n_layer, block_size, dropout)
model.load_state_dict(checkpoint['model_state_dict'])
model.to(device)
model.eval()

# Interactive Chat Mode
print("\n--- Ինտերակտիվ Զրույց (Interactive Chat) ---")
print("Գրեք 'quit' դուրս գալու համար:")
print("ԶԳՈՒՇԱՑՈՒՄ: Մոդելը ճանաչում է միայն այն տառերը, որոնք կան data/ թղթապանակում:\n")

while True:
    user_input = input("ՄԱՐԴ: ")
    if user_input.lower() == 'quit':
        break
    
    prompt = f"ՄԱՐԴ: {user_input}\nԱԻ:"
    
    try:
        context = torch.tensor([encode(prompt)], dtype=torch.long, device=device)
    except KeyError as e:
        print(f"[Սխալ: '{e.args[0]}' տառը գոյություն չունի մոդելի բառարանում: Փորձեք այլ բառ:]")
        continue
        
    # Գեներացնում ենք 100 նոր տառ
    generated_output = decode(model.generate(context, max_new_tokens=100)[0].tolist())
    
    # Առանձնացնում ենք ԱԻ-ի պատասխանը (կտրում ենք մեր գրած prompt-ը)
    response = generated_output[len(prompt):].split("ՄԱՐԴ:")[0] # stop if it starts writing for human
    print(f"ԱԻ:{response.strip()}\n")