import pandas as pd
import numpy as np
import gradio as gr
import matplotlib.pyplot as plt
import h3.api.basic_str as h3
import folium
from shapely.geometry import Polygon
import geopandas as gpd
import tempfile
import zipfile
import os
import glob
from shapely.ops import unary_union

# Define feature lists & descriptions
PRIMARY_FEATURES = ['OEM Critical Facilities', 'High Frequency Flooding', 'Low Frequency Flooding', 'METRO', 'HPW Critical Facilities', 'Other Critical Facilities']
SECONDARY_FEATURES = ['PCI', 'ADT', 'RPL_THEME_1', 'RPL_THEME_2', 'RPL_THEME_3', 'RPL_THEME_4', 'RPL_THEME_ALL']
ALL_FEATURES = PRIMARY_FEATURES + SECONDARY_FEATURES
EXTRA_COLUMNS = ['PCI Score', 'SURFACETYP', 'METRO_ROUTES', 'STREETNAME']

FEATURE_DESCRIPTIONS = {
    'OEM Critical Facilities': 'Distance to HFD, HPD, hospitals, and OEM facilities.',
    'HPW Critical Facilities': 'Distance to Houston Water facilities.',
    'Other Critical Facilities': 'Distance to NGOs, private entities, MSC, CC.',
    'High Frequency Flooding': 'Roadway flooding instance for 2, 5, 10 and 25 Year Storm events.',
    'Low Frequency Flooding': 'Roadway flooding instance for 50, 100 and 500 Year Storm events.',
    'METRO': 'Number of METRO bus routes.',
    'PCI': 'Pavement Condition Index.',
    'ADT': 'Annual Average Daily Traffic.',
    'RPL_THEME_1': 'Socioeconomic vulnerability (SVI theme 1).',
    'RPL_THEME_2': 'Household characteristics (SVI theme 2).',
    'RPL_THEME_3': 'Racial & ethnic minority (SVI theme 3).',
    'RPL_THEME_4': 'Housing & transportation (SVI theme 4).',
    'RPL_THEME_ALL': 'Overall social vulnerability (combined).'
}

# Load data including SurfaceTyp only once
def load_data(path="./MainInstances.xlsx"):
    df = pd.read_excel(path, dtype={'SURFACETYP': str, 'METRO_ROUTES': str})
    df.columns = df.columns.str.strip()
    
    # Ensure SurfaceTyp and METRO_ROUTES columns exist
    if 'SURFACETYP' not in df.columns:
        df['SURFACETYP'] = ''
    if 'METRO_ROUTES' not in df.columns:
        df['METRO_ROUTES'] = ''
    
    # Ensure other columns exist
    for col in ['GRID_ID'] + ALL_FEATURES + ['PCI Score', 'STREETNAME']:
        if col not in df.columns:
            df[col] = 0
            
    # Trim to needed columns
    cols = ['GRID_ID'] + ALL_FEATURES + ['PCI Score', 'SURFACETYP', 'METRO_ROUTES', 'STREETNAME']
    df = df[cols]
    
    # Convert numeric fields
    num_cols = ALL_FEATURES + ['PCI Score']
    df[num_cols] = df[num_cols].apply(pd.to_numeric, errors='coerce').fillna(0)
    
    # Clean SurfaceTyp and METRO_ROUTES
    df['METRO_ROUTES'] = df['METRO_ROUTES'].fillna('').astype(str).str.strip()
    df['SURFACETYP'] = df['SURFACETYP'].fillna('').astype(str).str.strip()
    
    return df

# H3 boundary to coords
def h3_poly_coords(h3_id):
    return [(lat, lon) for lat, lon in h3.cell_to_boundary(h3_id)]

# Compute weighted index
def calculate_index(df, features, weights):
    w = np.array(weights) / np.sum(weights)
    df_idx = df.copy()
    df_idx['Index'] = df_idx[features].values.dot(w)
    return df_idx

# Stats markdown
def stats_text(df_idx):
    s = df_idx['Index'].describe()
    return (
        "### Index Statistics\n"
        f"- Min: {s['min']:.4f}\n"
        f"- 25%: {s['25%']:.4f}\n"
        f"- Median: {s['50%']:.4f}\n"
        f"- 75%: {s['75%']:.4f}\n"
        f"- Max: {s['max']:.4f}"
    )

# Histogram
def make_histogram(df_idx):
    fig, ax = plt.subplots()
    ax.hist(df_idx['Index'], bins=30)
    ax.set_title('Index Distribution')
    ax.set_xlabel('Index')
    ax.set_ylabel('Frequency')
    plt.tight_layout()
    plt.close(fig)
    return fig

