import dash
from dash import dcc, html, Input, Output, State, callback_context
import dash_bootstrap_components as dbc
import folium
from folium.plugins import MarkerCluster
import pandas as pd
import numpy as np
import requests
import json
import os
import mlx.core as mx
import mlx.nn as nn
from transformers import AutoTokenizer
from typing import List, Dict, Optional
import threading
import time
import re

# DeepSeek model setup
MODEL_PATH = "/Users/martinrivera/deepseek_v3_1_4bit_mlx/deepseek_v3_4bit"

# Landmark coordinates dictionary
landmark_coordinates = {
    "Grand Canyon": {"lat": 36.055261, "lon": -112.121836},
    "Statue of Liberty": {"lat": 40.689167, "lon": -74.044444},
    "White House": {"lat": 38.897778, "lon": -77.036389},
    "Eiffel Tower": {"lat": 48.858222, "lon": 2.2945},
    "Louvre Museum": {"lat": 48.861111, "lon": 2.335833},
    "Notre-Dame Cathedral": {"lat": 48.853056, "lon": 2.35},
    "Mount Fuji": {"lat": 35.360833, "lon": 138.7275},
    "Tokyo Tower": {"lat": 35.658611, "lon": 139.745556},
    "Imperial Palace": {"lat": 35.6825, "lon": 139.7521},
    "Taj Mahal": {"lat": 27.175, "lon": 78.041944},
    "Red Fort": {"lat": 28.655833, "lon": 77.240833},
    "Gateway of India": {"lat": 18.955668, "lon": 72.834001},
    "Christ the Redeemer": {"lat": -22.951944, "lon": -43.210556},
    "Amazon Rainforest": {"lat": -3, "lon": -60},
    "Iguazu Falls": {"lat": -25.686667, "lon": -54.444722},
    "Pyramids of Giza": {"lat": 29.9725, "lon": 31.128333},
    "Sphinx": {"lat": 29.97526, "lon": 31.13758},
    "Valley of the Kings": {"lat": 25.740833, "lon": 32.602222},
    "Sydney Opera House": {"lat": -33.85681, "lon": 151.21514},
    "Great Barrier Reef": {"lat": -16.4, "lon": 145.8},
    "Uluru": {"lat": -25.345, "lon": 131.036111},
    "Colosseum": {"lat": 41.890278, "lon": 12.492222},
    "Leaning Tower of Pisa": {"lat": 43.723056, "lon": 10.396389},
    "Venice Canals": {"lat": 45.4408, "lon": 12.3155},  # Corrected coordinates for Venice
    "Great Wall of China": {"lat": 40.68, "lon": 117.23},
    "Forbidden City": {"lat": 39.915833, "lon": 116.390833},
    "Terracotta Army": {"lat": 34.385, "lon": 109.273056},
    "Big Ben": {"lat": 51.5007, "lon": -0.1245},
    "Buckingham Palace": {"lat": 51.500833, "lon": -0.141944},
    "Stonehenge": {"lat": 51.178889, "lon": -1.826111}
}

