# app.py
import os
import torch
import torch.nn as nn
from transformers import AutoImageProcessor, AutoModel
import torchvision.transforms as transforms
from PIL import Image
import numpy as np
import gradio as gr
import tempfile
import json
from datetime import datetime
import logging
from datasets import load_dataset
import matplotlib.pyplot as plt
import folium
from io import BytesIO
import base64

# Logging konfigürasyonu
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class S2NAIPGeoModel(nn.Module):
    """S2-NAIP Dataset için Jeo-Referanslama Modeli"""
    
    def __init__(self):
        super(S2NAIPGeoModel, self).__init__()
        
        # DINOv2 backbone
        self.dinov2 = AutoModel.from_pretrained("facebook/dinov2-base")
        
        # Regresyon başı - koordinat tahmini için
        self.regressor = nn.Sequential(
            nn.Linear(768, 512),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, 2)  # enlem, boylam
        )
        
    def forward(self, pixel_values):
        # Görüntü özelliklerini çıkar
        features = self.dinov2(pixel_values=pixel_values).last_hidden_state
        features = features.mean(dim=1)  # Global average pooling
        
        # Koordinat tahmini
        coordinates = self.regressor(features)
        return coordinates

class S2NAIPGeoSystem:
    """S2-NAIP Jeo-Referanslama Sistemi"""
    
    def __init__(self, model_path=None):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        logger.info(f"Cihaz: {self.device}")
        
        # Model ve processor
        self.processor = AutoImageProcessor.from_pretrained("facebook/dinov2-base")
        self.model = S2NAIPGeoModel().to(self.device)
        
        # Transformations
        self.transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225]
            )
        ])
        
        # Model yükleme
        if model_path and os.path.exists(model_path):
            self.model.load_state_dict(torch.load(model_path, map_location=self.device))
            logger.info("Model yüklendi")
        else:
            logger.info("Rastgele ağırlıklı model kullanılıyor")
        
        self.model.eval()
        
        # Dataset bilgileri
        self.dataset_info = {
            'name': 'allenai/s2-naip',
            'description': 'Sentinel-2 ve NAIP görüntü çiftleri',
            'resolution': '10m (Sentinel-2), 1m (NAIP)'
        }
    
    def predict(self, image):
        """Görüntüden koordinat tahmini"""
        try:
            # Görüntüyü yükle ve işle
            if isinstance(image, str):
                image = Image.open(image).convert('RGB')
            elif isinstance(image, np.ndarray):
                image = Image.fromarray(image.astype('uint8')).convert('RGB')
            
            # Transform uygula
            processed_image = self.transform(image).unsqueeze(0).to(self.device)
            
            with torch.no_grad():
                coordinates = self.model(processed_image)
                coords = coordinates.cpu().numpy()[0]
                
                # Gerçek koordinat aralığına uygun hale getir
                lat = float(coords[0] * 180 - 90)   # -90 ile 90 arası
                lon = float(coords[1] * 360 - 180)  # -180 ile 180 arası
                
                # Basit güven skoru (örnek)
                confidence = max(0.0, min(1.0, 1.0 - abs(coords[0]) - abs(coords[1])))
                
                result = {
                    'latitude': lat,
                    'longitude': lon,
                    'confidence': float(confidence),
                    'coordinates': [lat, lon],
                    'timestamp': datetime.now().isoformat(),
                    'dataset': self.dataset_info['name']
                }
                
                return result
                
        except Exception as e:
            logger.error(f"Tahmin hatası: {e}")
            return {
                'error': str(e),
                'latitude': 0.0,
                'longitude': 0.0,
                'confidence': 0.0
            }
    
    def load_sample_data(self, num_samples=5):
        """S2-NAIP datasetinden örnek veriler yükle"""
        try:
            dataset = load_dataset("allenai/s2-naip", split=f"train[:{num_samples}]")
            samples = []
            
            for i, item in enumerate(dataset):
                sample = {
                    'sentinel': item['sentinel'],
                    'naip': item['naip'],
                    'index': i,
                    'has_coords': 'lat' in item and 'lon' in item
                }
                
                if sample['has_coords']:
                    sample['lat'] = item['lat']
                    sample['lon'] = item['lon']
                
                samples.append(sample)
            
            logger.info(f"{len(samples)} örnek yüklendi")
            return samples
            
        except Exception as e:
            logger.error(f"Dataset yükleme hatası: {e}")
            return []

