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

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

os.environ["TORCH_BLAS_BACKEND"] = "hipblas"
os.environ["PYTORCH_JIT"] = "0"
os.environ["OMP_NUM_THREADS"] = "4"
os.environ["MKL_DISABLE_FAST_MM"] = "1"
os.environ["ATEN_CPU_CAPABILITY"] = "avx"
os.environ["TORCH_CPU_CAPABILITY"] = "avx"


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"""
    config = TernaryConfig()
    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
        print("📦 Using safetensors loader...")
        state_dict = load_file(checkpoint_path, device=str(device))
    else:
        print("📦 Using torch.load...")
        state_dict = torch.load(checkpoint_path, map_location=device, weights_only=True)
    
    # Clean state dict
    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"]
    
    # Remove freqs_cis if present
    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("✅ Model successfully loaded from SafeTensor / PT checkpoint.")
    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):
    
    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.clone()
    
    print("\nJiRack: ", end="", flush=True)
    printed_text_len = 0

    for _ in range(max_new_tokens):
        with torch.autocast(device_type="cuda" if device.type == "cuda" else "cpu", 
                           dtype=torch.float16):
            logits, _ = model(generated)
        
        next_token_logits = logits[:, -1, :].clone().to(torch.float32)
        
        # Repetition penalty
        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)
        
        # Top-p
        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)
        
        # Streaming decode
        full_reply = tokenizer.decode(generated[0][input_ids.shape[1]:], skip_special_tokens=True)
        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)
        
        if next_token.item() in [tokenizer.eos_token_id, 128009]:
            break
    
    print("\n" + "-"*30)


def main():
    print("--- 🔍 Checking Hardware Acceleration ---")
    cuda_available = torch.cuda.is_available()
    print(f"PyTorch CUDA/ROCm Available: {cuda_available}")
    
    if cuda_available:
        device = torch.device("cuda")
        print(f"Device: {torch.cuda.get_device_name(0)}")
    else:
        device = torch.device("cpu")
        print("⚠️ Using CPU")

    CHECKPOINT = "model.safetensors"   # ← измени на свой safetensors файл
    
    tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B-Instruct",clean_up_tokenization_spaces=False)
    
    try:
        model = load_jirack_model(CHECKPOINT, device)
    except Exception as e:
        print(f"❌ Loading Error: {e}")
        return

    print("\n" + "="*60)
    print(f"💬 JiRack 1B TERNARY CHAT (SafeTensor mode)")
    print("Type 'exit' to quit")
    print("="*60 + "\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,
                         temperature=0.7, top_p=0.9, repetition_penalty=1.3)
            print(f"(Gen Time: {time.time() - start_time:.2f}s)\n")
            
        except KeyboardInterrupt:
            print("\nStopped.")
            break
        except Exception as e:
            print(f"Error: {e}")


if __name__ == "__main__":
    main()