import numpy as np
from gpu_arch import Chip
from ai import AIAccelerator
from custom_vram import CustomVRAM
from PIL import Image
import requests

def test_ai_integration():
    print("\n--- Testing AI Integration ---")

    # Test 1: Model Loading (Florence-2 model)
    print("\nTest 1: Model Loading (Florence-2)")
    try:
        # Initialize a Chip for model loading
        chip_for_loading = Chip(chip_id=0, vram_size_gb=10)
        ai_accelerator_for_loading = chip_for_loading.ai_accelerator

        # Load BLIP-2 Large model and processor using Hugging Face Auto classes
        from transformers import Blip2ForConditionalGeneration, Blip2Processor
        model_id = "Salesforce/blip2-flan-t5-xxl"
        model = Blip2ForConditionalGeneration.from_pretrained(model_id, device_map="auto")
        processor = Blip2Processor.from_pretrained(model_id)

        ai_accelerator_for_loading.load_model(model_id, model, processor)
        print(f"Model '{model_id}' loaded successfully on chip 0.")
        assert ai_accelerator_for_loading.has_model(model_id), "Model not found in registry after loading."

    except Exception as e:
        print(f"Model loading test failed: {e}")
        return
    # Test 2: Multi-Chip Inference (on all images in sample_task folder)
    print("\nTest 2: Multi-Chip Inference (Florence-2, all images in sample_task)")
    import os
    num_chips = 1 # You can increase this if you want to test with more chips
    chips = []
    ai_accelerators = []

    try:
        # Initialize multiple chips and their AI accelerators
        for i in range(num_chips):
            chip = Chip(chip_id=i, vram_size_gb=1)
            chips.append(chip)
            ai_accelerators.append(chip.ai_accelerator)
            ai_accelerators[i].load_model(model_id, model, processor)
            print(f"Model '{model_id}' loaded successfully on chip {i}.")

        # Get all image files in sample_task folder
        image_folder = os.path.join(os.path.dirname(__file__), '..', 'sample_task')
        image_files = [f for f in os.listdir(image_folder) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif'))]
        image_files.sort()
        if not image_files:
            print("No images found in sample_task folder.")
            return

        # Perform inference on each image using all chips
        import time
        for img_name in image_files:
            img_path = os.path.join(image_folder, img_name)
            raw_image = Image.open(img_path).convert('RGB')
            print(f"\nRunning inference for image: {img_name}")
            for i, accelerator in enumerate(ai_accelerators):
                print(f"Performing inference on chip {i}...")
                start_time = time.time()
                result = accelerator.inference(model_id, raw_image)
                elapsed = time.time() - start_time
                print(f"Inference result from chip {i} on {img_name}: {result}")
                print(f"Inference time for chip {i} on {img_name}: {elapsed:.3f} seconds")
                assert result is not None, f"Inference returned None for chip {i} on {img_name}."
                assert isinstance(result, str), f"Inference result from chip {i} on {img_name} is not a string."
        print("Multi-chip inference test on all images successful.")

    except Exception as e:
        print(f"Multi-chip inference test failed: {e}")
        return
        return

    # Test 3: Matrix Operations (using CustomVRAM) - still on a single chip
    # print("\nTest 3: Matrix Operations (using CustomVRAM)")
    # try:
    #     matrix_a = np.array([[1, 2], [3, 4]], dtype=np.float32)
    #     matrix_b = np.array([[5, 6], [7, 8]], dtype=np.float32)

    #     matrix_a_id = ai_accelerator_for_loading.load_matrix(matrix_a, "matrix_A")
    #     matrix_b_id = ai_accelerator_for_loading.load_matrix(matrix_b, "matrix_B")

    #     result_matrix_id = ai_accelerator_for_loading.matrix_multiply(matrix_a_id, matrix_b_id, "result_C")
    #     result_matrix = ai_accelerator_for_loading.get_matrix(result_matrix_id)

    #     print(f"Matrix A:\n{matrix_a}")
    #     print(f"Matrix B:\n{matrix_b}")
    #     print(f"Result Matrix C:\n{result_matrix}")

    #     expected_result = np.dot(matrix_a, matrix_b)
    #     assert np.array_equal(result_matrix, expected_result), "Matrix multiplication result incorrect."
    #     print("Matrix operations test successful.")

    # except Exception as e:
    #     print(f"Matrix operations test failed: {e}")
    #     return

    print("\n--- All AI Integration Tests Completed ---")

if __name__ == "__main__":
    test_ai_integration()