import ibis
from ibis import _
import streamlit as st

import altair as alt



st.set_page_config(layout="wide",
                   page_title="TPL LandVote",
                   page_icon=":globe:")

'''
# LandVote Prototype

An experimental platform for visualizing data on ballot measures for conservation, based on data from <https://landvote.org/> curated by the Trust for Public Land.

'''

dark_orange = 'rgba(171, 86, 1, 1)'   # dark orange - min value 
light_orange = 'rgba(243, 211, 177, 1)' # light orange 
grey = 'rgba(211, 211, 211, 1)'        # grey
light_green = 'rgba(195, 219, 195, 1)' # light green 
dark_green = 'rgba(65, 125, 65, 1)'    # dark green  - max value 

dem_blue = "#1b46c2"
rep_red = "#E81B23"

## Chatbot
import os
import pandas as pd
import matplotlib.pyplot as plt
from pandasai.llm.openai import OpenAI
from pandasai import Agent
from pandasai.responses.streamlit_response import StreamlitResponse

llm = OpenAI(api_token=st.secrets["OPENAI_API_KEY"])
df1 = pd.read_csv("data.csv")

agent = Agent(
    [df1],
    config={"verbose": True, "response_parser": StreamlitResponse, "llm": llm},
)



with st.sidebar:


    '''
    ## Data Assistant (experimental)

    Ask questions about the landvote data, like:

    - What are the top states for approved conservation funds?
    - Plot the total funds spent in conservation each year.
    - What city has approved the most funds in a single measure? What was the description of that vote?
    - Which state has had largest number measures fail? What is that as a fraction of it's total measures?
    '''
    
    prompt = st.chat_input("Ask about the data")
    if prompt:
        with st.spinner():
            resp = agent.chat(prompt)
            if os.path.isfile('exports/charts/temp_chart.png'):
                im = plt.imread('exports/charts/temp_chart.png')
                st.image(im)
                os.remove('exports/charts/temp_chart.png')
            st.write(resp)
        


        
# year = st.slider("Select a year", min_value=1988, max_value=2024, value=2022, step=2)
year = st.slider("Select a year", min_value=1988, max_value=2024, value=2022, step=1)


url = "https://huggingface.co/datasets/boettiger-lab/landvote/resolve/main/votes.pmtiles"
parties = "https://huggingface.co/datasets/boettiger-lab/landvote/resolve/main/votes.parquet"

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

party = (con
            .read_parquet(parties)
            .cast({"geometry": "geometry"})
            )

def get_passes(party):
    df = (party
        .filter(_.year >= 2000)
        .group_by("year", "party")
        .aggregate(total=_.count(),
                   passes = (_.Status.isin(["Pass", "Pass*"]).sum())
                  )
        .mutate(
            percent_passed=(_.passes / _.total).round(2),
            color=ibis.case()
                .when(_.party == "DEMOCRAT", ibis.literal("#2e4a93"))
                .else_(ibis.literal("#E81B23"))
                .end()
        )
    )
    
    df = df.to_pandas()
    return df


def percent_chart(df_passes):
    alt.themes.enable('fivethirtyeight')

    chart = alt.Chart(df_passes).mark_line(strokeWidth=3).encode(
    x=alt.X('year:N', title='Year'),
    y=alt.Y('percent_passed:Q', title='Percent Passed'),
    color=alt.Color('party:N',  # Map 'party' to color
                    scale=alt.Scale(domain=["DEMOCRAT", "REPUBLICAN"], 
                                    range=["#1b46c2", "#E81B23"]),
                    legend=alt.Legend(title="Party")
                   )
    ).properties(
        title='% of Measures Passed'
    )
    
    st.altair_chart(chart, use_container_width=True)


def funding_chart(party):
    df = (party
        .mutate(amount=_.amount.replace('$', '').replace(',', '').cast('float64'))
        .filter(_.Status.isin(["Pass", "Pass*"]))
        .group_by("year")
        .aggregate(total_funding=_.amount.sum())
        .order_by("year")
        .mutate(cumulative_funding=_.total_funding.cumsum()/1e9)
        .execute()
        )
  
    chart = alt.Chart(df).mark_line(strokeWidth=3).encode(
        x=alt.X('year:N', title='Year'),
        y=alt.Y('cumulative_funding:Q', title='Billions of Dollars'),
    ).properties(
        title='Cumulative Funding'
    )
    
    st.altair_chart(chart, use_container_width=True)