class GeoVisualizer:
    """Jeo-görselleştirme sınıfı"""
    
    def __init__(self):
        self.map_style = 'openstreetmap'
    
    def create_map(self, predictions, center=(39, 35), zoom=4):
        """Folium haritası oluştur"""
        try:
            m = folium.Map(location=center, zoom_start=zoom, tiles=self.map_style)
            
            for i, pred in enumerate(predictions):
                if 'error' in pred:
                    continue
                    
                lat, lon = pred['latitude'], pred['longitude']
                confidence = pred.get('confidence', 0.5)
                
                # Güven skoruna göre renk
                color = 'red' if confidence < 0.3 else 'orange' if confidence < 0.7 else 'green'
                
                popup_text = f"""
                <b>Tahmin {i+1}</b><br>
                <b>Koordinatlar:</b> {lat:.4f}, {lon:.4f}<br>
                <b>Güven:</b> {confidence:.2%}
                """
                
                folium.Marker(
                    [lat, lon],
                    popup=folium.Popup(popup_text, max_width=300),
                    tooltip=f"Güven: {confidence:.2%}",
                    icon=folium.Icon(color=color, icon='info-sign')
                ).add_to(m)
            
            # Haritayı HTML olarak kaydet
            with tempfile.NamedTemporaryFile(suffix='.html', delete=False) as tmp:
                m.save(tmp.name)
                return tmp.name
                
        except Exception as e:
            logger.error(f"Harita oluşturma hatası: {e}")
            return None
    
    def plot_prediction_comparison(self, predictions):
        """Tahmin karşılaştırması grafiği"""
        try:
            fig, ax = plt.subplots(figsize=(10, 6))
            
            latitudes = [p['latitude'] for p in predictions if 'error' not in p]
            longitudes = [p['longitude'] for p in predictions if 'error' not in p]
            confidences = [p.get('confidence', 0) for p in predictions if 'error' not in p]
            
            scatter = ax.scatter(longitudes, latitudes, c=confidences, 
                               cmap='viridis', s=100, alpha=0.7)
            
            ax.set_xlabel('Boylam')
            ax.set_ylabel('Enlem')
            ax.set_title('Koordinat Tahminleri ve Güven Skorları')
            ax.grid(True, alpha=0.3)
            
            # Renk barı ekle
            plt.colorbar(scatter, ax=ax, label='Güven Skoru')
            
            # Geçici dosyaya kaydet
            with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp:
                plt.savefig(tmp.name, dpi=150, bbox_inches='tight')
                plt.close()
                return tmp.name
                
        except Exception as e:
            logger.error(f"Grafik oluşturma hatası: {e}")
            return None

# Ana uygulama sınıfı
class S2NAIPGeoApp:
    def __init__(self):
        self.geo_system = S2NAIPGeoSystem()
        self.visualizer = GeoVisualizer()
        self.predictions_history = []
        
        # Örnek verileri yükle
        self.sample_data = self.geo_system.load_sample_data(3)
        logger.info("S2-NAIP Jeo-Referanslama Uygulaması başlatıldı")
    
    def process_image(self, image):
        """Tek görüntü işleme"""
        result = self.geo_system.predict(image)
        
        if 'error' not in result:
            self.predictions_history.append(result)
            
            # Harita oluştur
            map_path = self.visualizer.create_map([result])
            
            return result, map_path
        else:
            return result, None
    
    def process_batch(self, files):
        """Toplu işleme"""
        results = []
        
        for file in files:
            try:
                result = self.geo_system.predict(file.name)
                result['filename'] = os.path.basename(file.name)
                results.append(result)
            except Exception as e:
                results.append({
                    'filename': os.path.basename(file.name),
                    'error': str(e)
                })
        
        # Başarılı tahminler
        successful = [r for r in results if 'error' not in r]
        
        if successful:
            map_path = self.visualizer.create_map(successful)
            plot_path = self.visualizer.plot_prediction_comparison(successful)
        else:
            map_path = None
            plot_path = None
        
        batch_result = {
            'results': results,
            'summary': {
                'toplam_gorsel': len(files),
                'basarili_tahmin': len(successful),
                'basarisiz_tahmin': len(results) - len(successful),
                'ortalama_guven': np.mean([r.get('confidence', 0) for r in successful]) if successful else 0
            }
        }
        
        self.predictions_history.extend(successful)
        
        return batch_result, map_path, plot_path
    
    def get_sample_images(self):
        """Örnek görüntüleri getir"""
        samples = []
        for sample in self.sample_data:
            samples.append({
                'sentinel': sample['sentinel'],
                'naip': sample['naip'],
                'coordinates': f"Enlem: {sample.get('lat', 'Bilinmiyor')}, Boylam: {sample.get('lon', 'Bilinmiyor')}"
            })
        return samples

