from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
from matplotlib.patches import Patch
from src.operations import remap_mask, create_overlay, to_binary_clothes, morphological_operations
from src.utils import denormalize_image


Path("assets").mkdir(parents=True, exist_ok=True)

CLASS_NAMES = [
    "Background",
    "Upper-Clothes",
    "Skirt",
    "Pants",
    "Dress"
]

CLASS_COLORS = [
    "#000000",  # Background
    "#FF0000",  # Upper-Clothes
    "#00FF00",  # Skirt
    "#0000FF",  # Pants
    "#FFFF00",  # Dress
]

SEGMENTATION_CMAP = ListedColormap(CLASS_COLORS)

def colorize_mask(mask):
    """
    Convert class-index mask into RGB colored mask.
    """

    mask = np.asarray(mask).astype(np.uint8)

    colored_mask = np.zeros(
        (*mask.shape, 3),
        dtype=np.uint8
    )

    colors = [
        (0, 0, 0),        # Background
        (255, 0, 0),      # Upper-Clothes
        (0, 255, 0),      # Skirt
        (0, 0, 255),      # Pants
        (255, 255, 0),    # Dress
    ]

    for class_id, color in enumerate(colors):
        colored_mask[mask == class_id] = color

    return colored_mask

def add_segmentation_legend(fig):

    legend_elements = [
        Patch(
            facecolor=CLASS_COLORS[i],
            label=CLASS_NAMES[i]
        )
        for i in range(len(CLASS_NAMES))
    ]

    fig.legend(
        handles=legend_elements,
        loc="lower center",
        ncol=5,
        frameon=False
    )
    
def visualize_batch(images, masks, num_samples = 4, save_path="assets/visualization_batch.png"):
    num_samples = min(num_samples, len(images))
    
    fig, axes = plt.subplots(2, num_samples, figsize = (4 * num_samples, 8))
    
    if num_samples == 1:
        axes = np.expand_dims(axes, axis = 1)
    
    for i in range(num_samples):
        image = images[i].permute(1, 2, 0).cpu().numpy()
        image = denormalize_image(image)

        axes[0, i].imshow(image)
        axes[0, i].set_title("Image")
        axes[0, i].axis("off")

        mask = masks[i].cpu().numpy()
        axes[1, i].imshow(mask, cmap = SEGMENTATION_CMAP, vmin = 0, vmax = 4)
        axes[1, i].set_title("Mask")
        axes[1, i].axis("off")
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches="tight")
    plt.close()


def visualize_preprocessing_sample(image, mask, label_map, alpha=0.5, save_path="assets/preprocessing.png"):
    if hasattr(image, "permute"):
        image = image.permute(1, 2, 0).cpu().numpy()

    if hasattr(mask, "numpy"):
        mask = mask.cpu().numpy()

    # If image is uint8 -> convert to [0, 1]
    if image.dtype == np.uint8:
        image_display = image / 255.0
    else:
        image_display = np.clip(image, 0, 1)

    remapped_mask = remap_mask(mask, label_map)

    overlay = create_overlay((image_display * 255).astype(np.uint8), remapped_mask, alpha=alpha)

    fig, axes = plt.subplots(1, 4, figsize=(16, 4))

    axes[0].imshow(image_display)
    axes[0].set_title("Original Image")
    axes[0].axis("off")

    axes[1].imshow(mask, cmap=SEGMENTATION_CMAP, vmin=0, vmax=4)
    axes[1].set_title("Original Mask")
    axes[1].axis("off")

    axes[2].imshow(remapped_mask, cmap=SEGMENTATION_CMAP, vmin=0, vmax=4)
    axes[2].set_title("Remapped Mask")
    axes[2].axis("off")

    axes[3].imshow(overlay)
    axes[3].set_title("Overlay")
    axes[3].axis("off")

    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches="tight")
    plt.close()


