def initialize_parcels(parcel_map, cluster_labels):
    parcel_dict = {}
    for i in range(parcel_map.shape[0]):
        for j in range(parcel_map.shape[1]):
            cluster_id = parcel_map[i, j]
            land_type = cluster_labels.get(cluster_id, "Unknown")

            parcel_dict[(i, j)] = {
                "land_type": land_type,
                "forage": None,
                "health": 1.0,
                "degraded": False,
                "cattle_grazing": {
                    "conservative": land_type == "Productive Grass",
                    "moderate": land_type in ["Productive Grass", "Pasture/Desert"],
                    "aggressive": land_type not in ["Water"]
                },
                "elk_grazing": {
                    "default": land_type in ["Riparian Sensitive Zone", "Productive Grass"]
                }
            }
    return parcel_dict

def get_land_forage_rates():
    return {
        'Productive Grass': 1.0,
        'Pasture/Desert': 0.4,
        'Riparian Sensitive Zone': 1.2,
        'Rocky Area': 0.2,
        'Water': 0.0
    }


def assign_initial_forage(parcel_dict, land_forage_rates):
    for parcel in parcel_dict.values():
        rate = land_forage_rates.get(parcel["land_type"], 0.0)
        parcel["forage"] = rate * 100  # initial AUMs

import numpy as np
import numpy as np
import matplotlib.pyplot as plt
from collections import Counter
from matplotlib.colors import ListedColormap
import matplotlib.patches as mpatches

def simulate_period(parcel_dict, grazing_strategy="moderate", cattle_stocking_rate=500, elk_pressure=300):
    print(f"\n🟢 Running simulation using cattle grazing strategy: **{grazing_strategy.upper()}**")

    # 1. Reset all forage to simulate winter die-off
    for parcel in parcel_dict.values():
        parcel["forage"] = 0

    # 2. Count eligible parcels for cattle grazing
    total_grazed_parcels = sum(
        1 for parcel in parcel_dict.values() if parcel["cattle_grazing"][grazing_strategy]
    )

    if total_grazed_parcels == 0:
        print("⚠️ No parcels match the selected grazing strategy.")
        return

    # 3. Calculate AUM pressure per parcel
    cattle_grazing_per_parcel = cattle_stocking_rate / total_grazed_parcels
    elk_grazing_per_parcel = elk_pressure / len(parcel_dict)  # uniform elk pressure

    # 4. Simulate forage growth + grazing effects
    for parcel in parcel_dict.values():
        if not parcel["cattle_grazing"][grazing_strategy]:
            continue

        # Forage regrowth (8-month season with weather factor)
        base_growth = {
            'Productive Grass': 1.0,
            'Pasture/Desert': 0.4,
            'Riparian Sensitive Zone': 1.2,
            'Rocky Area': 0.2,
            'Water': 0.0
        }.get(parcel["land_type"], 0.0)

        weather = np.random.normal(1.0, 0.2)
        growth = base_growth * weather * 8
        parcel["forage"] = growth  # no carryover between years

        # Total grazing pressure
        total_grazing = cattle_grazing_per_parcel + elk_grazing_per_parcel

        # Degradation or recovery
        if total_grazing > parcel["forage"]:
            parcel["degraded"] = True
            parcel["health"] -= 0.1
        else:
            parcel["health"] += 0.02

        # Subtract forage
        parcel["forage"] = max(parcel["forage"] - total_grazing, 0)


def get_forage_map(parcel_dict, n_rows, n_cols):
    return np.array([[parcel_dict[(i, j)]["forage"] for j in range(n_cols)] for i in range(n_rows)])


def plot_forage_map(forage_map, title="Parcel Forage Levels"):
    plt.figure(figsize=(8, 6))
    plt.imshow(forage_map, cmap='YlGn', origin='upper')
    plt.colorbar(label="Forage AUMs")
    plt.title(title)
    plt.axis('off')
    plt.tight_layout()
    plt.show()

def run_full_simulation(parcel_map, cluster_labels, n_rows, n_cols, strategy="moderate"):
    parcel_dict = initialize_parcels(parcel_map, cluster_labels)
    land_forage_rates = get_land_forage_rates()
    assign_initial_forage(parcel_dict, land_forage_rates)
    simulate_period(parcel_dict, grazing_strategy=strategy)
    return parcel_dict