import math
import numpy as np
import matplotlib.pyplot as plt
from skimage.io import imread
from skimage.morphology import binary_opening, binary_closing, remove_small_objects, dilation, disk
from skimage.segmentation import watershed, relabel_sequential, find_boundaries
from skimage.measure import label
from scipy.ndimage import distance_transform_edt as edt

# -------------------------
# 0) I/O helpers
# -------------------------
def load_mask(path: str) -> np.ndarray:
    m = imread(path)
    if m.ndim == 3:
        m = m[..., 0]
    return m.astype(np.int32)

# -------------------------
# 1) Semantic preprocessing (light morphology)
# -------------------------
def preprocess_semantic(mask: np.ndarray) -> dict:
    axon   = (mask == 6)
    myelin = (mask == 7)
    pas    = (mask == 5)
    abn    = (mask == 1)

    # light cleanup (same spirit as your pipeline)
    axon = binary_opening(axon, disk(1))
    axon = remove_small_objects(axon, 80)

    myelin = binary_opening(myelin, disk(1))
    myelin = binary_closing(myelin, disk(1))
    myelin = remove_small_objects(myelin, 60)

    pas = remove_small_objects(pas, 10)
    abn = remove_small_objects(abn, 10)

    return dict(axon=axon, myelin=myelin, pas=pas, abn=abn)

# -------------------------
# 2) Axon instances (myelin-guided watershed)
# -------------------------
def myelin_guided_axon_instances(axon: np.ndarray, myelin: np.ndarray,
                                barrier_dilate_radius: int = 2) -> np.ndarray:
    barrier = dilation(myelin, disk(max(1, barrier_dilate_radius)))
    ax_core = axon & (~barrier)

    dist = edt(ax_core)

    cc = label(ax_core.astype(np.uint8), connectivity=1)
    markers = np.zeros_like(ax_core, dtype=np.int32)
    mid = 0
    for cid in range(1, cc.max() + 1):
        comp = (cc == cid)
        if int(comp.sum()) < 80:
            continue
        yx = np.argmax(dist * comp)
        y, x = np.unravel_index(yx, dist.shape)
        mid += 1
        markers[y, x] = mid

    inst = watershed(-dist, markers=markers, mask=ax_core)
    inst, _, _ = relabel_sequential(inst)
    return inst.astype(np.int32)

# -------------------------
# 3) Merge axon pieces by PAS/abn adjacency (core label field)
#    This reproduces your "merge_axons_by_pas_abn" idea.
# -------------------------
def merge_axons_by_pas_abn(axon_labels: np.ndarray,
                           pas: np.ndarray,
                           abn: np.ndarray,
                           myelin: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    ax = (axon_labels > 0)
    domain = (ax | pas | abn) & (~myelin)

    comp = label(domain.astype(np.uint8), connectivity=1)

    axon_labels_merged = np.zeros_like(axon_labels, dtype=np.int32)
    core_labels = np.zeros_like(axon_labels, dtype=np.int32)

    new_id = 0
    for cid in range(1, comp.max() + 1):
        core_mask = (comp == cid)
        ax_mask = core_mask & ax
        if not ax_mask.any():
            continue
        new_id += 1
        axon_labels_merged[ax_mask] = new_id
        core_labels[core_mask] = new_id

    return axon_labels_merged, core_labels

# -------------------------
# 4) Myelin partition (seeded geodesic assignment via watershed)
#    Seeds come from a dilated rim around each axon core.
# -------------------------
def partition_myelin(axon_labels: np.ndarray,
                     myelin: np.ndarray,
                     core_labels: np.ndarray,
                     barrier_dilate_radius: int = 2) -> np.ndarray:
    tgt = myelin.astype(bool)
    if not tgt.any():
        return np.zeros_like(axon_labels, dtype=np.int32)

    markers = np.zeros_like(axon_labels, dtype=np.int32)
    gap = max(1, barrier_dilate_radius + 1)

    for a in np.unique(axon_labels[axon_labels > 0]):
        base = (core_labels == a)
        rim = dilation(base, disk(gap))
        seeds = rim & tgt
        if seeds.any():
            markers[seeds] = int(a)

    if not (markers > 0).any():
        return np.zeros_like(axon_labels, dtype=np.int32)

    # zero "height field" just means: grow by shortest distance (geodesic)
    zero_field = np.zeros_like(axon_labels, dtype=np.uint8)
    part = watershed(zero_field, markers=markers, mask=tgt)
    return part.astype(np.int32)

# -------------------------
# 5) PAS / abn assignment (directly via core labels)
# -------------------------
def assign_by_core(bin_mask: np.ndarray, core_labels: np.ndarray) -> np.ndarray:
    return np.where(bin_mask, core_labels, 0).astype(np.int32)

# -------------------------
# 6) Minimal visualization
# -------------------------
def color_instances(label_im: np.ndarray, seed: int = 0) -> np.ndarray:
    rng = np.random.default_rng(seed)
    labs = [l for l in np.unique(label_im) if l != 0]
    out = np.zeros((*label_im.shape, 3), np.uint8)
    for lab in labs:
        out[label_im == lab] = rng.integers(30, 255, size=3, dtype=np.uint8)
    return out

def visualize(mask_path: str):
    mask = load_mask(mask_path)
    sem = preprocess_semantic(mask)

    ax0 = myelin_guided_axon_instances(sem["axon"], sem["myelin"], barrier_dilate_radius=2)
    ax, core = merge_axons_by_pas_abn(ax0, sem["pas"], sem["abn"], sem["myelin"])

    my_part  = partition_myelin(ax, sem["myelin"], core, barrier_dilate_radius=2)
    pas_part = assign_by_core(sem["pas"], core)
    abn_part = assign_by_core(sem["abn"], core)

    fig, axs = plt.subplots(2, 2, figsize=(12, 10))
    axs = axs.ravel()

    axs[0].imshow(color_instances(ax, seed=1))
    axs[0].contour(find_boundaries(ax > 0, mode="outer"), colors="w", linewidths=0.4)
    axs[0].set_title("Axon instances")
    axs[0].axis("off")

    axs[1].imshow(color_instances(my_part, seed=2))
    axs[1].set_title("Myelin partition (assigned per axon)")
    axs[1].axis("off")

    axs[2].imshow(color_instances(pas_part, seed=3))
    axs[2].set_title("PAS assignment (by core)")
    axs[2].axis("off")

    axs[3].imshow(color_instances(abn_part, seed=4))
    axs[3].set_title("Abnormality assignment (by core)")
    axs[3].axis("off")

    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    MASK_PATH = "path/to/your/merged_mask.png"
    visualize(MASK_PATH)