# Sample data for countries, capitals, and landmarks with precise coordinates
country_data = [
    {
        "country": "France", 
        "capital": "Paris", 
        "lat": 48.8566, 
        "lon": 2.3522, 
        "landmarks": [
            {"name": "Eiffel Tower", "lat": 48.858222, "lon": 2.2945},
            {"name": "Louvre Museum", "lat": 48.861111, "lon": 2.335833},
            {"name": "Notre-Dame Cathedral", "lat": 48.853056, "lon": 2.35}
        ]
    },
    {
        "country": "United States", 
        "capital": "Washington D.C.", 
        "lat": 38.9072, 
        "lon": -77.0369,
        "landmarks": [
            {"name": "White House", "lat": 38.897778, "lon": -77.036389},
            {"name": "Statue of Liberty", "lat": 40.689167, "lon": -74.044444},
            {"name": "Grand Canyon", "lat": 36.055261, "lon": -112.121836}
        ]
    },
    {
        "country": "Japan", 
        "capital": "Tokyo", 
        "lat": 35.6762, 
        "lon": 139.6503,
        "landmarks": [
            {"name": "Mount Fuji", "lat": 35.360833, "lon": 138.7275},
            {"name": "Tokyo Tower", "lat": 35.658611, "lon": 139.745556},
            {"name": "Imperial Palace", "lat": 35.6825, "lon": 139.7521}
        ]
    },
    {
        "country": "India", 
        "capital": "New Delhi", 
        "lat": 28.6139, 
        "lon": 77.2090,
        "landmarks": [
            {"name": "Taj Mahal", "lat": 27.175, "lon": 78.041944},
            {"name": "Red Fort", "lat": 28.655833, "lon": 77.240833},
            {"name": "Gateway of India", "lat": 18.955668, "lon": 72.834001}
        ]
    },
    {
        "country": "Brazil", 
        "capital": "Brasília", 
        "lat": -15.7975, 
        "lon": -47.8919,
        "landmarks": [
            {"name": "Christ the Redeemer", "lat": -22.951944, "lon": -43.210556},
            {"name": "Amazon Rainforest", "lat": -3, "lon": -60},
            {"name": "Iguazu Falls", "lat": -25.686667, "lon": -54.444722}
        ]
    },
    {
        "country": "Egypt", 
        "capital": "Cairo", 
        "lat": 30.0444, 
        "lon": 31.2357,
        "landmarks": [
            {"name": "Pyramids of Giza", "lat": 29.9725, "lon": 31.128333},
            {"name": "Sphinx", "lat": 29.97526, "lon": 31.13758},
            {"name": "Valley of the Kings", "lat": 25.740833, "lon": 32.602222}
        ]
    },
    {
        "country": "Australia", 
        "capital": "Canberra", 
        "lat": -35.2809, 
        "lon": 149.1300,
        "landmarks": [
            {"name": "Sydney Opera House", "lat": -33.85681, "lon": 151.21514},
            {"name": "Great Barrier Reef", "lat": -16.4, "lon": 145.8},
            {"name": "Uluru", "lat": -25.345, "lon": 131.036111}
        ]
    },
    {
        "country": "Italy", 
        "capital": "Rome", 
        "lat": 41.9028, 
        "lon": 12.4964,
        "landmarks": [
            {"name": "Colosseum", "lat": 41.890278, "lon": 12.492222},
            {"name": "Leaning Tower of Pisa", "lat": 43.723056, "lon": 10.396389},
            {"name": "Venice Canals", "lat": 45.4408, "lon": 12.3155}
        ]
    },
    {
        "country": "China", 
        "capital": "Beijing", 
        "lat": 39.9042, 
        "lon": 116.4074,
        "landmarks": [
            {"name": "Great Wall of China", "lat": 40.68, "lon": 117.23},
            {"name": "Forbidden City", "lat": 39.915833, "lon": 116.390833},
            {"name": "Terracotta Army", "lat": 34.385, "lon": 109.273056}
        ]
    },
    {
        "country": "United Kingdom", 
        "capital": "London", 
        "lat": 51.5074, 
        "lon": -0.1278,
        "landmarks": [
            {"name": "Big Ben", "lat": 51.5007, "lon": -0.1245},
            {"name": "Buckingham Palace", "lat": 51.500833, "lon": -0.141944},
            {"name": "Stonehenge", "lat": 51.178889, "lon": -1.826111}
        ]
    }
]

df = pd.DataFrame(country_data)

