cluster_to_class = {
    0: "desert",
    1: "pasture",
    2: "riaprain",
    3: "sensitive riparian",
    4: "wetland",
    5: "water"
}


import numpy as np
from collections import deque

def identify_zones(parcel_map, connectivity="queen"):
    """
    Identifies contiguous zones in a 2D parcel map using connected component labeling.

    Parameters:
        parcel_map (np.ndarray): 2D array where each value is a cluster ID (e.g., land type).
        connectivity (str): "rook" (4-way) or "queen" (8-way) connectivity.

    Returns:
        zone_map (np.ndarray): Same shape as parcel_map, each zone gets a unique integer ID.
        zone_to_cluster (dict): Maps each zone ID to its underlying cluster ID.
    """
    n_rows, n_cols = parcel_map.shape
    zone_map = -1 * np.ones_like(parcel_map, dtype=int)
    visited = np.zeros_like(parcel_map, dtype=bool)
    zone_id = 0

    if connectivity == "queen":
        directions = [(-1, -1), (-1, 0), (-1, 1),
                      (0, -1),          (0, 1),
                      (1, -1),  (1, 0), (1, 1)]
    else:  # "rook"
        directions = [(-1, 0), (1, 0), (0, -1), (0, 1)]

    for i in range(n_rows):
        for j in range(n_cols):
            if visited[i, j]:
                continue
            cluster_id = parcel_map[i, j]
            queue = deque([(i, j)])
            while queue:
                x, y = queue.popleft()
                if visited[x, y] or parcel_map[x, y] != cluster_id:
                    continue
                visited[x, y] = True
                zone_map[x, y] = zone_id
                for dx, dy in directions:
                    nx, ny = x + dx, y + dy
                    if 0 <= nx < n_rows and 0 <= ny < n_cols and not visited[nx, ny]:
                        if parcel_map[nx, ny] == cluster_id:
                            queue.append((nx, ny))
            zone_id += 1

    # Optional: Map zone_id → cluster_id
    zone_to_cluster = {}
    for zid in range(zone_id):
        indices = np.argwhere(zone_map == zid)
        if len(indices) > 0:
            i, j = indices[0]
            zone_to_cluster[zid] = parcel_map[i, j]

    return zone_map, zone_to_cluster

import matplotlib.pyplot as plt
import numpy as np

def plot_labeled_zones(zone_map, zone_labels, zone_to_cluster, save_path="labeled_zones.png"):
    """
    Plots a zone map with human-readable labels and cluster-based custom colors.

    Colors (by cluster ID):
        0: tan
        1: green
        2: rose
        3: red
        4: purple
        5: blue
    """
    # Custom color map for cluster IDs (NOT zone IDs)
    cluster_colors = {
        0: "#D2B48C",  # tan
        1: "#228B22",  # green
        2: "#FF66CC",  # rose
        3: "#FF0000",  # red
        4: "#800080",  # purple
        5: "#1E90FF",  # blue
    }

    n_rows, n_cols = zone_map.shape
    rgb_map = np.zeros((n_rows, n_cols, 3))

    # Map each parcel to its cluster color
    for i in range(n_rows):
        for j in range(n_cols):
            zone_id = zone_map[i, j]
            cluster_id = zone_to_cluster.get(zone_id, 0)
            hex_color = cluster_colors.get(cluster_id, "#AAAAAA")  # fallback = gray
            rgb = tuple(int(hex_color.lstrip('#')[k:k+2], 16)/255 for k in (0, 2, 4))
            rgb_map[i, j] = rgb

    fig, ax = plt.subplots(figsize=(8, 6))
    ax.imshow(rgb_map)

    for zone_id, label in zone_labels.items():
        positions = np.argwhere(zone_map == zone_id)
        if len(positions) == 0:
            continue
        center_i, center_j = positions.mean(axis=0)
        cluster = zone_to_cluster.get(zone_id, "?")
        label_text = f"{label} ({cluster})"
        ax.text(center_j, center_i, label_text, color="white", ha="center", va="center",
                fontsize=9, fontweight="bold",
                bbox=dict(facecolor="black", alpha=0.5, boxstyle="round,pad=0.3"))

    ax.set_title("Labeled Grazing & Riparian Zones (Custom Colors)")
    ax.axis('off')
    plt.tight_layout()
    plt.savefig(save_path)
    plt.close(fig)


