import streamlit as st
import streamlit.components.v1 as components
import base64
import leafmap.maplibregl as leafmap
import altair as alt
import ibis
from ibis import _
import ibis.selectors as s


from typing import Optional
def to_streamlit(
    self,
    width: Optional[int] = None,
    height: Optional[int] = 600,
    scrolling: Optional[bool] = False,
    **kwargs,
    ):

    try:
        import streamlit.components.v1 as components
        import base64

        raw_html = self.to_html().encode("utf-8")
        raw_html = base64.b64encode(raw_html).decode()
        return components.iframe(
            f"data:text/html;base64,{raw_html}",
            width=width,
            height=height,
            scrolling=scrolling,
            **kwargs,
        )

    except Exception as e:
        raise Exception(e)


# ca_pmtiles = "https://huggingface.co/datasets/boettiger-lab/ca-30x30/resolve/main/ca2024-30m-tippe.pmtiles"
# ca_parquet = "https://huggingface.co/datasets/boettiger-lab/ca-30x30/resolve/main/ca2024-30m.parquet"
ca_pmtiles = "https://huggingface.co/datasets/boettiger-lab/ca-30x30/resolve/main/cpad-stats.pmtiles"
ca_parquet = "https://huggingface.co/datasets/boettiger-lab/ca-30x30/resolve/main/cpad-stats.parquet"



ca_area_acres = 1.014e8 #acres 
style_choice = "GAP Status Code"

con = ibis.duckdb.connect(extensions=["spatial"])

ca = (con
      .read_parquet(ca_parquet)
      .cast({"geom": "geometry"})
      .cast({"crop_expansion": "int64"})
     )

private_access_color = "#DE881E" # orange 
public_access_color = "#3388ff" # blue
tribal_color = "#BF40BF" # purple
mixed_color = "#005a00" # green
year2023_color = "#26542C" # green
year2024_color = "#F3AB3D" # orange 

federal_color = "#529642" # green
state_color = "#A1B03D" # light green
local_color = "#365591" # blue
special_color = "#529642" # brown
private_color = "#7A3F1A" # brown
joint_color = "#DAB0AE" # pink
county_color = "#BFD76B" # green
city_color = "#BDC368" #green
hoa_color = "#A89BBC" # purple
nonprofit_color =  "#D77031" #orange


from functools import reduce

def get_summary(ca, combined_filter, column, colors=None):
    df = ca.filter(combined_filter)

    # if colors is not None and not colors.empty: # df used for chart. 
    #     df = df.filter(_.reGAP <3) # only show gaps 1 and 2 for chart.
        
    df = (df
        .group_by(*column)  # unpack the list for grouping
        .aggregate(percent_protected=100 * _.Acres.sum() / ca_area_acres,
                   mean_richness = (_.richness * _.Acres).sum() / _.Acres.sum(),
                   mean_rsr = (_.rsr * _.Acres).sum() / _.Acres.sum(),
                   all_rsr =  (_.all_species_rwr * _.Acres).sum() / _.Acres.sum(),
                   all_richness =  (_.all_species_richness * _.Acres).sum() / _.Acres.sum(),
                   irrecoverable_carbon = (_.irrecoverable_carbon * _.Acres).sum() / _.Acres.sum(),
                   manageable_carbon = (_.manageable_carbon * _.Acres).sum() / _.Acres.sum(),
                   carbon_lost = (_.deforest_carbon * _.Acres).sum() / _.Acres.sum(),
                   human_impact =  (_.human_impact * _.Acres).sum() / _.Acres.sum(),
                   svi =  (_.svi * _.Acres).sum() / _.Acres.sum(),
                   svi_socioeconomic_status =  (_.svi_socioeconomic_status * _.Acres).sum() / _.Acres.sum(),
                   svi_household_char =  (_.svi_household_char * _.Acres).sum() / _.Acres.sum(),
                   svi_racial_ethnic_minority =  (_.svi_racial_ethnic_minority * _.Acres).sum() / _.Acres.sum(),
                   svi_housing_transit =  (_.svi_housing_transit * _.Acres).sum() / _.Acres.sum(),

                   # biodiversity_intactness_loss = (_.biodiversity_intactness_loss * _.Acres).sum() / _.Acres.sum(),
                   # crop_reduction = (_.crop_reduction * _.Acres).sum() / _.Acres.sum(),
                   # crop_expansion = (_.crop_expansion * _.Acres).sum() / _.Acres.sum(),
                   # forest_loss = (_.forest_integrity_loss * _.Acres).sum() / _.Acres.sum(),
                  )
        .mutate(percent_protected=_.percent_protected.round(1))
         )

    if colors is not None and not colors.empty: #
        df = df.inner_join(colors, column) # chart colors 
             
    df = df.cast({col: "string" for col in column})
    df = df.to_pandas()
    return df