# Create index map
def create_index_map(df_idx, features):
    sel = df_idx[df_idx['Index'] >= df_idx['Index'].quantile(0.5)]
    center = (29.76, -95.37)
    if not sel.empty:
        center = h3.cell_to_latlng(sel.iloc[0]['GRID_ID'])
    m = folium.Map(location=center, zoom_start=10)
    for _, row in sel.iterrows():
        idx = row['Index']
        fill_opacity = 0.6 if 3 <= idx <= 4 else 0
        norm = (idx - sel['Index'].min()) / (sel['Index'].max() - sel['Index'].min()) if sel['Index'].max() > sel['Index'].min() else 0.5
        color = f"rgb({int(255*norm)},0,{int(255*(1-norm))})"
        popup = "<br>".join(
            [f"{feat}: {row[feat]:.2f}" for feat in features] +
            [f"{col}: {row[col]}" for col in EXTRA_COLUMNS] +
            [f"Index: {idx:.2f}"]
        )
        folium.Polygon(
            locations=h3_poly_coords(row['GRID_ID']),
            color='black', fill=True, fill_color=color, fill_opacity=fill_opacity,
            popup=popup
        ).add_to(m)
    return m._repr_html_()

# Clustering logic
def find_clusters(hex_ids, min_size):
    clusters, visited = [], set()
    for h in hex_ids:
        if h in visited: continue
        cluster, stack = {h}, [h]
        visited.add(h)
        while stack:
            cur = stack.pop()
            for nb in h3.grid_disk(cur, 1):
                if nb in hex_ids and nb not in visited:
                    visited.add(nb)
                    cluster.add(nb)
                    stack.append(nb)
        if len(cluster) >= min_size:
            clusters.append(cluster)
    return clusters

# Generate clusters map/table and shapefile
def create_cluster_outputs(df_idx, threshold, min_size, pci_path="./PCI.json"):
    # 1) Filter by threshold
    filt = df_idx[df_idx['Index'] >= threshold]
    if filt.empty:
        return "No hexes above threshold.", None, None, None

    # 2) Find clusters
    clusters = find_clusters(filt['GRID_ID'].tolist(), min_size)
    if not clusters:
        return "No clusters found.", None, None, None

    # 3) Flatten all hex IDs
    ids = set().union(*clusters)

    # 4) Build table of cluster hexes
    table = df_idx[df_idx['GRID_ID'].isin(ids)][['GRID_ID', 'Index'] + EXTRA_COLUMNS]

    # 5) Create Folium map centered on first hex
    center = h3.cell_to_latlng(next(iter(ids)))
    m = folium.Map(location=center, zoom_start=12)

    # 6) Draw each hexagon with popup, and collect Shapely polygons
    hex_polys = []
    for _, row in table.iterrows():
        h3_id = row['GRID_ID']
        coords = h3_poly_coords(h3_id)
        # HTML popup
        popup_html = (
            f"<b>Hex:</b> {h3_id}<br>"
            f"<b>Index:</b> {row['Index']:.2f}<br>"
            + "<br>".join(f"{col}: {row[col]}" for col in EXTRA_COLUMNS)
        )
        folium.Polygon(
            locations=coords,
            color='blue', fill=True, fill_opacity=0.5,
            popup=folium.Popup(popup_html, max_width=300)
        ).add_to(m)
        # Note: Shapely expects (lon, lat)
        hex_polys.append(Polygon([(lon, lat) for lat, lon in coords]))

    # 7) Union hex polygons for the PCI intersection
    cluster_union = unary_union(hex_polys)

    # 8) Load PCI GeoJSON, set & transform CRS
    if os.path.exists(pci_path):
        pci_gdf = gpd.read_file(pci_path)
        pci_gdf = pci_gdf.set_crs("EPSG:2278", allow_override=True).to_crs(epsg=4326)

        # 9) Filter line segments that intersect the cluster_union
        lines = pci_gdf[pci_gdf.intersects(cluster_union)]

        if not lines.empty:
            # 10) Serialize any Timestamp fields to string
            for col in lines.columns:
                if pd.api.types.is_datetime64_any_dtype(lines[col]):
                    lines[col] = lines[col].astype(str)

            # 11) Add PCI lines as a GeoJson layer
            folium.GeoJson(
                data=lines.to_json(),
                name="PCI Lines",
                popup=folium.GeoJsonPopup(
                    fields=[c for c in lines.columns if c != 'geometry'],
                    aliases=[c.replace('_', ' ').title() for c in lines.columns if c != 'geometry'],
                    localize=True
                ),
                style_function=lambda feat: {'color': 'red', 'weight': 2}
            ).add_to(m)

    # 12) Add layer control so overlays are visible/toggleable
    folium.LayerControl().add_to(m)

    # 13) Generate the Shapefile ZIP
    tmp = tempfile.mkdtemp()
    base = os.path.join(tmp, 'focus_area')
    
    # Crucial Fix: Shapely needs (lon, lat) while h3 provides (lat, lon)
    shp_polys = [Polygon([(lon, lat) for lat, lon in h3.cell_to_boundary(h)]) for h in table['GRID_ID']]
    gdf = gpd.GeoDataFrame(table, geometry=shp_polys, crs='EPSG:4326')
    gdf.to_file(base + '.shp')
    
    zip_path = base + '.zip'
    with zipfile.ZipFile(zip_path, 'w') as zf:
        for f in glob.glob(base + '.*'):
            if not f.endswith('.zip'):
                zf.write(f, arcname=os.path.basename(f))

    # Return the zip_path to the Gradio file component instead of None
    return "Clusters, PCI overlay, and Shapefile ready!", m._repr_html_(), table, zip_path

