app.py

import os
import gradio as gr
import numpy as np
from openai import OpenAI
from Sim_Setup_Fcns import (
    load_and_crop_image, cluster_image, build_parcel_map,
    get_cluster_labels, get_land_colors, plot_parcel_map_to_file
)
from Sim_Engine import run_full_simulation
from feedback_fcns import (
    summarize_initial_conditions, plot_forage_map_to_file,
    elk_feedback, usfs_feedback, simulate_and_summarize, full_response
)

from zone_utils import (
    identify_zones, plot_labeled_zones,
    assign_zone_labels, save_zone_info_to_excel,override_zone_id_and_label
)
from R6_global import sim_state
#sim_state = SimState()  # ✅ Create it ONCE here, globally
#if sim_state.parcel_dict is None:
#    from Sim_Engine import initialize_parcels, assign_initial_forage, get_land_forage_rates
#    sim_state.parcel_dict = initialize_parcels(parcel_map, cluster_labels)
#    assign_initial_forage(sim_state.parcel_dict, get_land_forage_rates())

# === Setup on Launch ===
img = load_and_crop_image("Carson_map.png")
clustered_img = cluster_image(img)
parcel_map, n_rows, n_cols = build_parcel_map(clustered_img)
cluster_labels = get_cluster_labels()
land_colors = get_land_colors()
plot_parcel_map_to_file(parcel_map, cluster_labels, land_colors, save_path="clustered_map.png")

# === Zoning ===

# 1. Identify contiguous zones
zone_map, zone_to_cluster = identify_zones(parcel_map, connectivity="queen")

# 2. Assign human-readable labels (before override)
zone_labels = assign_zone_labels(zone_to_cluster)
# === Manual override for mislabeled riparian zone ===
# First, update the zone label once
for zid, lbl in zone_labels.items():
    if lbl == "A" and zone_to_cluster[zid] == 1:
        zone_labels[zid] = "Riparian A1"
    if lbl == "M" and zone_to_cluster[zid] == 1:
         zone_labels[zid] = "Riparian A2"
# Then update all matching parcels
for i in range(n_rows):
    for j in range(n_cols):
        zone_id = zone_map[i, j]
        if zone_labels.get(zone_id) == "Riparian A1":
            parcel_map[i, j] = 2
        if zone_labels.get(zone_id) == "Riparian A2":
            parcel_map[i, j] = 2

#       


# ⬇️ Add this block right after the override
zone_to_cluster = {}
for zone_id in np.unique(zone_map):
    indices = np.argwhere(zone_map == zone_id)
    if len(indices) > 0:
        i, j = indices[0]
        zone_to_cluster[zone_id] = parcel_map[i, j]

# 6. Plot labeled zones after override and mapping
plot_labeled_zones(zone_map, zone_labels, zone_to_cluster, save_path="zones_labeled.png")

# 5. Define cluster-to-class mapping (should stay after override)
cluster_to_class = {
    0: "desert",
    1: "pasture",
    2: "riparain",
    3: "sensitive riparian",
    4: "wetland",
    5: "water"
}

# 7. Save zone info to Excel
zone_excel_path = "zone_info.xlsx"
save_zone_info_to_excel(
    parcel_map, zone_map, zone_labels, zone_to_cluster, cluster_to_class,
    save_path=zone_excel_path
)


client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))

