"""
Fix OpenCLIP GPU Dependencies and Complete Isolation from PaddleOCR
Create separate virtual environment with GPU-enabled PyTorch for OpenCLIP
"""

import os
import sys
import subprocess
import tempfile
import platform

def check_system_gpu():
    """Check system GPU capabilities"""
    print("🔍 CHECKING SYSTEM GPU CAPABILITIES")
    print("=" * 40)
    
    try:
        # Check NVIDIA GPU
        result = subprocess.run(['nvidia-smi'], capture_output=True, text=True)
        if result.returncode == 0:
            print("✅ NVIDIA GPU detected")
            # Parse driver version
            for line in result.stdout.split('\n'):
                if 'Driver Version' in line:
                    print(f"   Driver: {line.strip()}")
                if 'CUDA Version' in line:
                    print(f"   CUDA: {line.strip()}")
        else:
            print("❌ No NVIDIA GPU detected or nvidia-smi not available")
            return False
        
        # Check PyTorch CUDA in main environment
        try:
            import torch
            print(f"✅ Main PyTorch CUDA: {torch.cuda.is_available()}")
            if torch.cuda.is_available():
                print(f"   CUDA Version: {torch.version.cuda}")
                print(f"   GPU Count: {torch.cuda.device_count()}")
                print(f"   Current GPU: {torch.cuda.get_device_name(0)}")
        except Exception as e:
            print(f"⚠️ Main PyTorch check failed: {e}")
        
        return True
        
    except Exception as e:
        print(f"❌ GPU check failed: {e}")
        return False

def create_openclip_gpu_environment():
    """Create isolated OpenCLIP environment with GPU support"""
    print("\n🛠️ CREATING OPENCLIP GPU ENVIRONMENT")
    print("=" * 40)
    
    venv_path = "openclip_gpu_env"
    
    # Remove existing environment if it exists
    if os.path.exists(venv_path):
        print(f"🗑️ Removing existing environment: {venv_path}")
        import shutil
        shutil.rmtree(venv_path)
    
    try:
        # Create virtual environment
        print(f"📦 Creating virtual environment: {venv_path}")
        result = subprocess.run([
            sys.executable, '-m', 'venv', venv_path,
            '--system-site-packages'  # Use system CUDA libraries
        ], capture_output=True, text=True)
        
        if result.returncode != 0:
            print(f"❌ Virtual environment creation failed: {result.stderr}")
            return False
        
        print("✅ Virtual environment created successfully")
        
        # Get virtual environment Python path
        if platform.system() == "Windows":
            venv_python = os.path.join(venv_path, "Scripts", "python.exe")
        else:
            venv_python = os.path.join(venv_path, "bin", "python")
        
        # Install GPU-enabled PyTorch
        print("📥 Installing GPU-enabled PyTorch...")
        
        # Use PyTorch with CUDA 11.8 (compatible with most systems)
        torch_install_cmd = [
            venv_python, '-m', 'pip', 'install',
            'torch==2.0.1+cu118', 'torchvision==0.15.2+cu118',
            '--index-url', 'https://download.pytorch.org/whl/cu118'
        ]
        
        result = subprocess.run(torch_install_cmd, capture_output=True, text=True, timeout=300)
        if result.returncode != 0:
            print(f"❌ PyTorch installation failed: {result.stderr}")
            return False
        
        print("✅ PyTorch with CUDA 11.8 installed")
        
        # Install OpenCLIP
        print("📥 Installing OpenCLIP...")
        openclip_cmd = [venv_python, '-m', 'pip', 'install', 'open-clip-torch']
        result = subprocess.run(openclip_cmd, capture_output=True, text=True, timeout=300)
        if result.returncode != 0:
            print(f"❌ OpenCLIP installation failed: {result.stderr}")
            return False
        
        print("✅ OpenCLIP installed")
        
        # Install additional dependencies
        print("📥 Installing additional dependencies...")
        deps_cmd = [venv_python, '-m', 'pip', 'install', 'Pillow', 'requests']
        result = subprocess.run(deps_cmd, capture_output=True, text=True, timeout=300)
        if result.returncode != 0:
            print(f"❌ Dependencies installation failed: {result.stderr}")
            return False
        
        print("✅ Additional dependencies installed")
        
        # Test the environment
        print("🧪 Testing OpenCLIP GPU environment...")
        test_script = f"""
import sys
import torch
import open_clip

print("=== OPENCLIP GPU ENVIRONMENT TEST ===")
print(f"Python: {{sys.version}}")
print(f"PyTorch version: {{torch.__version__}}")
print(f"CUDA available: {{torch.cuda.is_available()}}")
print(f"CUDA version: {{torch.version.cuda}}")
print(f"GPU count: {{torch.cuda.device_count()}}")

if torch.cuda.is_available():
    for i in range(torch.cuda.device_count()):
        print(f"GPU {{i}}: {{torch.cuda.get_device_name(i)}}")

# Test OpenCLIP model loading
try:
    model, _, preprocess = open_clip.create_model_and_transforms('ViT-B-32', pretrained='laion2b_s34b_b79k')
    print("✅ OpenCLIP model loaded successfully")
    
    if torch.cuda.is_available():
        model = model.cuda()
        print("✅ OpenCLIP model moved to GPU")
    else:
        print("⚠️ OpenCLIP model running on CPU")
        
except Exception as e:
    print(f"❌ OpenCLIP test failed: {{e}}")
    sys.exit(1)
"""
        
        with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
            f.write(test_script)
            test_file = f.name
        
        result = subprocess.run([venv_python, test_file], capture_output=True, text=True, timeout=60)
        os.unlink(test_file)
        
        if result.returncode == 0:
            print("🎉 OPENCLIP GPU ENVIRONMENT SETUP SUCCESSFUL!")
            print(result.stdout)
            return True
        else:
            print("❌ OpenCLIP environment test failed")
            print(result.stderr)
            return False
            
    except Exception as e:
        print(f"❌ Environment creation failed: {e}")
        return False

