# -*- coding: utf-8 -*-
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import leafmap.foliumap as leafmap
import streamlit as st
import altair as alt
import ibis
from ibis import _
import ibis.selectors as s

st.set_page_config(layout="wide", page_title="Protected Areas Explorer", page_icon=":globe:")

#pad_pmtiles = "https://data.source.coop/cboettig/pad-us-3/pad-stats.pmtiles"
#parquet = "https://data.source.coop/cboettig/pad-us-3/pad-stats.parquet"
#pad_pmtiles = "https://huggingface.co/spaces/boettiger-lab/pad-us/resolve/main/pad-stats.pmtiles"
#parquet = "https://huggingface.co/spaces/boettiger-lab/pad-us/resolve/main/pad-stats.parquet"
pad_pmtiles = "https://huggingface.co/datasets/boettiger-lab/pad-us-3/resolve/main/pad-stats.pmtiles"
parquet = "https://huggingface.co/datasets/boettiger-lab/pad-us-3/resolve/main/pad-stats.parquet"


# some default color variables, consider user palette via st.color_picker()
private_color = "#DE881E" # orange #"#850101" # red
tribal_color = "#BF40BF" # purple
mixed_color = "#005a00" # green
public_color = "#3388ff" # blue

# default color breaks, consider tool via st.slider()
low = 2
high = 3
alpha = .5
style_choice = "Manager Type"
us_lower_48_area_m2 = 7.8e+12


## Helper functions

#@st.cache_resource
def ibis_connection(parquet):
    return ibis.read_parquet(parquet)
pad_data = ibis_connection(parquet)

#@st.cache_data()
def summary_table(column, colors):
    df = (pad_data
        .rename(area = "area_square_meters")
        .group_by(_[column])
        .aggregate(hectares_protected =  (_.area.sum() / 10000).round(),
                   percent_protected =  100 * _.area.sum() / us_lower_48_area_m2,
                   mean_richness = (_.richness * _.area).sum() / _.area.sum(),
                   mean_rsr = (_.rsr * _.area).sum() / _.area.sum(),
                   carbon_lost = (_.deforest_carbon * _.area).sum() / _.area.sum(),
                   crop_expansion = (_.crop_expansion * _.area).sum() / _.area.sum(),
                   human_impact =  (_.human_impact * _.area).sum() / _.area.sum(),
                   )
        .mutate(percent_protected = _.percent_protected.round(1))
        .inner_join(colors, column)
        )
    df = df.to_pandas()
    df[column] = df[column].astype(str)
    return df

def bar_chart(df, x, y):
    chart = alt.Chart(df).mark_bar().encode(
        x=x,
        y=y,
        color=alt.Color('color').scale(None)
    ).properties(width="container", height=300)
    return chart


def area_plot(df, column):
    base = alt.Chart(df).encode(
        alt.Theta("percent_protected:Q").stack(True),
    )
    pie = ( base
           .mark_arc(innerRadius= 40, outerRadius=70)
           .encode(alt.Color("color:N").scale(None).legend(None),
                   tooltip=['percent_protected', 'hectares_protected', column])
    )
    text = ( base
            .mark_text(radius=60, size=12, color="white")
            .encode(text = column + ":N")
    )
    plot = pie # pie + text
    return plot.properties(width=180, height=180)

def pad_style(paint, alpha):
    return {
    "version": 8,
    "sources": {
        "pad": {
            "type": "vector",
            "url": "pmtiles://" + pad_pmtiles,
            "attribution": "US PAD v3"}},
    "layers": [{
            "id": "public",
            "source": "pad",
            "source-layer": "pad-stats",
            "type": "fill",
            "paint": {
                "fill-color": paint,
                "fill-opacity": alpha
            }
        }]}




'''
# US Protected Area Database Explorer

'''


m = leafmap.Map(center=[35, -100], zoom=4, layers_control=True)

custom_style = '''
"blue"
'''


sample_q = '''(
ibis.read_parquet(parquet).
mutate(area = _.area_square_meters).
group_by(_.gap_code).
aggregate(percent_protected =  100 * _.area.sum() / us_lower_48_area_m2,
        mean_richness = (_.richness * _.area).sum() / _.area.sum(),
        mean_rsr = (_.rsr * _.area).sum() / _.area.sum()
        ).
mutate(percent_protected = _.percent_protected.round())
)
'''