import matplotlib.pyplot as plt
import numpy as np

def plot_labeled_zonesold3am(zone_map, zone_labels, zone_to_cluster, save_path="labeled_zones.png"):
    """
    Plots a zone map with fixed colors based on cluster class (e.g., pasture = green, desert = tan),
    and appends the cluster ID to each label (e.g., "A (1)").
    
    Parameters:
        zone_map (np.ndarray): 2D array of zone IDs (integers).
        zone_labels (dict): Mapping from zone_id to human-readable label (str).
        zone_to_cluster (dict): Mapping from zone_id to cluster/group ID (int).
        save_path (str): File path to save the plotted image.
    """

    # === Fixed colors by cluster ID ===
    cluster_color_map = {
        0: "#d2b48c",  # Desert – tan
        1: "#228B22",  # Pasture – green
        2: "#87CEEB",  # Water – light blue
        3: "#FF69B4",  # Riparian – pink
        4: "#8B0000",  # Sensitive riparian – dark red
        5: "#9370DB",  # Town – purple
    }

    # === Build a color image based on cluster color ===
    rows, cols = zone_map.shape
    color_image = np.zeros((rows, cols, 3))

    for i in range(rows):
        for j in range(cols):
            zone_id = zone_map[i, j]
            cluster_id = zone_to_cluster.get(zone_id, 0)
            hex_color = cluster_color_map.get(cluster_id, "#888888")  # default gray
            rgb = tuple(int(hex_color.lstrip("#")[k:k+2], 16)/255.0 for k in (0, 2, 4))
            color_image[i, j] = rgb

    # === Plot map ===
    fig, ax = plt.subplots(figsize=(8, 6))
    ax.imshow(color_image)

    for zone_id in sorted(np.unique(zone_map)):
        if zone_id not in zone_labels:
            continue
        label = zone_labels[zone_id]
        group = zone_to_cluster.get(zone_id, "?")
        label_text = f"{label} ({group})"
        positions = np.argwhere(zone_map == zone_id)
        if len(positions) == 0:
            continue
        center_i, center_j = positions.mean(axis=0)
        ax.text(center_j, center_i, label_text, color="white", ha="center", va="center",
                fontsize=9, fontweight="bold",
                bbox=dict(facecolor="black", alpha=0.5, boxstyle="round,pad=0.3"))

    ax.set_title("Labeled Grazing & Riparian Zones")
    ax.axis('off')
    plt.tight_layout()
    plt.savefig(save_path)
    plt.close(fig)



import matplotlib.pyplot as plt
import numpy as np

def plot_labeled_zonesold(zone_map, zone_labels, zone_to_cluster, save_path="labeled_zones.png"):
    """
    Plots a zone map with human-readable labels (e.g., "A", "B", "Riparian A"),
    and appends the cluster ID to each label (e.g., "A (0)").
    
    Parameters:
        zone_map (np.ndarray): 2D array of zone IDs (integers).
        zone_labels (dict): Mapping from zone_id to human-readable label (str).
        zone_to_cluster (dict): Mapping from zone_id to cluster/group ID (int).
        save_path (str): File path to save the plotted image.
    """
    unique_ids = sorted(np.unique(zone_map))
    color_map = plt.get_cmap('tab20', len(unique_ids))

    fig, ax = plt.subplots(figsize=(8, 6))
    cax = ax.imshow(zone_map, cmap=color_map, vmin=0, vmax=len(unique_ids) - 1)

    # Add label with group ID
    for zone_id in unique_ids:
        if zone_id not in zone_labels:
            continue
        label = zone_labels[zone_id]
        group = zone_to_cluster.get(zone_id, "?")
        label_text = f"{label} ({group})"  # append group ID

        positions = np.argwhere(zone_map == zone_id)
        if len(positions) == 0:
            continue
        center_i, center_j = positions.mean(axis=0)
        ax.text(center_j, center_i, label_text, color="white", ha="center", va="center",
                fontsize=9, fontweight="bold",
                bbox=dict(facecolor="black", alpha=0.5, boxstyle="round,pad=0.3"))

    ax.set_title("Labeled Grazing & Riparian Zones")
    ax.axis('off')
    plt.tight_layout()
    plt.savefig(save_path)
    plt.close(fig)