def summary_table(column, colors, filter_cols, filter_vals,colorby_vals):
    filters = [] 
    
    if filter_cols and filter_vals: #if a filter is selected, add to list of filters 
        for filter_col, filter_val in zip(filter_cols, filter_vals):
            if len(filter_val) > 1:
                filters.append(getattr(_, filter_col).isin(filter_val))
            else:
                filters.append(getattr(_, filter_col) == filter_val[0])

    if column not in filter_cols: #show color_by variable in table
        filter_cols.append(column)
        filters.append(getattr(_, column).isin(colorby_vals[column]))  
        
    combined_filter = reduce(lambda x, y: x & y, filters)

    df = get_summary(ca, combined_filter, [column], colors) # df used for charts 
    df_tab = get_summary(ca, combined_filter, filter_cols, colors = None) #df used for printed table
    return df, df_tab


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


    
def get_pmtiles_style(paint, alpha, cols, values):
    filters = []
    
    for col, val in zip(cols, values):
        filter_condition = ["match", ["get", col], val, True, False]
        filters.append(filter_condition)
    combined_filter = ["all"] + filters

    return {
        "version": 8,
        "sources": {
            "ca": {
                "type": "vector",
                "url": "pmtiles://" + ca_pmtiles,
            }
        },
        "layers": [{
            "id": "ca30x30",
            "source": "ca",
            "source-layer": "layer",
            # "source-layer": "ca202430m",
            # "source-layer": "ca2024",
            "type": "fill",
            "filter": combined_filter,  # Use the combined filter
            "paint": {
                "fill-color": paint,
                "fill-opacity": alpha
            }
        }]
    }



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 getButtons(style_options, style_choice, default_gap=None):
    column = style_options[style_choice]['property']
    opts = [style[0] for style in style_options[style_choice]['stops']]   
    default_gap = default_gap or {}  
    buttons = {
        name: st.checkbox(f"{name}", value=default_gap.get(name, False), key=column + str(name))
        for name in opts
    }
    filter_choice = [key for key, value in buttons.items() if value]  # return only selected
    d = {}
    d[column] = filter_choice
    return d

default_gap = {
    1: True,  
    2: True,  
}

def getColorVals(style_options, style_choice): #adding "color by" values to table
    column = style_options[style_choice]['property']
    opts = [style[0] for style in style_options[style_choice]['stops']]   
    d = {}
    d[column] = opts
    return d

manager = {
            'property': 'manager_type',
            'type': 'categorical',
            'stops': [
                ['Federal', federal_color],
                ['State', state_color],
                ['Non Profit', nonprofit_color],
                ['Special District', special_color],
                ['Unknown', "grey"],
                ['County', county_color],
                ['City', city_color],
                ['Joint', joint_color],
                ['Tribal', tribal_color],
                ['Private', private_color],
                ['Home Owners Association', hoa_color]
            ]
            }


easement = {
            'property': 'Easement',
            'type': 'categorical',
            'stops': [
                [0, public_access_color],
                [1, private_access_color]
            ]
            }


