import os
import random
import numpy as np
import matplotlib.pyplot as plt
from openai import OpenAI
from Sim_Engine import simulate_period

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

def summarize_initial_conditions(n_rows, n_cols):
    num_parcels = n_rows * n_cols
    avg_forage = round(random.uniform(6.5, 8.0), 1)
    degraded_pct = round(random.uniform(0.0, 2.0), 1)
    riparian_health = random.choice(["excellent", "moderate", "fragile"])
    elk_corridor_status = random.choice([
        "completely intact and lightly used",
        "intact but under slight pressure from cattle movement",
        "showing signs of fragmentation near key crossings"
    ])
    rainfall_outlook = random.choice(["normal", "below average", "above average"])

    return (
        f"There are {num_parcels} parcels in total. Grazing has not yet occurred.\n"
        f"Average forage availability is {avg_forage} AUMs per parcel, with about {degraded_pct}% of land already degraded due to prior conditions.\n"
        f"Riparian zone condition is {riparian_health}, and the elk movement corridor is {elk_corridor_status}.\n"
        f"The seasonal rainfall outlook is {rainfall_outlook}."
    )

def plot_forage_map_to_file(parcel_dict, n_rows, n_cols, title="Forage Map", save_path="forage_map.png"):
    forage_map = np.array([
        [parcel_dict[(i, j)]["forage"] for j in range(n_cols)]
        for i in range(n_rows)
    ])
    fig, ax = plt.subplots(figsize=(8, 6))
    cax = ax.imshow(forage_map, cmap='YlGn', origin='upper')
    fig.colorbar(cax, label="Forage AUMs")
    ax.set_title(title)
    ax.axis('off')
    plt.tight_layout()
    plt.savefig(save_path)
    plt.close(fig)

def elk_feedback(plan_choice, current_summary):
    prompt = f"""
You represent a coalition of elk-related interests: conservationists, hunting advocates, and the hospitality/lodging industry.
A student has selected the **'{plan_choice}'** cattle grazing strategy. Below are the current ecological conditions:
-----
{current_summary}
-----
Please do the following:
- Explicitly choose **one elk management strategy** from the list:
  - **Preserve**: strict elk protections, no hunting, unrestricted movement
  - **Cooperate**: shared use corridor, some riparian restrictions, sustainable elk population
  - **Exploit**: prioritize hunting/tourism, tolerate reduced elk numbers and access
- Reflect each group's view briefly, but unify the final position.
- Justify your strategy choice based on the above ecological indicators.
"""
    response = client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.8
    )
    return response.choices[0].message.content

def usfs_feedback(plan_choice, student_essay, current_summary):
    # Extract the AUM line from the summary
    aum_line = ""
    for line in current_summary.split("\n"):
        if "Average forage availability" in line:
            aum_line = line.strip()
            break

    prompt = f"""
You are a USFS land management agent evaluating a student’s cattle grazing proposal.

The student selected the **'{plan_choice}'** strategy and submitted this justification:
-----
{student_essay}
-----

Here are the current rangeland conditions:
-----
{current_summary}
-----

🚨 **Must Include Block**:
You must include this sentence exactly in your response:
→ "{aum_line}"

Then provide your evaluation:
- Comment on forage, degradation, riparian and corridor health
- State whether the plan is ecologically sound
- Suggest improvements if needed
- Keep the tone professional, clear, and grounded in the data
"""

    response = client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.7
    )
    return response.choices[0].message.content

def simulate_and_summarizeold(plan_choice, round_counter, parcel_dict, parcel_map, cluster_labels, n_rows, n_cols, elk_pressure):
    # Interpret strategy
    strategy_map = {
        "conservative": "conservative",
        "normal": "moderate",
        "aggressive": "aggressive"
    }
    strategy = strategy_map.get(plan_choice.lower(), "moderate")

    # ✅ Reuse incoming parcel_dict — do not reset
    simulate_period(parcel_dict, grazing_strategy=strategy, elk_pressure=elk_pressure)

    # Extract summary
    forage_vals = [p["forage"] for p in parcel_dict.values()]
    avg_forage = sum(forage_vals) / len(forage_vals)
    degraded_pct = 100 * sum(p["health"] < 0.5 for p in parcel_dict.values()) / len(parcel_dict)

    summary = (
        f"After Round {round_counter} with the '{plan_choice}' plan:\n"
        f"Avg forage: {avg_forage:.1f} AUMs, Parcels with low health (<0.5): {degraded_pct:.1f}%"
    )

    # Generate map visuals
    from Sim_Engine import get_forage_map, get_health_map, plot_forage_map, plot_health_map
    forage_map = get_forage_map(parcel_dict, n_rows, n_cols)
    health_map = get_health_map(parcel_dict, n_rows, n_cols)
    plot_forage_map(forage_map, title=f"Forage Map (Round {round_counter})", save_path="forage_map.png")
    plot_health_map(health_map, title=f"Health Map (Round {round_counter})", save_path="health_map.png")

    return summary, "forage_map.png", "health_map.png", round_counter + 1, parcel_dict

def simulate_and_summarize(plan_choice, round_counter, parcel_map, cluster_labels, n_rows, n_cols, run_full_simulation, history):

    strategy_map = {
        "conservative": "conservative",
        "normal": "moderate",
        "aggressive": "aggressive"
    }
    strategy = strategy_map.get(plan_choice.lower())
    parcel_dict = run_full_simulation(parcel_map, cluster_labels, n_rows, n_cols, strategy=strategy)

    plot_forage_map_to_file(parcel_dict, n_rows, n_cols, title=f"Round {round_counter} Forage Map")

    # ✅ ADD THIS BLOCK
    from Sim_Engine import get_health_map, plot_health_map
    health_map = get_health_map(parcel_dict, n_rows, n_cols)
    plot_health_map(health_map, title=f"Round {round_counter} Health Map", save_path="health_map.png")

    
    forage_vals = [p["forage"] for p in parcel_dict.values()]
    avg_forage = sum(forage_vals) / len(forage_vals)
    degraded_pct = 100 * sum(p["degraded"] for p in parcel_dict.values()) / len(parcel_dict)

    summary = (
        f"After Round {round_counter} with the '{plan_choice}' plan:\n"
        f"Avg forage: {avg_forage:.1f} AUMs, Degraded parcels: {degraded_pct:.1f}%"
    )
    return summary, "forage_map.png", "health_map.png", round_counter + 1


def full_response(plan_choice, essay_text, current_summary):
    elk_resp = elk_feedback(plan_choice, current_summary)
    usfs_resp = usfs_feedback(plan_choice, essay_text, current_summary)
    return elk_resp, usfs_resp