"""
Create Pareto front visualizations for M1 Mirror optimization data.
Shows relationship between geometric parameters and 60/20 WFE performance.
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
import matplotlib.patches as mpatches

# Set style for publication-quality plots
plt.rcParams.update({
    'font.family': 'sans-serif',
    'font.sans-serif': ['Arial', 'Helvetica', 'DejaVu Sans'],
    'font.size': 11,
    'axes.titlesize': 16,
    'axes.labelsize': 13,
    'xtick.labelsize': 11,
    'ytick.labelsize': 11,
    'legend.fontsize': 10,
    'figure.dpi': 150,
    'savefig.dpi': 150,
    'axes.spines.top': False,
    'axes.spines.right': False,
})

# Load data
df = pd.read_csv(r'c:\Users\antoi\Atomizer\studies\m1_mirror_all_trials_export.csv')

print("=== Data Overview ===")
print(f"Total rows: {len(df)}")
print(f"\nColumn data availability:")
for col in df.columns:
    non_null = df[col].notna().sum()
    if non_null > 0:
        print(f"  {col}: {non_null} ({100*non_null/len(df):.1f}%)")

print(f"\nStudies: {df['study'].unique()}")

# Filter for rows with the key parameters
thickness_col = 'center_thickness'
angle_col = 'blank_backface_angle'
wfe_col = 'rel_filtered_rms_60_vs_20'

print(f"\n=== Key columns ===")
print(f"center_thickness non-null: {df[thickness_col].notna().sum()}")
print(f"blank_backface_angle non-null: {df[angle_col].notna().sum()}")
print(f"rel_filtered_rms_60_vs_20 non-null: {df[wfe_col].notna().sum()}")

# Create filtered dataset with valid WFE values
df_valid = df[df[wfe_col].notna()].copy()
print(f"\nRows with valid WFE data (before outlier removal): {len(df_valid)}")

if len(df_valid) == 0:
    print("No valid WFE data found!")
    exit()

# Remove outliers - WFE values above 1000 are clearly failed simulations
WFE_THRESHOLD = 100  # Reasonable upper bound for WFE ratio
df_valid = df_valid[df_valid[wfe_col] < WFE_THRESHOLD].copy()
print(f"Rows with valid WFE data (after outlier removal, WFE < {WFE_THRESHOLD}): {len(df_valid)}")

# Show ranges
print(f"\n=== Value ranges (clean data) ===")
if df_valid[thickness_col].notna().any():
    print(f"center_thickness: {df_valid[thickness_col].min():.2f} - {df_valid[thickness_col].max():.2f} mm")
if df_valid[angle_col].notna().any():
    print(f"blank_backface_angle: {df_valid[angle_col].min():.2f} - {df_valid[angle_col].max():.2f}°")
print(f"rel_filtered_rms_60_vs_20: {df_valid[wfe_col].min():.4f} - {df_valid[wfe_col].max():.4f}")

# Also check mass
if 'mass_kg' in df_valid.columns and df_valid['mass_kg'].notna().any():
    print(f"mass_kg: {df_valid['mass_kg'].min():.2f} - {df_valid['mass_kg'].max():.2f} kg")


def compute_pareto_front(x, y, minimize_x=True, minimize_y=True):
    """
    Compute Pareto front indices.
    Returns indices of points on the Pareto front.
    """
    # Create array of points
    points = np.column_stack([x, y])
    n_points = len(points)

    # Adjust for minimization/maximization
    if not minimize_x:
        points[:, 0] = -points[:, 0]
    if not minimize_y:
        points[:, 1] = -points[:, 1]

    # Find Pareto front
    pareto_mask = np.ones(n_points, dtype=bool)

    for i in range(n_points):
        if pareto_mask[i]:
            # Check if any other point dominates point i
            for j in range(n_points):
                if i != j and pareto_mask[j]:
                    # j dominates i if j is <= in all objectives and < in at least one
                    if (points[j, 0] <= points[i, 0] and points[j, 1] <= points[i, 1] and
                        (points[j, 0] < points[i, 0] or points[j, 1] < points[i, 1])):
                        pareto_mask[i] = False
                        break

    return np.where(pareto_mask)[0]


def create_pareto_plot(df_plot, x_col, y_col, x_label, y_label, title, filename,
                       minimize_x=True, minimize_y=True, color_by=None, color_label=None):
    """Create a publication-quality Pareto front plot."""

    # Filter valid data
    mask = df_plot[x_col].notna() & df_plot[y_col].notna()
    df_clean = df_plot[mask].copy()

    if len(df_clean) < 2:
        print(f"Not enough data for {title}")
        return

    x = df_clean[x_col].values
    y = df_clean[y_col].values

    # Compute Pareto front
    pareto_idx = compute_pareto_front(x, y, minimize_x, minimize_y)

    # Sort Pareto points by x for line drawing
    pareto_points = np.column_stack([x[pareto_idx], y[pareto_idx]])
    sort_idx = np.argsort(pareto_points[:, 0])
    pareto_sorted = pareto_points[sort_idx]

    # Create figure with professional styling
    fig, ax = plt.subplots(figsize=(12, 8))
    fig.patch.set_facecolor('white')

    # Professional color palette
    bg_color = '#f8f9fa'
    grid_color = '#dee2e6'
    point_color = '#6c757d'
    pareto_color = '#dc3545'
    pareto_fill = '#ffc107'

    ax.set_facecolor(bg_color)

    # Color scheme - use mass as color if available
    if color_by is not None and color_by in df_clean.columns and df_clean[color_by].notna().sum() > 10:
        # Only use color if we have enough colored points
        color_mask = df_clean[color_by].notna()
        colors = df_clean.loc[color_mask, color_by].values

        # Plot non-colored points in gray
        ax.scatter(x[~color_mask.values], y[~color_mask.values],
                  c=point_color, alpha=0.3, s=40,
                  edgecolors='white', linewidth=0.3, zorder=2)

        # Plot colored points
        scatter = ax.scatter(x[color_mask.values], y[color_mask.values],
                            c=colors, cmap='plasma', alpha=0.7, s=60,
                            edgecolors='white', linewidth=0.5, zorder=2)
        cbar = plt.colorbar(scatter, ax=ax, pad=0.02, shrink=0.8)
        cbar.set_label(color_label or color_by, fontsize=12, fontweight='bold')
        cbar.ax.tick_params(labelsize=10)
    else:
        ax.scatter(x, y, c=point_color, alpha=0.4, s=50,
                  edgecolors='white', linewidth=0.3, zorder=2, label='Design candidates')

    # Draw Pareto front fill area (visual emphasis)
    if len(pareto_sorted) > 1:
        # Smooth interpolation for the Pareto front line
        from scipy.interpolate import interp1d
        if len(pareto_sorted) >= 4:
            # Use cubic interpolation for smooth curve
            try:
                f = interp1d(pareto_sorted[:, 0], pareto_sorted[:, 1], kind='cubic')
                x_smooth = np.linspace(pareto_sorted[:, 0].min(), pareto_sorted[:, 0].max(), 100)
                y_smooth = f(x_smooth)
                ax.plot(x_smooth, y_smooth, color=pareto_color, linewidth=3, alpha=0.9, zorder=3)
            except:
                ax.plot(pareto_sorted[:, 0], pareto_sorted[:, 1], color=pareto_color,
                       linewidth=3, alpha=0.9, zorder=3)
        else:
            ax.plot(pareto_sorted[:, 0], pareto_sorted[:, 1], color=pareto_color,
                   linewidth=3, alpha=0.9, zorder=3)

    # Plot Pareto front points with emphasis
    ax.scatter(x[pareto_idx], y[pareto_idx], c=pareto_fill, s=180,
              edgecolors=pareto_color, linewidth=2.5, zorder=5,
              label=f'Pareto optimal ({len(pareto_idx)} designs)')

    # Styling
    ax.set_xlabel(x_label, fontsize=14, fontweight='bold', labelpad=12)
    ax.set_ylabel(y_label, fontsize=14, fontweight='bold', labelpad=12)
    ax.set_title(title, fontsize=18, fontweight='bold', pad=20, color='#212529')

    # Refined grid
    ax.grid(True, alpha=0.5, linestyle='-', linewidth=0.5, color=grid_color)
    ax.set_axisbelow(True)

    # Add minor grid
    ax.minorticks_on()
    ax.grid(True, which='minor', alpha=0.2, linestyle=':', linewidth=0.3, color=grid_color)

    # Legend with professional styling
    legend = ax.legend(loc='upper right', fontsize=11, framealpha=0.95,
                      edgecolor=grid_color, fancybox=True, shadow=True)

    # Add annotation for best point
    if minimize_y:
        best_idx = pareto_idx[np.argmin(y[pareto_idx])]
    else:
        best_idx = pareto_idx[np.argmax(y[pareto_idx])]

    # Professional annotation box - position dynamically based on data location
    # Determine best quadrant for annotation
    x_range = x.max() - x.min()
    y_range = y.max() - y.min()
    x_mid = x.min() + x_range / 2
    y_mid = y.min() + y_range / 2

    # Place annotation away from the best point
    if x[best_idx] < x_mid:
        x_offset = 50
    else:
        x_offset = -120
    if y[best_idx] < y_mid:
        y_offset = 50
    else:
        y_offset = -60

    ax.annotate(f'Best WFE: {y[best_idx]:.2f}\n{x_label.split()[0]}: {x[best_idx]:.1f}',
                xy=(x[best_idx], y[best_idx]),
                xytext=(x_offset, y_offset), textcoords='offset points',
                fontsize=11, fontweight='bold',
                bbox=dict(boxstyle='round,pad=0.6', facecolor='white',
                         edgecolor=pareto_color, linewidth=2, alpha=0.95),
                arrowprops=dict(arrowstyle='->', connectionstyle='arc3,rad=0.2',
                               color=pareto_color, lw=2))

    # Statistics box in bottom left
    stats_text = f'Total designs explored: {len(df_clean):,}\nPareto optimal: {len(pareto_idx)}'
    ax.text(0.02, 0.02, stats_text, transform=ax.transAxes, fontsize=10,
            verticalalignment='bottom',
            bbox=dict(boxstyle='round,pad=0.5', facecolor='white',
                     edgecolor=grid_color, alpha=0.9))

    # Adjust spines
    for spine in ax.spines.values():
        spine.set_color(grid_color)
        spine.set_linewidth(1.5)

    plt.tight_layout()
    plt.savefig(filename, dpi=200, bbox_inches='tight', facecolor='white',
               edgecolor='none', pad_inches=0.2)
    plt.close()
    print(f"Saved: {filename}")

    return pareto_sorted, pareto_idx


# Create plots
output_dir = r'c:\Users\antoi\Atomizer\studies'

# 1. Blank Thickness vs 60/20 WFE
print("\n--- Creating Blank Thickness vs WFE plot ---")
if df_valid[thickness_col].notna().any():
    result = create_pareto_plot(
        df_valid,
        x_col=thickness_col,
        y_col=wfe_col,
        x_label='Blank Thickness (mm)',
        y_label='60/20 WFE (Relative RMS)',
        title='M1 Mirror Optimization\nBlank Thickness vs Wavefront Error',
        filename=f'{output_dir}\\pareto_thickness_vs_wfe.png',
        minimize_x=False,  # Thinner may be desirable
        minimize_y=True,   # Lower WFE is better
        color_by='mass_kg' if 'mass_kg' in df_valid.columns else None,
        color_label='Mass (kg)'
    )
else:
    print("No thickness data available")

# 2. Blank Backface Angle vs 60/20 WFE
print("\n--- Creating Backface Angle vs WFE plot ---")
if df_valid[angle_col].notna().any():
    result = create_pareto_plot(
        df_valid,
        x_col=angle_col,
        y_col=wfe_col,
        x_label='Blank Backface Angle (degrees)',
        y_label='60/20 WFE (Relative RMS)',
        title='M1 Mirror Optimization\nBackface Angle vs Wavefront Error',
        filename=f'{output_dir}\\pareto_angle_vs_wfe.png',
        minimize_x=False,
        minimize_y=True,
        color_by='mass_kg' if 'mass_kg' in df_valid.columns else None,
        color_label='Mass (kg)'
    )
else:
    print("No backface angle data available")

# 3. Combined 2D Design Space plot
print("\n--- Creating Design Space plot ---")
if df_valid[thickness_col].notna().any() and df_valid[angle_col].notna().any():
    mask = df_valid[thickness_col].notna() & df_valid[angle_col].notna()
    df_both = df_valid[mask].copy()

    if len(df_both) > 0:
        fig, ax = plt.subplots(figsize=(12, 9))
        fig.patch.set_facecolor('white')

        bg_color = '#f8f9fa'
        grid_color = '#dee2e6'
        ax.set_facecolor(bg_color)

        # Use a perceptually uniform colormap
        scatter = ax.scatter(
            df_both[thickness_col],
            df_both[angle_col],
            c=df_both[wfe_col],
            cmap='RdYlGn_r',  # Red=bad (high WFE), Green=good (low WFE)
            s=100,
            alpha=0.8,
            edgecolors='white',
            linewidth=0.5,
            vmin=df_both[wfe_col].quantile(0.05),
            vmax=df_both[wfe_col].quantile(0.95)
        )

        cbar = plt.colorbar(scatter, ax=ax, pad=0.02, shrink=0.85)
        cbar.set_label('60/20 WFE (Relative RMS)\nLower = Better Performance',
                      fontsize=12, fontweight='bold')
        cbar.ax.tick_params(labelsize=10)

        ax.set_xlabel('Blank Thickness (mm)', fontsize=14, fontweight='bold', labelpad=12)
        ax.set_ylabel('Blank Backface Angle (degrees)', fontsize=14, fontweight='bold', labelpad=12)
        ax.set_title('M1 Mirror Design Space Exploration\nGeometric Parameters vs Optical Performance',
                    fontsize=18, fontweight='bold', pad=20)

        ax.grid(True, alpha=0.5, color=grid_color)
        ax.minorticks_on()
        ax.grid(True, which='minor', alpha=0.2, linestyle=':', color=grid_color)

        # Mark best point with star
        best_idx = df_both[wfe_col].idxmin()
        best_row = df_both.loc[best_idx]
        ax.scatter(best_row[thickness_col], best_row[angle_col],
                  c='#ffc107', s=400, marker='*', edgecolors='#dc3545', linewidth=3,
                  zorder=5, label=f'Best Design (WFE={best_row[wfe_col]:.2f})')

        # Add annotation for best point - position in upper left to avoid overlap
        ax.annotate(f'Best Design\nThickness: {best_row[thickness_col]:.1f}mm\nAngle: {best_row[angle_col]:.2f}°\nWFE: {best_row[wfe_col]:.2f}',
                   xy=(best_row[thickness_col], best_row[angle_col]),
                   xytext=(-100, 60), textcoords='offset points',
                   fontsize=10, fontweight='bold',
                   bbox=dict(boxstyle='round,pad=0.6', facecolor='white',
                            edgecolor='#dc3545', linewidth=2, alpha=0.95),
                   arrowprops=dict(arrowstyle='->', connectionstyle='arc3,rad=0.3',
                                  color='#dc3545', lw=2))

        ax.legend(loc='upper right', fontsize=11, framealpha=0.95, fancybox=True, shadow=True)

        # Stats
        stats_text = f'Designs evaluated: {len(df_both):,}'
        ax.text(0.02, 0.02, stats_text, transform=ax.transAxes, fontsize=10,
               verticalalignment='bottom',
               bbox=dict(boxstyle='round,pad=0.5', facecolor='white',
                        edgecolor=grid_color, alpha=0.9))

        for spine in ax.spines.values():
            spine.set_color(grid_color)
            spine.set_linewidth(1.5)

        plt.tight_layout()
        plt.savefig(f'{output_dir}\\design_space_wfe.png',
                   dpi=200, bbox_inches='tight', facecolor='white', pad_inches=0.2)
        plt.close()
        print(f"Saved: design_space_wfe.png")
else:
    print("Not enough data for combined design space plot")

print("\n" + "="*60)
print("PARETO VISUALIZATION COMPLETE")
print("="*60)
print(f"\nOutput files saved to: {output_dir}")
print("\nFiles created:")
print("  1. pareto_thickness_vs_wfe.png - Thickness vs WFE Pareto front")
print("  2. pareto_angle_vs_wfe.png - Backface Angle vs WFE Pareto front")
print("  3. design_space_wfe.png - Combined design space heatmap")