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 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
from streamlit_plotly_events import plotly_events
import warnings
from folium import plugins
warnings.filterwarnings("ignore")


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


@st.cache_data
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):
    files = [
        "data/GUST_hrrr_2021_all_max.h5",
        "data/GUST_hrrr_2022_all_max.h5",
        "data/GUST_hrrr_2023_all_max.h5",
        "data/GUST_hrrr_202401_202409_max.h5",
    ]

    all_data = []
    all_dates = []
    for f in files:
        with h5py.File(f, 'r') as f:
            data = f['GUST'][:, 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)*2.23694

    dates_mat = np.concatenate(all_dates)

    data_actual = np.array([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)
    data_max_2 = data_max_2

    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 lat_lon_to_row_col(lat, lon):
    crs_dic = pickle.load(open('data/hrrr_crs.pkl', 'rb'))
    transform = crs_dic['affine']
    trans_rtma = crs_dic['proj_4326']
    lon_rtma, lat_rtma = trans_rtma.transform(lon, lat)

    row, col = rasterio.transform.rowcol(transform, lon_rtma, lat_rtma)
    row, col = int(row), int(col)
    return row, col


def map_folium(lat, lon, files_dates_selected, actual_point, max_point):
    popup_content = f"""
    <div style="font-size: 7pt; width: 100px; height: 100px;">
        <b>{address}</b><br>
        <b>Gust: {actual_point:,.2f} MPH</b><br>
        <b>Max: {max_point:,.2f} MPH</b><br>
    </div>
    """
    # Create a base map
    m = folium.Map(location=[lat, lon], zoom_start=6)
    folium.Marker(location=[lat, lon],

                  popup=folium.Popup(popup_content, max_width=400)


                  ).add_to(m)

    # Define the image bounds (SW and NE corners)
    image_bounds = [[21.081227294744885, -134.02385805744265],
                    [52.62502506520796, -60.98015065638055]]

    # 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, vmax=75)
    # Add the color legend to the map
    colormap_hail.caption = 'Legend: Wind Gust (MPH)'
    colormap_hail.add_to(m)

    plugins.Fullscreen().add_to(m)

    return m


def get_all_data(lat, lon, radius, start_date, end_date):
    #Geocode and get Data

    row, col = lat_lon_to_row_col(lat, lon)

    df_data, max_values = get_data(row, col, radius)

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

    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>Gust: {row["Actual"]}<br>Max: {row["Max"]}', axis=1)

    ))

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

    return fig, df_data


#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(2021, 7, 1))
within_days = st.sidebar.selectbox('Days Within', (90, 180, 365))
circle_radius = st.sidebar.selectbox('Box Radius (Miles)', (5, 10, 25))
interactive_map = st.sidebar.radio(
    'Interactive Map (Lower Res)', (False, True))


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

result = geocode(address)
lat, lon = result.values[0]
radius = int(np.ceil(circle_radius*1.6/2.5))

fig, df_data = get_all_data(lat, lon, radius, start_date, end_date)

_, actual_point, max_point = df_data.query(f"Date=='{date_focus}'").values[0]

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

with col1:
    st.title('Gust')
    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('GUST MAP')

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

   
        _, actual_point, max_point = df_data.query(
            f"Date=='{selected_data['Date']}'").values[0]
        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, actual_point, max_point)
        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, actual_point, max_point)
        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)