# === Gradio App ===
with gr.Blocks() as demo:
    gr.Markdown("# AGEC 3052 — Grazing Strategy Simulation")
    gr.Image(value="clustered_map.png", label="Initial 25×25 Parcel Layout")
    gr.Image(value="zones_labeled.png", label="Labeled Pasture & Riparian Zones")
    # ✅ Downloadable Excel
    gr.File(value=zone_excel_path, label="Download Zone Info (Excel)")
    
    plan = gr.Radio(["Conservative", "Normal", "Aggressive"], label="Grazing Plan")
    essay = gr.Textbox(lines=8, label="Your Essay Justifying the Plan")

    elk_output = gr.Textbox(label="Elk Stakeholder Feedback")
    usfs_output = gr.Textbox(label="USFS Feedback")
    sim_output = gr.Textbox(label="Simulation Results", lines=2)
    sim_image = gr.Image(label="Forage Map After Simulation", type="filepath")
    health_image = gr.Image(label="Health Map After Simulation", type="filepath")

    round_counter = gr.State(value=1)
    history = gr.State(value=[summarize_initial_conditions(n_rows, n_cols)])

    from R6_global import sim_state
#   sim_state = SimState()  # ✅ Define this at the top of app.py

    parcel_map_state = gr.State(value=parcel_map)
    cluster_labels_state = gr.State(value=cluster_labels)
    n_rows_state = gr.State(value=n_rows)
    n_cols_state = gr.State(value=n_cols)

    def submit_handler(plan_choice, essay_text, parcel_map, cluster_labels, n_rows, n_cols, history_val):
        from R6_global import sim_state
        if sim_state.parcel_dict is None:
            from Sim_Engine import initialize_parcels, assign_initial_forage, get_land_forage_rates
            sim_state.parcel_dict = initialize_parcels(parcel_map, cluster_labels)
            assign_initial_forage(sim_state.parcel_dict, get_land_forage_rates())

        parcel_dict = sim_state.parcel_dict
        round_counter = sim_state.round_counter

        return full_response(plan_choice, essay_text, history_val[-1])


#    def sim_handler(plan_choice, round_val, history_val):
#        summary, map_path, health_path, new_round = simulate_and_summarize(
#            plan_choice, round_val, parcel_map, cluster_labels, n_rows, n_cols,
#            run_full_simulation, history_val
#        )
#        history_val.append(summary)
    def sim_handler(plan_choice, round_val, history_val):
        from R6_global import sim_state
#        sim_state = SimState()

        print(f"DEBUG: sim_state.parcel_dict is {type(sim_state.parcel_dict)}")

        # Step 1: Initialize parcel_dict if it's missing
        if sim_state.parcel_dict is None:
                from Sim_Engine import initialize_parcels, assign_initial_forage, get_land_forage_rates
                print("DEBUG: Initializing parcel_dict inside sim_handler")

                sim_state.parcel_dict = initialize_parcels(parcel_map, cluster_labels)
                assign_initial_forage(
                        sim_state.parcel_dict,
                        get_land_forage_rates()
                )

        # Ensure it's initialized
        assert sim_state.parcel_dict is not None, "🚨 sim_handler: parcel_dict is STILL None after attempted init"

        # Step 2: Run the simulation
        summary, map_path, health_path, new_round, updated_parcel_dict = simulate_and_summarize(
            plan_choice, round_val, parcel_map, cluster_labels, n_rows, n_cols,
            run_full_simulation, history_val
        )
        sim_state.parcel_dict = updated_parcel_dict  # ✅ Store the updated parcel dict


        # Step 3: Manually increment round and return
        sim_state.round_counter += 1
        print(f"DEBUG: Updated round = {sim_state.round_counter}")
        history_val.append(summary)
        return summary, map_path, health_path, sim_state.round_counter, history_val

    
    submit_btn = gr.Button("Submit Grazing Plan")
    sim_btn = gr.Button("Run Simulation")

#    submit_btn.click(fn=submit_handler, inputs=[plan, essay, history], outputs=[elk_output, usfs_output])
    submit_btn.click(
        fn=submit_handler,
        inputs=[plan, essay, parcel_map_state, cluster_labels_state, n_rows_state, n_cols_state, history],
        outputs=[elk_output, usfs_output]
    )


    sim_btn.click(fn=sim_handler, inputs=[plan, round_counter, history], outputs=[sim_output, sim_image, health_image, round_counter, history])

demo.launch()