import gradio as gr
from ultralyticsplus import YOLO, render_result
import cv2
import time
import torch

# --------------------------
# System Checks & Optimization
# --------------------------
print("\n" + "="*40)
print("System Configuration Check:")
print(f"PyTorch Version: {torch.__version__}")
print(f"CUDA Available: {torch.cuda.is_available()}")
print(f"CUDA Device Count: {torch.cuda.device_count()}")
if torch.cuda.is_available():
    print(f"Using GPU: {torch.cuda.get_device_name(0)}")
else:
    print("Using CPU - For better performance, consider using a GPU environment")
print("="*40 + "\n")

# --------------------------
# Model Configuration
# --------------------------
# Load model with performance optimizations
model = YOLO('foduucom/plant-leaf-detection-and-classification')

# Configure model parameters
model_params = {
    'conf': 0.25,
    'iou': 0.45,
    'imgsz': 640,
    'device': 'cuda' if torch.cuda.is_available() else 'cpu',
    'half': True if torch.cuda.is_available() else False  # FP16 acceleration
}
model.overrides.update(model_params)

# Warmup model with dummy input
print("Performing model warmup...")
dummy_input = torch.randn(1, 3, 640, 640).to(model_params['device'])
if model_params['half']:
    dummy_input = dummy_input.half()
model.predict(dummy_input, verbose=False)
print("Model warmup complete!\n")

# --------------------------
# Image Processing Pipeline
# --------------------------
def preprocess_image(image):
    """Optimized image preprocessing"""
    # Convert RGB to BGR
    img = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    
    # Resize maintaining aspect ratio
    max_size = 1280
    h, w = img.shape[:2]
    scale = min(max_size/h, max_size/w)
    img = cv2.resize(img, (int(w*scale), int(h*scale)), 
                    interpolation=cv2.INTER_LINEAR)
    return img

# --------------------------
# Detection Function
# --------------------------
def detect_leaves(image):
    try:
        start_time = time.time()
        
        # Step 1: Preprocessing
        preprocess_start = time.time()
        img = preprocess_image(image)
        print(f"Preprocessing time: {time.time() - preprocess_start:.2f}s")
        
        # Step 2: Prediction
        predict_start = time.time()
        results = model.predict(
            source=img,
            verbose=False,
            stream=False,  # Disable streaming mode
            augment=False   # Disable TTA for speed
        )
        print(f"Prediction time: {time.time() - predict_start:.2f}s")
        
        # Step 3: Postprocessing
        postprocess_start = time.time()
        num_leaves = len(results[0].boxes)
        rendered_img = render_result(model=model, image=img, result=results[0])
        rendered_img = cv2.cvtColor(rendered_img, cv2.COLOR_BGR2RGB)
        print(f"Postprocessing time: {time.time() - postprocess_start:.2f}s")
        
        total_time = time.time() - start_time
        print(f"\nTotal processing time: {total_time:.2f}s")
        print(f"Detected leaves: {num_leaves}")
        print("-"*50)
        
        return rendered_img, num_leaves
    
    except Exception as e:
        print(f"Error processing image: {str(e)}")
        return None, 0

# --------------------------
# Gradio Interface
# --------------------------
with gr.Blocks(title="Leaf Detection", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🍃 Real-Time Plant Leaf Detection")
    gr.Markdown("Upload a plant image to analyze leaf count and health")
    
    with gr.Row():
        input_image = gr.Image(label="Input Image", type="numpy")
        output_image = gr.Image(label="Detection Results", width=600)
    
    with gr.Row():
        leaf_count = gr.Number(label="Detected Leaves", precision=0)
        process_btn = gr.Button("Analyze Image", variant="primary")
    
    progress = gr.Textbox(label="Processing Status", visible=True)
    
    process_btn.click(
        fn=detect_leaves,
        inputs=[input_image],
        outputs=[output_image, leaf_count]
    )

if __name__ == "__main__":
    demo.launch(
        server_port=7860,
        show_error=True,
        share=False
    )