import gradio as gr
from ultralytics import YOLO
import cv2
import numpy as np
import torch
import wikipedia
from PIL import Image
import pandas as pd
import os
import uuid
from datetime import datetime
from h3 import h3

# Initialize models
model = YOLO("yolov8n.pt")  # You can change this to a fine-tuned model
midas = torch.hub.load("intel-isl/MiDaS", "MiDaS_small")
midas.to("cpu").eval()
midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms").small_transform

# CSV path
csv_file = "tree_measurements.csv"
if not os.path.exists(csv_file):
    pd.DataFrame(columns=[
        "Timestamp", "Latitude", "Longitude", "H3_Index", 
        "Estimated_Height", "Species", "Image_File"
    ]).to_csv(csv_file, index=False)

# Convert lat/lon to H3 index
def latlon_to_h3(lat, lon, resolution=9):
    return h3.geo_to_h3(lat, lon, resolution)

# Main processing function
def process_tree_image(image, latitude, longitude):
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    img_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)

    # Object detection
    results = model(img_cv)
    detections = results[0].boxes.data.cpu().numpy()

    if len(detections) == 0:
        return "No tree detected", image

    # Crop the first detected tree
    x1, y1, x2, y2, conf, cls = detections[0]
    tree_crop = img_cv[int(y1):int(y2), int(x1):int(x2)]

    # Depth estimation
    input_tensor = midas_transforms(Image.fromarray(cv2.cvtColor(tree_crop, cv2.COLOR_BGR2RGB))).unsqueeze(0)
    with torch.no_grad():
        depth = midas(input_tensor).squeeze().cpu().numpy()
    approx_height = round(np.max(depth) - np.min(depth), 2)

    # Get H3 location index
    h3_index = latlon_to_h3(latitude, longitude)

    # Wikipedia species info
    try:
        species_guess = wikipedia.summary("tree", sentences=1)
    except:
        species_guess = "Could not retrieve species info"

    # Save image with unique name
    image_id = f"tree_{uuid.uuid4().hex[:8]}.png"
    image.save(image_id)

    # Append to CSV
    new_data = pd.DataFrame([{
        "Timestamp": timestamp,
        "Latitude": latitude,
        "Longitude": longitude,
        "H3_Index": h3_index,
        "Estimated_Height": approx_height,
        "Species": species_guess,
        "Image_File": image_id
    }])
    new_data.to_csv(csv_file, mode='a', index=False, header=False)

    return (
        f"📏 Estimated Tree Height: {approx_height} meters\n"
        f"🌍 H3 Location Index: {h3_index}\n"
        f"🌿 Species Info: {species_guess}",
        image
    )

# Gradio UI
interface = gr.Interface(
    fn=process_tree_image,
    inputs=[
        gr.Image(type="pil", label="📸 Capture Tree Image"),
        gr.Number(label="🌐 Latitude"),
        gr.Number(label="🌐 Longitude")
    ],
    outputs=[
        gr.Textbox(label="📊 Tree Info"),
        gr.Image(label="🖼️ Captured Image")
    ],
    title="🌳 Tree Height & Species Estimator",
    description="Use your webcam or upload an image of a tree. Enter your GPS location to get height, species info, and H3 geolocation. All data is saved to CSV."
)

interface.launch()