import plotly.graph_objects as go
from folium.raster_layers import ImageOverlay
import re
import glob
import pickle
import h5py
import rasterio
import streamlit as st
import branca.colormap as cm
import folium
import numpy as np
import pandas as pd
from geopy.extra.rate_limiter import RateLimiter
from geopy.geocoders import Nominatim
from streamlit_plotly_events import plotly_events
import warnings
warnings.filterwarnings("ignore")

#
@st.cache_data
def convert_df(df):
    return df.to_csv(index=0).encode('utf-8')


def geocode(address):
    try:
        address2 = address.replace(' ', '+').replace(',', '%2C')
        df = pd.read_json(
            f'https://geocoding.geo.census.gov/geocoder/locations/onelineaddress?address={address2}&benchmark=2020&format=json')
        results = df.iloc[:1, 0][0][0]['coordinates']
        lat, lon = results['y'], results['x']
    except:
        geolocator = Nominatim(user_agent="GTA Lookup")
        geocode = RateLimiter(geolocator.geocode, min_delay_seconds=1)
        location = geolocator.geocode(address)
        lat, lon = location.latitude, location.longitude
    return pd.DataFrame({'Lat': lat, 'Lon': lon}, index=[0])

@st.cache_data
def get_data(row, col, radius=8):
    files = [
        "data/2023_hail.h5",
        "data/2022_hail.h5",
        "data/2021_hail.h5",
        "data/2020_hail.h5"
    ]

    all_data = []
    all_dates = []
    for i in files:
        with h5py.File(i, 'r') as f:
            data = f['hail'][:, row-radius:row +
                             radius+1, col-radius:col+radius+1]
            dates = f['dates'][:]
            all_data.append(data)
            all_dates.append(dates)

    data_mat = np.concatenate(all_data)
    data_mat = np.where(data_mat < 0, 0, data_mat)*0.0393701
    dates_mat = np.concatenate(all_dates)

    data_actual = [i[radius, radius] for i in data_mat]
    data_max = np.max(data_mat, axis=(1, 2))
    data_max_2 = np.max(data_mat, axis=0)

    df = pd.DataFrame({'Date': dates_mat,
                       'Actual': data_actual,
                      'Max': data_max})

    df['Date'] = pd.to_datetime(df['Date'], format='%Y%m%d')
    # df['Date']=df['Date']+pd.Timedelta(days=1)

    return df, data_max_2

@st.cache_data
def map_folium(lat, lon, files_dates_selected):

    # Create a base map
    m = folium.Map(location=[lat, lon], zoom_start=6)
    folium.Marker(location=[lat, lon], popup=address).add_to(m)

    # Define the image bounds (SW and NE corners)
    image_bounds = [[20.0000010001429, -129.99999999985712],
                    [54.9999999998571, -60.00000200014287]]

    # Add ImageOverlays for each image

    overlay = ImageOverlay(image=files_dates_selected, bounds=image_bounds,
                           opacity=.75,
                           mercator_project=False)

    overlay.add_to(m)
    colormap_hail = cm.LinearColormap(
        colors=['blue', 'lightblue', 'pink', 'red'], vmin=0.01, vmax=2)
    # Add the color legend to the map
    colormap_hail.caption = 'Legend: Hail (Inches)'
    colormap_hail.add_to(m)
    return m




#Set up 2 Columns
st.set_page_config(layout="wide")
col1, col2 = st.columns((2))


#Input Values
address = st.sidebar.text_input("Address", "Dallas, TX")
date_focus = st.sidebar.date_input("Date",  pd.Timestamp(2023, 7, 1))
within_days = st.sidebar.selectbox('Days Within', (90, 180, 365))
# start_date = st.sidebar.date_input("Start Date",  pd.Timestamp(2023, 1, 1))
# end_date = st.sidebar.date_input("End Date",  pd.Timestamp(2023, 12, 31))

start_date = date_focus+pd.Timedelta(days=-within_days)
end_date = date_focus+pd.Timedelta(days=within_days)

date_range = pd.date_range(start=start_date, end=end_date).strftime('%Y%m%d')

circle_radius = st.sidebar.selectbox('Box Radius (Miles)', (5, 10, 25))


#Geocode and get Data
result = geocode(address)
lat, lon = result.values[0]


crs_dic = pickle.load(open('data/mrms_hail_crs.pkl', 'rb'))
transform = crs_dic['affine']

row, col = rasterio.transform.rowcol(transform, lon, lat)
row, col = int(row), int(col)

radius = int(np.ceil(circle_radius*1.6))

# Get Data
df_data, max_values = get_data(row, col, radius)

df_data = df_data.query(f"'{start_date}'<=Date<='{end_date}'")
df_data['Max'] = df_data['Max'].round(3)
df_data['Actual'] = df_data['Actual'].round(3)


fig = go.Figure()

# Add bars for actual values
fig.add_trace(go.Bar(
    x=df_data['Date'],
    y=df_data['Actual'],
    name='Actual Value',
    marker_color='#2D5986',
    hoverinfo='text',  # Show text information only
    text=df_data.apply(
        lambda row: f'Date: {row["Date"].date()}<br>Hail: {row["Actual"]}<br>Max: {row["Max"]}', axis=1)

))


# Update layout
fig.update_layout(
    title='',
    xaxis_title='Date',
    yaxis_title='Hail (Inches)',
    barmode='group'
)


files = glob.glob(r'webp/**/*.webp', recursive=True)

with col1:
    st.title('Hail')
    try:

        selected_points = plotly_events(fig)
        csv = convert_df(df_data)
        st.download_button(
            label="Download data as CSV",
            data=csv,
            file_name='data.csv',
            mime='text/csv')
    except:
        pass


with col2:
    st.title('Hail Mesh')

    if selected_points:
        # Extract the details of the first selected point
        selected_index = selected_points[0]['pointIndex']
        selected_data = df_data.iloc[selected_index]

        # # Display the selected point details
        # st.write("You selected the following point:")
        # st.write(selected_data)

        files_dates_selected = [i for i in files if selected_data['Date'].strftime(
            '%Y%m%d') in re.search(r'(\d{8})', i).group()][0]
        m = map_folium(lat, lon, files_dates_selected)
        m.save("map_new.html")

        st.write('Date: ' + selected_data['Date'].strftime('%m-%d-%Y'))
        st.components.v1.html(
            open("map_new.html", 'r').read(), height=500, width=500)

    else:
        files_dates_selected = [i for i in files if date_focus.strftime(
            '%Y%m%d') in re.search(r'(\d{8})', i).group()][0]
        st.write('Date: ' + date_focus.strftime('%m-%d-%Y'))
        m = map_folium(lat, lon, files_dates_selected)
        m.save("map_new.html")
        st.components.v1.html(
            open("map_new.html", 'r').read(), height=500, width=500)


st.markdown(""" <style>
#MainMenu {visibility: hidden;}
footer {visibility: hidden;}
</style> """, unsafe_allow_html=True)