# Build Gradio app
def create_app():
    # If load_data fails during initial boot without excel file, we can bypass or load a dummy
    try:
        df = load_data()
    except FileNotFoundError:
        df = pd.DataFrame() 

    with gr.Blocks() as app:
        gr.Markdown("# Focus Area Identification by creating a Risk Index")
        gr.Markdown("Build a custom risk index for Houston thoroughfare hexes by assigning feature weights.")
        href = "https://www.mermaidchart.com/raw/7431d57c-8e4d-47ec-bb22-ec64c34f5d03?theme=light&version=v0.1&format=svg"
        gr.Markdown(f"[This is a link to view the behind the scenes flow chart]({href})")

        # Feature Weights
        gr.Markdown("### Step 1: Assign weights to features")
        init = pd.DataFrame({
            "Feature": ALL_FEATURES,
            "Description": [FEATURE_DESCRIPTIONS[f] for f in ALL_FEATURES],
            "Weight": [0]*len(ALL_FEATURES)
        })
        wf = gr.Dataframe(value=init, headers=["Feature","Description","Weight"], datatype=["str","str","number"], interactive=True)
        btn1 = gr.Button("Calculate Index")
        stat_msg = gr.Markdown()
        out_stats = gr.Markdown()
        out_hist = gr.Plot()
        out_map = gr.HTML()

        def calc(wt):
            try:
                df_wt = pd.DataFrame(wt, columns=["Feature","Description","Weight"]) if isinstance(wt, list) else wt
                df_wt['Weight'] = pd.to_numeric(df_wt['Weight'], errors='coerce').fillna(0)
                sel = df_wt[df_wt['Weight']>0]
                if not np.isclose(sel['Weight'].sum(), 100):
                    raise ValueError(f"Weights sum to {sel['Weight'].sum():.2f}, must be 100.")
                df_idx = calculate_index(df, sel['Feature'].tolist(), sel['Weight'].tolist())
                return "Risk Index created!", stats_text(df_idx), make_histogram(df_idx), create_index_map(df_idx, sel['Feature'].tolist())
            except Exception as e:
                return f"⚠️ Error calculating index: {e}", None, None, None

        btn1.click(calc, inputs=[wf], outputs=[stat_msg, out_stats, out_hist, out_map])

        # Clustering
        gr.Markdown("### Step 2: Prioritized Focus Area Creation")
        thr = gr.Number(label="Index Threshold", value=3)
        ms = gr.Number(label="Minimum Cluster Size", value=3)
        btn2 = gr.Button("Create Clusters")
        cl_msg = gr.Markdown()
        cl_map = gr.HTML()
        cl_tbl = gr.Dataframe(label="Focus Area Hexes Data")
        cl_zip = gr.File(label="Download Focus Area Shapefile Zip")

        def cluster(wt, t, msize):
            try:
                df_wt = pd.DataFrame(wt, columns=["Feature","Description","Weight"]) if isinstance(wt, list) else wt
                df_wt['Weight'] = pd.to_numeric(df_wt['Weight'], errors='coerce').fillna(0)
                sel = df_wt[df_wt['Weight']>0]
                if not np.isclose(sel['Weight'].sum(), 100):
                    raise ValueError(f"Weights sum to {sel['Weight'].sum():.2f}, must be 100.")
                df_idx = calculate_index(df, sel['Feature'].tolist(), sel['Weight'].tolist())
                return create_cluster_outputs(df_idx, t, int(msize))
            except Exception as e:
                return f"⚠️ Error clustering: {e}", None, None, None

        btn2.click(cluster, inputs=[wf, thr, ms], outputs=[cl_msg, cl_map, cl_tbl, cl_zip])

    return app

if __name__ == '__main__':
    app = create_app()
    app.launch()