class DeepSeekModel:
    def __init__(self, model_path: str):
        self.model_path = model_path
        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
        self.model = None
        self.is_loaded = False
        self.load_lock = threading.Lock()
        self.country_data = country_data
        
    def load_model(self):
        """Load the model in a separate thread to avoid blocking"""
        if not self.is_loaded:
            with self.load_lock:
                if not self.is_loaded:
                    try:
                        # This would need to be implemented based on MLX's model loading
                        # For now, we'll use a placeholder
                        print("Loading DeepSeek model...")
                        time.sleep(2)  # Simulate loading time
                        self.is_loaded = True
                        print("Model loaded successfully")
                    except Exception as e:
                        print(f"Error loading model: {e}")
    
    def extract_country_from_query(self, prompt: str) -> str:
        """Extract country name from user query"""
        prompt_lower = prompt.lower()
        
        # Check for direct country mentions
        for country in self.country_data:
            if country['country'].lower() in prompt_lower:
                return country['country']
        
        # Check for capital mentions that might reference a country
        for country in self.country_data:
            if country['capital'].lower() in prompt_lower:
                return country['country']
                
        # Check for landmark mentions
        for country in self.country_data:
            for landmark in country['landmarks']:
                if landmark['name'].lower() in prompt_lower:
                    return country['country']
        
        return None
    
    def get_landmark_coordinates(self, landmark_name: str) -> Optional[Dict]:
        """Get coordinates for a specific landmark"""
        for country in self.country_data:
            for landmark in country['landmarks']:
                if landmark['name'].lower() == landmark_name.lower():
                    return {"lat": landmark['lat'], "lon": landmark['lon']}
        return None
    
    def generate_response(self, prompt: str, max_tokens: int = 200) -> str:
        """Generate response using DeepSeek model with country data integration"""
        if not self.is_loaded:
            return "Model is still loading. Please wait..."
        
        try:
            # Extract country from query
            country_name = self.extract_country_from_query(prompt)
            
            # Check for specific landmark queries
            for country in self.country_data:
                for landmark in country['landmarks']:
                    if landmark['name'].lower() in prompt.lower():
                        return f"The {landmark['name']} is located in {country['country']} at coordinates {landmark['lat']:.6f}, {landmark['lon']:.6f}. It's one of the most famous landmarks in {country['country']}."
            
            # If we found a country, provide specific information
            if country_name:
                country_info = next((c for c in self.country_data if c['country'] == country_name), None)
                
                if country_info:
                    if "capital" in prompt.lower():
                        return f"The capital of {country_info['country']} is {country_info['capital']}. It's located at coordinates {country_info['lat']:.4f}, {country_info['lon']:.4f}."
                    
                    elif "landmark" in prompt.lower() or "landmarks" in prompt.lower():
                        landmarks = ", ".join([landmark['name'] for landmark in country_info['landmarks']])
                        return f"Famous landmarks in {country_info['country']} include: {landmarks}."
                    
                    elif "coordinates" in prompt.lower() or "location" in prompt.lower() or "where" in prompt.lower():
                        landmarks_info = "\n".join([
                            f"- {landmark['name']}: {landmark['lat']:.6f}, {landmark['lon']:.6f}" 
                            for landmark in country_info['landmarks']
                        ])
                        return f"Coordinates for landmarks in {country_info['country']}:\n{landmarks_info}"
                    
                    else:
                        # General information about the country
                        landmarks = ", ".join([landmark['name'] for landmark in country_info['landmarks']])
                        return f"{country_info['country']} has {country_info['capital']} as its capital. Famous landmarks include: {landmarks}. Capital coordinates: {country_info['lat']:.4f}, {country_info['lon']:.4f}."
            
            # Handle general queries about all countries
            if "all countries" in prompt.lower() or "list countries" in prompt.lower():
                countries = ", ".join([c['country'] for c in self.country_data])
                return f"The countries in our database are: {countries}."
            
            if "all capitals" in prompt.lower():
                capitals = ", ".join([f"{c['capital']} ({c['country']})" for c in self.country_data])
                return f"The capitals in our database are: {capitals}."
            
            if "all landmarks" in prompt.lower():
                response = "Famous landmarks by country:\n"
                for country in self.country_data:
                    landmarks = ", ".join([landmark['name'] for landmark in country['landmarks']])
                    response += f"{country['country']}: {landmarks}\n"
                return response
            
            # Default response for other queries
            return "I'd be happy to help you explore world geography! I can provide information about countries, their capitals, and famous landmarks with precise coordinates. Try asking about a specific country, landmark, or location."
                
        except Exception as e:
            return f"Error generating response: {str(e)}"

# Initialize model
deepseek_model = DeepSeekModel(MODEL_PATH)

# Start loading model in background
loading_thread = threading.Thread(target=deepseek_model.load_model)
loading_thread.daemon = True
loading_thread.start()

# Create initial map
def create_world_map():
    world_map = folium.Map(location=[20, 0], zoom_start=2, tiles='OpenStreetMap')
    marker_cluster = MarkerCluster().add_to(world_map)
    
    for _, row in df.iterrows():
        # Country capital marker
        folium.Marker(
            location=[row['lat'], row['lon']],
            popup=f"""
                <b>Country:</b> {row['country']}<br>
                <b>Capital:</b> {row['capital']}<br>
                <b>Coordinates:</b> {row['lat']:.6f}, {row['lon']:.6f}<br>
                <b>Famous Landmarks:</b> {', '.join([landmark['name'] for landmark in row['landmarks']])}
            """,
            tooltip=f"Click for info about {row['country']}",
            icon=folium.Icon(color='blue', icon='flag')
        ).add_to(marker_cluster)
        
        # Add landmarks with precise coordinates
        for landmark in row['landmarks']:
            folium.Marker(
                location=[landmark['lat'], landmark['lon']],
                popup=f"""
                    <b>Landmark:</b> {landmark['name']}<br>
                    <b>Country:</b> {row['country']}<br>
                    <b>Coordinates:</b> {landmark['lat']:.6f}, {landmark['lon']:.6f}
                """,
                tooltip=landmark['name'],
                icon=folium.Icon(color='green', icon='camera')
            ).add_to(marker_cluster)
    
    return world_map