## Protected Area polygon color codes 
manager = {
            'property': 'manager_type',
            'type': 'categorical',
            'stops': [
                ['Federal', "darkblue"],
                ['State', public_color],
                ['Local Government', "lightblue"],
                ['Regional Agency Special District', "darkgreen"],
                ['Unknown', "grey"],
                ['Joint', "green"],
                ['American Indian Lands', tribal_color],
                ['Private', "darkred"],
                ['Non-Governmental Organization', "orange"]
            ]
            }

easement = {
            'property': 'category',
            'type': 'categorical',
            'stops': [
                ['Fee', public_color],
                ['Easement', private_color],
                ['Proclamation', tribal_color]
            ]
            }

access = {
    'property': 'public_access',
    'type': 'categorical',
    'stops': [
        ['Open Access', public_color],
        ['Closed', private_color],
        ['Unknown', "grey"],
        ['Restricted Access', tribal_color]
    ]
}

gap = {
        'property': 'gap_code',
        'type': 'categorical',
        'stops': [
            ["1", "#26633d"],
            ["2", "#879647"],
            ["3", "#BBBBBB"],
            ["4", "#F8F8F8"]
        ]
        }

iucn = {
            'property': 'iucn_category',
            'type': 'categorical',
            'stops': [
                ["Ia: Strict nature reserves", "#4B0082"],
                ["Ib: Wilderness areas", "#663399"],
                ["II: National park", "#7B68EE"],
                ["III: Natural monument or feature", "#9370DB"],
                ["IV: Habitat / species management", "#8A2BE2"],
                ["V: Protected landscape / seascape", "#9932CC"],
                ["VI: Protected area with sustainable use of natural resources", "#9400D3"],
                ["Other Conservation Area", "#DDA0DD"],
                ["Unassigned", "#F8F8F8"],
            ]
        }

style_options = {
                    "GAP Status Code": gap,
                    "IUCN Status Code": iucn,
                    "Manager Type": manager,
                    "Fee/Easement": easement,
                    "Public Access": access,
                    # "Mean Richness": richness,
                    # "Mean RSR": rsr,
                    # "custom": eval(custom)
                    }

code_ex='''
m.add_cog_layer("https://data.source.coop/vizzuality/lg-land-carbon-data/natcrop_expansion_100m_cog.tif",
                palette="oranges", name="Cropland Expansion", transparent_bg=True, opacity = 0.7, zoom_to_layer=False)
'''



justice40 = "https://data.source.coop/cboettig/justice40/disadvantaged-communities.pmtiles"
justice40_fill = {
        'property': 'Disadvan',
        'type': 'categorical',
        'stops': [
            [0, "rgba(255, 255, 255, 0)"],
            [1, "rgba(0, 0, 139, 1)"]]}
justice40_style = {
    "version": 8,
    "sources": {
        "source1": {
            "type": "vector",
            "url": "pmtiles://" + justice40,
            "attribution": "Justice40"}
    },
    "layers": [{
            "id": "layer1",
            "source": "source1",
            "source-layer": "DisadvantagedCommunitiesCEJST",
            "type": "fill",
            "paint": {"fill-color": justice40_fill, "fill-opacity": 0.6}}]
}




# +
## Map controls sidebar

