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
)

# === 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)])

    def submit_handler(plan_choice, essay_text, history_val):
        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)
        return summary, map_path, health_path, new_round, 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])
    sim_btn.click(fn=sim_handler, inputs=[plan, round_counter, history], outputs=[sim_output, sim_image, health_image, round_counter, history])

demo.launch()