def visualize_random_samples(dataset, label_map, num_samples=4, alpha=0.5, save_path="assets/random_samples.png"):
    num_samples = min(num_samples, len(dataset))
    indices = np.random.choice(len(dataset), size=num_samples, replace=False)

    fig, axes = plt.subplots(num_samples, 4, figsize=(16, 4 * num_samples))
    if num_samples == 1:
        axes = np.expand_dims(axes, axis=0)

    for row, idx in enumerate(indices):
        sample = dataset[idx]
        image = sample["pixel_values"].convert("RGB")
        image = np.array(image, dtype=np.uint8)
        mask = np.array(sample["label"], dtype=np.uint8)
        remapped_mask = remap_mask(mask, label_map)

        overlay = create_overlay(image, remapped_mask, alpha=alpha)

        axes[row, 0].imshow(image)
        axes[row, 0].set_title(f"Image #{idx}")
        axes[row, 0].axis("off")
        
        axes[row, 1].imshow(mask, cmap="tab10")
        axes[row, 1].set_title("Original Mask")
        axes[row, 1].axis("off")
        
        axes[row, 2].imshow(remapped_mask, cmap="tab10", vmin=0, vmax=4)
        axes[row, 2].set_title("Remapped Mask")
        axes[row, 2].axis("off")

        axes[row, 3].imshow(overlay)
        axes[row, 3].set_title("Overlay")
        axes[row, 3].axis("off")

    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches="tight")
    plt.close()


def visualize_prediction(
    image,
    gt_mask,
    pred_mask,
    opening_kernel=3,
    closing_kernel=11,
    opening_iterations=1,
    closing_iterations=1,
    alpha=0.5,
    save_path="assets/prediction.png"
):
    # Convert image to numpy
    if hasattr(image, "permute"):
        image = image.permute(1, 2, 0).cpu().numpy()

    # Convert masks to numpy
    if hasattr(gt_mask, "numpy"):
        gt_mask = gt_mask.cpu().numpy()

    if hasattr(pred_mask, "numpy"):
        pred_mask = pred_mask.cpu().numpy()

    # If image is normalized, convert it back to display range
    if image.dtype != np.uint8:
        image = np.clip(image, 0, 1)
        image = (image * 255).astype(np.uint8)

    # Ground Truth
    gt_mask = gt_mask.astype(np.uint8)

    # Prediction -> Binary
    pred_binary = to_binary_clothes(pred_mask)

    # Prediction -> Morphological Operations
    final_prediction = morphological_operations(
        pred_binary,
        opening_kernel=opening_kernel,
        closing_kernel=closing_kernel,
        opening_iterations=opening_iterations,
        closing_iterations=closing_iterations
    )

    # Overlay
    overlay = create_overlay(image, final_prediction, alpha=alpha)

    # Visualization
    fig, axes = plt.subplots(1, 4, figsize=(16, 4))

    axes[0].imshow(image)
    axes[0].set_title("Original Image")
    axes[0].axis("off")

    axes[1].imshow(gt_mask, cmap="tab10", vmin=0, vmax=4)
    axes[1].set_title("Ground Truth")
    axes[1].axis("off")

    axes[2].imshow(final_prediction, cmap="gray", vmin=0, vmax=1)
    axes[2].set_title("Final Prediction")
    axes[2].axis("off")

    axes[3].imshow(overlay)
    axes[3].set_title("Overlay")
    axes[3].axis("off")

    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches="tight")
    plt.close()
    

def visualize_test_predictions(
    images,
    gt_masks,
    pred_masks,
    num_samples=4,
    opening_kernel=3,
    closing_kernel=11,
    opening_iterations=1,
    closing_iterations=1,
    alpha=0.5,
    save_path="assets/test_predictions.png"
):

    num_samples = min(num_samples, len(images))

    fig, axes = plt.subplots(num_samples, 4, figsize=(16, 4 * num_samples))

    if num_samples == 1:
        axes = np.expand_dims(axes, axis=0)

    for row in range(num_samples):

        image = images[row]
        gt_mask = gt_masks[row]
        pred_mask = pred_masks[row]

        # Convert image
        if hasattr(image, "permute"):
            image = image.permute(1, 2, 0).cpu().numpy()

        if image.dtype != np.uint8:
            image = np.clip(image, 0, 1)
            image = (image * 255).astype(np.uint8)

        # Convert masks
        if hasattr(gt_mask, "numpy"):
            gt_mask = gt_mask.cpu().numpy()

        if hasattr(pred_mask, "numpy"):
            pred_mask = pred_mask.cpu().numpy()

        gt_mask = gt_mask.astype(np.uint8)
        pred_mask = pred_mask.astype(np.uint8)

        # Binary prediction
        pred_binary = to_binary_clothes(pred_mask)

        # Morphological post-processing
        final_prediction = morphological_operations(
            pred_binary,
            opening_kernel=opening_kernel,
            closing_kernel=closing_kernel,
            opening_iterations=opening_iterations,
            closing_iterations=closing_iterations
        )

        # Overlay
        overlay = create_overlay(image, final_prediction, alpha=alpha)

        # Original image
        axes[row, 0].imshow(image)
        axes[row, 0].set_title("Original Image")
        axes[row, 0].axis("off")

        # Ground truth
        axes[row, 1].imshow(gt_mask, cmap=SEGMENTATION_CMAP, vmin=0, vmax=4)
        axes[row, 1].set_title("Ground Truth")
        axes[row, 1].axis("off")

        # Final prediction
        axes[row, 2].imshow(final_prediction, cmap="gray", vmin=0, vmax=1)
        axes[row, 2].set_title("Final Prediction")
        axes[row, 2].axis("off")

        # Overlay
        axes[row, 3].imshow(overlay)
        axes[row, 3].set_title("Overlay")
        axes[row, 3].axis("off")

    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches="tight")
    plt.close()
    