year = {
            'property': 'established',
            'type': 'categorical',
            'stops': [
                [2023, year2023_color],
                [2024, year2024_color]
            ]
            }


access = {
    'property': 'access_type',
    'type': 'categorical',
    'stops': [
        ['Open Access', public_access_color],
        ['No Public Access', private_access_color],
        ['Unknown Access', "grey"],
        ['Restricted Access', tribal_color]
    ]
}



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


# area_type = {
#         'property': 'type',
#         'type': 'categorical',
#         'stops': [
#             ["Land", "green"],
#             ["Water", "blue"]
#         ]
#         }


style_options = {
                "Year": year,
                "GAP Status Code": gap,
                "Manager Type": manager,
                "Easement": easement,
                "Public Access": access,
                # "Type": area_type
                }



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",
            "filter": ["match", ["get", "StateName"], "California", True, False],
            "type": "fill",
            "paint": {"fill-color": justice40_fill, "fill-opacity": 0.6}}]
}



    
sv_pmtiles = "https://data.source.coop/cboettig/social-vulnerability/svi2020_us_county.pmtiles"
def get_sv_style(column):
    return  {
            "layers": [
                {
                    "id": "SVI",
                    "source": "Social Vulnerability Index",
                    "source-layer": "SVI2020_US_county",
                    "filter": ["match", ["get", "STATE"], "California", True, False],
                    "type": "fill",
                    "paint": {
                        "fill-color": 
                            ["interpolate", ["linear"], ["get", column],
                    0, "#FFE6EE",
                    1, "#850101"] 
                            
                        }
                    }
            ]
    }



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

'''
# CA 30X30 Prototype

'''

m = leafmap.Map(style="positron")
filters = {}

