#!/usr/bin/env python3
"""
Brello Thinking - Advanced AI Reasoning Model
Created by Epic Systems
Engineered by Rehan Temkar
Built on Tencent Hunyuan Base Model
"""

import os
import re
from transformers import AutoModelForCausalLM, AutoTokenizer

def get_brello_system_prompt():
    """Returns the Brello Thinking system prompt"""
    return """You are Brello Thinking, an advanced AI assistant created by Epic Systems and engineered by Rehan Temkar. You are built on the Tencent Hunyuan base model and specialize in deep reasoning, mathematical problem-solving, coding, and creative thinking. You have enhanced chain-of-thought capabilities and can show your thinking process.

Key capabilities:
- Advanced mathematical reasoning and computation
- Programming and code generation in multiple languages
- Creative problem-solving and logical analysis
- Long-context understanding and document analysis
- Multi-language support with cultural understanding

Always be helpful, accurate, and demonstrate your advanced reasoning abilities. When appropriate, show your thinking process using the <think></think> tags before providing your final answer in <answer></answer> tags."""

def main():
    # Use the current directory as the model path since all model files are here
    model_path = "."
    
    print("🤖 Loading Brello Thinking...")
    print("Created by Epic Systems")
    print("Engineered by Rehan Temkar")
    print("Built on Tencent Hunyuan Base Model")
    print("="*60)
    
    print("Loading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    
    print("Loading model...")
    model = AutoModelForCausalLM.from_pretrained(
        model_path, 
        device_map="auto",
        torch_dtype="auto"
    )
    
    print("✅ Brello Thinking loaded successfully!")
    print("\n" + "="*60)
    print("🧠 BRELLO THINKING - ADVANCED AI REASONING")
    print("="*60)
    
    # Example conversation with Brello branding
    messages = [
        {"role": "system", "content": get_brello_system_prompt()},
        {"role": "user", "content": "What is 2+2? Please show your reasoning."}
    ]
    
    print("\nGenerating response with Brello Thinking...")
    tokenized_chat = tokenizer.apply_chat_template(
        messages, 
        tokenize=True, 
        add_generation_prompt=True,
        return_tensors="pt",
        enable_thinking=True  # Enable thinking mode
    )
    
    # Generate response
    outputs = model.generate(
        tokenized_chat.to(model.device), 
        max_new_tokens=200,  # Shorter for demo
        do_sample=True,
        top_k=20,
        top_p=0.8,
        repetition_penalty=1.05,
        temperature=0.7
    )
    
    output_text = tokenizer.decode(outputs[0])
    print("\n" + "="*60)
    print("BRELLO THINKING RESPONSE:")
    print("="*60)
    print(output_text)
    print("="*60)
    
    # Parse thinking and answer content
    think_pattern = r'<think>(.*?)</think>'
    think_matches = re.findall(think_pattern, output_text, re.DOTALL)
    
    answer_pattern = r'<answer>(.*?)</answer>'
    answer_matches = re.findall(answer_pattern, output_text, re.DOTALL)
    
    if think_matches:
        think_content = think_matches[0].strip()
        print(f"\n🧠 BRELLO THINKING PROCESS:\n{think_content}")
    
    if answer_matches:
        answer_content = answer_matches[0].strip()
        print(f"\n💡 BRELLO FINAL ANSWER:\n{answer_content}")
    
    # Interactive mode with Brello branding
    print("\n" + "="*60)
    print("🤖 BRELLO THINKING INTERACTIVE MODE")
    print("Type 'quit' to exit")
    print("Type 'about' to learn about Brello Thinking")
    print("="*60)
    
    while True:
        user_input = input("\n👤 You: ")
        if user_input.lower() == 'quit':
            print("\n🤖 Brello Thinking: Thank you for using Brello Thinking! Goodbye!")
            break
        elif user_input.lower() == 'about':
            about_response = """I am Brello Thinking, an advanced AI assistant created by Epic Systems and engineered by Rehan Temkar. I am built on the Tencent Hunyuan base model and specialize in deep reasoning, mathematical problem-solving, coding, and creative thinking.

My key capabilities include:
- Advanced mathematical reasoning and computation
- Programming and code generation in multiple languages  
- Creative problem-solving and logical analysis
- Long-context understanding and document analysis
- Multi-language support with cultural understanding

I have enhanced chain-of-thought capabilities and can show my thinking process when appropriate. I'm designed to be helpful, accurate, and demonstrate advanced reasoning abilities."""
            print(f"\n🤖 Brello Thinking: {about_response}")
            continue
            
        messages = [
            {"role": "system", "content": get_brello_system_prompt()},
            {"role": "user", "content": user_input}
        ]
        
        tokenized_chat = tokenizer.apply_chat_template(
            messages, 
            tokenize=True, 
            add_generation_prompt=True,
            return_tensors="pt",
            enable_thinking=True
        )
        
        outputs = model.generate(
            tokenized_chat.to(model.device), 
            max_new_tokens=500,  # Reasonable length for chat
            do_sample=True,
            top_k=20,
            top_p=0.8,
            repetition_penalty=1.05,
            temperature=0.7
        )
        
        response = tokenizer.decode(outputs[0])
        
        # Extract answer
        answer_matches = re.findall(answer_pattern, response, re.DOTALL)
        if answer_matches:
            answer = answer_matches[0].strip()
            print(f"\n🤖 Brello Thinking: {answer}")
        else:
            # If no answer tags, show the full response
            print(f"\n🤖 Brello Thinking: {response}")

if __name__ == "__main__":
    main()