#!/usr/bin/env python3
"""
Advanced Test Suite for Rax 4.0 Chat - Enterprise Edition
Developed by RaxCore Technologies
"""

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import time
import json

class RaxTester:
    def __init__(self, model_path="."):
        """Initialize Rax 4.0 Chat model for testing"""
        print("🚀 Initializing Rax 4.0 Chat - Enterprise Edition")
        print("=" * 60)
        
        self.model_path = model_path
        self.load_model()
        
    def load_model(self):
        """Load Rax 4.0 model and tokenizer"""
        print("📦 Loading Rax 4.0 Chat model...")
        start_time = time.time()
        
        try:
            self.tokenizer = AutoTokenizer.from_pretrained(self.model_path)
            self.model = AutoModelForCausalLM.from_pretrained(
                self.model_path,
                torch_dtype=torch.bfloat16,
                device_map="auto",
                trust_remote_code=True
            )
            
            load_time = time.time() - start_time
            print(f"✅ Model loaded successfully in {load_time:.2f} seconds")
            print(f"🧠 Model: {self.model.config._name_or_path}")
            print(f"🔢 Parameters: ~{self.model.num_parameters() / 1e9:.1f}B")
            print(f"💾 Device: {next(self.model.parameters()).device}")
            
        except Exception as e:
            print(f"❌ Error loading model: {e}")
            raise
    
    def generate_response(self, messages, max_tokens=512, temperature=0.7):
        """Generate response using Rax 4.0"""
        try:
            # Apply chat template
            input_text = self.tokenizer.apply_chat_template(
                messages, 
                tokenize=False, 
                add_generation_prompt=True
            )
            
            inputs = self.tokenizer(input_text, return_tensors="pt")
            
            # Move inputs to model device
            inputs = {k: v.to(next(self.model.parameters()).device) for k, v in inputs.items()}
            
            # Generate with timing
            start_time = time.time()
            
            with torch.no_grad():
                outputs = self.model.generate(
                    **inputs,
                    max_new_tokens=max_tokens,
                    temperature=temperature,
                    do_sample=True,
                    top_p=0.9,
                    repetition_penalty=1.1,
                    pad_token_id=self.tokenizer.eos_token_id
                )
            
            generation_time = time.time() - start_time
            
            # Decode response
            response = self.tokenizer.decode(
                outputs[0][inputs['input_ids'].shape[1]:], 
                skip_special_tokens=True
            )
            
            # Calculate tokens per second
            tokens_generated = len(outputs[0]) - len(inputs['input_ids'][0])
            tokens_per_second = tokens_generated / generation_time if generation_time > 0 else 0
            
            return {
                'response': response,
                'generation_time': generation_time,
                'tokens_generated': tokens_generated,
                'tokens_per_second': tokens_per_second
            }
            
        except Exception as e:
            print(f"❌ Error generating response: {e}")
            return None
    
    def test_basic_conversation(self):
        """Test basic conversational capabilities"""
        print("\n🗣️  Testing Basic Conversation")
        print("-" * 40)
        
        messages = [
            {"role": "system", "content": "You are Rax 4.0, the most advanced AI assistant created by RaxCore. You excel at complex reasoning, coding, and multilingual communication."},
            {"role": "user", "content": "Hello! Can you tell me about yourself and what makes you special?"}
        ]
        
        result = self.generate_response(messages, max_tokens=256)
        
        if result:
            print(f"👤 User: {messages[1]['content']}")
            print(f"🤖 Rax 4.0: {result['response']}")
            print(f"⚡ Generation: {result['generation_time']:.2f}s ({result['tokens_per_second']:.1f} tokens/s)")
            return True
        return False
    
    def test_coding_capabilities(self):
        """Test advanced coding capabilities"""
        print("\n💻 Testing Coding Capabilities")
        print("-" * 40)
        
        messages = [
            {"role": "system", "content": "You are Rax 4.0, an expert programming assistant created by RaxCore."},
            {"role": "user", "content": "Write a Python function to implement a binary search algorithm with detailed comments."}
        ]
        
        result = self.generate_response(messages, max_tokens=512)
        
        if result:
            print(f"👤 User: {messages[1]['content']}")
            print(f"🤖 Rax 4.0: {result['response']}")
            print(f"⚡ Generation: {result['generation_time']:.2f}s ({result['tokens_per_second']:.1f} tokens/s)")
            return True
        return False
    
    def test_reasoning_capabilities(self):
        """Test advanced reasoning and problem-solving"""
        print("\n🧠 Testing Reasoning Capabilities")
        print("-" * 40)
        
        messages = [
            {"role": "system", "content": "You are Rax 4.0, an advanced AI with superior reasoning capabilities created by RaxCore."},
            {"role": "user", "content": "Explain the concept of quantum entanglement and its potential applications in quantum computing. Then solve this logic puzzle: If all roses are flowers, and some flowers fade quickly, can we conclude that some roses fade quickly?"}
        ]
        
        result = self.generate_response(messages, max_tokens=768)
        
        if result:
            print(f"👤 User: {messages[1]['content']}")
            print(f"🤖 Rax 4.0: {result['response']}")
            print(f"⚡ Generation: {result['generation_time']:.2f}s ({result['tokens_per_second']:.1f} tokens/s)")
            return True
        return False
    
    def test_multilingual_capabilities(self):
        """Test multilingual communication"""
        print("\n🌍 Testing Multilingual Capabilities")
        print("-" * 40)
        
        messages = [
            {"role": "system", "content": "You are Rax 4.0, a multilingual AI assistant created by RaxCore with native-level proficiency in multiple languages."},
            {"role": "user", "content": "Please respond in French: Explain the importance of artificial intelligence in modern business, then translate your response to Spanish."}
        ]
        
        result = self.generate_response(messages, max_tokens=512)
        
        if result:
            print(f"👤 User: {messages[1]['content']}")
            print(f"🤖 Rax 4.0: {result['response']}")
            print(f"⚡ Generation: {result['generation_time']:.2f}s ({result['tokens_per_second']:.1f} tokens/s)")
            return True
        return False
    
    def test_enterprise_scenario(self):
        """Test enterprise-grade business scenario"""
        print("\n🏢 Testing Enterprise Scenario")
        print("-" * 40)
        
        messages = [
            {"role": "system", "content": "You are Rax 4.0, an enterprise-grade AI assistant created by RaxCore for business applications."},
            {"role": "user", "content": "I'm the CEO of a fintech startup. Analyze the current AI market trends, identify 3 key opportunities for our company, and create a brief strategic plan with implementation timeline."}
        ]
        
        result = self.generate_response(messages, max_tokens=1024)
        
        if result:
            print(f"👤 User: {messages[1]['content']}")
            print(f"🤖 Rax 4.0: {result['response']}")
            print(f"⚡ Generation: {result['generation_time']:.2f}s ({result['tokens_per_second']:.1f} tokens/s)")
            return True
        return False
    
    def run_comprehensive_test(self):
        """Run comprehensive test suite"""
        print("🧪 Starting Comprehensive Rax 4.0 Test Suite")
        print("=" * 60)
        
        tests = [
            ("Basic Conversation", self.test_basic_conversation),
            ("Coding Capabilities", self.test_coding_capabilities),
            ("Reasoning Capabilities", self.test_reasoning_capabilities),
            ("Multilingual Capabilities", self.test_multilingual_capabilities),
            ("Enterprise Scenario", self.test_enterprise_scenario)
        ]
        
        results = []
        total_time = 0
        
        for test_name, test_func in tests:
            print(f"\n🔍 Running: {test_name}")
            start_time = time.time()
            
            try:
                success = test_func()
                test_time = time.time() - start_time
                total_time += test_time
                
                results.append({
                    'test': test_name,
                    'success': success,
                    'time': test_time
                })
                
                status = "✅ PASSED" if success else "❌ FAILED"
                print(f"Status: {status} ({test_time:.2f}s)")
                
            except Exception as e:
                test_time = time.time() - start_time
                total_time += test_time
                
                results.append({
                    'test': test_name,
                    'success': False,
                    'time': test_time,
                    'error': str(e)
                })
                
                print(f"Status: ❌ FAILED - {e} ({test_time:.2f}s)")
        
        # Print summary
        print("\n" + "=" * 60)
        print("📊 TEST SUMMARY")
        print("=" * 60)
        
        passed = sum(1 for r in results if r['success'])
        total = len(results)
        
        print(f"Tests Passed: {passed}/{total}")
        print(f"Success Rate: {(passed/total)*100:.1f}%")
        print(f"Total Time: {total_time:.2f}s")
        print(f"Average Time per Test: {total_time/total:.2f}s")
        
        print("\n📋 Detailed Results:")
        for result in results:
            status = "✅" if result['success'] else "❌"
            print(f"{status} {result['test']}: {result['time']:.2f}s")
            if 'error' in result:
                print(f"   Error: {result['error']}")
        
        print("\n🎉 Rax 4.0 Chat testing completed!")
        print("🌟 Developed by RaxCore - Premier AI Innovation Company")
        
        return results

def main():
    """Main test execution"""
    try:
        # Initialize tester
        tester = RaxTester()
        
        # Run comprehensive tests
        results = tester.run_comprehensive_test()
        
        # Save results
        with open('test_results.json', 'w') as f:
            json.dump(results, f, indent=2)
        
        print(f"\n💾 Test results saved to: test_results.json")
        
    except Exception as e:
        print(f"❌ Test execution failed: {e}")
        return False
    
    return True

if __name__ == "__main__":
    main()