plot_cluster_validation.py

#!/usr/bin/env python3
"""
Plot cluster validation metrics (elbow and silhouette) from parallel clustering results.

This script reads the k*_results.txt files generated by the parallel voxel clustering
workflow and creates clean elbow and silhouette plots for cluster validation.

Usage:
    python plot_cluster_validation.py --results_dir /path/to/results --output_dir /path/to/plots
    python plot_cluster_validation.py --results_dir derivatives/voxel_clustering
"""

import argparse
import os
import os.path as op
import re
import glob
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt


def get_parser():
    """Command line argument parser"""
    p = argparse.ArgumentParser(description="Plot cluster validation metrics")
    p.add_argument(
        "--results_dir",
        required=True,
        help="Directory containing k*_results.txt files",
    )
    p.add_argument(
        "--output_dir",
        help="Directory to save plots",
    )
    p.add_argument(
        "--k_min",
        type=int,
        default=2,
        help="Minimum k value to plot",
    )
    p.add_argument(
        "--k_max",
        type=int,
        default=10,
        help="Maximum k value to plot",
    )
    p.add_argument(
        "--figsize",
        nargs=2,
        type=float,
        default=[12, 5],
        help="Figure size (width height)",
    )
    p.add_argument(
        "--dpi",
        type=int,
        default=300,
        help="Figure DPI for saved plots",
    )
    return p


def parse_results_file(filepath):
    """Parse a k*_results.txt file to extract metrics"""
    metrics = {}

    try:
        with open(filepath, "r") as f:
            content = f.read()

        # Extract k value from filename
        k_match = re.search(r"k(\d+)_results\.txt", op.basename(filepath))
        if k_match:
            k_value = int(k_match.group(1))
            metrics["k"] = k_value

        # Extract method
        method_match = re.search(r"Method:\s*(\w+)", content)
        if method_match:
            metrics["method"] = method_match.group(1)

        # Extract silhouette score
        sil_match = re.search(r"Silhouette score:\s*([\d\.-]+)", content)
        if sil_match:
            metrics["silhouette"] = float(sil_match.group(1))

        # Extract gap statistic
        gap_match = re.search(r"gap_statistic:\s*([\d\.-]+)", content)
        if gap_match:
            metrics["gap_statistic"] = float(gap_match.group(1))

        # Extract gap standard error
        gap_std_match = re.search(r"gap_std:\s*([\d\.-]+)", content)
        if gap_std_match:
            metrics["gap_std"] = float(gap_std_match.group(1))

        # Extract group sizes
        groups_match = re.search(r"Participants per group:\s*\[(.*?)\]", content)
        if groups_match:
            group_sizes = [int(x.strip()) for x in groups_match.group(1).split()]
            metrics["group_sizes"] = group_sizes

    except Exception as e:
        print(f"Warning: Could not parse {filepath}: {e}")

    return metrics


def collect_all_results(results_dir, k_min=2, k_max=10):
    """Collect results from all k*_results.txt files or DataFrame CSV"""
    print(f"Collecting results from: {results_dir}")

    # First try to load from validation DataFrame CSV
    csv_file = op.join(results_dir, "cluster_validation_metrics.csv")
    if op.exists(csv_file):
        print(f"Found validation DataFrame: {csv_file}")
        try:
            df = pd.read_csv(csv_file)
            # Filter by k range
            df = df[(df["k"] >= k_min) & (df["k"] <= k_max)]

            # Convert to the expected format
            all_metrics = []
            for _, row in df.iterrows():
                metrics = {
                    "k": int(row["k"]),
                    "method": "hierarchical",
                    "silhouette": row.get("silhouette_score", np.nan),
                    "gap_statistic": row.get("gap_statistic", np.nan),
                    "gap_std": row.get("gap_std", np.nan),
                }
                all_metrics.append(metrics)

                sil_score = metrics.get("silhouette", "N/A")
                gap_val = metrics.get("gap_statistic", "N/A")
                print(f"  k={metrics['k']}: silhouette={sil_score}, gap={gap_val}")

            print(f"Loaded {len(all_metrics)} results from DataFrame")
            return sorted(all_metrics, key=lambda x: x["k"])

        except Exception as e:
            print(f"Error reading DataFrame CSV: {e}")
            print("Falling back to individual k*_results.txt files...")

    # Fallback to original method with individual files
    # Find all k*_results.txt files
    pattern = op.join(results_dir, "k*_results.txt")
    result_files = glob.glob(pattern)

    if not result_files:
        raise FileNotFoundError(
            f"No k*_results.txt files or validation CSV found in {results_dir}"
        )

    print(f"Found {len(result_files)} result files")

    all_metrics = []
    for filepath in result_files:
        metrics = parse_results_file(filepath)
        if metrics and "k" in metrics:
            # Filter by k range
            if k_min <= metrics["k"] <= k_max:
                all_metrics.append(metrics)
                sil_score = metrics.get("silhouette", "N/A")
                gap_val = metrics.get("gap_statistic", "N/A")
                print(f"  k={metrics['k']}: silhouette={sil_score}, gap={gap_val}")

    # Sort by k value
    all_metrics.sort(key=lambda x: x["k"])
    return all_metrics


