import json
import geopandas as gpd
import folium
from folium import FeatureGroup, LayerControl
import numpy as np

# ==========================================
# 1. Load Configuration
# ==========================================
config_path = "config/cvi_scoring_simple.json"

try:
    with open(config_path, "r") as f:
        config = json.load(f)
except FileNotFoundError:
    print(f"⚠️ Config file not found at {config_path}.")
    raise

# Extract the "Master" Palette (1-5 -> Green-Red) for general use
default_palette = {
    k: v["color_hex"] 
    for k, v in config["meta"]["default_palette"].items()
}

# ==========================================
# 2. Helper: Color Lookup Logic
# ==========================================
def get_layer_style(feature, layer_type, config, default_palette):
    """
    Determines the color based on the '_score' column found in the GeoJSON
    and the definitions in the config JSON.
    """
    props = feature.get("properties", {})
    
    # Default fallback color (Gray for missing data/NaN)
    color = "#999999"
    opacity = 0.8
    
    try:
        score_val = None
        
        # --- A. LAND COVER (Unique Config Colors) ---
        if layer_type == "land_cover":
            # The file has 'land_cover_score' (1.0 to 5.0)
            val = props.get("land_cover_score")
            
            # Check if valid number (not None or NaN)
            if val is not None and not np.isnan(val):
                # Convert 1.0 -> "1"
                key = str(int(val))
                # Use specific colors defined in config["land_cover"]["classes"]
                if key in config["land_cover"]["classes"]:
                    color = config["land_cover"]["classes"][key]["color_hex"]
            else:
                # If land cover is NaN (like in your sample), make it faint gray
                opacity = 0.3

        # --- B. PHYSICAL VARIABLES (Slope, Erosion, Elevation) ---
        elif layer_type in ["slope", "erosion", "elevation"]:
            # The files have 'slope_score', 'erosion_score', etc.
            col_name = f"{layer_type}_score"
            val = props.get(col_name)
            
            if val is not None and not np.isnan(val):
                # Convert 1.0 -> "1"
                key = str(int(val))
                # Map 1-5 score to the Default Traffic Light Palette
                color = default_palette.get(key, color)

        # --- C. FINAL CVI ---
        elif layer_type == "cvi":
            # The file uses 'CVI_equal_class' (1, 2, 3, 4, 5)
            val = props.get("CVI_equal_class")
            
            if val is not None and not np.isnan(val):
                key = str(int(val))
                color = default_palette.get(key, color)

    except Exception as e:
        # Keep default gray if parsing fails
        pass

    return {
        "color": color,
        "weight": 4,
        "opacity": opacity
    }

# ==========================================
# 3. Define Layers to Load
# ==========================================
layers_info = [
    {
        "name": "1. Land Cover",
        "type": "land_cover",
        "file": "./output_data/transects_with_land_cover.geojson"
    },
    {
        "name": "2. Slope",
        "type": "slope",
        "file": "./output_data/transects_with_slope.geojson"
    },
    {
        "name": "3. Erosion",
        "type": "erosion",
        "file": "./output_data/transects_with_erosion.geojson"
    },
    {
        "name": "4. Elevation",
        "type": "elevation",
        "file": "./output_data/transects_with_elevation.geojson"
    },
    {
        "name": "5. Final CVI (Equal)",
        "type": "cvi",
        "file": "./output_data/transects_with_cvi_equal.geojson"
    }
]

# ==========================================
# 4. Initialize Map
# ==========================================
# Load CVI first to center the map
gdf_center = gpd.read_file("./output_data/transects_with_cvi_equal.geojson")
center = gdf_center.unary_union.centroid
m = folium.Map(location=[center.y, center.x], zoom_start=11, tiles="CartoDB positron")

# ==========================================
# 5. Add Layers
# ==========================================
for layer in layers_info:
    try:
        gdf = gpd.read_file(layer["file"])
        
        # Only show CVI by default
        is_visible = (layer["type"] == "cvi")
        
        fg = FeatureGroup(name=layer["name"], show=is_visible)
        
        # Dynamically get all columns for tooltip (except geometry)
        tooltip_cols = [c for c in gdf.columns if c != 'geometry']
        
        folium.GeoJson(
            gdf,
            style_function=lambda feature, ltype=layer["type"]: get_layer_style(
                feature, ltype, config, default_palette
            ),
            tooltip=folium.GeoJsonTooltip(fields=tooltip_cols, localize=True)
        ).add_to(fg)
        
        fg.add_to(m)
        print(f"Added layer: {layer['name']}")
        
    except Exception as e:
        print(f"Could not load {layer['file']}: {e}")

# ==========================================
# 6. Add Legend
# ==========================================
legend_items_html = ""
sorted_keys = sorted(config["meta"]["default_palette"].keys())

for key in sorted_keys:
    item = config["meta"]["default_palette"][key]
    color = item["color_hex"]
    
    label = f"Score {key}" 
    if key == "1": label += " (Low Vuln.)"
    if key == "5": label += " (High Vuln.)"
    
    legend_items_html += f"""
    <div style="margin-bottom: 4px;">
        <span style="background:{color};width:12px;height:12px;display:inline-block;border:1px solid #999;margin-right:5px;"></span>
        {label}
    </div>
    """

legend_html = f"""
<div style="
     position: fixed;
     bottom: 30px;
     left: 30px;
     width: 160px;
     z-index: 9999;
     background: white;
     padding: 10px;
     border-radius: 6px;
     border: 2px solid rgba(0,0,0,0.2);
     font-size: 13px;
     font-family: sans-serif;">
     <b>CVI Scores</b><hr style="margin: 5px 0;">
     {legend_items_html}
     <hr style="margin: 5px 0;">
     <small><i>Gray = No Data</i></small>
</div>
"""

m.get_root().html.add_child(folium.Element(legend_html))
folium.LayerControl().add_to(m)

m