import streamlit as st
import geopandas as gpd
import pydeck as pdk
import pandas as pd
from branca import colormap as cm
import pathlib
import os
import requests
import json
from shapely.geometry import shape
import matplotlib.pyplot as plt
import io
import base64
import numpy as np

st.set_page_config(layout="wide")

# Data source
data_links = {
    "Tokyo": "https://raw.githubusercontent.com/kunifujiwara/data/master/frac_veg/FRAC_VEG_Tokyo.geojson",
    "Kanagawa": "https://raw.githubusercontent.com/kunifujiwara/data/master/frac_veg/FRAC_VEG_Kanagawa.geojson",
    "Chiba": "https://raw.githubusercontent.com/kunifujiwara/data/master/frac_veg/FRAC_VEG_Chiba.geojson",
    "Saitama": "https://raw.githubusercontent.com/kunifujiwara/data/master/frac_veg/FRAC_VEG_Saitama.geojson",
}

@st.cache_data
def get_geom_data(prefecture):
    response = requests.get(data_links[prefecture])
    if response.status_code == 200:
        geojson_data = json.loads(response.content)
        gdf = gpd.GeoDataFrame.from_features(geojson_data['features'])
        return gdf
    else:
        st.error(f"Failed to fetch data for {prefecture}. Status code: {response.status_code}")
        return None

def get_data_columns(gdf):
    return gdf.select_dtypes(include=[float, int]).columns.tolist()

def create_colormap_legend(vmin, vmax, cmap, n_colors, attribute_name):
    fig, ax = plt.subplots(figsize=(6, 0.8))
    fig.subplots_adjust(bottom=0.5)

    norm = plt.Normalize(vmin=vmin, vmax=vmax)
    cbar = fig.colorbar(plt.cm.ScalarMappable(norm=norm, cmap=cmap), cax=ax, orientation='horizontal')
    cbar.set_ticks([vmin, (vmin+vmax)/2, vmax])
    cbar.set_ticklabels([f'{vmin:.2f}', f'{(vmin+vmax)/2:.2f}', f'{vmax:.2f}'])
    ax.set_title(f"{attribute_name}", fontsize=10, pad=10)
    
    buf = io.BytesIO()
    plt.savefig(buf, format='png', dpi=300, bbox_inches='tight')
    plt.close(fig)
    return base64.b64encode(buf.getvalue()).decode()

def calculate_zoom_level(bbox):
    lon_range = bbox[2] - bbox[0]
    lat_range = bbox[3] - bbox[1]
    max_range = max(lon_range, lat_range)
    zoom = int(np.log2(360 / max_range)) - 1
    return min(max(1, zoom), 20)  # Clamp zoom between 1 and 20

def app():
    st.title("Japan Vegetation Cover Fraction")
    st.markdown(
        """**Introduction:** This interactive dashboard visualizes Japan Fractional Vegetation Cover at town block levels."""
    )

    prefecture = st.selectbox("Prefecture", ["Tokyo", "Kanagawa", "Chiba", "Saitama"])
    
    gdf = get_geom_data(prefecture)

    if gdf is None:
        st.error("Failed to load data. Please try again later.")
        return

    # City filter
    cities = sorted(gdf['CITY_NAME_x'].unique().tolist())
    selected_cities = st.multiselect("Select cities to display", cities, default=[])
    
    # Filter GeoDataFrame based on selected cities
    if selected_cities:
        gdf_filtered = gdf[gdf['CITY_NAME_x'].isin(selected_cities)]
    else:
        gdf_filtered = gdf

    attributes = get_data_columns(gdf_filtered)
    selected_attribute = st.selectbox("Select attribute to visualize", attributes)

    col1, col2, col3, col4, col5, col6 = st.columns(6)
    with col1:
        n_colors = st.slider("Number of colors", min_value=2, max_value=20, value=8)
    with col2:
        alpha = st.slider("Fill opacity", min_value=0.0, max_value=1.0, value=0.8, step=0.1)
    with col3:
        vmin = st.number_input("Min value", value=float(gdf_filtered[selected_attribute].min()), step=0.1)
    with col4:
        vmax = st.number_input("Max value", value=float(gdf_filtered[selected_attribute].max()), step=0.1)
    with col5:
        show_3d = st.checkbox("Show 3D view", value=False)
    with col6:
        if show_3d:
            elev_scale = st.slider("Elevation scale", min_value=1, max_value=10000, value=1, step=10)
        else:
            elev_scale = 1

    # Create color scale
    cmap = plt.get_cmap('Greens')
    colors = [cmap(i / (n_colors - 1)) for i in range(n_colors)]
    hex_colors = ['#%02x%02x%02x' % (int(r * 255), int(g * 255), int(b * 255)) for r, g, b, _ in colors]

    color_scale = cm.LinearColormap(colors=hex_colors, vmin=vmin, vmax=vmax)
    
    gdf_filtered['color'] = gdf_filtered[selected_attribute].apply(lambda x: color_scale(x))
    gdf_filtered['color'] = gdf_filtered['color'].apply(lambda x: [int(x[1:3], 16), int(x[3:5], 16), int(x[5:7], 16)] + [int(alpha * 255)])

    layer = pdk.Layer(
        "GeoJsonLayer",
        gdf_filtered.__geo_interface__,
        pickable=True,
        opacity=alpha,
        stroked=True,
        filled=True,
        extruded=show_3d,
        wireframe=True,
        get_elevation=f"properties.{selected_attribute}" if show_3d else None,
        elevation_scale=elev_scale if show_3d else 1,
        get_fill_color="properties.color",
        get_line_color=[0, 0, 0],
        get_line_width=2,
        line_width_min_pixels=1,
    )

    # Calculate bounding box and zoom level
    bbox = gdf_filtered.total_bounds
    zoom = calculate_zoom_level(bbox)

    view_state = pdk.ViewState(
        latitude=(bbox[1] + bbox[3]) / 2,
        longitude=(bbox[0] + bbox[2]) / 2,
        zoom=zoom,
        pitch=45 if show_3d else 0,
    )

    r = pdk.Deck(
        layers=[layer],
        initial_view_state=view_state,
        map_style="mapbox://styles/mapbox/light-v9",
        tooltip={"text": "{CITY_NAME_x}\n{" + selected_attribute + "}"}
    )

    st.pydeck_chart(r)

    # Create and display color scale legend
    legend_img = create_colormap_legend(vmin, vmax, cmap, n_colors, selected_attribute)

    # Create three columns, with the middle one being 30% width
    left_spacer, center_col, right_spacer = st.columns([1, 1, 1])
    
    # Display the legend in the middle column
    with right_spacer:
        st.image(f"data:image/png;base64,{legend_img}", caption="Color Scale", use_column_width=True)

    if st.checkbox("Show raw data"):
        st.write(gdf_filtered[['CITY_NAME_x', selected_attribute]])

app()