def assign_zone_labels(zone_to_cluster):
    """
    Assigns human-readable labels to each zone based on its cluster type.
    Riparian zones are labeled like 'Riparian A', 'Riparian B', etc.
    Other zones are labeled 'A', 'B', etc. by group type.

    Returns:
        zone_labels (dict): zone_id → label string
    """
    label_counts = {}  # track how many zones per type
    zone_labels = {}

    for zone_id, cluster_id in zone_to_cluster.items():
        land_class = cluster_to_class.get(cluster_id, f"cluster{cluster_id}")
        count = label_counts.get(land_class, 0)
        suffix = chr(65 + count)  # A, B, C...
        if "riparian" in land_class:
            label = f"{land_class.title()} {suffix}"
        else:
            label = f"{suffix}"
        zone_labels[zone_id] = label
        label_counts[land_class] = count + 1

    return zone_labels


def assign_zone_labels_old(zone_to_cluster, cluster_to_class):
    """
    Creates a dictionary mapping zone_id → human-readable labels (e.g., "A", "Riparian B").

    Parameters:
        zone_to_cluster (dict): zone_id → cluster ID
        cluster_to_class (dict): cluster ID → class name (e.g., "pasture", "riparian")
    
    Returns:
        dict: zone_id → human-friendly label
    """
    label_counts = {}  # Track how many zones per class
    zone_labels = {}

    for zone_id, cluster_id in zone_to_cluster.items():
        zone_class = cluster_to_class.get(cluster_id, "Unknown")
        count = label_counts.get(zone_class, 0)

        # Generate a label like "Riparian A", "Pasture B", etc.
        letter = chr(ord('A') + count)
        label = f"{zone_class.capitalize()} {letter}" if zone_class != "Unknown" else f"Zone {zone_id}"

        zone_labels[zone_id] = label
        label_counts[zone_class] = count + 1

    return zone_labels

import numpy as np
import pandas as pd

def save_zone_info_to_excel(parcel_map, zone_map, zone_labels, zone_to_cluster, cluster_to_class, save_path="zone_details.xlsx"):
    """
    Save detailed zone information to an Excel file.

    Parameters:
        parcel_map (np.ndarray): 2D array with cluster IDs.
        zone_map (np.ndarray): 2D array with zone IDs.
        zone_labels (dict): zone_id → human-friendly label (e.g., "A", "Riparian B")
        zone_to_cluster (dict): zone_id → cluster ID
        cluster_to_class (dict): cluster ID → land class string (e.g., "pasture", "riparian")
        save_path (str): File path to save Excel.
    """
    rows, cols = parcel_map.shape
    data = []

    for i in range(rows):
        for j in range(cols):
            zone_id = zone_map[i, j]
            cluster_id = parcel_map[i, j]
            label = zone_labels.get(zone_id, "Unknown")
            zone_class = cluster_to_class.get(cluster_id, "Unknown")
            data.append({
                "Row": i,
                "Col": j,
                "Zone ID": zone_id,
                "Zone Label": label,
                "Cluster ID": cluster_id,
                "Zone Class": zone_class
            })

    df = pd.DataFrame(data)
    df.to_excel(save_path, index=False)
    print(f"✅ Zone info saved to {save_path}")


