#!/usr/bin/env python3
"""
cuDNN Automated Installer for CUDA 12.9
This script attempts to download and install cuDNN for GPU acceleration
"""

import os
import sys
import urllib.request
import zipfile
import tempfile
import shutil

def main():
    print("🚀 Starting cuDNN Automated Installation for CUDA 12.9")
    print("=" * 60)
    
    # Check CUDA installation
    cuda_path = r'C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.9'
    if not os.path.exists(cuda_path):
        print(f"❌ CUDA 12.9 not found at: {cuda_path}")
        print("💡 Please install CUDA 12.9 first")
        return False
    
    print(f"✅ Found CUDA 12.9 at: {cuda_path}")
    
    # Check if cuDNN already exists
    cudnn_dll_path = os.path.join(cuda_path, 'bin', 'cudnn64_8.dll')
    if os.path.exists(cudnn_dll_path):
        print(f"✅ cuDNN already installed at: {cudnn_dll_path}")
        return True
    
    print("❌ cuDNN not found. Starting automated installation...")
    
    # Known cuDNN download URLs (these may change, we'll try multiple sources)
    cudnn_urls = [
        # Try to find a direct download link (these are placeholders)
        "https://developer.download.nvidia.com/compute/redist/cudnn/v8.9.7/local_installers/12.x/cudnn-windows-x86_64-8.9.7.29_cuda12-archive.zip",
    ]
    
    # Create temporary directory
    temp_dir = tempfile.mkdtemp()
    
    try:
        for url in cudnn_urls:
            print(f"\n📥 Attempting to download cuDNN from: {url}")
            try:
                # Download cuDNN
                download_path = os.path.join(temp_dir, "cudnn.zip")
                urllib.request.urlretrieve(url, download_path)
                
                print("✅ Download successful. Extracting...")
                
                # Extract cuDNN
                with zipfile.ZipFile(download_path, 'r') as zip_ref:
                    zip_ref.extractall(temp_dir)
                
                # Find extracted cuDNN directory
                cudnn_extracted = None
                for item in os.listdir(temp_dir):
                    item_path = os.path.join(temp_dir, item)
                    if os.path.isdir(item_path) and 'cudnn' in item.lower():
                        cudnn_extracted = item_path
                        break
                
                if not cudnn_extracted:
                    print("❌ Could not find cuDNN in extracted files")
                    continue
                
                print(f"✅ Extracted cuDNN to: {cudnn_extracted}")
                
                # Copy cuDNN files to CUDA directory
                print("\n📁 Copying cuDNN files to CUDA directory...")
                
                # Copy bin files
                cudnn_bin = os.path.join(cudnn_extracted, 'bin')
                if os.path.exists(cudnn_bin):
                    for file in os.listdir(cudnn_bin):
                        if file.startswith('cudnn'):
                            src = os.path.join(cudnn_bin, file)
                            dst = os.path.join(cuda_path, 'bin', file)
                            shutil.copy2(src, dst)
                            print(f"  ✅ Copied: {file} -> {dst}")
                
                # Copy include files
                cudnn_include = os.path.join(cudnn_extracted, 'include')
                if os.path.exists(cudnn_include):
                    for file in os.listdir(cudnn_include):
                        if file.startswith('cudnn'):
                            src = os.path.join(cudnn_include, file)
                            dst = os.path.join(cuda_path, 'include', file)
                            shutil.copy2(src, dst)
                            print(f"  ✅ Copied: {file} -> {dst}")
                
                # Copy lib files
                cudnn_lib = os.path.join(cudnn_extracted, 'lib')
                if os.path.exists(cudnn_lib):
                    for file in os.listdir(cudnn_lib):
                        if file.startswith('cudnn'):
                            src = os.path.join(cudnn_lib, file)
                            dst = os.path.join(cuda_path, 'lib', 'x64', file)
                            shutil.copy2(src, dst)
                            print(f"  ✅ Copied: {file} -> {dst}")
                
                # Verify installation
                if os.path.exists(cudnn_dll_path):
                    print(f"\n🎉 cuDNN successfully installed at: {cudnn_dll_path}")
                    return True
                else:
                    print("❌ cuDNN installation verification failed")
                    continue
                    
            except Exception as e:
                print(f"❌ Download/Installation failed: {e}")
                continue
        
        # If all URLs failed, provide manual instructions
        print("\n" + "=" * 60)
        print("🔧 MANUAL INSTALLATION REQUIRED")
        print("=" * 60)
        print("Automated download failed. Please install cuDNN manually:")
        print("\nSteps:")
        print("1. Visit: https://developer.nvidia.com/cudnn")
        print("2. Login with NVIDIA Developer account (free)")
        print("3. Download cuDNN for CUDA 12.x")
        print("4. Extract the zip file")
        print("5. Copy these files to CUDA directory:")
        print(f"   - bin\\cudnn64_8.dll -> {cuda_path}\\bin\\")
        print(f"   - include\\cudnn*.h -> {cuda_path}\\include\\")
        print(f"   - lib\\cudnn*.lib -> {cuda_path}\\lib\\x64\\")
        print("\n6. Add CUDA to PATH:")
        print(f"   {cuda_path}\\bin")
        
        return False
        
    finally:
        # Clean up temporary directory
        if os.path.exists(temp_dir):
            shutil.rmtree(temp_dir)

def test_gpu_after_install():
    print("\n" + "=" * 60)
    print("🧪 Testing GPU Acceleration After Installation")
    print("=" * 60)
    
    try:
        import paddle
        print(f"✅ PaddlePaddle version: {paddle.__version__}")
        print(f"🏁 CUDA compiled: {paddle.is_compiled_with_cuda()}")
        print(f"🏁 GPU count: {paddle.device.cuda.device_count()}")
        
        if paddle.device.cuda.device_count() > 0:
            print(f"🏁 Current GPU device: {paddle.device.cuda.get_device_name()}")
            
            # Test GPU tensor operations
            paddle.set_device('gpu:0')
            x = paddle.to_tensor([1.0, 2.0, 3.0])
            y = paddle.to_tensor([4.0, 5.0, 6.0])
            z = x + y
            print(f"✅ GPU tensor operation successful: {z.numpy()}")
            print(f"✅ Tensor is on GPU: {z.place.is_gpu_place()}")
            print("🎉 PaddlePaddle GPU is working!")
            return True
        else:
            print("❌ No GPU devices detected")
            return False
            
    except Exception as e:
        print(f"❌ GPU test failed: {e}")
        return False

if __name__ == "__main__":
    # First try to install cuDNN
    success = main()
    
    # Then test GPU
    if success:
        test_gpu_after_install()
    else:
        print("\n❌ cuDNN installation failed. GPU acceleration unavailable.")
        print("💡 The system will continue to work in CPU mode.")