def plot_training_curves(history):
    epochs = range(1, len(history["train_loss"]) + 1)

    # Loss Curves
    plt.figure(figsize=(8, 5))
    plt.plot(epochs, history["train_loss"], label="Train Loss")
    plt.plot(epochs, history["val_loss"], label="Validation Loss")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.title("Training and Validation Loss")
    plt.legend()
    plt.grid(True)

    plt.tight_layout()
    plt.savefig("assets/training_loss.png", dpi=150, bbox_inches="tight")
    plt.close()

    # Loss Components
    plt.figure(figsize=(8, 5))
    plt.plot(epochs, history["train_ce"], label="Train CE")
    plt.plot(epochs, history["val_ce"], label="Validation CE")
    plt.plot(epochs, history["train_dice_loss"], label="Train Dice Loss")
    plt.plot(epochs, history["val_dice_loss"], label="Validation Dice Loss")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.title("Loss Components")
    plt.legend()
    plt.grid(True)

    plt.tight_layout()
    plt.savefig("assets/training_components.png", dpi=150, bbox_inches="tight")
    plt.close()

    # Segmentation Metrics
    plt.figure(figsize=(8, 5))
    plt.plot(epochs, history["val_miou"], label="mIoU")
    plt.plot(epochs, history["val_miou_clothes"], label="mIoU Clothes")
    plt.plot(epochs, history["val_dice"], label="Dice")
    plt.plot(epochs, history["val_dice_clothes"], label="Dice Clothes")
    plt.xlabel("Epoch")
    plt.ylabel("Score")
    plt.title("Validation Segmentation Metrics")
    plt.legend()
    plt.grid(True)

    plt.tight_layout()
    plt.savefig("assets/validation_metrics.png", dpi=150, bbox_inches="tight")
    plt.close()
    
def visualize_inference(image, pred_mask, final_mask, alpha=0.5, save_path="assets/inference_result.png"):
    if hasattr(image, "permute"):
        image = image.permute(1, 2, 0).cpu().numpy()

    if image.dtype != np.uint8:
        image = np.clip(image, 0, 1)
        image = (image * 255).astype(np.uint8)

    if hasattr(pred_mask, "numpy"):
        pred_mask = pred_mask.cpu().numpy()

    if hasattr(final_mask, "numpy"):
        final_mask = final_mask.cpu().numpy()

    pred_mask = pred_mask.astype(np.uint8)
    final_mask = final_mask.astype(np.uint8)

    overlay = create_overlay(image, final_mask, alpha=alpha)
    fig, axes = plt.subplots(1, 4, figsize=(16, 4))
    # Original Image
    axes[0].imshow(image)
    axes[0].set_title("Original Image")
    axes[0].axis("off")

    # Raw Prediction
    colored_pred = colorize_mask(pred_mask)
    axes[1].imshow(colored_pred)
    axes[1].set_title("Raw Prediction")
    axes[1].axis("off")

    # Final Prediction
    axes[2].imshow(final_mask, cmap="gray", vmin=0, vmax=1)
    axes[2].set_title("Final Prediction")
    axes[2].axis("off")

    # Overlay
    axes[3].imshow(overlay)
    axes[3].set_title("Overlay")
    axes[3].axis("off")

    add_segmentation_legend(fig)
    plt.tight_layout(rect = [0, 0.08, 1, 1])
    plt.savefig(save_path, dpi=150, bbox_inches="tight")
    plt.close()