import gradio as gr
from transformers import pipeline
from PIL import Image
import requests
from datetime import datetime, timedelta

# -----------------------------
# 1. Модельдерді ГЛОБАЛЬДЫ ЖҮКТЕМЕЙМІЗ (Lazy Loading)
# -----------------------------
clf = None  # Image Classification (Trash-Net)
det = None  # Object Detection (DETR)
seg = None  # Image Segmentation (SegFormer)
dep = None  # Depth Estimation (DPT-Hybrid) - ЖАҢА МОДЕЛЬ!

# -----------------------------
# 2. Қазақша аударма сөздіктері (КЕҢЕЙТІЛГЕН)
# -----------------------------

# prithivMLmods/Trash-Net үшін
trash_labels = {
    "plastic": "Пластик",
    "metal": "Металл",
    "paper": "Қағаз",
    "glass": "Шыны (Әйнек)",
    "cardboard": "Картон",
    "organic": "Органикалық қалдық",
    "other": "Басқа"
}

# facebook/detr-resnet-50 (COCO) үшін
object_labels = {
    "person": "адам", "bottle": "бөтелке", "chair": "орындық", "table": "үстел", 
    "cat": "мысық", "dog": "ит", "cup": "кесе", "couch": "диван", "book": "кітап",
    "tv": "теледидар", "laptop": "ноутбук", "clock": "сағат", "car": "көлік",
    "potted plant": "гүл құмы", "traffic light": "бағдаршам", "umbrella": "қолшатыр",
}
segment_labels = {
    "wall": "қабырға", "building": "ғимарат", "sky": "аспан", "floor": "еден", "tree": "ағаш",
    "ceiling": "төбе", "road": "жол", "bed": "төсек", "windowpane": "терезе", "grass": "шөп",
    "cabinet": "шкаф", "sidewalk": "жаяужол", "person": "адам", "door": "есік",
    "table": "үстел", "mountain": "тау", "plant": "өсімдік", "chair": "орындық", "car": "көлік",
    "water": "су", "sofa": "диван", "shelf": "сөре", "house": "үй", "sea": "теңіз",
}
WMO_WEATHER_CODES = {
    0: "Ашық", 1: "Аздап бұлтты", 2: "Бұлыңғыр", 3: "Бұлтты", 45: "Тұман",
    51: "Жеңіл сіркіреу", 61: "Жеңіл жаңбыр", 63: "Орташа жаңбыр", 65: "Қатты жаңбыр", 
    71: "Жеңіл қар", 73: "Орташа қар", 75: "Қатты қар", 80: "Жеңіл нөсер", 
    95: "Найзағай", 99: "Қатты қарлы найзағай",
}

# -----------------------------
# 3. Ауа райын алу (Open-Meteo)
# -----------------------------
def get_weather():
    try:
        latitude = 48.9711
        longitude = 90.0500
        url = (
            f"https://api.open-meteo.com/v1/forecast?latitude={latitude}&longitude={longitude}"
            "&daily=weathercode,temperature_2m_max,temperature_2m_min&forecast_days=7&timezone=Asia/Ulaanbaatar"
        )
        response = requests.get(url, timeout=10).json()
        
        output = "### 🌤 Баян-Өлгий ауа райы\n\n"
        output += "**7 күндік болжам:**\n"
        warnings = []
        
        for i in range(7):
            date = response["daily"]["time"][i]
            max_temp = response["daily"]["temperature_2m_max"][i]
            min_temp = response["daily"]["temperature_2m_min"][i]
            code = response["daily"]["weathercode"][i]
            desc = WMO_WEATHER_CODES.get(code, "Белгісіз")
            
            output += f"* **{date}**: 🌡 {min_temp}°C / {max_temp}°C, {desc}\n"
            
            if min_temp <= -25:
                warnings.append(f"{date}: Қатты аяз күтілуде ({min_temp}°C)!")

        if warnings:
            output += "\n**⚠️ Ерекше ескертулер:**\n"
            for w in warnings:
                output += f"* {w}\n"
        
        return output

    except Exception:
        return "Ауа райын алу кезінде қате болды."