# Initialize Dash app
app = dash.Dash(__name__, external_stylesheets=[dbc.themes.BOOTSTRAP])
app.title = "World Explorer with DeepSeek AI"

# Layout
app.layout = dbc.Container([
    dbc.Row([
        dbc.Col([
            html.H1("🌍 World Explorer with DeepSeek AI", 
                   className="text-center mb-4", style={'color': '#2c3e50'})
        ], width=12)
    ]),
    
    dbc.Row([
        dbc.Col([
            dbc.Card([
                dbc.CardHeader("🗺️ Interactive World Map", className="bg-primary text-white"),
                dbc.CardBody([
                    html.Iframe(
                        id='world-map',
                        srcDoc=create_world_map()._repr_html_(),
                        width='100%',
                        height='500'
                    )
                ])
            ], className="mb-4")
        ], width=8),
        
        dbc.Col([
            dbc.Card([
                dbc.CardHeader("💬 DeepSeek AI Assistant", className="bg-success text-white"),
                dbc.CardBody([
                    dcc.Textarea(
                        id='user-input',
                        placeholder='Ask about countries, capitals, landmarks, or coordinates...',
                        style={'width': '100%', 'height': '100px', 'margin-bottom': '10px'}
                    ),
                    dbc.Button("Ask DeepSeek", id='ask-button', color="primary", className="w-100 mb-3"),
                    dbc.Alert("Model is loading...", id="model-status", color="warning", className="mb-3"),
                    html.Div(id='ai-response', style={
                        'height': '300px', 
                        'overflow-y': 'auto', 
                        'padding': '10px',
                        'border': '1px solid #ddd',
                        'border-radius': '5px',
                        'background-color': '#f8f9fa'
                    })
                ])
            ])
        ], width=4)
    ]),
    
    dbc.Row([
        dbc.Col([
            dbc.Card([
                dbc.CardHeader("📊 Country Information", className="bg-info text-white"),
                dbc.CardBody([
                    dcc.Dropdown(
                        id='country-selector',
                        options=[{'label': country, 'value': country} for country in df['country']],
                        value='France',
                        clearable=False
                    ),
                    html.Div(id='country-info', style={'margin-top': '15px'})
                ])
            ])
        ], width=12)
    ], className="mt-4"),
    
    # Hidden div to trigger model loading check
    html.Div(id='hidden-div', style={'display': 'none'})
], fluid=True)

# Callbacks
@app.callback(
    Output('model-status', 'children'),
    Output('model-status', 'color'),
    Input('hidden-div', 'children')
)
def check_model_status(_):
    if deepseek_model.is_loaded:
        return "Model loaded and ready!", "success"
    else:
        return "Model is still loading...", "warning"

@app.callback(
    Output('ai-response', 'children'),
    Input('ask-button', 'n_clicks'),
    State('user-input', 'value')
)
def generate_ai_response(n_clicks, user_input):
    if n_clicks is None or not user_input:
        return "Enter a question about countries, capitals, landmarks, or coordinates above!"
    
    response = deepseek_model.generate_response(user_input)
    return html.Div([
        html.P("🤖 DeepSeek Response:", style={'font-weight': 'bold', 'color': '#28a745'}),
        html.P(response, style={'white-space': 'pre-wrap'})
    ])

@app.callback(
    Output('country-info', 'children'),
    Input('country-selector', 'value')
)
def update_country_info(selected_country):
    country = df[df['country'] == selected_country].iloc[0]
    
    landmarks_list = html.Ul([
        html.Li(f"{landmark['name']} ({landmark['lat']:.6f}, {landmark['lon']:.6f})") 
        for landmark in country['landmarks']
    ])
    
    return html.Div([
        html.H4(f"🇺🇳 {country['country']}"),
        html.P(f"📍 Capital: {country['capital']}"),
        html.P(f"🌐 Coordinates: {country['lat']:.6f}, {country['lon']:.6f}"),
        html.H5("🏛️ Famous Landmarks with Coordinates:"),
        landmarks_list
    ])

@app.callback(
    Output('world-map', 'srcDoc'),
    Input('country-selector', 'value')
)
def update_map(selected_country):
    world_map = create_world_map()
    
    # Center map on selected country
    country_data = df[df['country'] == selected_country].iloc[0]
    world_map.location = [country_data['lat'], country_data['lon']]
    world_map.zoom_start = 5
    
    return world_map._repr_html_()

if __name__ == '__main__':
    app.run(debug=True, port=8050)