with st.sidebar:

    if st.toggle("Protected Areas", True):
        
        style_choice = st.radio("Color by:", style_options)
        alpha = st.slider("transparency", 0.0, 1.0, 0.5)
        style = pad_style(style_options[style_choice], alpha)
        m.add_pmtiles(pad_pmtiles, name="Protected Areas (PAD-US-3)", style=style, overlay=True, show=True, zoom_to_layer=False)
        ## Add legend based on selected style?
        # m.add_legend(legend_dict=legend_dict)

    with st.expander("🦜 Biodiversity"):
        if st.toggle("Species Richness", False):
            m.add_tile_layer(url="https://data.source.coop/cboettig/mobi/tiles/red/species-richness-all/{z}/{x}/{y}.png",
                            name="MOBI Species Richness",
                            attribution="NatureServe",
                            opacity=0.9
                            )
        
        if st.toggle("Range-Size Rarity"):
            m.add_tile_layer(url="https://data.source.coop/cboettig/mobi/tiles/green/range-size-rarity-all/{z}/{x}/{y}.png",
                            name="MOBI Range-Size Rarity",
                            attribution="NatureServe",
                            opacity=0.9
                            )
            #m.add_cog_layer("https://data.source.coop/cboettig/mobi/range-size-rarity-all/RSR_All.tif",
            #                palette="greens", name="Range-Size Rarity", transparent_bg=True, opacity = 0.9, zoom_to_layer=False)
    with st.expander("⛅ Carbon & Climate"):
        if st.toggle("Carbon Lost (2002-2022)"):
            m.add_cog_layer("https://data.source.coop/vizzuality/lg-land-carbon-data/deforest_carbon_100m_cog.tif",
                            palette="reds", name="Carbon Lost (2002-2022)", transparent_bg=True, opacity = 0.8, zoom_to_layer=False)
        
        if st.toggle("Irrecoverable Carbon"):
            m.add_cog_layer("https://data.source.coop/cboettig/carbon/cogs/irrecoverable_c_total_2018.tif",
                            palette="purples", name="Irrecoverable Carbon", transparent_bg=True, opacity = 0.8, zoom_to_layer=False)

        if st.toggle("Manageable Carbon"):
            m.add_cog_layer("https://data.source.coop/cboettig/carbon/cogs/manageable_c_total_2018.tif",
                            palette="greens", name="Manageable Carbon", transparent_bg=True, opacity = 0.8, zoom_to_layer=False)

    with st.expander("🚜 Human Impacts"):

        if st.toggle("Human Impact"):
            hi="https://data.source.coop/vizzuality/hfp-100/hfp_2021_100m_v1-2_cog.tif"
            m.add_cog_layer(hi, palette="purples", name="Human Impact", transparent_bg=True, opacity = 0.8, zoom_to_layer=False)

        if st.toggle("cropland expansion"):
            m.add_cog_layer("https://data.source.coop/vizzuality/lg-land-carbon-data/natcrop_expansion_100m_cog.tif",
                            palette="greens", name="cropland expansion", transparent_bg=True, opacity = 0.8, zoom_to_layer=False)

        if st.toggle("Biodiversity Intactness Loss"):
            m.add_cog_layer("https://data.source.coop/vizzuality/lg-land-carbon-data/natcrop_bii_100m_cog.tif",
                            palette="reds", name="biodiversity intactness loss", transparent_bg=True, opacity = 0.8, zoom_to_layer=False)
    
        if st.toggle("Forest Integrity Loss"):
            m.add_cog_layer("https://data.source.coop/vizzuality/lg-land-carbon-data/natcrop_fii_100m_cog.tif",
                            palette="reds", name="forest integrity loss", transparent_bg=True, opacity = 0.8, zoom_to_layer=False)

    with st.expander("⚖ Equity and Justice"):
        if st.toggle("Climate & Economic Justice"):
            m.add_pmtiles(justice40, name="Climate & Economic Justice", style = justice40_style, overlay=True, show=True, zoom_to_layer=False)

    with st.expander("🔥 Wildfire"):
        # Fire Polygons, USGS
        if st.toggle("Fire boundaries"):
            usgs = "https://data.source.coop/cboettig/fire/usgs-mtbs.pmtiles"
            fire_style = {
                "version": 8,
                "sources": {
                    "source1": {
                        "type": "vector",
                        "url": "pmtiles://" + usgs,
                        "attribution": "USGS"}},
                "layers": [{
                        "id": "usgs",
                        "source": "source1",
                        "source-layer": "mtbs_perims_DD",
                        "type": "fill",
                        "paint": {"fill-color": "#FFA500", "fill-opacity": 0.4}}]}
            m.add_pmtiles(usgs, name="Fire", style=fire_style, overlay=True, show=True, zoom_to_layer=False)


    with st.expander("💻 Custom code"):
        if st.toggle("Custom map layers"):
            
            code = st.text_area(label = "leafmap code:",
                                value = code_ex, 
                                height = 100)
            eval(compile(code, "<string>", "exec"))

#    "## Boundaries"
#    boundaries = st.radio("Boundaries:",
#                          ["None",
#                           "State Boundaries",
#                           "County Boundaries",
#                           "Congressional District",  
#                           "custom"]
#    )

    with st.expander("🗺 Basemaps"):
        # radio selector would make more sense
        if st.toggle("Topography"):
            m.add_basemap("Esri.WorldShadedRelief")
        if st.toggle("Satellite"):
            m.add_basemap("Esri.WorldImagery")
        
   
    
