from branca.element import Template, MacroElement
from folium.raster_layers import ImageOverlay
import re
import glob
import altair as alt
import pickle
import h5py
import rasterio
import streamlit as st
import os
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

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])


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 f in files:
        with h5py.File(f, '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


def map_folium(lat, lon,files_dates_selected, within_days ):
    
    # Create a base map
    m = folium.Map(location=[lat, lon], zoom_start=5)
    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
    dates = []
    for f in files_dates_selected:
        overlay = ImageOverlay(image=f, bounds=image_bounds,
                               opacity=.75,
                               mercator_project=False)
        filename = os.path.basename(f)
        date_str = re.search(r'(\d{8})', filename).group()
        formatted_date = f"{date_str[:4]}-{date_str[4:6]}-{date_str[6:8]}"
        dates.append(formatted_date)
        overlay.add_to(m)
    

    # HTML template for the slider control with dates
    template_1 = '{% macro html(this, kwargs) %}' + f"""
    

    
    <div id="slider-control" style="position: fixed; top: 50px; left: 50px; z-index: 9999; background-color: white; padding: 10px; border-radius: 5px; box-shadow: 0px 0px 5px rgba(0, 0, 0, 0.3);">
        <label for="image-slider">Select Date:</label>
        <input type="range" min="0" max="{len(dates)-1}" value="{within_days}" class="slider" id="image-slider" style="width: 150px;" oninput="updateFromSlider(this.value)">
        <input type="text" id="date-input" placeholder="YYYY-MM-DD" oninput="updateFromInput(this.value)">
        <span id="slider-value">{dates[within_days]}</span>
    </div>
    <script>"""
    
    template_2 = f"""
    var dates = {dates};"""
    
    template_3 = """
    var currentIndex = 0;
    
    function updateImage(index) {
        index = Math.round(index);  // Ensure the index is an integer
        // Update the displayed date
        document.getElementById('slider-value').innerHTML = dates[index];
        document.getElementById('date-input').value = dates[index];
        
        // Hide all images
        document.querySelectorAll('.leaflet-image-layer').forEach(function(layer) {
            layer.style.display = 'none';
        });
        
        // Show the current image
        document.querySelectorAll('.leaflet-image-layer')[index].style.display = 'block';
        
        currentIndex = index;
    }
    
    function updateFromSlider(value) {
        updateImage(parseFloat(value));
    }
    
    function updateFromInput(inputDate) {
        var index = dates.indexOf(inputDate);
        if (index !== -1) {
            document.getElementById('image-slider').value = index;
            updateImage(index);
        } else {
            alert('Invalid date. Please enter a date in the format YYYY-MM-DD that exists in the dataset.');
        }
    }
    
    // Initially show only the first image
    document.addEventListener('DOMContentLoaded', function() {
        document.querySelectorAll('.leaflet-image-layer').forEach(function(layer, index) {
            layer.style.display = index === 0 ? 'block' : 'none';
        });
    });
    </script>
    {% endmacro %}
    """
    template = template_1+template_2+template_3
    
    
    # Add the custom control to the map
    macro = MacroElement()
    macro._template = Template(template)
    m.get_root().add_child(macro)
    
    
    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", "123 Main Street, Dallas, TX 75126")
date_focus = st.sidebar.date_input("Date",  pd.Timestamp(2023, 7, 1))
within_days = st.sidebar.selectbox('Days Within', (30, 90))
# 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))

zoom_dic = {5: 12, 10: 11, 25: 10}
zoom = zoom_dic[circle_radius]

#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)

# center=row,col
radius = int(np.ceil(circle_radius*1.6))
crop_coords = col-radius, row-radius, col+radius+1, row+radius+1

files = glob.glob(r'webp/**/*.webp', recursive=True)
files_dates_selected = [i for i in files if any(
    i for j in date_range if str(j) in re.search(r'(\d{8})', i).group())]




# 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)


# Create the bar chart
fig = alt.Chart(df_data).mark_bar(size=3, color='red').encode(
    x='Date:T',  # Temporal data type
    y='Actual:Q',  # Quantitative data type
    color='Actual:Q',  # Color based on Actual values
    tooltip=[  # Adding tooltips
        alt.Tooltip('Date:T', title='Date'),
        alt.Tooltip('Actual:Q', title='Actual Value'),
        alt.Tooltip('Max:Q', title=f'Max Value with {circle_radius} Miles')
    ]
).configure(
    view=alt.ViewConfig(
        strokeOpacity=0  # No border around the chart
    )
).configure_axis(
    grid=False  # Disable grid lines
).configure_legend(
    fillColor='transparent',  # Ensure no legend is shown
    strokeColor='transparent'
)




with col1:
    st.title(f'Hail')
    try:
        st.altair_chart(fig, use_container_width=True)
        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 'is_first_run' not in st.session_state:
         # First run
         st.session_state.is_first_run = True
         st.components.v1.html(open("data/map.html", 'r').read(), height=500, width=500)
    else:
         with st.spinner("Loading... Please wait, it's gonna be great..."):
             # st_folium(m, height=500)
             # Not the first run; create a new map
             m=map_folium(lat, lon,files_dates_selected, within_days )
             m.save("map_new.html")
             st.components.v1.html(open("map_new.html", 'r').read(), height=500, width=500)

# st.bokeh_chart(hv.render(nice_plot*points_lat_lon, backend='bokeh'),use_container_width=True)


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