import os
import json
import numpy as np
import torch
from PIL import Image, ImageDraw
import gradio as gr
from openai import OpenAI
from geopy.geocoders import Nominatim
from staticmap import StaticMap, CircleMarker, Polygon
from diffusers import ControlNetModel, StableDiffusionControlNetInpaintPipeline
import spaces
import logging
from pydantic import BaseModel, ValidationError, Field
from typing import List, Union

# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Initialize APIs
openai_client = OpenAI(api_key=os.environ['OPENAI_API_KEY'])
geolocator = Nominatim(user_agent="geoapi")

# Define Pydantic models for GeoJSON validation
class Geometry(BaseModel):
    type: str
    coordinates: Union[List[float], List[List[float]]]

class Feature(BaseModel):
    type: str = "Feature"
    properties: dict
    geometry: Geometry

class FeatureCollection(BaseModel):
    type: str = "FeatureCollection"
    features: List[Feature]

# Function to fetch coordinates
@spaces.GPU
def get_geo_coordinates(location_name):
    try:
        location = geolocator.geocode(location_name)
        if location:
            return [location.longitude, location.latitude]
        return None
    except Exception as e:
        logger.error(f"Error fetching coordinates for {location_name}: {e}")
        return None

# Function to process OpenAI chat response
@spaces.GPU
def process_openai_response(query):
    response = openai_client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {
                "role": "system",
                "content": """
You are an assistant that generates structured JSON output for geographical queries. Your task is to generate a JSON object containing information about geographical features and their representation based on the user's query. Follow these rules:

1. The JSON should always have the following structure:
    {
        "input": "<user's query>",
        "output": {
            "answer": "<concise text answering the query>",
            "feature_representation": {
                "type": "<one of: Point, LineString, Polygon, MultiPoint, MultiLineString, MultiPolygon, GeometryCollection>",
                "cities": ["<list of city names>"],
                "properties": {
                    "description": "<a prompt for a diffusion model describing the geographical feature>"
                }
            }
        }
    }

2. For the `type` field in `feature_representation`:
    - Use "Point" for single city queries.
    - Use "MultiPoint" for queries involving multiple cities not forming a line or area.
    - Use "LineString" for queries about paths between two or more cities.
    - Use "Polygon" for queries about areas formed by three or more cities.

3. For the `cities` field:
    - List the names of cities mentioned in the query in the order they appear.
    - If no cities are mentioned, leave the array empty.

4. For the `properties.description` field:
    - Describe the geographical feature in a creative way, suitable for generating an image with a diffusion model.

### Example Input:
"Mark a triangular area using New York, Boston, and Philadelphia."

### Example Output:
{
    "input": "Mark a triangular area using New York, Boston, and Philadelphia.",
    "output": {
        "answer": "The cities New York, Boston, and Philadelphia form a triangle.",
        "feature_representation": {
            "type": "Polygon",
            "cities": ["New York", "Boston", "Philadelphia"],
            "properties": {
                "description": "A satellite image of a triangular area formed by New York, Boston, and Philadelphia, with green fields and urban regions, 4k resolution, highly detailed."
            }
        }
    }
}

Generate similar JSON for the following query:
"""
            },
            {
                "role": "user",
                "content": query
            }
        ],
        temperature=1,
        max_tokens=2048,
        top_p=1,
        frequency_penalty=0,
        presence_penalty=0,
        response_format={"type": "json_object"}
    )
    return json.loads(response.choices[0].message.content)

# Generate GeoJSON from OpenAI response
@spaces.GPU
def generate_geojson(response):
    #log response
    logger.info(f"OpenAI response: {response}")
    feature_type = response['output']['feature_representation']['type']
    city_names = response['output']['feature_representation']['cities']
    properties = response['output']['feature_representation']['properties']

    coordinates = []

    # Fetch coordinates for cities
    for city in city_names:
        try:
            coord = get_geo_coordinates(city)  # Function to fetch city coordinates
            if coord:
                coordinates.append(coord)
            else:
                logger.warning(f"Coordinates not found for city: {city}")
        except Exception as e:
            logger.error(f"Error fetching coordinates for {city}: {e}")

    if feature_type == "Polygon":
        if len(coordinates) < 3:
            raise ValueError("Polygon requires at least 3 coordinates.")
        # Close the polygon by appending the first point at the end
        coordinates.append(coordinates[0])
        coordinates = [coordinates]  # Nest coordinates for Polygon

    # Create the GeoJSON object
    geojson_data = {
        "type": "FeatureCollection",
        "features": [
            {
                "type": "Feature",
                "properties": properties,
                "geometry": {
                    "type": feature_type,
                    "coordinates": coordinates,
                },
            }
        ],
    }

    # Validate the GeoJSON
    try:
        validated_geojson = FeatureCollection(**geojson_data)
        return validated_geojson.dict()
    except ValidationError as e:
        logger.error(f"Invalid GeoJSON data: {e}")
        raise ValueError("Generated GeoJSON is invalid.")

