import torch
import torch.nn.functional as F
from transformers import AutoTokenizer
import os
import sys
import time
import warnings

# Suppress the clean_up_tokenization warning
warnings.filterwarnings("ignore", message=".*Ignoring clean_up_tokenization_spaces.*")

try:
    from JiRackTernaryPyTorch_1b import TernaryTransformer1B, TernaryConfig
except ImportError:
    print("❌ Error: File JiRackTernaryPyTorch_1b.py not found")
    sys.exit(1)

def load_jirack_model(checkpoint_path, device):
    """Загрузка модели с поддержкой как Safetensors, так и стандартных PT чекпоинтов"""
    config = TernaryConfig()
    # Инициализируем модель сразу в float16
    model = TernaryTransformer1B(config).to(device=device, dtype=torch.float16)

    print(f"--- 🚀 Loading checkpoint: {checkpoint_path} on {device} ---")

    if checkpoint_path.endswith(".safetensors"):
        from safetensors.torch import load_file
        state_dict = load_file(checkpoint_path, device=str(device))
    else:
        state_dict = torch.load(checkpoint_path, map_location=device)

    if isinstance(state_dict, dict) and "model_state_dict" in state_dict:
        state_dict = state_dict["model_state_dict"]
    elif isinstance(state_dict, dict) and "model" in state_dict:
        state_dict = state_dict["model"]

    state_dict = {k: v for k, v in state_dict.items() if "freqs_cis" not in k}

    model.load_state_dict(state_dict, strict=False)
    model.eval()
    print("✅ Weights successfully loaded.")
    return model

@torch.no_grad()
def chat_generate(model, tokenizer, prompt, device, max_new_tokens=128, temperature=0.5, top_p=0.9, repetition_penalty=1.3):
    """JiRack chat with llama style generation loop"""

    formatted_prompt = f"<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\n{prompt}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"

    input_ids = tokenizer.encode(formatted_prompt, return_tensors="pt").to(device)
    generated = input_ids

    print("\nJiRack: ", end="", flush=True)

    # We track what we've printed so far to handle streaming cleanly without leaking special tokens
    printed_text_len = 0

    for _ in range(max_new_tokens):
        device_type = "cuda" if device.type == "cuda" else "cpu"
        with torch.autocast(device_type=device_type, dtype=torch.float16):
            logits, _ = model(generated)

        next_token_logits = logits[:, -1, :].clone().to(torch.float32)

        for token_id in set(generated[0].tolist()):
            val = next_token_logits[0, token_id]
            if val > 0:
                next_token_logits[0, token_id] = val / repetition_penalty
            else:
                next_token_logits[0, token_id] = val * repetition_penalty

        next_token_logits = next_token_logits / max(temperature, 1e-6)

        sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True, dim=-1)
        probs = F.softmax(sorted_logits, dim=-1)
        cumulative_probs = torch.cumsum(probs, dim=-1)

        sorted_indices_to_remove = cumulative_probs > top_p
        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
        sorted_indices_to_remove[..., 0] = False

        indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
        next_token_logits = torch.where(indices_to_remove, torch.tensor(float('-inf'), device=device), next_token_logits)

        probs = F.softmax(next_token_logits, dim=-1)
        next_token = torch.multinomial(probs, num_samples=1)

        generated = torch.cat([generated, next_token], dim=-1)

        # Decode the whole generated sequence slice omitting the prompt to safely skip special tokens
        full_reply = tokenizer.decode(generated[0][input_ids.shape[1]:], skip_special_tokens=True)
        
        # Stream out only the newly decoded characters
        if len(full_reply) > printed_text_len:
            new_text = full_reply[printed_text_len:]
            print(new_text, end="", flush=True)
            printed_text_len = len(full_reply)

        # Stop condition check
        if next_token.item() in [tokenizer.eos_token_id, 128009]:
            break

    print("\n" + "-"*30)

def main():
    # Diagnostic Check for AMD GPU (ROCm)
    print("--- 🔍 Checking Hardware Acceleration ---")
    cuda_available = torch.cuda.is_available()
    print(f"PyTorch CUDA/ROCm Available: {cuda_available}")
    if cuda_available:
        print(f"Device Name: {torch.cuda.get_device_name(0)}")
        device = torch.device("cuda")
    else:
        print("⚠️ WARNING: GPU not found by PyTorch. Falling back to CPU.")
        print("If using AMD MI50, verify HIP/ROCm environment variables (e.g., ROCM_PATH or HIP_VISIBLE_DEVICES).")
        device = torch.device("cpu")
    
    CHECKPOINT = "jiarck_pro_1b_model.pt" 

    print(f"\n--- 📚 Loading Tokenizer on {device} ---")
    tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B-Instruct")

    try:
        model = load_jirack_model(CHECKPOINT, device)
    except Exception as e:
        print(f"❌ Loading Error: {e}")
        return

    print("\n" + "="*50)
    print(f"💬 JiRack 1B TERNARY CHAT MODE ({device.type.upper()}-MODE)")
    print("Type 'exit' to quit")
    print("="*50 + "\n")

    while True:
        try:
            user_input = input("User: ")
            if user_input.lower() in ["exit", "quit", "выход"]:
                break
            if not user_input.strip(): 
                continue

            start_time = time.time()
            chat_generate(model, tokenizer, user_input, device)
            print(f"(Gen Time: {time.time() - start_time:.2f}s)")
        except KeyboardInterrupt:
            print("\nStopped.")
            break

if __name__ == "__main__":
    main()