# -----------------------------
# 4. Анализ функциясы (4 Модельді жүктеу осы жерде)
# -----------------------------
def analyze_image(image):
    global clf, det, seg, dep
    
    # Lazy loading: Модельдерді тек бірінші рет қана жүктейміз
    try:
        if clf is None:
            clf = pipeline("image-classification", model="prithivMLmods/Trash-Net")
        if det is None:
            det = pipeline("object-detection", model="facebook/detr-resnet-50")
        if seg is None:
            seg = pipeline("image-segmentation", model="nvidia/segformer-b0-finetuned-ade-512-512")
        # ЖАҢА МОДЕЛЬДІ ЖҮКТЕУ: Depth Estimation
        if dep is None:
            dep = pipeline("depth-estimation", model="Intel/dpt-hybrid-midas")
    except Exception as e:
        return f"Модельдерді жүктеу кезінде қате шықты. Қате: {e}"

    # --- 1. Қоқыс классификациясы ---
    clf_result = clf(image)
    clf_outputs = [f"{trash_labels.get(res['label'], res['label'])} ({round(res['score'] * 100, 2)}%)" for res in clf_result]

    # --- 2. Объект табу ---
    det_result = det(image)
    objects = [object_labels.get(obj['label'], obj['label']) for obj in det_result]
    unique_objects = list(set(objects))

    # --- 3. Сегментация ---
    seg_result = seg(image)
    seg_labels_kz = list({segment_labels.get(s['label'], s['label']) for s in seg_result})
    
    # --- 4. Тереңдікті бағалау (Нәтижесін сурет етіп қайтарамыз) ---
    dep_result = dep(image)
    # dep_result[0] - бұл тереңдік картасының суреті (PIL Image)

    # --- Ауа райы ---
    weather_info = get_weather()

    # Нәтиже
    output_text = (
        f"### 🗑️ Қоқыс түрі (Trash-Net)\n**{', '.join(clf_outputs)}**\n\n---\n\n"
        f"### 🔍 Табылған нысандар (DETR)\n{', '.join(unique_objects) if unique_objects else 'Нысандар табылған жоқ.'}\n\n---\n\n"
        f"### 🧩 Суреттегі аймақтар (SegFormer)\n{', '.join(seg_labels_kz) if seg_labels_kz else 'Аймақтар табылған жоқ.'}\n\n"
    )
    
    # Бір функциядан бірнеше нәтиже қайтару үшін
    return output_text, dep_result[0] 

# -----------------------------
# 5. Интерфейс (Gradio)
# -----------------------------
css = """
body { background-image: url('https://uploads8.wikiart.org/images/vincent-van-gogh/starry-night.jpg'); background-size: cover; background-repeat: no-repeat; background-attachment: fixed; }
.gradio-container, .gr-panel, .gr-box { background-color: rgba(255, 255, 255, 0.9) !important; border-radius: 15px !important; }
h1.title { color: white !important; text-shadow: 2px 2px 4px #000000; text-align: center; }
"""

with gr.Blocks(title="SmartWaste-YOLOv8n", css=css) as demo:
    gr.Markdown("<h1 class='title'>♻️ SmartWaste — Ақылды Қоқыс + Ауа райы Анализ жүйесі</h1>")
    gr.Markdown("<p style='text-align:center;'>**Ескерту:** Бұл қосымша жедел жадты үнемдеу үшін модельдерді **'АНАЛИЗ ЖАСАУ'** батырмасын бірінші басқанда жүктейді. Алғашқы анализ **өте** ұзаққа созылуы мүмкін.</p>")

    with gr.Row():
        image_input = gr.Image(label="Суретті таңдаңыз немесе камера қосыңыз", type="pil", sources=["upload", "webcam", "clipboard"])
        depth_output = gr.Image(label="4. Тереңдік Картасы (Depth Estimation)", type="pil", interactive=False)
        
    with gr.Row():
        analyze_btn = gr.Button("🔎 АНАЛИЗ ЖАСАУ", variant="primary")
        output_text = gr.Markdown(label="Нәтижелер (Классификация, Объект, Сегментация)")
        
    # analyze_image функциясы екі нәтиже қайтарады: мәтін және сурет
    analyze_btn.click(fn=analyze_image, inputs=image_input, outputs=[output_text, depth_output])
    
    # Ауа райын жеке көрсету (қосымша іске қосылғанда бірден жүктеледі)
    gr.Markdown(get_weather())


demo.launch()