"""
Farm Human Recognition API - Gradio Interface
YOLO and pose estimation models for farm worker detection
"""

import gradio as gr
import torch
import cv2
import numpy as np
from PIL import Image
import json
import base64
import io
import time
from typing import List, Dict, Any

# Import models
try:
    from transformers import YolosImageProcessor, YolosForObjectDetection
    from transformers import AutoImageProcessor, AutoModelForObjectDetection
    MODELS_AVAILABLE = True
except ImportError:
    MODELS_AVAILABLE = False

class HumanRecognitionAPI:
    def __init__(self):
        self.models = {}
        self.processors = {}
        self.model_configs = {
            "yolos_tiny": "hustvl/yolos-tiny",
            "yolos_small": "hustvl/yolos-small",
            "yolos_base": "hustvl/yolos-base"
        }
        
        # Human activity classes relevant to farming
        self.farm_activities = {
            "harvesting": ["picking", "collecting", "gathering", "harvesting"],
            "planting": ["sowing", "planting", "seeding"],
            "maintenance": ["pruning", "watering", "fertilizing", "weeding"],
            "inspection": ["examining", "checking", "monitoring", "inspecting"],
            "operation": ["driving", "operating", "machinery", "equipment"],
            "general": ["working", "standing", "walking", "person"]
        }
        
        if MODELS_AVAILABLE:
            self.load_models()
    
    def load_models(self):
        """Load human detection models"""
        for model_key, model_name in self.model_configs.items():
            try:
                print(f"Loading {model_name}...")
                processor = YolosImageProcessor.from_pretrained(model_name)
                model = YolosForObjectDetection.from_pretrained(model_name)
                
                self.processors[model_key] = processor
                self.models[model_key] = model
                print(f"✅ {model_name} loaded successfully")
            except Exception as e:
                print(f"❌ Failed to load {model_name}: {e}")
    
    def detect_humans(self, image: Image.Image, model_key: str = "yolos_small") -> Dict[str, Any]:
        """Detect humans and analyze farm activities"""
        if not MODELS_AVAILABLE or model_key not in self.models:
            return {"error": "Model not available"}
        
        start_time = time.time()
        
        try:
            # Preprocess image
            processor = self.processors[model_key]
            model = self.models[model_key]
            
            inputs = processor(images=image, return_tensors="pt")
            
            # Run inference
            with torch.no_grad():
                outputs = model(**inputs)
            
            # Post-process results
            target_sizes = torch.tensor([image.size[::-1]])
            results = processor.post_process_object_detection(
                outputs, threshold=0.5, target_sizes=target_sizes
            )[0]
            
            # Filter for human detections
            human_detections = []
            for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
                class_name = model.config.id2label[label.item()].lower()
                if "person" in class_name and score > 0.5:
                    human_detections.append({
                        "class": "person",
                        "confidence": float(score),
                        "bbox": [float(x) for x in box],
                        "area": float((box[2] - box[0]) * (box[3] - box[1])),
                        "activity": self.infer_activity(box, image.size)
                    })
            
            # Analyze safety and productivity
            safety_analysis = self.analyze_safety(human_detections, image.size)
            productivity_metrics = self.calculate_productivity_metrics(human_detections)
            
            processing_time = time.time() - start_time
            
            return {
                "humans_detected": len(human_detections),
                "detections": human_detections,
                "safety_analysis": safety_analysis,
                "productivity_metrics": productivity_metrics,
                "processing_time": round(processing_time, 2),
                "model_used": model_key
            }
            
        except Exception as e:
            return {"error": str(e)}
    
    def infer_activity(self, bbox: List[float], image_size: tuple) -> str:
        """Infer farm activity from bounding box characteristics"""
        x1, y1, x2, y2 = bbox
        width = x2 - x1
        height = y2 - y1
        
        # Simple activity inference based on pose characteristics
        aspect_ratio = width / height
        relative_size = (width * height) / (image_size[0] * image_size[1])
        
        if aspect_ratio > 1.2:  # Wide bounding box
            return "operating_equipment"
        elif relative_size > 0.1:  # Large person in frame
            return "close_work"
        elif y2 > image_size[1] * 0.8:  # Person near bottom
            return "ground_work"
        else:
            return "general_activity"
    
    def analyze_safety(self, detections: List[Dict], image_size: tuple) -> Dict[str, Any]:
        """Analyze workplace safety factors"""
        if not detections:
            return {"status": "no_workers", "score": 1.0}
        
        safety_score = 1.0
        concerns = []
        
        # Check worker density
        workers_per_area = len(detections) / (image_size[0] * image_size[1] / 1000000)  # per megapixel
        if workers_per_area > 5:
            safety_score -= 0.2
            concerns.append("High worker density - ensure adequate spacing")
        
        # Check for workers near equipment (simplified check)
        for detection in detections:
            if detection["activity"] == "operating_equipment":
                # Check if other workers are nearby
                nearby_workers = sum(1 for d in detections 
                                   if d != detection and self.calculate_distance(d["bbox"], detection["bbox"]) < 100)
                if nearby_workers > 0:
                    safety_score -= 0.3
                    concerns.append("Workers detected near operating equipment")
        
        return {
            "status": "safe" if safety_score > 0.7 else "caution" if safety_score > 0.4 else "unsafe",
            "score": max(0.0, safety_score),
            "concerns": concerns,
            "workers_detected": len(detections)
        }
    
    def calculate_distance(self, bbox1: List[float], bbox2: List[float]) -> float:
        """Calculate distance between bounding box centers"""
        center1 = [(bbox1[0] + bbox1[2]) / 2, (bbox1[1] + bbox1[3]) / 2]
        center2 = [(bbox2[0] + bbox2[2]) / 2, (bbox2[1] + bbox2[3]) / 2]
        return ((center1[0] - center2[0]) ** 2 + (center1[1] - center2[1]) ** 2) ** 0.5
    
    def calculate_productivity_metrics(self, detections: List[Dict]) -> Dict[str, Any]:
        """Calculate farm productivity metrics"""
        if not detections:
            return {"active_workers": 0, "productivity_score": 0.0}
        
        activity_counts = {}
        for detection in detections:
            activity = detection["activity"]
            activity_counts[activity] = activity_counts.get(activity, 0) + 1
        
        # Simple productivity scoring
        productive_activities = ["close_work", "ground_work", "operating_equipment"]
        productive_workers = sum(activity_counts.get(activity, 0) for activity in productive_activities)
        productivity_score = productive_workers / len(detections) if detections else 0
        
        return {
            "active_workers": len(detections),
            "productivity_score": round(productivity_score, 2),
            "activity_breakdown": activity_counts,
            "recommendations": self.generate_productivity_recommendations(activity_counts)
        }
    
    def generate_productivity_recommendations(self, activity_counts: Dict[str, int]) -> List[str]:
        """Generate productivity improvement recommendations"""
        recommendations = []
        total_workers = sum(activity_counts.values())
        
        if activity_counts.get("general_activity", 0) > total_workers * 0.3:
            recommendations.append("Consider assigning specific tasks to idle workers")
        
        if activity_counts.get("operating_equipment", 0) > 1:
            recommendations.append("Multiple equipment operators detected - ensure coordination")
        
        if total_workers > 10:
            recommendations.append("Large workforce detected - consider team organization")
        
        return recommendations[:3]  # Limit to 3 recommendations
    
    def draw_detections(self, image: Image.Image, detections: List[Dict]) -> Image.Image:
        """Draw bounding boxes and labels on image"""
        img_array = np.array(image)
        
        for detection in detections:
            bbox = detection["bbox"]
            x1, y1, x2, y2 = [int(coord) for coord in bbox]
            
            # Draw bounding box
            cv2.rectangle(img_array, (x1, y1), (x2, y2), (0, 255, 0), 2)
            
            # Draw label
            label = f"Worker {detection['confidence']:.2f}"
            cv2.putText(img_array, label, (x1, y1-10), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
        
        return Image.fromarray(img_array)

# Initialize API
api = HumanRecognitionAPI()

def predict_humans(image, model_choice):
    """Gradio prediction function"""
    if image is None:
        return None, "Please upload an image"
    
    # Convert to PIL Image
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    
    # Run human detection
    results = api.detect_humans(image, model_choice)
    
    if "error" in results:
        return None, f"Error: {results['error']}"
    
    # Create visualization
    annotated_image = api.draw_detections(image, results["detections"])
    
    # Format results text
    safety = results["safety_analysis"]
    productivity = results["productivity_metrics"]
    
    safety_emoji = "🟢" if safety["status"] == "safe" else "🟡" if safety["status"] == "caution" else "🔴"
    
    results_text = f"""
👥 **Farm Worker Analysis**

{safety_emoji} **Safety Status**: {safety['status'].title()} (Score: {safety['score']:.1%})
👷 **Workers Detected**: {results['humans_detected']}
📊 **Productivity Score**: {productivity['productivity_score']:.1%}

**🛡️ Safety Analysis**:
"""
    
    if safety["concerns"]:
        for concern in safety["concerns"]:
            results_text += f"\n⚠️ {concern}"
    else:
        results_text += "\n✅ No immediate safety concerns detected"
    
    results_text += f"\n\n**📈 Productivity Metrics**:"
    if productivity["activity_breakdown"]:
        for activity, count in productivity["activity_breakdown"].items():
            results_text += f"\n• {activity.replace('_', ' ').title()}: {count} workers"
    
    if productivity["recommendations"]:
        results_text += f"\n\n**💡 Recommendations**:"
        for rec in productivity["recommendations"]:
            results_text += f"\n• {rec}"
    
    return annotated_image, results_text

# Gradio Interface
with gr.Blocks(title="👥 Farm Human Recognition API") as app:
    gr.Markdown("# 👥 Farm Human Recognition API")
    gr.Markdown("AI-powered farm worker detection, safety analysis, and productivity assessment")
    
    with gr.Tab("👷 Worker Detection"):
        with gr.Row():
            with gr.Column():
                image_input = gr.Image(type="pil", label="Upload Farm Image")
                model_choice = gr.Dropdown(
                    choices=["yolos_tiny", "yolos_small", "yolos_base"],
                    value="yolos_small",
                    label="Select Model"
                )
                detect_btn = gr.Button("🔍 Detect Workers", variant="primary")
            
            with gr.Column():
                output_image = gr.Image(label="Worker Detection Results")
                results_text = gr.Textbox(label="Analysis Results", lines=20)
        
        detect_btn.click(
            predict_humans,
            inputs=[image_input, model_choice],
            outputs=[output_image, results_text]
        )
    
    with gr.Tab("📡 API Documentation"):
        gr.Markdown("""
## 🚀 API Endpoint

**POST** `/api/predict`

### Request Format
```json
{
  "data": ["<base64_image>", "<model_choice>"]
}
```

### Model Options
- **yolos_tiny**: Fastest processing, basic accuracy
- **yolos_small**: Balanced performance (recommended)
- **yolos_base**: Highest accuracy, slower processing

### Response Format
```json
{
  "humans_detected": 3,
  "detections": [
    {
      "class": "person",
      "confidence": 0.92,
      "bbox": [120, 45, 180, 200],
      "activity": "ground_work"
    }
  ],
  "safety_analysis": {
    "status": "safe",
    "score": 0.85,
    "concerns": []
  },
  "productivity_metrics": {
    "active_workers": 3,
    "productivity_score": 0.75,
    "activity_breakdown": {
      "ground_work": 2,
      "operating_equipment": 1
    }
  }
}
```

### Activity Types
- **ground_work**: Workers performing field operations
- **close_work**: Detailed inspection or harvesting
- **operating_equipment**: Machinery operation
- **general_activity**: General farm activities
        """)

if __name__ == "__main__":
    app.launch()