# Gradio arayüzü oluştur
def create_interface():
    app = S2NAIPGeoApp()
    
    with gr.Blocks(title="🌍 S2-NAIP Jeo-Referanslama Sistemi", theme=gr.themes.Soft()) as demo:
        gr.Markdown("""
        # 🌍 S2-NAIP Jeo-Referanslama Sistemi
        **Sentinel-2 ve NAIP görüntüleri için AI destekli koordinat tahmini**
        
        Bu sistem, `allenai/s2-naip` dataseti kullanılarak geliştirilmiştir.
        """)
        
        with gr.Tab("📍 Tek Görüntü Analizi"):
            with gr.Row():
                with gr.Column():
                    image_input = gr.Image(
                        type="filepath", 
                        label="Uydu Görüntüsü Yükle",
                        height=300
                    )
                    predict_btn = gr.Button("Koordinatları Tahmin Et", variant="primary")
                    
                with gr.Column():
                    output_json = gr.JSON(label="Tahmin Sonuçları")
                    output_map = gr.HTML(label="Harita Görünümü")
            
            # Örnekler
            gr.Markdown("### Örnek Görüntüler")
            with gr.Row():
                for i, sample in enumerate(app.sample_data):
                    with gr.Column():
                        gr.Image(
                            value=sample['sentinel'],
                            label=f"Sentinel-2 Örnek {i+1}",
                            height=150
                        )
                        gr.Textbox(
                            value=f"Koordinatlar: {sample.get('lat', 'Bilinmiyor')}, {sample.get('lon', 'Bilinmiyor')}",
                            label="Gerçek Koordinatlar"
                        )
            
            predict_btn.click(
                fn=app.process_image,
                inputs=image_input,
                outputs=[output_json, output_map]
            )
        
        with gr.Tab("📊 Toplu İşleme"):
            with gr.Row():
                with gr.Column():
                    batch_files = gr.File(
                        file_count="multiple",
                        file_types=[".jpg", ".jpeg", ".png", ".tiff"],
                        label="Toplu Görüntü Yükle"
                    )
                    batch_btn = gr.Button("Toplu İşle", variant="primary")
                
                with gr.Column():
                    batch_output = gr.JSON(label="Toplu Sonuçlar")
                    batch_map = gr.HTML(label="Toplu Harita")
                    batch_plot = gr.Image(label="Karşılaştırma Grafiği")
            
            batch_btn.click(
                fn=app.process_batch,
                inputs=batch_files,
                outputs=[batch_output, batch_map, batch_plot]
            )
        
        with gr.Tab("ℹ️ Dataset Bilgisi"):
            gr.Markdown("""
            ## S2-NAIP Dataset Hakkında
            
            **Dataset:** `allenai/s2-naip`
            
            **Açıklama:** 
            - Sentinel-2 (10m çözünürlük) ve NAIP (1m çözünürlük) görüntü çiftleri
            - ABD genelinde çeşitli lokasyonlar
            - Her örnek için koordinat bilgisi içerir
            
            **Özellikler:**
            - Multi-spektral Sentinel-2 görüntüleri
            - Yüksek çözünürlüklü NAIP görüntüleri  
            - Coğrafi koordinat metadata'sı
            
            **Kullanım:** Bu dataset, uydu görüntülerinden koordinat tahmini için kullanılmaktadır.
            """)
            
            # Dataset istatistikleri
            gr.Markdown("""
            ### Model Bilgisi
            - **Temel Model:** DINOv2 Base
            - **Input Çözünürlük:** 224x224
            - **Çıktı:** Enlem, Boylam koordinatları
            - **Güven Skoru:** Tahmin kalitesi göstergesi
            """)
        
        # Uyarı
        gr.Markdown("""
        ---
        ⚠️ **Not:** Bu demo için rastgele ağırlıklar kullanılmaktadır. Gerçek tahminler için modelin eğitilmesi gerekmektedir.
        """)
    
    return demo

if __name__ == "__main__":
    # Gradio arayüzünü başlat
    demo = create_interface()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )