import os
import sys
import dash
from dash import dcc, html, dash_table, callback, Input, Output, State
import dash_bootstrap_components as dbc
import pandas as pd
from datetime import datetime
import numpy as np
import plotly.express as px
import plotly.graph_objects as go
from geopy.extra.rate_limiter import RateLimiter
from geopy.geocoders import Nominatim
from dash.exceptions import PreventUpdate
from vincenty import vincenty
import duckdb
import requests
import urllib
from dotenv import load_dotenv
import time
from functools import wraps
import glob


# Load environment variables
load_dotenv()

# Initialize the Dash app
app = dash.Dash(
    __name__,
    external_stylesheets=[dbc.themes.BOOTSTRAP],
    suppress_callback_exceptions=True
)
app.title = "Hail Damage Analyzer"
server = app.server

# Cache functions
def simple_cache(expire_seconds=300):
    def decorator(func):
        cache = {}
        @wraps(func)
        def wrapper(*args, **kwargs):
            key = (func.__name__, args, frozenset(kwargs.items()))
            current_time = time.time()
            if key in cache:
                result, timestamp = cache[key]
                if current_time - timestamp < expire_seconds:
                    return result
            result = func(*args, **kwargs)
            cache[key] = (result, current_time)
            return result
        return wrapper
    return decorator

@simple_cache(expire_seconds=300)
def duck_sql(sql_code):
    con = duckdb.connect()
    con.execute("PRAGMA threads=2")
    con.execute("PRAGMA enable_object_cache")
    return con.execute(sql_code).df()

@simple_cache(expire_seconds=300)
def get_data(lat, lon, date_str):
    data_dir = r"C:/Users/aammann/OneDrive - Great American Insurance Group/Documents/Python Scripts/hail_data"
    parquet_files = glob.glob(f"{data_dir}/hail_*.parquet")
    print("Parquet files found:", parquet_files)
    if not parquet_files:
        raise ValueError("No parquet files found in the specified directory")
    
    file_paths = ", ".join([f"'{file}'" for file in parquet_files])
    lat_min, lat_max = lat-1, lat+1
    lon_min, lon_max = lon-1, lon+1
    
    code = f"""
    SELECT 
        "#ZTIME" as "Date_utc", 
        LON, 
        LAT, 
        MAXSIZE 
    FROM read_parquet([{file_paths}], hive_partitioning=1)
    WHERE 
        LAT BETWEEN {lat_min} AND {lat_max}
        AND LON BETWEEN {lon_min} AND {lon_max}
        AND "#ZTIME" <= '{date_str}'
    """
    return duck_sql(code)

def distance(x):
    left_coords = (x[0], x[1])  # LAT, LON
    right_coords = (x[2], x[3])  # Lat_address, Lon_address
    return vincenty(left_coords, right_coords, miles=True)

def geocode(address):
    try:
        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[0, 0]['results'].iloc[0]['coordinates']
            return results['y'], results['x']
        except:
            geolocator = Nominatim(user_agent="HailDamageAnalyzer")
            geocode = RateLimiter(geolocator.geocode, min_delay_seconds=1)
            location = geolocator.geocode(address)
            if location:
                return location.latitude, location.longitude
            raise Exception("Geocoding failed")
    except:
        try:
            geocode_key = os.getenv("GEOCODE_KEY")
            if not geocode_key:
                raise Exception("Geocode API key not found")
            address_encoded = urllib.parse.quote(address)
            url = f'https://api.geocod.io/v1.7/geocode?q={address_encoded}&api_key={geocode_key}'
            response = requests.get(url, verify=False)
            response.raise_for_status()
            json_response = response.json()
            return json_response['results'][0]['location']['lat'], json_response['results'][0]['location']['lng']
        except Exception as e:
            print(f"Geocoding error: {str(e)}")
            raise Exception("Could not geocode address. Please try again with a different address.")

# Layout
app.layout = html.Div([
    dcc.Store(id="filtered-data-store"),
    dcc.Download(id="download-dataframe-csv"),
    dbc.Button("Download Data as CSV", id="btn-download-csv", color="secondary", className="mb-3"),
    
    dbc.Container([
        dbc.Row([
            dbc.Col([
                html.H1("Hail Damage Analyzer", className="text-center my-4"),
                html.P("Analyze historical hail data", className="text-center text-muted"),
                html.Hr()
            ], width=12)
        ]),
        
        dbc.Row([
            dbc.Col([
                html.Div([
                    html.H5("Search Parameters", className="mb-3"),
                    dbc.Form([
                        dbc.Label("Address"),
                        dbc.Input(id="address-input", type="text", placeholder="Enter address", value="Dallas, TX", className="mb-3"),
                        dbc.Label("Maximum Date"),
                        dcc.DatePickerSingle(
                            id='date-picker',
                            min_date_allowed=datetime(2010, 1, 1),
                            max_date_allowed=datetime(2025, 7, 5),
                            date=datetime(2025, 7, 5),
                            className="mb-3 w-100"
                        ),
                        dbc.Label("Show Data Within"),
                        dcc.Dropdown(
                            id='distance-dropdown',
                            options=[
                                {'label': 'All Distances', 'value': 'all'},
                                {'label': 'Within 1 Mile', 'value': '1'},
                                {'label': 'Within 3 Miles', 'value': '3'},
                                {'label': 'Within 5 Miles', 'value': '5'},
                                {'label': 'Within 10 Miles', 'value': '10'}
                            ],
                            value='all',
                            className="mb-4"
                        ),
                        dbc.Button("Search", id="search-button", color="primary", className="w-100 mb-3")
                    ]),
                    html.Div(id="summary-cards", className="mt-4")
                ], className="p-3 bg-light rounded-3")
            ], md=4),
            
            dbc.Col([
                dbc.Row([
                    dbc.Col([
                        dbc.Card([
                            dbc.CardHeader("Hail Data Overview"),
                            dbc.CardBody([
                                dcc.Loading(
                                    id="loading-hail-data",
                                    type="circle",
                                    children=[
                                        html.Div(id="hail-data-table"),
                                        html.Div(id="map-container", className="mt-4")
                                    ]
                                )
                            ])
                        ])
                    ])
                ]),
                dbc.Row([
                    dbc.Col([
                        dbc.Card([
                            dbc.CardHeader("Hail Size Over Time"),
                            dbc.CardBody([
                                dcc.Loading(
                                    id="loading-hail-chart",
                                    type="circle",
                                    children=[
                                        dcc.Graph(id="hail-size-chart")
                                    ]
                                )
                            ])
                        ], className="mt-4")
                    ])
                ])
            ], md=8)
        ]),
        
        html.Div(id="intermediate-data", style={"display": "none"}),
        dbc.Row([
            dbc.Col([
                html.Hr(),
                html.P("© 2025 Hail Damage Analyzer", className="text-center text-muted small")
            ])
        ], className="mt-4")
    ], fluid=True)
])

# Main callback
@app.callback(
    [Output("intermediate-data", "children"),
     Output("summary-cards", "children"),
     Output("hail-data-table", "children"),
     Output("map-container", "children"),
     Output("hail-size-chart", "figure"),
     Output("filtered-data-store", "data")],
    [Input("search-button", "n_clicks"),
     Input("address-input", "n_submit")],
    [State("address-input", "value"),
     State("date-picker", "date"),
     State("distance-dropdown", "value")]
)
def update_all(n_clicks, n_submit, address, date_str, distance_filter):
    print("Update all callback triggered")  # Debug
    ctx = dash.callback_context
    if not ctx.triggered:
        raise PreventUpdate
    
    try:
        lat, lon = geocode(address)
        date_obj = datetime.strptime(date_str.split('T')[0], '%Y-%m-%d')
        date_formatted = date_obj.strftime('%Y%m%d')
        df = get_data(lat, lon, date_formatted)
        
        if df.empty:
            error_alert = dbc.Alert("No hail data found for this location and date range.", color="warning")
            return dash.no_update, error_alert, "", "", {}, []
        
        df["Lat_address"] = lat
        df["Lon_address"] = lon
        df['Miles to Hail'] = [
            distance(i) for i in df[['LAT', 'LON', 'Lat_address', 'Lon_address']].values
        ]
        df['MAXSIZE'] = df['MAXSIZE'].round(2)
        
        if distance_filter != 'all':
            max_distance = float(distance_filter)
            df = df[df['Miles to Hail'] <= max_distance]
        
        max_size = df['MAXSIZE'].max()
        last_date = df['Date_utc'].max()
        total_events = len(df)
        
        summary_cards = dbc.Row([
            dbc.Col([
                dbc.Card([
                    dbc.CardBody([
                        html.H6("Max Hail Size (in)", className="card-title"),
                        html.H3(f"{max_size:.1f}", className="text-center")
                    ])
                ], className="text-center")
            ], md=4, className="mb-3"),
            dbc.Col([
                dbc.Card([
                    dbc.CardBody([
                        html.H6("Last Hail Event", className="card-title"),
                        html.H3(last_date, className="text-center")
                    ])
                ], className="text-center")
            ], md=4, className="mb-3"),
            dbc.Col([
                dbc.Card([
                    dbc.CardBody([
                        html.H6("Total Events", className="card-title"),
                        html.H3(f"{total_events}", className="text-center")
                    ])
                ], className="text-center")
            ], md=4, className="mb-3")
        ])
        
        df_display = df[['Date_utc', 'MAXSIZE', 'Miles to Hail']].copy()
        df_display['Miles to Hail'] = df_display['Miles to Hail'].round(2)
        df_display = df_display.rename(columns={
            'Date_utc': 'Date',
            'MAXSIZE': 'Hail Size (in)',
            'Miles to Hail': 'Distance (miles)'
        })
        
        data_table = dash_table.DataTable(
            id='hail-data-table',
            columns=[{"name": i, "id": i} for i in df_display.columns],
            data=df_display.to_dict('records'),
            page_size=10,
            style_table={'overflowX': 'auto'},
            style_cell={
                'textAlign': 'left',
                'padding': '8px',
                'minWidth': '50px', 'width': '100px', 'maxWidth': '180px',
                'whiteSpace': 'normal'
            },
            style_header={
                'backgroundColor': 'rgb(230, 230, 230)',
                'fontWeight': 'bold'
            },
            style_data_conditional=[
                {
                    'if': {
                        'filter_query': '{Hail Size (in)} >= 2',
                        'column_id': 'Hail Size (in)'
                    },
                    'backgroundColor': '#ffcccc',
                    'fontWeight': 'bold'
                }
            ]
        )

        map_fig = go.Figure()
        for _, row in df.iterrows():
            size = row['MAXSIZE']
            map_fig.add_trace(
                go.Scattermapbox(
                    lon=[row['LON']],
                    lat=[row['LAT']],
                    mode='markers',
                    marker=go.scattermapbox.Marker(
                        size=size * 3,
                        color='red',
                        opacity=0.7
                    ),
                    text=f"Size: {size} in Date: {row['Date_utc']}",
                    hoverinfo='text',
                    showlegend=False
                )
            )

        if not df.empty:
            center_lat = df['Lat_address'].iloc[0]
            center_lon = df['Lon_address'].iloc[0]
            map_fig.add_trace(
                go.Scattermapbox(
                    lon=[center_lon],
                    lat=[center_lat],
                    mode='markers',
                    marker=go.scattermapbox.Marker(
                        size=14,
                        color='blue',
                        symbol='star'
                    ),
                    text=f"Your Location: {address}",
                    hoverinfo='text',
                    showlegend=False
                )
            )

        map_fig.update_layout(
            mapbox_style="open-street-map",
            mapbox=dict(
                center=dict(lat=center_lat, lon=center_lon),
                zoom=10
            ),
            margin={"r":0, "t":0, "l":0, "b":0},
            height=400
        )

        df_chart = df.copy()
        df_chart['Date'] = pd.to_datetime(df_chart['Date_utc'])
        df_chart = df_chart.sort_values('Date')

        chart_fig = px.scatter(
            df_chart,
            x='Date',
            y='MAXSIZE',
            color='Miles to Hail',
            size='MAXSIZE',
            hover_data=['Miles to Hail'],
            title='Hail Size Over Time',
            labels={
                'MAXSIZE': 'Hail Size (in)',
                'Miles to Hail': 'Distance (miles)'
            }
        )

        chart_fig.update_traces(
            marker=dict(
                line=dict(width=1, color='DarkSlateGrey'),
                opacity=0.7
            ),
            selector=dict(mode='markers')
        )

        chart_fig.update_layout(
            xaxis_title='Date',
            yaxis_title='Hail Size (in)',
            plot_bgcolor='rgba(0,0,0,0.02)',
            paper_bgcolor='white',
            hovermode='closest'
        )

        intermediate_data = df.to_json(date_format='iso', orient='split')
        map_figure = dcc.Graph(figure=map_fig)
        chart_figure = chart_fig
        store_data = df.to_dict('records')
        print("Store data populated:", store_data[:2])

        return (
            intermediate_data,
            summary_cards,
            data_table,
            map_figure,
            chart_figure,
            store_data
        )

    except Exception as e:
        error_alert = dbc.Alert(f"Error: {str(e)}", color="danger")
        return dash.no_update, error_alert, "", "", {}, []

from dash import callback_context

@callback(
    Output("download-dataframe-csv", "data"),
    [Input("btn-download-csv", "n_clicks")],
    [State("filtered-data-store", "data")],
    prevent_initial_call=True
)
def download_csv(n_clicks, data):
    if not n_clicks or not data:
        return dash.no_update

    df = pd.DataFrame(data)
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"hail_data_export_{timestamp}.csv"
    csv_string = df.to_csv(index=False, encoding='utf-8')
    return dict(content=csv_string, filename=filename)


if __name__ == '__main__':
    print("🚀 Dash app is running! Open this link in your browser:")
    print("👉 http://localhost:7860/")
    app.run(debug=True, host='0.0.0.0', port=7860)