with st.sidebar:
    color_choice = st.radio("Color by:", style_options)      
    colorby_vals = getColorVals(style_options, color_choice)
    
    alpha = st.slider("transparency", 0.0, 1.0, 0.5)
    st.divider()
    "Filters:"
    for label in style_options:
        with st.expander(label): 
             
            if label == "GAP Status Code": # gap code 1 and 2 are on by default
                opts = getButtons(style_options, label, default_gap)
            else:
                opts = getButtons(style_options, label)
            filters.update(opts)
            
        selected = {k: v for k, v in filters.items() if v} #get selected filters
        if selected: 
            filter_cols = list(selected.keys())
            filter_vals = list(selected.values())
        else: 
            filter_cols = []
            filter_vals = []
            
    style = get_pmtiles_style(style_options[color_choice], alpha, filter_cols, filter_vals)
    m.add_pmtiles(ca_pmtiles, style=style, visible=True, name="CA", opacity=alpha, tooltip=True)

    st.divider()
    "Data Layers:"

    with st.expander("🦜 Biodiversity"):
        alpha_bio = st.slider("transparency", 0.0, 1.0, 0.4, key = "biodiversity")

        show_richness = st.toggle("Species Richness", False)
        if show_richness:
            m.add_tile_layer(
                url = "https://huggingface.co/datasets/boettiger-lab/ca-30x30/resolve/main/species-richness-ca/{z}/{x}/{y}.png",
                name="MOBI Species Richness",
                opacity=alpha_bio
            )
            
        show_rsr = st.toggle("Range-Size Rarity")
        if show_rsr:           
            m.add_tile_layer(
                url="https://huggingface.co/datasets/boettiger-lab/ca-30x30/resolve/main/range-size-rarity/{z}/{x}/{y}.png",
                name="MOBI Range-Size Rarity",
                opacity=alpha_bio)
         
    with st.expander("⛅ Carbon & Climate"):
        alpha_climate = st.slider("transparency", 0.0, 1.0, 0.3, key = "climate")

        show_irrecoverable_carbon = st.toggle("Irrecoverable Carbon")
        if show_irrecoverable_carbon:
            m.add_cog_layer("https://huggingface.co/datasets/boettiger-lab/ca-30x30/resolve/main/ca_irrecoverable_c_2018_cog.tif",
                            palette="reds", name="Irrecoverable Carbon", transparent_bg=True, opacity = alpha_climate, fit_bounds=False, bidx=[1])

        show_manageable_carbon = st.toggle("Manageable Carbon")
        if show_manageable_carbon:
           m.add_cog_layer("https://huggingface.co/datasets/boettiger-lab/ca-30x30/resolve/main/ca_manageable_c_2018_cog.tif",
                            palette="purples", name="Manageable Carbon", transparent_bg=True, opacity = alpha_climate, fit_bounds=False)
          
    with st.expander("Climate and Economic Justice"):
        alpha_justice40 = st.slider("transparency", 0.0, 1.0, 0.3, key = "social justice")

        show_justice40 = st.toggle("Justice40")
        if show_justice40:
            m.add_pmtiles(justice40, style=justice40_style, visible=True, name="Justice40", opacity=alpha_justice40, tooltip=False, fit_bounds = False)

    with st.expander("Social Vulnerability"):
        alpha_justice = st.slider("transparency", 0.0, 1.0, 0.3, key = "SVI")

        show_sv = st.toggle("Social Vulnerability Index (SVI)")
        if show_sv:
            m.add_pmtiles(sv_pmtiles, style = get_sv_style("RPL_THEMES") ,visible=True, opacity=alpha_justice, tooltip=False, fit_bounds = False)

        show_sv_socio = st.toggle("Socioeconomic Status")
        if show_sv_socio:
            m.add_pmtiles(sv_pmtiles, style = get_sv_style("RPL_THEME1") ,visible=True, opacity=alpha_justice, tooltip=False, fit_bounds = False)
            
        show_sv_household = st.toggle("Household Characteristics")
        if show_sv_household:
            m.add_pmtiles(sv_pmtiles, style = get_sv_style("RPL_THEME2") ,visible=True, opacity=alpha_justice, tooltip=False, fit_bounds = False)

        show_sv_minority = st.toggle("Racial & Ethnic Minority Status")
        if show_sv_minority:
            m.add_pmtiles(sv_pmtiles, style = get_sv_style("RPL_THEME3") ,visible=True, opacity=alpha_justice, tooltip=False, fit_bounds = False)

        show_sv_housing = st.toggle("Housing Type & Transportation")
        if show_sv_housing:
            m.add_pmtiles(sv_pmtiles, style = get_sv_style("RPL_THEME4") ,visible=True, opacity=alpha_justice, tooltip=False, fit_bounds = False)

            
    with st.expander("🚜 Human Impacts"):
        alpha_hi = st.slider("transparency", 0.0, 1.0, 0.5, key = "hi")

        show_carbon_lost = st.toggle("Carbon Lost (2002-2022)")
        if show_carbon_lost:
            m.add_tile_layer(
                url="https://huggingface.co/datasets/boettiger-lab/ca-30x30/resolve/main/deforest-carbon-ca/{z}/{x}/{y}.png",
                name="Carbon Lost (2002-2022)",
                opacity = alpha_hi)

        show_human_impact = st.toggle("Human Impact")
        if show_human_impact:
            m.add_cog_layer(
                url = "https://huggingface.co/datasets/boettiger-lab/ca-30x30/resolve/main/ca_human_impact_cog.tif", name="Human Impact", transparent_bg=True, opacity = alpha_hi, fit_bounds=False)



select_column = {
                "Year": "established",
                "GAP Status Code": "reGAP",
                "Manager Type": "manager_type",
                "Easement": "Easement",
                "Public Access": "access_type",
                # "Type": "type",
                }


column = select_column[color_choice]

select_colors = {
                "Year": year["stops"],
                "GAP Status Code": gap["stops"],
                "Manager Type": manager["stops"],
                "Easement": easement["stops"],
                "Public Access": access["stops"],
                # "Type": area_type["stops"]
                }