def plot_validation_metrics(all_metrics, output_dir, figsize=(12, 5), dpi=300):
    """Create elbow and silhouette validation plots"""
    if not all_metrics:
        raise ValueError("No valid metrics found to plot")

    print("Creating validation plots...")

    # Extract data
    k_values = [m["k"] for m in all_metrics]
    silhouette_scores = [m.get("silhouette", np.nan) for m in all_metrics]

    # Check if we have silhouette scores
    has_silhouette = not all(np.isnan(silhouette_scores))

    # Create figure
    if has_silhouette:
        fig, axes = plt.subplots(1, 2, figsize=figsize)
    else:
        fig, axes = plt.subplots(1, 1, figsize=(figsize[0] / 2, figsize[1]))
        axes = [axes]

    # Plot 1: Silhouette Analysis
    if has_silhouette:
        # Remove NaN values for plotting
        valid_indices = ~np.isnan(silhouette_scores)
        valid_k = np.array(k_values)[valid_indices]
        valid_sil = np.array(silhouette_scores)[valid_indices]

        axes[0].plot(valid_k, valid_sil, "o-", linewidth=2, markersize=6, color="blue")
        axes[0].set_xlabel("Number of Clusters (k)")
        axes[0].set_ylabel("Silhouette Score")
        axes[0].set_title("Silhouette Analysis")
        axes[0].grid(True, alpha=0.3)
        axes[0].set_xticks(valid_k)

        # Highlight best k
        if len(valid_sil) > 0:
            best_idx = np.argmax(valid_sil)
            best_k = valid_k[best_idx]
            best_score = valid_sil[best_idx]
            axes[0].axvline(best_k, color="red", linestyle="--", alpha=0.7)

            # Position annotation based on plot area
            y_range = axes[0].get_ylim()
            y_offset = (y_range[1] - y_range[0]) * 0.05  # 5% of y-range

            axes[0].annotate(
                f"Best k={best_k}\nScore={best_score:.3f}",
                xy=(best_k, best_score),
                xytext=(best_k + 0.5, best_score + y_offset),
                arrowprops=dict(arrowstyle="->", color="red"),
                bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.8),
                fontsize=10,
            )

    # Plot 2: Gap Statistic
    if len(axes) > 1:
        gap_values = [m.get("gap_statistic", np.nan) for m in all_metrics]
        gap_stds = [m.get("gap_std", np.nan) for m in all_metrics]
        has_gap = not all(np.isnan(gap_values))

        if has_gap:
            valid_indices = ~np.isnan(gap_values)
            valid_k = np.array(k_values)[valid_indices]
            valid_gap = np.array(gap_values)[valid_indices]
            valid_gap_std = np.array(gap_stds)[valid_indices]

            axes[1].errorbar(
                valid_k,
                valid_gap,
                yerr=valid_gap_std,
                fmt="o-",
                linewidth=2,
                markersize=6,
                color="green",
                capsize=5,
            )
            axes[1].set_xlabel("Number of Clusters (k)")
            axes[1].set_ylabel("Gap Statistic")
            axes[1].set_title("Gap Statistic Analysis")
            axes[1].grid(True, alpha=0.3)
            axes[1].set_xticks(valid_k)
        else:
            # Plot group sizes as bar chart instead
            axes[1].set_title("Group Sizes by k")
            axes[1].set_xlabel("Number of Clusters (k)")
            axes[1].set_ylabel("Participants per Group")

            # Create grouped bar chart
            width = 0.8 / max(k_values) if k_values else 0.1
            for i, metrics in enumerate(all_metrics):
                k = metrics["k"]
                group_sizes = metrics.get("group_sizes", [])
                if group_sizes:
                    positions = [
                        k + j * width - width * len(group_sizes) / 2
                        for j in range(len(group_sizes))
                    ]
                    axes[1].bar(
                        positions,
                        group_sizes,
                        width=width,
                        alpha=0.7,
                        label=f"k={k}" if len(all_metrics) <= 5 else None,
                    )

            if len(all_metrics) <= 5:
                axes[1].legend()
            axes[1].grid(True, alpha=0.3)

    plt.tight_layout()

    # Save plot
    os.makedirs(output_dir, exist_ok=True)
    output_file = op.join(output_dir, "cluster_validation_summary.png")
    plt.savefig(output_file, dpi=dpi, bbox_inches="tight")
    print(f"Validation plot saved to: {output_file}")

    plt.show()

    return output_file


def main():
    """Main function"""
    args = get_parser().parse_args()

    # Set default output directory
    if args.output_dir is None:
        default_dir = op.join("derivatives", "hierarchical_clustering", "figures")
        args.output_dir = default_dir

    print("Cluster Validation Plotting Tool")
    print("=" * 40)
    print(f"Results directory: {args.results_dir}")
    print(f"Output directory: {args.output_dir}")
    print(f"K range: {args.k_min} to {args.k_max}")

    try:
        # Collect all results
        all_metrics = collect_all_results(args.results_dir, args.k_min, args.k_max)

        if not all_metrics:
            print("No valid results found!")
            return 1

        # Create plots
        plot_validation_metrics(
            all_metrics, args.output_dir, figsize=args.figsize, dpi=args.dpi
        )

        print("\nSUCCESS!")
        print(f"Processed {len(all_metrics)} cluster solutions")
        print(f"Plots saved to: {args.output_dir}")

        return 0

    except Exception as e:
        print(f"ERROR: {e}")
        return 1


if __name__ == "__main__":
    exit(main())