#!/usr/bin/env python3
"""
Test multiple model loading strategies to find what works.
Each strategy is tested independently.
"""

import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfig

def check_gpu_memory():
    """Check memory usage on all GPUs."""
    print("   Memory Usage:")
    for i in range(torch.cuda.device_count()):
        mem = torch.cuda.memory_allocated(i) / 1e9
        total = torch.cuda.get_device_properties(i).total_memory / 1e9
        print(f"   GPU {i}: {mem:.2f} GB / {total:.2f} GB")
    
    gpu0_mem = torch.cuda.memory_allocated(0) / 1e9
    gpu1_mem = torch.cuda.memory_allocated(1) / 1e9
    
    # Determine pattern
    if abs(gpu0_mem - gpu1_mem) < 2.0:  # Within 2GB
        if gpu0_mem < 15.0:
            return "DISTRIBUTED"
        else:
            return "DUPLICATE"
    else:
        if gpu0_mem > gpu1_mem:
            return f"GPU0_ONLY ({gpu0_mem:.1f}GB)"
        else:
            return f"GPU1_ONLY ({gpu1_mem:.1f}GB)"

def test_strategy_1():
    """Test 1: bf16 model + BnB 4-bit (ON-THE-FLY quantization)"""
    print("\n" + "=" * 80)
    print("TEST 1: bf16 model + BnB 4-bit (ON-THE-FLY)")
    print("=" * 80)
    
    try:
        print("   Loading bf16 model with BnB 4-bit...")
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16,
        )
        model = AutoModelForCausalLM.from_pretrained(
            "/data/models/Ornith-1.0-35B",  # ← bf16 model
            quantization_config=bnb_config,
            device_map="auto",
            trust_remote_code=True,
            low_cpu_mem_usage=True,
        )
        print("   ✓ Model loaded successfully")
        
        pattern = check_gpu_memory()
        print(f"\n   Pattern: {pattern}")
        return True, pattern
    except Exception as e:
        print(f"\n   ✗ FAILED: {e}")
        return False, str(e)

def test_strategy_2():
    """Test 2: bf16 model + BnB 4-bit (alternative config)"""
    print("\n" + "=" * 80)
    print("TEST 2: bf16 model + BnB 4-bit (alt config)")
    print("=" * 80)
    
    try:
        torch.cuda.empty_cache()
        print("   Loading bf16 model with BnB 4-bit...")
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_use_double_quant=True,
        )
        model = AutoModelForCausalLM.from_pretrained(
            "/data/models/Ornith-1.0-35B",  # ← bf16 model
            quantization_config=bnb_config,
            device_map="auto",
            trust_remote_code=True,
            low_cpu_mem_usage=True,
        )
        print("   ✓ Model loaded successfully")
        
        pattern = check_gpu_memory()
        print(f"\n   Pattern: {pattern}")
        return True, pattern
    except Exception as e:
        print(f"\n   ✗ FAILED: {e}")
        return False, str(e)

def test_strategy_3():
    """Test 3: device_map with explicit GPU assignment"""
    print("\n" + "=" * 80)
    print("TEST 3: device_map with explicit GPU assignment")
    print("=" * 80)
    
    try:
        torch.cuda.empty_cache()
        print("   Loading model with explicit device_map...")
        
        # Get model config to determine layers
        from transformers import AutoConfig
        config = AutoConfig.from_pretrained("/data/models/Ornith-1.0-35B", trust_remote_code=True)
        num_layers = config.num_hidden_layers
        
        # Split layers: first half on GPU 0, second half on GPU 1
        device_map = {}
        for i in range(num_layers):
            if i < num_layers // 2:
                device_map[f"model.layers.{i}"] = 0
            else:
                device_map[f"model.layers.{i}"] = 1
        
        # Embeddings and norm on GPU 0
        device_map["model.embed_tokens"] = 0
        device_map["model.norm"] = 0
        device_map["lm_head"] = 0
        
        print(f"   Created device_map with {len(device_map)} entries")
        model = AutoModelForCausalLM.from_pretrained(
            "/data/models/Ornith-1.0-35B",
            device_map=device_map,
            torch_dtype=torch.bfloat16,
            trust_remote_code=True,
            low_cpu_mem_usage=True,
        )
        print("   ✓ Model loaded successfully")
        
        pattern = check_gpu_memory()
        print(f"\n   Pattern: {pattern}")
        return True, pattern
    except Exception as e:
        print(f"\n   ✗ FAILED: {e}")
        return False, str(e)

