import requests
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
from io import BytesIO
import math


def fetch_tile(x, y, z):
    """
    Fetches an RGB-encoded elevation tile from a tile server.
    Args:
        x (int): X-coordinate of the tile.
        y (int): Y-coordinate of the tile.
        z (int): Zoom level of the tile.
    Returns:
        Image: An RGB image of the elevation tile.
    """
    url = f"https://elevation-tiles-prod.s3.amazonaws.com/terrarium/{z}/{x}/{y}.png"
    response = requests.get(url)
    if response.status_code == 200:
        return Image.open(BytesIO(response.content))
    else:
        raise ValueError(f"Failed to fetch tile: {url}")


def rgb_to_elevation(rgb):
    """
    Convert an RGB value to an elevation value using Terrarium encoding.
    Args:
        rgb (tuple): RGB tuple.
    Returns:
        float: Elevation in meters.
    """
    r, g, b = rgb
    elevation = (r * 256 + g + b / 256) - 32768
    return elevation


def tile_to_elevation_array(tile_image):
    """
    Converts an RGB tile image to a numpy array of elevation values.
    Args:
        tile_image (Image): An RGB image representing a tile.
    Returns:
        numpy.ndarray: A 2D array of elevation values.
    """
    rgb_array = np.array(tile_image)
    elevation_array = np.apply_along_axis(rgb_to_elevation, 2, rgb_array)
    return elevation_array


def plot_elevation_hsb(elevation_array):
    """
    Plots elevation data using an HSB scientific color spectrum.
    Args:
        elevation_array (numpy.ndarray): A 2D array of elevation values.
    """
    max_elev = np.max(elevation_array)
    min_elev = np.min(elevation_array)

    if max_elev == min_elev:
        print("No elevation data available for this tile.")
        return

    # Lop off top 45 in hue space and normalize elevation for color mapping
    scale = 210 / (max_elev - min_elev)
    elevation_normalized = (max_elev - elevation_array) * scale

    # Use Hue-Saturation-Brightness mapping
    hsv_image = np.zeros((*elevation_array.shape, 3))
    hsv_image[..., 0] = elevation_normalized / 360.0  # Hue varies with elevation, normalized to [0, 1]
    hsv_image[..., 1] = 1.0  # Full saturation
    hsv_image[..., 2] = 1.0  # Full brightness

    # Convert HSV to RGB for visualization
    rgb_image = plt.cm.hsv(hsv_image[..., 0])[:, :, :3]  # Only take RGB, ignore alpha

    # Plot the image
    plt.figure(figsize=(10, 10))
    plt.imshow(rgb_image)
    plt.axis('off')
    plt.title("Elevation Visualization with HSB Spectrum")
    plt.show()


def tile_zxy_from_lat_long_zoom(lat, long, z):
    """
    Converts latitude, longitude, and zoom level to tile indices.
    Args:
        lat (float): Latitude in degrees.
        long (float): Longitude in degrees.
        z (int): Zoom level.
    Returns:
        tuple: Tile indices (z, x, y).
    """
    lat_rad = math.radians(lat)
    n = 2 ** z
    x_tile = int(n * ((long + 180) / 360))
    y_tile = int(n * (1 - (math.log(math.tan(lat_rad) + (1 / math.cos(lat_rad))) / math.pi)) / 2)
    return z, x_tile, y_tile


def render_riudoso_elevation():
    """
    Fetches and renders the elevation data for Ruidoso, NM using a scientific HSB color spectrum.
    """
    # Latitude, longitude, and zoom level for Ruidoso, NM
    lat, long, z = 33.339077, -105.678695, 11
    z, x, y = tile_zxy_from_lat_long_zoom(lat, long, z)
    tile_image = fetch_tile(x, y, z)
    elevation_array = tile_to_elevation_array(tile_image)
    plot_elevation_hsb(elevation_array)


if __name__ == "__main__":
    render_riudoso_elevation()