def override_zone_id_and_label(zone_map, zone_labels, from_id, from_label, to_id, to_label):
    """
    Reassigns all parcels with a given zone ID and label to a new zone ID and label.
    
    Parameters:
        zone_map (np.ndarray): 2D array with zone IDs.
        zone_labels (dict): zone_id → label.
        from_id (int): Zone ID to search for.
        from_label (str): Must match the current label for that zone.
        to_id (int): Zone ID to assign.
        to_label (str): New label for the new zone ID.
    
    Returns:
        zone_map (np.ndarray): Updated zone map.
        zone_labels (dict): Updated labels.
    """
    # Step 1: Confirm the label matches
    if zone_labels.get(from_id) != from_label:
        print(f"❌ Mismatch: Zone {from_id} label is '{zone_labels.get(from_id)}', not '{from_label}'")
        return zone_map, zone_labels

    # Step 2: Loop through all parcels
    for i in range(zone_map.shape[0]):
        for j in range(zone_map.shape[1]):
            if zone_map[i, j] == from_id:
                zone_map[i, j] = to_id  # Override ID

    # Step 3: Update label dictionary
    zone_labels[to_id] = to_label
    if from_id not in zone_map:
        zone_labels.pop(from_id, None)

    print(f"✅ Overrode zone {from_id} ('{from_label}') → zone {to_id} ('{to_label}')")
    return zone_map, zone_labels



def remap_zone_id_and_label(zone_map, zone_labels, old_zone_id, old_label, new_zone_id, new_label):
    """
    For all parcels where zone_id == old_zone_id and label == old_label:
    → Change zone_id to new_zone_id and label to new_label.

    Args:
        zone_map (np.ndarray): 2D map of zone IDs.
        zone_labels (dict): zone_id → label.
        old_zone_id (int)
        old_label (str)
        new_zone_id (int)
        new_label (str)

    Returns:
        zone_map, zone_labels
    """
    # Only proceed if the label for old_zone_id matches
    if zone_labels.get(old_zone_id) != old_label:
        print(f"❌ Skipping: Label mismatch for zone {old_zone_id}")
        return zone_map, zone_labels

    # Go through every (i,j) and remap matching zones
    rows, cols = zone_map.shape
    for i in range(rows):
        for j in range(cols):
            if zone_map[i, j] == old_zone_id:
                zone_map[i, j] = new_zone_id

    # Update the label dictionary
    zone_labels[new_zone_id] = new_label
    if old_zone_id in zone_labels:
        del zone_labels[old_zone_id]

    print(f"✅ Reassigned zone {old_zone_id} ('{old_label}') → {new_zone_id} ('{new_label}')")
    return zone_map, zone_labels


import matplotlib.pyplot as plt
import numpy as np

def plot_labeled_zones_old(zone_map, zone_labels, save_path="labeled_zones.png"):
    """
    Plots a zone map with human-readable labels (e.g., "A", "B", "Riparian A").
    
    Parameters:
        zone_map (np.ndarray): 2D array of integers where each unique int is a zone ID.
        zone_labels (dict): Mapping from zone_id (int) to human label (str).
        save_path (str): File path to save the plotted image.
    """
    unique_ids = sorted(np.unique(zone_map))
    color_map = plt.get_cmap('tab20', len(unique_ids))

    fig, ax = plt.subplots(figsize=(8, 6))
    cax = ax.imshow(zone_map, cmap=color_map, vmin=0, vmax=len(unique_ids) - 1)

    # Add human-readable labels at zone centers
    for zone_id in unique_ids:
        if zone_id not in zone_labels:
            continue
        label = zone_labels[zone_id]
        positions = np.argwhere(zone_map == zone_id)
        if len(positions) == 0:
            continue
        center_i, center_j = positions.mean(axis=0)
        ax.text(center_j, center_i, label, color="white", ha="center", va="center",
                fontsize=9, fontweight="bold", bbox=dict(facecolor="black", alpha=0.5, boxstyle="round,pad=0.3"))

    ax.set_title("Labeled Grazing & Riparian Zones")
    ax.axis('off')
    plt.tight_layout()
    plt.savefig(save_path)
    plt.close(fig)