def test_strategy_4():
    """Test 4: Load to CPU, then move to GPU manually"""
    print("\n" + "=" * 80)
    print("TEST 4: Load to CPU, then move to GPU")
    print("=" * 80)
    
    try:
        torch.cuda.empty_cache()
        print("   Loading bf16 model to CPU...")
        model = AutoModelForCausalLM.from_pretrained(
            "/data/models/Ornith-1.0-35B",
            device_map="cpu",
            torch_dtype=torch.bfloat16,
            trust_remote_code=True,
            low_cpu_mem_usage=True,
        )
        print("   ✓ Model loaded to CPU")
        
        # Count params on CPU
        cpu_params = sum(p.numel() for p in model.parameters() if p.device.type == 'cpu')
        print(f"   CPU parameters: {cpu_params / 1e9:.2f}B")
        
        # Move to GPU 0
        print("\n   Moving to GPU 0...")
        model = model.to("cuda:0")
        pattern = check_gpu_memory()
        print(f"   Pattern after move to GPU 0: {pattern}")
        
        # Move to GPU 1
        print("\n   Moving to GPU 1...")
        model = model.to("cuda:1")
        pattern = check_gpu_memory()
        print(f"   Pattern after move to GPU 1: {pattern}")
        
        return True, "LOADED_TO_CPU_THEN_GPU"
    except Exception as e:
        print(f"\n   ✗ FAILED: {e}")
        return False, str(e)

def test_strategy_5():
    """Test 5: Sequential layer loading (manual distribution)"""
    print("\n" + "=" * 80)
    print("TEST 5: Sequential layer loading (manual distribution)")
    print("=" * 80)
    
    try:
        torch.cuda.empty_cache()
        print("   Loading model layer by layer...")
        
        # This is a simplified version - in reality would need more complex logic
        # For now, just test if we can load to one GPU
        print("   Loading bf16 to GPU 0 only...")
        model = AutoModelForCausalLM.from_pretrained(
            "/data/models/Ornith-1.0-35B",
            device_map={"": 0},
            torch_dtype=torch.bfloat16,
            trust_remote_code=True,
            low_cpu_mem_usage=True,
        )
        print("   ✓ Model loaded to GPU 0")
        
        pattern = check_gpu_memory()
        print(f"\n   Pattern: {pattern}")
        
        # Now try GPU 1
        torch.cuda.empty_cache()
        print("\n   Loading bf16 to GPU 1 only...")
        model = AutoModelForCausalLM.from_pretrained(
            "/data/models/Ornith-1.0-35B",
            device_map={"": 1},
            torch_dtype=torch.bfloat16,
            trust_remote_code=True,
            low_cpu_mem_usage=True,
        )
        print("   ✓ Model loaded to GPU 1")
        
        pattern = check_gpu_memory()
        print(f"   Pattern: {pattern}")
        
        return True, "SEQUENTIAL_LOAD"
    except Exception as e:
        print(f"\n   ✗ FAILED: {e}")
        return False, str(e)

if __name__ == "__main__":
    print("=" * 80)
    print("Testing multiple model loading strategies")
    print("=" * 80)
    
    # Check GPU availability
    print(f"\n1. GPU Check:")
    print(f"   CUDA available: {torch.cuda.is_available()}")
    print(f"   GPU count: {torch.cuda.device_count()}")
    for i in range(torch.cuda.device_count()):
        props = torch.cuda.get_device_properties(i)
        print(f"   GPU {i}: {props.name} ({props.total_memory / 1e9:.2f} GB)")
    
    # Run all tests
    results = []
    
    tests = [
        ("Test 1: device_map=auto (no BnB)", test_strategy_1),
        ("Test 2: device_map=auto + BnB 4-bit", test_strategy_2),
        ("Test 3: Explicit device_map", test_strategy_3),
        ("Test 4: Load to CPU then GPU", test_strategy_4),
        ("Test 5: Sequential layer loading", test_strategy_5),
    ]
    
    for name, test_func in tests:
        try:
            success, pattern = test_func()
            results.append((name, success, pattern))
        except Exception as e:
            print(f"\n   ✗ Test crashed: {e}")
            results.append((name, False, str(e)))
        
        # Clear GPU memory between tests
        torch.cuda.empty_cache()
    
    # Summary
    print("\n" + "=" * 80)
    print("SUMMARY")
    print("=" * 80)
    
    for name, success, pattern in results:
        status = "✓ PASS" if success else "✗ FAIL"
        print(f"{status}: {name}")
        print(f"         Pattern: {pattern}")
    
    # Find working strategies
    working = [name for name, success, _ in results if success]
    if working:
        print(f"\n✓ {len(working)} strategy/strategies work:")
        for w in working:
            print(f"  - {w}")
    else:
        print("\n✗ No strategies work!")