# Generate static map image
@spaces.GPU
def generate_static_map(geojson_data, invisible=False):
    m = StaticMap(600, 600)
    logger.info(f"GeoJSON data: {geojson_data}")

    for feature in geojson_data["features"]:
        geom_type = feature["geometry"]["type"]
        coords = feature["geometry"]["coordinates"]

        if geom_type == "Point":
            m.add_marker(CircleMarker((coords[0][0], coords[0][1]), '#1C00ff00' if invisible else 'blue', 1000))
        elif geom_type in ["MultiPoint", "LineString"]:
            for coord in coords:
                m.add_marker(CircleMarker((coord[0], coord[1]), '#1C00ff00' if invisible else 'blue', 1000))
        elif geom_type in ["Polygon", "MultiPolygon"]:
            for polygon in coords:
                m.add_polygon(Polygon([(c[0], c[1]) for c in polygon], '#1C00ff00' if invisible else 'blue', 3))

    return m.render()

# ControlNet pipeline setup
controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_inpaint", torch_dtype=torch.float16)
pipeline = StableDiffusionControlNetInpaintPipeline.from_pretrained(
    "stable-diffusion-v1-5/stable-diffusion-inpainting", controlnet=controlnet, torch_dtype=torch.float16
)
pipeline.to('cuda')

@spaces.GPU
def make_inpaint_condition(init_image, mask_image):
    init_image = np.array(init_image.convert("RGB")).astype(np.float32) / 255.0
    mask_image = np.array(mask_image.convert("L")).astype(np.float32) / 255.0

    assert init_image.shape[0:1] == mask_image.shape[0:1], "image and image_mask must have the same image size"
    init_image[mask_image > 0.5] = -1.0  # set as masked pixel
    init_image = np.expand_dims(init_image, 0).transpose(0, 3, 1, 2)
    init_image = torch.from_numpy(init_image)
    return init_image

@spaces.GPU
def generate_satellite_image(init_image, mask_image, prompt):
    control_image = make_inpaint_condition(init_image, mask_image)
    result = pipeline(
        prompt=prompt, 
        image=init_image, 
        mask_image=mask_image, 
        control_image=control_image,
        strength=0.42,
        guidance_scale=62
    )
    return result.images[0]

# Gradio UI
@spaces.GPU
def handle_query(query):
    response = process_openai_response(query)
    geojson_data = generate_geojson(response)

    map_image = generate_static_map(geojson_data)
    empty_map_image = generate_static_map(geojson_data, invisible=True)

    difference = np.abs(np.array(map_image.convert("RGB")) - np.array(empty_map_image.convert("RGB")))
    threshold = 10
    mask = (np.sum(difference, axis=-1) > threshold).astype(np.uint8) * 255

    mask_image = Image.fromarray(mask, mode="L")
    satellite_image = generate_satellite_image(
        empty_map_image, mask_image, response['output']['feature_representation']['properties']['description']
    )

    return map_image, empty_map_image, satellite_image, mask_image, response

def update_query(selected_query):
    return selected_query

query_options = [
    "Area covering south asian subcontinent",
    "Due to considerable rainfall in the up- and mid- stream areas of Kala Oya, the Rajanganaya reservoir is now spilling at a rate of 17,000 cubic feet per second, the department said."
]

with gr.Blocks() as demo:
    with gr.Row():
        selected_query = gr.Dropdown(label="Select Query", choices=query_options, value=query_options[-1])
        query_input = gr.Textbox(label="Enter Query", value=query_options[-1])
        selected_query.change(update_query, inputs=selected_query, outputs=query_input)
        submit_btn = gr.Button("Submit")
    with gr.Row():
        map_output = gr.Image(label="Map Visualization")
        empty_map_output = gr.Image(label="Empty Visualization")
    with gr.Row():
        satellite_output = gr.Image(label="Generated Satellite Image")
        mask_output = gr.Image(label="Mask")
        image_prompt = gr.Textbox(label="Image Prompt Used")
    submit_btn.click(handle_query, inputs=[query_input], outputs=[map_output, empty_map_output, satellite_output, mask_output, image_prompt])

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