def update_image_classifier_for_gpu():
    """Update the image classifier to use the new GPU environment"""
    print("\n🔄 UPDATING IMAGE CLASSIFIER FOR GPU")
    print("=" * 40)
    
    # Update the isolated_image_classifier.py to use the new GPU environment
    new_classifier_code = '''
"""
Isolated Image Classifier with GPU Support
Runs in separate virtual environment with GPU-enabled PyTorch
"""

import os
import sys
import subprocess
import tempfile
import logging
from pathlib import Path
from typing import List, Dict, Any

logger = logging.getLogger(__name__)

class IsolatedImageClassifier:
    """Image classifier with GPU support in isolated environment"""
    
    def __init__(self):
        self.available = False
        self._check_availability()
    
    def _check_availability(self):
        """Check if OpenCLIP GPU environment is available"""
        try:
            # Use the GPU virtual environment Python executable
            venv_python = "openclip_gpu_env\\\\Scripts\\\\python.exe" if os.name == 'nt' else "openclip_gpu_env/bin/python"
            
            if not os.path.exists(venv_python):
                raise RuntimeError(f"GPU virtual environment not found: {venv_python}")
            
            # Test OpenCLIP and GPU availability
            test_script = """
import torch
import open_clip

print(f"GPU_AVAILABLE:{torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU_COUNT:{torch.cuda.device_count()}")
    print(f"GPU_NAME:{torch.cuda.get_device_name(0)}")

# Test model loading
model, _, _ = open_clip.create_model_and_transforms('ViT-B-32', pretrained='laion2b_s34b_b79k')
print("MODEL_LOADED:SUCCESS")
"""
            result = subprocess.run([venv_python, '-c', test_script], 
                                  capture_output=True, text=True, timeout=30)
            
            if "MODEL_LOADED:SUCCESS" in result.stdout:
                self.available = True
                # Parse GPU info
                gpu_available = "GPU_AVAILABLE:True" in result.stdout
                if gpu_available:
                    logger.info("✅ OpenCLIP GPU environment initialized with GPU support")
                else:
                    logger.warning("⚠️ OpenCLIP GPU environment initialized but running on CPU")
            else:
                error_msg = result.stderr or result.stdout or "Unknown error"
                raise RuntimeError(f"OpenCLIP GPU check failed: {error_msg}")
                
        except Exception as e:
            logger.error(f"OpenCLIP GPU availability check failed: {e}")
            raise RuntimeError(f"OpenCLIP GPU availability check failed: {e}")
    
    def classify_image(self, image_path: str, top_k: int = 5) -> List[Dict[str, Any]]:
        """
        Classify image using isolated GPU environment
        
        Args:
            image_path: Path to image file
            top_k: Number of top predictions to return
            
        Returns:
            List of classification results with confidence scores
        """
        if not self.available:
            return [{"label": "classification_unavailable", "confidence": 0.0}]
        
        try:
            venv_python = "openclip_gpu_env\\\\Scripts\\\\python.exe" if os.name == 'nt' else "openclip_gpu_env/bin/python"
            
            classification_script = f"""
import open_clip
import torch
from PIL import Image

try:
    # Load model
    model, _, processor = open_clip.create_model_and_transforms(
        model_name="ViT-B-32",
        pretrained="laion2b_s34b_b79k"
    )
    
    # Move to GPU if available
    if torch.cuda.is_available():
        model = model.cuda()
        print("USING_GPU:YES")
    else:
        print("USING_GPU:NO")
    
    # Load and process image
    image = Image.open(r"{image_path}").convert("RGB")
    image_tensor = processor(image).unsqueeze(0)
    
    if torch.cuda.is_available():
        image_tensor = image_tensor.cuda()
    
    # Get predictions
    with torch.no_grad():
        image_features = model.encode_image(image_tensor)
        image_features /= image_features.norm(dim=-1, keepdim=True)
        
        # Common labels for document images
        text_labels = [
            "a photo of a bee", "a photo of a flower", "a photo of a person", 
            "a photo of a document", "a photo of a chart", "a photo of a diagram",
            "a photo of a table", "a photo of a graph", "a photo of a logo",
            "a photo of a signature", "a photo of a stamp", "a photo of a barcode",
            "a photo of a QR code", "a photo of a screenshot", "a photo of a landscape",
            "a photo of an animal", "a photo of a building", "a photo of a vehicle",
            "a photo of food", "a photo of clothing", "a photo of electronics",
            "a photo of furniture", "a photo of nature", "a photo of art",
            "a photo of text", "a photo of numbers", "a photo of symbols"
        ]
        
        # Encode text labels
        text_tokens = open_clip.tokenize(text_labels)
        if torch.cuda.is_available():
            text_tokens = text_tokens.cuda()
        
        text_features = model.encode_text(text_tokens)
        text_features /= text_features.norm(dim=-1, keepdim=True)
        
        # Calculate similarity
        similarity = (100.0 * image_features @ text_features.T).softmax(dim=-1)
        values, indices = similarity[0].topk({top_k})
        
        for value, index in zip(values, indices):
            label = text_labels[index]
            confidence = float(value)
            print(f"CLASSIFICATION_RESULT:{{label}}|{{confidence}}")
            
except Exception as e:
    print(f"CLASSIFICATION_ERROR:{{e}}")
"""
            result = subprocess.run([venv_python, '-c', classification_script], 
                                  capture_output=True, text=True, timeout=30)
            
            results = []
            gpu_used = False
            
            for line in result.stdout.split('\\n'):
                if line.startswith("USING_GPU:YES"):
                    gpu_used = True
                if line.startswith("CLASSIFICATION_RESULT:"):
                    parts = line.replace("CLASSIFICATION_RESULT:", "").split("|")
                    if len(parts) == 2:
                        try:
                            label = parts[0]
                            confidence = float(parts[1])
                            results.append({{"label": label, "confidence": confidence, "gpu_used": gpu_used}})
                        except:
                            continue
            
            if results:
                if gpu_used:
                    logger.info(f"✅ Classification completed with GPU on {image_path}")
                else:
                    logger.warning(f"⚠️ Classification completed on CPU for {image_path}")
                return results
            else:
                logger.error("No results from classification")
                return [{{"label": "classification_failed", "confidence": 0.0, "gpu_used": False}}]
                
        except Exception as e:
            logger.error(f"Classification failed: {e}")
            return [{{"label": "classification_error", "confidence": 0.0, "gpu_used": False}}]


# Singleton instance
_classifier_instance = None

def get_isolated_classifier() -> IsolatedImageClassifier:
    """Get singleton isolated image classifier instance"""
    global _classifier_instance
    if _classifier_instance is None:
        _classifier_instance = IsolatedImageClassifier()
    return _classifier_instance


def test_isolated_classifier():
    """Test function for isolated image classifier"""
    classifier = get_isolated_classifier()
    
    if classifier.available:
        print("✅ Isolated image classifier is available")
        # Test with a sample image if available
        test_images = ["test_bee_image.png", "sample_image.jpg"]
        for test_image in test_images:
            if os.path.exists(test_image):
                results = classifier.classify_image(test_image)
                print(f"Classification results for {test_image}:")
                for result in results:
                    print(f"  {result['label']}: {result['confidence']:.4f} (GPU: {result.get('gpu_used', False)})")
                break
        else:
            print("⚠️ No test images found for classification")
    else:
        print("❌ Isolated image classifier is not available")


if __name__ == "__main__":
    test_isolated_classifier()
'''
    
    with open('isolated_image_classifier_gpu.py', 'w') as f:
        f.write(new_classifier_code)
    
    print("✅ Updated image classifier created: isolated_image_classifier_gpu.py")
    return True