colors = (ibis
          .memtable(select_colors[color_choice], columns = [column, "color"])
          .to_pandas()
         )

main = st.container()
df,df_tab = summary_table(column, colors, filter_cols, filter_vals, colorby_vals)

with main:
    map_col, stats_col = st.columns([2,1])
    with map_col:
        to_streamlit(m, height=700)
        st.dataframe(df_tab, use_container_width = True)
    svi = bar_chart(df, column, 'svi')
    svi_socioeconomic_status = bar_chart(df, column, 'svi_socioeconomic_status')
    svi_household_char = bar_chart(df, column, 'svi_household_char')
    svi_racial_ethnic_minority = bar_chart(df, column, 'svi_racial_ethnic_minority')
    svi_housing_transit = bar_chart(df, column, 'svi_housing_transit')

    richness_chart = bar_chart(df, column, 'mean_richness')
    rsr_chart = bar_chart(df, column, 'mean_rsr')
    # all_rsr = bar_chart(df, column, 'all_rsr')
    # all_richness = bar_chart(df, column, 'all_richness')
    irrecoverable_carbon = bar_chart(df, column, 'irrecoverable_carbon')
    manageable_carbon = bar_chart(df, column, 'manageable_carbon')
    carbon_lost = bar_chart(df, column, 'carbon_lost')
    human_impact = bar_chart(df, column, 'human_impact')

    # crop_expansion = bar_chart(df, column, 'crop_expansion')
    # biodiversity_intactness_loss = bar_chart(df, column, biodiversity_intactness_loss')
    # crop_reduction = bar_chart(df, column, 'crop_reduction')
    # forest_loss = bar_chart(df, column, 'forest_loss')
    




    total_percent = df.percent_protected.sum().round(1)
    with stats_col:
        with st.container():
            f"{total_percent}% CA Covered"
            st.altair_chart(area_plot(df, column), use_container_width=True)
            

        with st.container():

            if show_richness:
                "Species Richness"
                st.altair_chart(richness_chart, use_container_width=True)

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

            if show_irrecoverable_carbon:
                "Irrecoverable Carbon"
                st.altair_chart(irrecoverable_carbon, use_container_width=True)

            if show_manageable_carbon:
                "Manageable Carbon"
                st.altair_chart(manageable_carbon, use_container_width=True)

            if show_sv:
                "Social Vulnerability Index"
                st.altair_chart(svi, use_container_width=True)
                
            if show_sv_socio:
                "SVI - Socioeconomic Status"
                st.altair_chart(svi_socioeconomic_status, use_container_width=True)
            
            if show_sv_household:
                "SVI - Household Characteristics"
                st.altair_chart(svi_household_char, use_container_width=True)
            
            if show_sv_minority:
                "SVI - Racial and Ethnic Minority"
                st.altair_chart(svi_racial_ethnic_minority, use_container_width=True)
            
            if show_sv_housing:
                "SVI - Housing Type and Transit"
                st.altair_chart(svi_housing_transit, use_container_width=True)

            
            if show_carbon_lost:
                "Carbon Lost ('02-'22)"
                st.altair_chart(carbon_lost, use_container_width=True)

            if show_human_impact:
                "Human Impact"
                st.altair_chart(human_impact, use_container_width=True)




st.divider()
footer = st.container()



'''
## Credits
Authors: Cassie Buhler & Carl Boettiger, UC Berkeley
License: BSD-2-clause

### Data sources
- California Protected Areas Database by CA Nature. Data: https://www.californianature.ca.gov/datasets/CAnature::30x30-conserved-areas-terrestrial-2024/about. License: Public Domain

- Climate and Economic Justice Screening Tool, US Council on Environmental Quality, Justice40, data: https://beta.source.coop/repositories/cboettig/justice40/description/, License: Public Domain

- CDC 2020 Social Vulnerability Index by US Census Track. Data: https://source.coop/repositories/cboettig/social-vulnerability/description. 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

- 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



'''