paint_fill = {
                "fill-color": [
                    "case",
                    # if passed, color green
                    ["==", ["get", "Status"], "Pass"],
                    [
                        "interpolate", ["linear"], [
                            "to-number", ["slice", ["get", "yes"], 0, -1]  # convert 'yes' string to number
                        ],
                        50, grey,    
                        55, light_green,     # higher yes % -> darker green
                        100, dark_green     # 100 is the max of data 
                    ],
                    # if failed, color orange
                    ["==", ["get", "Status"], "Fail"],
                    [
                        "interpolate", ["linear"], [
                            "to-number", ["slice", ["get", "yes"], 0, -1]  # convert 'yes' string to number
                        ],
                        0, dark_orange,     # higher yes % -> lighter orange 
                        50, light_orange,   
                        67, grey # 67 is the max of data.
                    ],
                    grey  # if no match 
                ]
            }


paint_extrusion = {
                "fill-extrusion-color": [
                    "case",
                    # if passed, color green
                    ["==", ["get", "Status"], "Pass"],
                    [
                        "interpolate", ["linear"], [
                            "to-number", ["slice", ["get", "yes"], 0, -1]  # convert 'yes' string to number
                        ],
                        50, grey,    
                        55, light_green,     # higher yes % -> darker green
                        100, dark_green     # 100 is the max of data 
                    ],
                    # if failed, color orange
                    ["==", ["get", "Status"], "Fail"],
                    [
                        "interpolate", ["linear"], [
                            "to-number", ["slice", ["get", "yes"], 0, -1]  # convert 'yes' string to number
                        ],
                        0, dark_orange,     # higher yes % -> lighter orange 
                        50, light_orange,   
                        67, grey # 67 is the max of data.
                    ],
                    grey  # if no match 
                ],
                "fill-extrusion-height": ["*", ["get", "log_amount"], 5000],
            }


def get_style_status(jurisdiction):
    if jurisdiction == "State":
        name = "state"
        label = "states"
        paint_type = paint_fill
        layer_type = "fill"

    elif jurisdiction == "County":
        name = "county"
        label = "counties"
        paint_type = paint_extrusion
        layer_type = "fill-extrusion"

    else: 
        name = "municipal"
        label = "cities"
        paint_type = paint_extrusion
        layer_type = "fill-extrusion"

    style = {
        "layers": [
            {
                "id": label,
                "source": name,
                "source-layer": name,
                "type": layer_type,
                "filter": [
                    "==",
                    ["get", "year"],
                    year,
                ],
                "paint": paint_type
            },
        ],
    }
    return style 


def get_style_party(jurisdiction):
    if jurisdiction == "State":
        name = "state"
        label = "states"
    elif jurisdiction == "County":
        name = "county"
        label = "counties"

    else: 
        name = "municipal"
        label = "cities"


    style_party = {
        "layers": [
            {
                "id": label,
                "source": name,
                "source-layer": name,
                "type": "fill",
                "filter": [
                    "==",
                    ["get", "year"],
                    year,
                ],
                "paint": {
                "fill-color": 
                    {
                    'property': 'party',
                    'type': 'categorical',
                    'stops': [
                        ["DEMOCRAT", dem_blue],
                        ["REPUBLICAN", rep_red],
                    ]
                    }
                }
            },
        ],
    }
    return style_party



style_options= ["Measure Status", "Political Party"]

color_choice = st.radio("Color by:", style_options)
            
import leafmap.maplibregl as leafmap
m = leafmap.Map(style="positron",   center=(-100, 40), zoom=3)

if color_choice == "Measure Status":

        #states are 2D and transparent, thus added separately.
    m.add_pmtiles(
        url,
        style=get_style_status("State"),
        visible=True,
        opacity=0.4,
        tooltip=True,
        fit_bounds=False
    )
    
    m.add_pmtiles(
        url,
        style=get_style_status("County"),
        visible=True,
        opacity=1.0,
        tooltip=True,
        fit_bounds=False
    )
    
    
    m.add_pmtiles(
        url,
        style=get_style_status("Municipal"),
        visible=True,
        opacity=1.0,
        tooltip=True,
        fit_bounds=False
    )
    
else:
    m.add_pmtiles(
    url,
    style=get_style_party("State"),
    visible=True,
    opacity=0.6,
    tooltip=True,
    fit_bounds=False
    )
    
    m.add_pmtiles(
        url,
        style=get_style_party("County"),
        visible=True,
        opacity=1.0,
        tooltip=True,
        fit_bounds=False
    )
    
    
    m.add_pmtiles(
        url,
        style=get_style_party("Municipal"),
        visible=True,
        opacity=1.0,
        tooltip=True,
        fit_bounds=False
    )
    
    
m.add_layer_control()

m.to_streamlit()

df_passes = get_passes(party)
percent_chart(df_passes)
funding_chart(party.filter(_.year >= 2000))


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