def main():
    """Main function to fix OpenCLIP GPU dependencies"""
    print("🚀 FIXING OPENCLIP GPU DEPENDENCIES")
    print("=" * 50)
    
    # Step 1: Check system GPU
    if not check_system_gpu():
        print("❌ System GPU check failed. Cannot proceed with GPU setup.")
        return False
    
    # Step 2: Create OpenCLIP GPU environment
    if not create_openclip_gpu_environment():
        print("❌ Failed to create OpenCLIP GPU environment")
        return False
    
    # Step 3: Update image classifier for GPU
    if not update_image_classifier_for_gpu():
        print("❌ Failed to update image classifier for GPU")
        return False
    
    print("\n🎉 OPENCLIP GPU DEPENDENCIES FIXED SUCCESSFULLY!")
    print("=" * 50)
    print("✅ New virtual environment: openclip_gpu_env")
    print("✅ GPU-enabled PyTorch installed")
    print("✅ OpenCLIP with GPU support configured")
    print("✅ Updated image classifier: isolated_image_classifier_gpu.py")
    print("\n📝 Next steps:")
    print("   1. Update your document processor to use isolated_image_classifier_gpu.py")
    print("   2. Run test.py to verify GPU usage for both PaddleOCR and OpenCLIP")
    
    return True

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)