# Map radio buttons to corresponding column:
select_column = {
                "GAP Status Code": "gap_code",
                "IUCN Status Code": "iucn_category",
                "Manager Type": "manager_type",
                "Fee/Easement": "category",
                "Public Access": "public_access",
                "Mean Richness": "gap_code",
                "Mean RSR": "gap_code",
                 "custom": "gap_code"}
column = select_column[style_choice]

# Map radio buttons to corresponding color-scheme:
select_colors = {
                "GAP Status Code": gap["stops"],
                "IUCN Status Code": iucn["stops"],
                "Manager Type": manager["stops"],
                "Fee/Easement": easement["stops"],
                "Public Access": access["stops"],
                "Mean Richness": gap["stops"],
                "Mean RSR": gap["stops"],
                "custom": gap["stops"]}
                
colors = (ibis
          .memtable(select_colors[style_choice], columns = [column, "color"])
          .to_pandas()
         )


main = st.container()

with main:
    map_col, stats_col = st.columns([2,1])

    with map_col:
        m.to_streamlit(height=700)

    df = summary_table(column, colors)
    total_percent = df.percent_protected.sum()
    richness_chart = bar_chart(df, column, 'mean_richness')
    rsr_chart = bar_chart(df, column, 'mean_rsr')
    carbon_lost = bar_chart(df, column, 'carbon_lost')
    crop_expansion = bar_chart(df, column, 'crop_expansion')
    human_impact = bar_chart(df, column, 'human_impact')

    with stats_col:
        with st.container():
            col1, col2, col3 = st.columns(3)
            with col1:
                f"{total_percent}% Continental US Covered"
                st.altair_chart(area_plot(df, column), use_container_width=False)

            with col2:
                "Species Richness"
                st.altair_chart(richness_chart, use_container_width=True)

            with col3:
                "Range-Size Rarity"
                st.altair_chart(rsr_chart, use_container_width=True)

        with st.container():
            col1b, col2b, col3b = st.columns(3)
            with col1b:
                "Carbon Lost ('02-'22)"
                st.altair_chart(carbon_lost, use_container_width=True)

            with col2b:
                "Crop expansion"
                st.altair_chart(crop_expansion, use_container_width=True)

            with col3b:
                "Human Impact"
                st.altair_chart(human_impact, use_container_width=True)


st.divider()

footer = st.container()


with footer:
    '''
    ## Custom queries

    Input custom python code below to interactively explore the data. 

    '''

    col2_1, col2_2 = st.columns(2)



    with col2_1:
        query = st.text_area(
            label = "Python code:",
            value = sample_q, 
            height = 300)

    with col2_2:
        "Output table:"
        df = eval(query)
        st.write(df.to_pandas())

    st.divider()

    '''
    ## Credits

    Author: Carl Boettiger, UC Berkeley
    License: BSD-2-clause

    ### Data sources

    - US Protected Areas Database v3 by USGS. Data: https://beta.source.coop/cboettig/us-pad-3. Citation: https://doi.org/10.5066/P9Q9LQ4B, License: Public Domain
    - Imperiled Species Richness and Range-Size-Rarity from NatureServe (2022). Data: https://beta.source.coop/repositories/cboettig/mobi. License CC-BY-NC-ND
    - Carbon-loss and farming impact by Vizzuality, on https://beta.source.coop/repositories/vizzuality/lg-land-carbon-data. Citation: https://doi.org/10.1101/2023.11.01.565036, License: CC-BY
    - Human Footprint by Vizzuality, on https://beta.source.coop/repositories/vizzuality/hfp-100.  Citation: https://doi.org/10.3389/frsen.2023.1130896, License: Public Domain
    - Fire polygons by USGS, reprocessed to PMTiles on https://beta.source.coop/cboettig/fire/. License: Public Domain
    - Irrecoverable Carbon from Conservation International, reprocessed to COG on https://beta.source.coop/cboettig/carbon, citation: https://doi.org/10.1038/s41893-021-00803-6, License: CC-BY-NC
    - Climate and Economic Justice Screening Tool, US Council on Environmental Quality, Justice40, data: https://beta.source.coop/repositories/cboettig/justice40/description/, License: Public Domain

    ### Software

    Proudly built with a free and Open Source software stack: Streamlit (reactive application), HuggingFace (application hosting), Source.Coop (data hosting),
    using cloud-native data serializations in COG, PMTiles, and GeoParquet. Coded in pure python using leafmap and duckdb. Map styling with [MapLibre](https://maplibre.org/).
    '''