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

st.set_page_config(layout="wide")

st.sidebar.info(
    """
    - Web App URL: <https://streamlit.gishub.org>
    - GitHub repository: <https://github.com/giswqs/streamlit-geospatial>
    """
)

st.sidebar.title("Contact")
st.sidebar.info(
    """
    Qiusheng Wu at [wetlands.io](https://wetlands.io) | [GitHub](https://github.com/giswqs) | [Twitter](https://twitter.com/giswqs) | [YouTube](https://www.youtube.com/c/QiushengWu) | [LinkedIn](https://www.linkedin.com/in/qiushengwu)
    """
)

STREAMLIT_STATIC_PATH = pathlib.Path(st.__path__[0]) / "static"
DOWNLOADS_PATH = STREAMLIT_STATIC_PATH / "downloads"
if not DOWNLOADS_PATH.is_dir():
    DOWNLOADS_PATH.mkdir()

# 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",
}

# Predefined list of color palettes
color_palettes = ['viridis', 'plasma', 'inferno', 'magma', 'cividis', 'Blues', 'Greens', 'Reds', 'Oranges', 'Purples']

@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)
        features = geojson_data['features']
        
        data = []
        for feature in features:
            properties = feature['properties']
            geometry = shape(feature['geometry'])
            properties['geometry'] = geometry
            data.append(properties)
        
        gdf = gpd.GeoDataFrame(data)
        gdf.set_geometry('geometry', inplace=True)
        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 select_non_null(gdf, col_name):
    return gdf[~gdf[col_name].isna()]

def select_null(gdf, col_name):
    return gdf[gdf[col_name].isna()]

def app():
    st.title("Japan Vegetation Cover Fraction")
    st.markdown(
        """**Introduction:** This interactive dashboard is designed for visualizing Japan Fractional Vegetation Cover at town block levels. 
        The data sources include [Vegetation Cover Fraction](https://zenodo.org/records/5553516) from a research project (https://doi.org/10.3130/aijt.28.521), 
        and [Cartographic Boundary Files](https://www.e-stat.go.jp/gis/statmap-search?page=1&type=2&aggregateUnitForBoundary=A&toukeiCode=00200521&toukeiYear=2015&serveyId=A002005212015&coordsys=1&format=shape&datum=2000) from Census of Japan 2015.
        """
    )

    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

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

    row2_col1, row2_col2, row2_col3, row2_col4, row2_col5, row2_col6 = st.columns(
        [0.6, 0.68, 0.7, 0.7, 1.5, 0.8]
    )

    with row2_col1:
        palette = st.selectbox("Color palette", color_palettes, index=color_palettes.index('Blues'))
    with row2_col2:
        n_colors = st.slider("Number of colors", min_value=2, max_value=20, value=8)
    with row2_col3:
        show_nodata = st.checkbox("Show nodata areas", value=True)
    with row2_col4:
        show_3d = st.checkbox("Show 3D view", value=False)
    with row2_col5:
        if show_3d:
            elev_scale = st.slider(
                "Elevation scale", min_value=1, max_value=1000000, value=1, step=10
            )
            with row2_col6:
                st.info("Press Ctrl and move the left mouse button.")
        else:
            elev_scale = 1

    color_scale = cm.LinearColormap(colors=palette, vmin=gdf[selected_attribute].min(), vmax=gdf[selected_attribute].max(), n=n_colors)
    gdf['color'] = gdf[selected_attribute].apply(lambda x: color_scale(x))
    gdf['color'] = gdf['color'].apply(lambda x: [int(x[1:3], 16), int(x[3:5], 16), int(x[5:7], 16)])

    gdf_null = select_null(gdf, selected_attribute)
    gdf = select_non_null(gdf, selected_attribute)

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

    if show_nodata:
        nodata_layer = pdk.Layer(
            "GeoJsonLayer",
            gdf_null.__geo_interface__,
            pickable=True,
            opacity=0.2,
            stroked=True,
            filled=True,
            extruded=False,
            get_fill_color=[200, 200, 200],
            get_line_color=[0, 0, 0],
            get_line_width=2,
            line_width_min_pixels=1,
        )
        layers = [layer, nodata_layer]
    else:
        layers = [layer]

    view_state = pdk.ViewState(
        latitude=gdf.geometry.centroid.y.mean(),
        longitude=gdf.geometry.centroid.x.mean(),
        zoom=9,
        pitch=45 if show_3d else 0,
    )

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

    st.pydeck_chart(r)

    st.write(color_scale)

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

app()