plotter.py

#!/usr/bin/env python3
"""
Plot results from squeezing.jl simulations.

Usage:
    python squeezing_plotter.py <filename1.jld2> [filename2.jld2 ...] [output.png]
    python squeezing_plotter.py [-vsL] [-instexp] [filename1.jld2 ...] [output.png]
    python squeezing_plotter.py  # Automatically finds files in data/ directory
    
    arguments:
    -vsL: Plot sigma_inf vs L instead of sigma(t) vs t
    -instexp: Plot instantaneous power-law exponent beta(t) 
"""

import h5py
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
import sys
import os
import glob
from pathlib import Path
try: # stupid matplotlib fix; may not be needed on your system 
    import matplotlib.cbook
    if not hasattr(matplotlib.cbook, "_Stack"):
        class _Stack(list):
            def push(self, item):
                self.append(item)
                return item
            def pop(self):
                return super().pop() if self else None
            def current(self):
                return self[-1] if self else None
        matplotlib.cbook._Stack = _Stack
except:
    pass

# Set matplotlib to use LaTeX for text rendering 
matplotlib.rcParams['text.usetex'] = True
matplotlib.rcParams['font.family'] = 'serif'
matplotlib.rcParams['font.serif'] = ['Times New Roman', 'Times', 'DejaVu Serif']
matplotlib.rcParams['font.size'] = 22
matplotlib.rcParams['axes.labelsize'] = 16
matplotlib.rcParams['axes.titlesize'] = 16
matplotlib.rcParams['xtick.labelsize'] = 12
matplotlib.rcParams['ytick.labelsize'] = 12
matplotlib.rcParams['legend.fontsize'] = 11
matplotlib.rcParams['figure.dpi'] = 75
matplotlib.rcParams['savefig.dpi'] = 150
matplotlib.rcParams['savefig.bbox'] = 'tight'

# Marker size options
matplotlib.rcParams['lines.markersize'] = 5  # Default marker size
matplotlib.rcParams['lines.markeredgewidth'] = 0.5  # Default marker edge width

# Plotting style options (global boolean variables)
plot_as_line = not False  # If True, plot data as lines instead of markers
plot_instantaneous_exponent = not False  # If True, plot instantaneous power-law exponent
linewidth = 2.75
b_exponent = 5  # Base for computing instantaneous exponent: beta(t) = log_b(sigma(t/b)/sigma(t))
cmap = plt.cm.coolwarm  # Colormap for plotting colors by order of input files

def read_jld2_file(filename):
    """
    Read a JLD2 file (which is an HDF5 file) and extract the data.
    
    Returns a dictionary with the data.
    """
    data = {}
    with h5py.File(filename, 'r') as f:
        # JLD2 files store data in a specific structure
        # The actual data is usually in a group or as attributes
        def extract_data(name, obj):
            if isinstance(obj, h5py.Dataset):
                # Get the dataset
                try:
                    # Try to read as array
                    arr = obj[:]
                    # Convert to appropriate type
                    if arr.dtype == 'object':
                        # Handle string arrays
                        if len(arr.shape) == 0:
                            data[name] = arr[()]
                        else:
                            data[name] = arr
                    else:
                        data[name] = arr
                except:
                    data[name] = obj[()]
            elif isinstance(obj, h5py.Group):
                # Recursively extract from groups
                for key in obj.keys():
                    extract_data(f"{name}/{key}" if name else key, obj[key])
        
        # Start extraction from root
        for key in f.keys():
            extract_data(key, f[key])
    
    return data

def load_jld2_simple(filename):
    """
    Simplified loader for JLD2 files.
    JLD2 files are HDF5 files with a specific structure.
    """
    data = {}
    with h5py.File(filename, 'r') as f:
        def extract_value(obj, path=""):
            """Recursively extract values from HDF5 structure."""
            if isinstance(obj, h5py.Dataset):
                try:
                    # Check shape first to avoid slicing errors on scalars
                    if obj.shape == ():
                        # Scalar value - use [()] to get the value directly
                        val = obj[()]
                        if isinstance(val, bytes):
                            return val.decode('utf-8', errors='ignore').strip('\x00')
                        elif isinstance(val, str):
                            return val
                        elif isinstance(val, (bool, np.bool_)):
                            return bool(val)
                        elif isinstance(val, (int, np.integer)):
                            return int(val)
                        elif isinstance(val, (float, np.floating)):
                            return float(val)
                        else:
                            return val
                    else:
                        # Array - safe to use [:]
                        arr = obj[:]
                        if arr.dtype.kind == 'S':  # String
                            # String array - try to decode
                            try:
                                if arr.size == 1:
                                    return arr.tobytes().decode('utf-8', errors='ignore').strip('\x00')
                                else:
                                    return np.array([a.tobytes().decode('utf-8', errors='ignore').strip('\x00') for a in arr])
                            except:
                                return np.array(arr)
                        elif arr.dtype.kind == 'U':  # Unicode
                            return np.array([str(a) for a in arr])
                        else:
                            return np.array(arr)
                except Exception as e:
                    # Don't print warning for every failed extraction to reduce noise
                    return None
            elif isinstance(obj, h5py.Group):
                # Group - recurse
                result = {}
                for key in obj.keys():
                    val = extract_value(obj[key], f"{path}/{key}" if path else key)
                    if val is not None:
                        result[key] = val
                return result if result else None
            return None
        
        # JLD2 stores data in _jld2/keys and _jld2/data
        if '_jld2' in f:
            jld2_group = f['_jld2']
            if 'keys' in jld2_group and 'data' in jld2_group:
                keys = jld2_group['keys']
                data_group = jld2_group['data']
                # Extract keys (they're stored as strings)
                key_names = []
                for i in range(len(keys)):
                    try:
                        key_str = keys[i].tobytes().decode('utf-8', errors='ignore').strip('\x00')
                        key_names.append(key_str)
                    except:
                        pass
                
                # Extract corresponding data
                for i, key_name in enumerate(key_names):
                    if str(i) in data_group:
                        val = extract_value(data_group[str(i)])
                        if val is not None:
                            data[key_name] = val
        
        # Also try direct access (JLD2 might store directly)
        for key in f.keys():
            if key != '_jld2':
                val = extract_value(f[key], key)
                if val is not None:
                    data[key] = val
    
    return data

def identify_rule(R_or, R_and):
    """
    Identify which pre-defined rule (R, F, or M) matches the given R_or and R_and.
    Returns the rule name as a string, or "unknown" if no match is found.
    """
    # Convert to lists of tuples for comparison
    def normalize_set(R):
        if R is None:
            return None
        # Convert to list of tuples, handling different array formats
        if isinstance(R, np.ndarray):
            if R.ndim == 2 and R.shape[1] == 2:
                # Array of shape (n, 2) - each row is a tuple
                result = sorted([tuple(int(R[i, j]) for j in range(2)) for i in range(len(R))])
                return result
            elif R.ndim == 1:
                # Might be a flattened array or object array
                # Try to reshape or handle as object array
                try:
                    # If it's an object array, each element might be a tuple
                    if R.dtype == object:
                        result = sorted([tuple(x) for x in R if x is not None])
                        return result
                    # Otherwise, might need special handling
                except:
                    pass
        elif isinstance(R, (list, tuple)):
            # Handle list of lists or list of tuples
            result = []
            for x in R:
                if isinstance(x, (list, tuple, np.ndarray)):
                    result.append(tuple(int(y) for y in x))
                else:
                    result.append(tuple(x))
            return sorted(result)
        return None
    
    R_or_norm = normalize_set(R_or)
    R_and_norm = normalize_set(R_and)
    
    if R_or_norm is None or R_and_norm is None:
        return "unknown"
    
    # Known rules
    R_or_R = sorted([(0, 1), (0, -1)])
    R_and_R = sorted([(1, 0), (-1, 0)])
    
    R_or_F = sorted([(1, 1), (0, -1)])
    R_and_F = sorted([(1, 0), (-1, -1)])
    
    R_or_M = sorted([(1, 0), (0, -1)])
    R_and_M = sorted([(1, 0), (0, 1)])
    
    if R_or_norm == R_or_R and R_and_norm == R_and_R:
        return "R"
    elif R_or_norm == R_or_F and R_and_norm == R_and_F:
        return "F"
    elif R_or_norm == R_or_M and R_and_norm == R_and_M:
        return "M"
    else:
        return "unknown"

def get_orientation(r, rotate):
    """
    Get orientation symbol based on (r, rotate) tuple.
    (T, T): ↓, (T, F): ←, (F, T): ↑, (F, F): →
    """
    # Convert to boolean if needed
    r_bool = bool(r) if not isinstance(r, bool) else r
    rotate_bool = bool(rotate) if not isinstance(rotate, bool) else rotate
    
    if r_bool and rotate_bool:
        return r'$\downarrow$'
    elif r_bool and not rotate_bool:
        return r'$\leftarrow$'
    elif not r_bool and rotate_bool:
        return r'$\uparrow$'
    else:  # not r_bool and not rotate_bool
        return r'$\rightarrow$'

def plot_dw_growth(filenames, output_filename=None, allhists=False):
    """
    Plot domain wall growth sigma(t) vs time for dw_growth mode.
    Can handle single file or list of files.
    
    Arguments:
    - filenames: List of JLD2 filenames or single filename
    - output_filename: Optional output filename for saving the plot
    - allhists: If True, plot all individual histories. If False, plot mean ± 1 std dev.
    """
    # Convert single filename to list
    if isinstance(filenames, str):
        filenames = [filenames]
    
    # Load data from all files
    all_data = []
    rule_name = None
    
    for filename in filenames:
        # Try multiple methods to read JLD2 file
        data = None
        
        # Method 1: Try jld2 Python package (if available)
        try:
            import jld2
            with jld2.open(filename, 'r') as f:
                data = dict(f)
            print(f"Successfully read {filename} using jld2 package")
        except ImportError:
            pass  # jld2 not available, try other methods
        except Exception as e:
            print(f"jld2 package failed: {e}")
        
        # Method 2: Try h5py with our custom loader
        if data is None:
            try:
                data = load_jld2_simple(filename)
                if data:
                    print(f"Successfully read {filename} using h5py")
            except Exception as e:
                print(f"h5py method failed: {e}")
        
        # Method 3: Try direct h5py access (simpler structure)
        if data is None or len(data) == 0:
            try:
                data = {}
                with h5py.File(filename, 'r') as f:
                    # Try to access variables directly
                    for key in ['mode', 'Lx', 'Ly', 'sigma_avg', 'r', 'rotate', 'R_or', 'R_and', 'histories']:
                        if key in f:
                            try:
                                obj = f[key]
                                if isinstance(obj, h5py.Dataset):
                                    # Check if it's a scalar
                                    if obj.shape == ():
                                        # Scalar value
                                        data[key] = obj[()]
                                    else:
                                        # Array
                                        data[key] = np.array(obj[:])
                                elif isinstance(obj, h5py.Group):
                                    # It's a group, try to extract recursively
                                    pass
                            except Exception as e:
                                # Silently skip if we can't read it
                                pass
                            
                    # Try to get kwargs
                    if 'kwargs' in f:
                        kwargs_group = f['kwargs']
                        if isinstance(kwargs_group, h5py.Group):
                            kwargs_dict = {}
                            for k in kwargs_group.keys():
                                try:
                                    v_obj = kwargs_group[k]
                                    if isinstance(v_obj, h5py.Dataset):
                                        if v_obj.shape == ():
                                            kwargs_dict[k] = v_obj[()]
                                        else:
                                            kwargs_dict[k] = np.array(v_obj[:])
                                except:
                                    pass
                            if kwargs_dict:
                                data['kwargs'] = kwargs_dict
            except Exception as e:
                print(f"Direct h5py access failed: {e}")
        
        if data is None or len(data) == 0:
            print(f"Could not read {filename}")
            print("Try installing jld2 package: pip install jld2")
            print("Or ensure the file is a valid JLD2/HDF5 file")
            continue
        
        # Determine mode from filename
        if '_dw_growth' in filename:
            mode = 'dw_growth'
        elif '_dw_roughness' in filename:
            mode = 'dw_roughness'
        else:
            mode = 'unknown'
        
        if mode != 'dw_growth':
            print(f"File {filename} is not in dw_growth mode (filename doesn't contain '_dw_growth'), skipping")
            continue
        
        # Get n_trials and print it
        # Try to extract from sigma_trials array shape/length
        n_trials = None
        
        # Method 1: Extract from sigma_trials shape (should be (time_steps, n_trials))
        sigma_trials = data.get('sigma_trials', None)
        if sigma_trials is not None:
            try:
                if hasattr(sigma_trials, 'shape') and len(sigma_trials.shape) >= 2:
                    n_trials = sigma_trials.shape[1]  # Second dimension is number of trials
                elif hasattr(sigma_trials, 'shape') and len(sigma_trials.shape) == 1:
                    # If it's a 1D array, might be structured differently
                    n_trials = len(sigma_trials)
                elif hasattr(sigma_trials, '__len__'):
                    # If it's a list, use length
                    n_trials = len(sigma_trials)
            except:
                pass
        
        # Method 2: Try using jld2 package if available (handles references properly)
        if n_trials is None:
            try:
                import jld2
                with jld2.open(filename, 'r') as f:
                    jld2_data = dict(f)
                    jld2_kwargs = jld2_data.get('kwargs', {})
                    if isinstance(jld2_kwargs, dict):
                        n_trials = jld2_kwargs.get('n_trials', None)
            except ImportError:
                pass
            except:
                pass
        
        # Method 3: Try from already loaded data kwargs
        if n_trials is None:
            kwargs = data.get('kwargs', {})
            if isinstance(kwargs, dict):
                n_trials = kwargs.get('n_trials', None)
        
        # Convert to int
        if n_trials is not None:
            if isinstance(n_trials, np.ndarray):
                if n_trials.size == 1:
                    n_trials = int(n_trials[()])
                else:
                    n_trials = None
            elif isinstance(n_trials, (int, np.integer)):
                n_trials = int(n_trials)
        
        # Print for each file
        if n_trials is not None:
            print(f"File {filename}: n_trials = {n_trials}")
        else:
            print(f"File {filename}: n_trials not found")
        
        sigma_avg = data.get('sigma_avg', None)
        if sigma_avg is None:
            print(f"No sigma_avg data found in {filename}, skipping")
            continue
        
        # Convert to numpy array if needed
        if not isinstance(sigma_avg, np.ndarray):
            sigma_avg = np.array(sigma_avg)
        
        # Flatten if needed
        if sigma_avg.ndim > 1:
            sigma_avg = sigma_avg.flatten()
        
        # Extract r and rotate
        r = data.get('r', False)
        rotate = data.get('rotate', False)
        # Convert to boolean if needed
        if isinstance(r, np.ndarray):
            r = bool(r[()]) if r.size == 1 else False
        if isinstance(rotate, np.ndarray):
            rotate = bool(rotate[()]) if rotate.size == 1 else False
        
        # Get Lx and Ly for threshold check
        Lx = data.get('Lx', None)
        Ly = data.get('Ly', None)
        if isinstance(Lx, np.ndarray):
            Lx = int(Lx[()]) if Lx.size == 1 else None
        if isinstance(Ly, np.ndarray):
            Ly = int(Ly[()]) if Ly.size == 1 else None
        
        # Check if sigma exceeds threshold
        if Lx is not None and Ly is not None:
            threshold = Ly / 2.0 if not rotate else Lx / 2.0
            max_sigma = np.max(sigma_avg[1:])  # Skip first point at t=0
            if max_sigma > threshold:
                print(f"WARNING: sigma(t) exceeds threshold in {filename}")
                print(f"  Maximum sigma = {max_sigma:.3f}, threshold = {threshold:.3f}")
                print(f"  (rotate = {rotate}, threshold is {'Lx' if rotate else 'Ly'}/2)")
        
        # Get rule - first try reading directly, then fall back to identifying from R_or/R_and
        file_rule = data.get('rule', None)
        if file_rule is None:
            # Fall back to identifying from R_or and R_and
            R_or = data.get('R_or', None)
            R_and = data.get('R_and', None)
            file_rule = identify_rule(R_or, R_and)
        else:
            # Convert to string if it's an array
            if isinstance(file_rule, np.ndarray):
                if file_rule.size == 1:
                    file_rule = str(file_rule[()])
                else:
                    # Try to decode as string if it's a byte array
                    try:
                        file_rule = file_rule.tobytes().decode('utf-8').strip('\x00')
                    except:
                        file_rule = str(file_rule[0]) if len(file_rule) > 0 else "unknown"
            else:
                file_rule = str(file_rule)
            
            # Clean up: extract just the rule letter (R, F, or M) if it's embedded in other text
            import re
            # First, try to extract just R, F, or M from the string
            match = re.search(r'([RFM])', str(file_rule).upper())
            if match:
                file_rule = match.group(1)
            # If it's already just R, F, or M, keep it
            elif str(file_rule).upper().strip() in ['R', 'F', 'M']:
                file_rule = str(file_rule).upper().strip()
            else:
                file_rule = "unknown"
        
        if rule_name is None:
            rule_name = file_rule
        elif rule_name != file_rule:
            print(f"Warning: Files have different rules ({rule_name} vs {file_rule})")
        
        # Also get R_or and R_and for storage
        R_or = data.get('R_or', None)
        R_and = data.get('R_and', None)
        
        # Get p from kwargs (defaults to 0.0)
        kwargs = data.get('kwargs', {})
        if isinstance(kwargs, dict):
            p = kwargs.get('p', 0.0)
        else:
            p = 0.0
        # Convert to float if needed
        if isinstance(p, np.ndarray):
            p = float(p[()]) if p.size == 1 else 0.0
        else:
            p = float(p) if p is not None else 0.0
        
        # Fallback: try to extract p from filename if not found in kwargs
        # Filename format: ..._dw_growthp0.02_... or ..._dw_growth_...
        if abs(p) < 1e-10:  # Check if p is effectively zero
            import re
            # Look for pattern: dw_roughnessp0.02 or dw_growthp0.02
            match = re.search(r'_dw_(?:roughness|growth)p([\d.]+)', filename)
            if match:
                try:
                    p = float(match.group(1))
                except:
                    p = 0.0
        
        # Get orientation
        orientation = get_orientation(r, rotate)
        
        # Get histories if available
        histories = data.get('histories', None)
        if histories is not None:
            # Convert to numpy array if needed
            if not isinstance(histories, np.ndarray):
                histories = np.array(histories)
            # Ensure it's 2D: (n_trials, time_steps)
            if histories.ndim == 1:
                # If 1D, it might be a single trial - reshape
                histories = histories.reshape(1, -1)
        else:
            print(f"No histories found in {filename}")
        
        # Store data for plotting
        all_data.append({
            'sigma_avg': sigma_avg,
            'r': r,
            'rotate': rotate,
            'orientation': orientation,
            'filename': filename,
            'rule': file_rule,
            'R_or': R_or,
            'R_and': R_and,
            'p': p,
            'histories': histories  # Can be None if not available
        })
    
    if len(all_data) == 0:
        print("No valid data to plot")
        return
    
    # Create figure with PRL-style aesthetics - make window smaller
    # Use global variable for instantaneous exponent plotting
    global plot_instantaneous_exponent
    if plot_instantaneous_exponent:
        # Create single plot for instantaneous exponent only
        fig, ax_exp = plt.subplots(figsize=(4., 4.), dpi=100)
        ax = None  # Don't plot sigma(t)
    else:
        fig, ax = plt.subplots(figsize=(4., 4.), dpi=100)  # Smaller window
        ax_exp = None
    
    # Color mapping for orientations (same as dw_roughness)
    orientation_colors = {
        r'$\rightarrow$': 'blue',
        r'$\uparrow$': 'red',
        r'$\leftarrow$': 'green',
        r'$\downarrow$': 'purple'
    }
    
    # Marker shapes for different p values (if any p != 0)
    marker_shapes = ['o', 's', '^', 'P']  # circle, square, triangle, plus
    p_values = sorted(set([item['p'] for item in all_data]))
    has_nonzero_p = any(p != 0.0 for p in p_values)
    
    # Create mapping from p to marker shape
    p_to_marker = {}
    if has_nonzero_p:
        # Only assign shapes to nonzero p values, zero p gets default 'o'
        nonzero_p = [p for p in p_values if p != 0.0]
        for i, p_val in enumerate(nonzero_p):
            if i < len(marker_shapes):
                p_to_marker[p_val] = marker_shapes[i]
        # Zero p gets circle
        if 0.0 in p_values:
            p_to_marker[0.0] = 'o'
    else:
        # All p are zero, use default circle
        p_to_marker = {0.0: 'o'}
    
    # Plot each curve with power law fit
    for idx, data_item in enumerate(all_data):
        sigma_avg = data_item['sigma_avg']
        orientation = data_item['orientation']
        p = data_item.get('p', 0.0)
        file_rule = data_item.get('rule', rule_name)  # Use stored rule or fallback to global
        
        # Get color for this orientation
        color = orientation_colors.get(orientation, 'gray')
        
        # Get marker shape for this p value
        marker = p_to_marker.get(p, 'o')
        
        # Create time array (skip first point at t=0)
        T_evolve = len(sigma_avg) - 1
        time = np.arange(1, T_evolve + 1)  # Skip t=0
        sigma = sigma_avg[1:]  # Skip first data point
        
        # Remove any zero or negative values for log scale
        valid = (sigma > 0) & (time > 0)
        time_plot = time[valid]
        sigma_plot = sigma[valid]
        
        if len(time_plot) < 2:
            continue
        
        # For rule M with uparrow orientation, only fit to first 100 data points
        if file_rule == "M" and orientation == r'$\uparrow$':
            n_fit_points = min(100, len(time_plot))
            time_fit = time_plot[:n_fit_points]
            sigma_fit = sigma_plot[:n_fit_points]
        else:
            time_fit = time_plot
            sigma_fit = sigma_plot
        
        if len(time_fit) < 2:
            continue
        
        # Fit power law: sigma ~ a * t^beta
        # On log-log scale: log(sigma) = log(a) + beta * log(t)
        log_t = np.log(time_fit)
        log_sigma = np.log(sigma_fit)
        
        # Remove any invalid values
        valid_fit = np.isfinite(log_t) & np.isfinite(log_sigma)
        if np.sum(valid_fit) < 2:
            continue
        
        log_t_fit = log_t[valid_fit]
        log_sigma_fit = log_sigma[valid_fit]
        
        # Linear regression: log(sigma) = beta * log(t) + log(a)
        # Using least squares: [log(t) ones] * [beta; log(a)] = log(sigma)
        X = np.vstack([log_t_fit, np.ones(len(log_t_fit))]).T
        coeffs = np.linalg.lstsq(X, log_sigma_fit, rcond=None)[0]
        beta = coeffs[0]  # Power law exponent
        log_a = coeffs[1]  # log of prefactor
        a = np.exp(log_a)  # Prefactor
        
        # Generate fit curve
        t_fit = np.logspace(np.log10(time_plot.min()), np.log10(time_plot.max()), 100)
        sigma_fit = a * (t_fit ** beta)
        
        # Format beta to 2 decimal places
        beta_str = f"{beta:.2f}"
        
        # Plot data points (only if not plotting instantaneous exponent)
        if not plot_instantaneous_exponent and ax is not None:
            # Extract orientation symbol from LaTeX string (e.g., r'$\downarrow$' -> '\downarrow')
            # Remove $ signs to get just the symbol
            orientation_symbol = orientation.replace('$', '').strip()
            
            # Format label - include p if any file has nonzero p
            if has_nonzero_p:
                p_str = f"{p:.2f}"
                label = rf'$\beta_{{{orientation_symbol}}} = {beta_str}$, $p = {p_str}$'
            else:
                label = rf'$\beta_{{{orientation_symbol}}} = {beta_str}$'
            
            # Use global variable for line plotting
            global plot_as_line
            
            if plot_as_line:
                # Plot as line instead of markers
                ax.loglog(time_plot, sigma_plot, '-', color=color, 
                         linewidth=linewidth, alpha=0.7, label=label)
            else:
                # Use marker shape based on p value
                ax.loglog(time_plot, sigma_plot, marker,
                          markerfacecolor=color, markeredgecolor='k', 
                          alpha=0.7, label=label, markersize=5)
            
            # Plot power law fit as dashed line with same color
            ax.loglog(t_fit, sigma_fit, '--', color=color, linewidth=linewidth, alpha=0.4)
            
            # Plot histories if available
            histories = data_item.get('histories', None)
            if histories is not None and isinstance(histories, np.ndarray):
                # Transpose if needed: histories should be (n_trials, time_steps) array
                if histories.ndim == 2:
                    # Check if we need to transpose (time_steps, n_trials) -> (n_trials, time_steps)
                    if histories.shape[0] < histories.shape[1]:
                        # Likely (n_trials, time_steps) - no transpose needed
                        histories_transposed = histories
                    else:
                        # Likely (time_steps, n_trials) - transpose
                        histories_transposed = np.transpose(histories)
                    
                    n_trials_hist, n_time = histories_transposed.shape
                    
                    # Extract histories, skipping first point at t=0
                    if n_time == len(sigma_avg):
                        # Histories include t=0
                        histories_plot = histories_transposed[:, 1:]
                    elif n_time == len(sigma_avg) - 1:
                        # Histories already skip t=0
                        histories_plot = histories_transposed[:, :]
                    else:
                        # Length mismatch, skip
                        histories_plot = None
                    
                    if histories_plot is not None and histories_plot.shape[1] == len(time_plot):
                        if allhists:
                            # Plot all individual histories as faint lines
                            for trial_idx in range(n_trials_hist):
                                trial_sigma = histories_plot[trial_idx, :]
                                # Filter out invalid values
                                valid_hist = (trial_sigma > 0) & (time_plot > 0)
                                if np.sum(valid_hist) >= 2:
                                    ax.loglog(time_plot[valid_hist], trial_sigma[valid_hist], 
                                             '-', color=color, linewidth=0.5, alpha=0.15)
                        else:
                            # Compute ±1 standard deviation over trials
                            # histories_plot is (n_trials, time_steps)
                            # Use sigma_plot (already plotted, which is sigma_avg) as the mean, compute std from histories
                            # Handle NaN values in histories (from padding)
                            std_sigma = np.nanstd(histories_plot, axis=0)
                            
                            # Compute ±1 std dev bounds around sigma_plot
                            sigma_upper = sigma_plot + std_sigma
                            sigma_lower = sigma_plot - std_sigma
                            
                            # Filter out invalid values (ensure all are positive for log scale and finite)
                            valid = (sigma_plot > 0) & (sigma_upper > 0) & (sigma_lower > 0) & (time_plot > 0) & \
                                    np.isfinite(sigma_upper) & np.isfinite(sigma_lower) & np.isfinite(std_sigma)
                            
                            if np.sum(valid) >= 2:
                                # Fill between ±1 std dev
                                ax.fill_between(time_plot[valid], sigma_lower[valid], sigma_upper[valid],
                                               color=color, alpha=0.15, linewidth=0)
        
        # Plot instantaneous power-law exponent if requested
        if plot_instantaneous_exponent and ax_exp is not None:
            # Label for exponent plot (no beta in label since we're plotting it)
            label_exp = orientation
            # Compute instantaneous exponent: beta(t) = log_b(sigma(t/b) / sigma(t))
            # Follow the exact approach from ca_plotter.py
            global b_exponent
            
            b = b_exponent
            ts = time_plot
            mt = sigma_plot  # mt is sigma in our case
            
            # Compress time and sigma arrays: compressed_ts[i] = ts[int(i/b)], compressed_mt[i] = mt[int(i/b)]
            compressed_ts = np.copy(ts)
            compressed_mt = np.copy(mt)
            for i in range(len(ts)):
                compressed_ts[i] = ts[int(i/b)]
                compressed_mt[i] = mt[int(i/b)]
            
            # Find unique times
            unique_times = len(np.unique(compressed_ts))
            avgd_comp_ts = np.zeros(unique_times)
            avgd_comp_mt = np.zeros(unique_times)
            
            # Average the exponent values
            for i in range(len(compressed_ts)):
                indx = i // b
                if indx < unique_times:
                    avgd_comp_ts[indx] = compressed_ts[i]
                    # Compute log_b(compressed_mt[i]/mt[i]) = log(compressed_mt[i]/mt[i]) / log(b)
                    # and average by dividing by b
                    if mt[i] > 0 and compressed_mt[i] > 0:
                        avgd_comp_mt[indx] += (np.log(compressed_mt[i]/mt[i]) / np.log(b)) / b
            
            # Plot instantaneous exponent
            # Plot avgd_comp_ts * b vs avgd_comp_mt (scale time back by b)
            # Use linear-linear scale instead of semilogx
            ax_exp.plot(avgd_comp_ts * b, -avgd_comp_mt, '-', 
                       color=color, linewidth=linewidth, alpha=0.7, label=label_exp)
            
            # Plot fitted beta as horizontal line
            ax_exp.axhline(y=beta, color=color, linestyle='--', 
                          linewidth=linewidth, alpha=0.4)
    
    # Add reference power-law with exponent 0.25 (only for sigma(t) plots, not instantaneous exponent)
    if not plot_instantaneous_exponent and ax is not None and len(all_data) > 0:
        # Find the overall time range from all datasets
        all_times = []
        for data_item in all_data:
            sigma_avg = data_item['sigma_avg']
            if len(sigma_avg) > 1:
                time = np.arange(1, len(sigma_avg))
                all_times.extend(time)
        
        reference_curve = False 
        reference_exponent = 0.25
        if reference_curve:
            # Generate reference power-law curve
            if len(all_times) > 0:
                t_min = min(all_times)
                t_max = max(all_times)
                t_ref = np.logspace(np.log10(t_min), np.log10(t_max), 100)
                # Scale to match typical data values - use middle sigma value
                # Find a typical sigma value from the first dataset
                first_sigma = all_data[0]['sigma_avg']
                if len(first_sigma) > 1:
                    # Use middle sigma value to scale (skip first point at t=0)
                    sigma_data = first_sigma[1:]
                    middle_idx = len(sigma_data) // 2
                    if middle_idx < len(sigma_data) and sigma_data[middle_idx] > 0:
                        sigma_scale = sigma_data[middle_idx]
                        # Use corresponding time value
                        t_scale = middle_idx + 1  # time starts at 1
                        sigma_ref = sigma_scale * ((t_ref / t_scale) ** reference_exponent)
                        ax.loglog(t_ref, sigma_ref, ':', color='gray', linewidth=linewidth*0.8, 
                                alpha=0.5, label=r'$t^{{reference_exponent}}$')
    
    # Set up plots based on mode
    if plot_instantaneous_exponent and ax_exp is not None:
        # Only plot instantaneous exponent
        ax_exp.set_xlabel(r'$t$')
        ax_exp.set_ylabel(r'$\beta(t)$')
        
        # Title with rule in mathsf font - just the rule symbol
        if rule_name and rule_name != "unknown":
            # Ensure rule_name is just a single character (R, F, or M)
            rule_clean = str(rule_name).strip().upper()
            if rule_clean in ['R', 'F', 'M']:
                title = r'$\mathsf{' + rule_clean + '}$'
            else:
                # Try to extract just the letter
                import re
                match = re.search(r'([RFM])', rule_clean)
                if match:
                    title = r'$\mathsf{' + match.group(1) + '}$'
                else:
                    title = ""
        else:
            # Try to extract rule name from filename as fallback
            # Filename format: R_dw_growth_... or data/R_dw_growth_...
            extracted_rule = None
            if len(all_data) > 0:
                import re
                filename = all_data[0].get('filename', '')
                # Extract rule from start of filename (before first underscore)
                match = re.search(r'^([RFM])_|/([RFM])_', filename)
                if match:
                    extracted_rule = match.group(1) or match.group(2)
            
            if extracted_rule:
                title = r'$\mathsf{' + extracted_rule + '}$'
            else:
                title = ""
        ax_exp.set_title(title)
        
        # Grid
        ax_exp.grid(True, alpha=0.3, linestyle='--', linewidth=0.5)
        ax_exp.set_axisbelow(True)
        
        # Legend
        if len(all_data) > 0:
            ax_exp.legend(loc='best', frameon=True, fancybox=False, edgecolor='black', framealpha=0.9)
    elif ax is not None:
        # Plot sigma(t) normally
        # Labels with LaTeX math (use rcParams for fontsize)
        ax.set_xlabel(r'$t$')
        
        # Check if "_maxheight" is in any filename
        has_maxheight = False
        if len(all_data) > 0:
            for data_item in all_data:
                filename = data_item.get('filename', '')
                if '_maxheight' in filename:
                    has_maxheight = True
                    break
        
        if has_maxheight:
            ax.set_ylabel(r'$\widetilde\sigma(t)$')
        else:
            ax.set_ylabel(r'$\sigma(t)$')
        
        # Title with rule in mathsf font - just the rule symbol
        if rule_name and rule_name != "unknown":
            # Ensure rule_name is just a single character (R, F, or M)
            rule_clean = str(rule_name).strip().upper()
            if rule_clean in ['R', 'F', 'M']:
                title = r'$\mathsf{' + rule_clean + '}$'
            else:
                # Try to extract just the letter
                import re
                match = re.search(r'([RFM])', rule_clean)
                if match:
                    title = r'$\mathsf{' + match.group(1) + '}$'
                else:
                    title = ""
        else:
            # Try to extract rule name from filename as fallback
            # Filename format: R_dw_growth_... or data/R_dw_growth_...
            extracted_rule = None
            if len(all_data) > 0:
                import re
                filename = all_data[0].get('filename', '')
                # Extract rule from start of filename (before first underscore)
                match = re.search(r'^([RFM])_|/([RFM])_', filename)
                if match:
                    extracted_rule = match.group(1) or match.group(2)
            
            if extracted_rule:
                title = r'$\mathsf{' + extracted_rule + '}$'
            else:
                title = ""
        ax.set_title(title)
        
        # Grid
        # ax.grid(True, which='both', alpha=0.3, linestyle='--', linewidth=0.5)
        # ax.set_axisbelow(True)
        
        # Legend
        if len(all_data) > 0:
            ax.legend(loc='best', frameon=True, fancybox=False, edgecolor='black', framealpha=0.9)
    
    # Tight layout
    plt.tight_layout()
    
    # Save or show
    if output_filename:
        plt.savefig(output_filename, dpi=300, bbox_inches='tight')
        print(f"Plot saved to {output_filename}")
    else:
        # Make the window smaller when displaying
        mngr = fig.canvas.manager
        if hasattr(mngr, 'window'):
            mngr.window.wm_geometry("400x400+100+100")  # width x height + x + y
        plt.show()
    
    plt.close()

def plot_dw_roughness(filenames, output_filename=None):
    """
    Plot L vs sigma_inf for dw_roughness mode files.
    Each file plots a single point where L = min(Lx, Ly) and sigma_inf is the estimated
    long-time steady-state value of sigma.
    """
    # Convert single filename to list
    if isinstance(filenames, str):
        filenames = [filenames]
    
    # Load data and extract L and sigma_inf for each file
    sigma_inf_list = []
    L_list = []
    orientations = []
    p_list = []
    rule_name = None
    
    for filename in filenames:
        # Try multiple methods to read JLD2 file
        data = None
        
        # Method 1: Try jld2 Python package (if available)
        try:
            import jld2
            with jld2.open(filename, 'r') as f:
                data = dict(f)
        except ImportError:
            pass
        except Exception as e:
            print(f"jld2 package failed: {e}")
        
        # Method 2: Try h5py with our custom loader
        if data is None:
            try:
                data = load_jld2_simple(filename)
            except Exception as e:
                print(f"h5py method failed: {e}")
        
        # Method 3: Try direct h5py access
        if data is None or len(data) == 0:
            try:
                data = {}
                with h5py.File(filename, 'r') as f:
                    for key in ['mode', 'Lx', 'Ly', 'sigma_avg', 'r', 'rotate', 'R_or', 'R_and']:
                        if key in f:
                            try:
                                obj = f[key]
                                if isinstance(obj, h5py.Dataset):
                                    if obj.shape == ():
                                        data[key] = obj[()]
                                    else:
                                        data[key] = np.array(obj[:])
                            except:
                                pass
                    # Try to get kwargs if it exists
                    if 'kwargs' in f:
                        kwargs_group = f['kwargs']
                        if isinstance(kwargs_group, h5py.Group):
                            kwargs_dict = {}
                            for k in kwargs_group.keys():
                                try:
                                    v_obj = kwargs_group[k]
                                    if isinstance(v_obj, h5py.Dataset):
                                        if v_obj.shape == ():
                                            kwargs_dict[k] = v_obj[()]
                                        else:
                                            kwargs_dict[k] = np.array(v_obj[:])
                                except:
                                    pass
                            if kwargs_dict:
                                data['kwargs'] = kwargs_dict
                    # Try to get roughness_params if it exists
                    if 'roughness_params' in f:
                        roughness_group = f['roughness_params']
                        if isinstance(roughness_group, h5py.Group):
                            for k in ['mean']:
                                if k in roughness_group:
                                    try:
                                        v_obj = roughness_group[k]
                                        if isinstance(v_obj, h5py.Dataset):
                                            if v_obj.shape == ():
                                                data[f'roughness_params_{k}'] = v_obj[()]
                                    except:
                                        pass
            except Exception as e:
                print(f"Direct h5py access failed: {e}")
        
        if data is None or len(data) == 0:
            print(f"Could not read {filename}, skipping")
            continue
        
        # Determine mode from filename
        if '_dw_growth' in filename:
            mode = 'dw_growth'
        elif '_dw_roughness' in filename:
            mode = 'dw_roughness'
        else:
            mode = 'unknown'
        
        if mode != 'dw_roughness':
            print(f"File {filename} is not in dw_roughness mode (filename doesn't contain '_dw_roughness'), skipping")
            continue
        
        # Get n_trials and print it
        # Try to extract from sigma_history_all_trials or sigma_trials array length
        n_trials = None
        
        # Method 1: Extract from sigma_history_all_trials length (if it's a list/array of trials)
        sigma_history_all_trials = data.get('sigma_history_all_trials', None)
        if sigma_history_all_trials is not None:
            try:
                if hasattr(sigma_history_all_trials, '__len__'):
                    n_trials = len(sigma_history_all_trials)
            except:
                pass
        
        # Method 2: Extract from sigma_trials length
        if n_trials is None:
            sigma_trials = data.get('sigma_trials', None)
            if sigma_trials is not None:
                try:
                    if hasattr(sigma_trials, '__len__'):
                        n_trials = len(sigma_trials)
                except:
                    pass
        
        # Method 3: Try using jld2 package if available (handles references properly)
        try:
            import jld2
            with jld2.open(filename, 'r') as f:
                jld2_data = dict(f)
                jld2_kwargs = jld2_data.get('kwargs', {})
                if isinstance(jld2_kwargs, dict):
                    n_trials = jld2_kwargs.get('n_trials', None)
        except ImportError:
            pass
        except:
            pass
        
        # If not found, try from already loaded data
        if n_trials is None:
            kwargs = data.get('kwargs', {})
            if isinstance(kwargs, dict):
                n_trials = kwargs.get('n_trials', None)
        
        # If still not found, try to extract from HDF5 object reference structure
        if n_trials is None:
            try:
                with h5py.File(filename, 'r') as f:
                    if 'kwargs' in f:
                        kwargs_obj = f['kwargs']
                        # Check if kwargs is a Group (direct access)
                        if isinstance(kwargs_obj, h5py.Group):
                            if 'n_trials' in kwargs_obj:
                                n_trials = kwargs_obj['n_trials'][()]
                        # Check if kwargs is a Dataset with object reference (JLD2 Dict storage)
                        elif isinstance(kwargs_obj, h5py.Dataset):
                            try:
                                ref = kwargs_obj[()]
                                if hasattr(ref, '__iter__') and len(ref) > 0:
                                    ref_obj = f[ref[0]]
                                    if isinstance(ref_obj, h5py.Dataset):
                                        arr = ref_obj[:]
                                        # The array contains references to key-value pairs
                                        # Try to extract: keys and values alternate or are paired
                                        kwargs_dict = {}
                                        for i, ref_item in enumerate(arr):
                                            try:
                                                resolved = f[ref_item]
                                                if isinstance(resolved, h5py.Dataset):
                                                    val = resolved[()]
                                                    # If it's a string, it's likely a key
                                                    if isinstance(val, bytes):
                                                        key_str = val.decode('utf-8', errors='ignore').strip('\x00')
                                                        # Next item should be the value
                                                        if i + 1 < len(arr):
                                                            try:
                                                                value_ref = arr[i + 1]
                                                                value_obj = f[value_ref]
                                                                if isinstance(value_obj, h5py.Dataset):
                                                                    value_val = value_obj[()]
                                                                    if isinstance(value_val, (int, np.integer, float, np.floating)):
                                                                        kwargs_dict[key_str] = int(value_val) if isinstance(value_val, (int, np.integer)) else float(value_val)
                                                            except:
                                                                pass
                                            except:
                                                pass
                                        if 'n_trials' in kwargs_dict:
                                            n_trials = kwargs_dict['n_trials']
                            except:
                                pass
            except:
                pass
        
        if n_trials is None:
            # Try to get from top level data
            n_trials = data.get('n_trials', None)
        
        # Convert to int if it's a numpy array/scalar
        if n_trials is not None:
            if isinstance(n_trials, np.ndarray):
                if n_trials.size == 1:
                    n_trials = int(n_trials[()])
                else:
                    n_trials = None
            elif isinstance(n_trials, (int, np.integer)):
                n_trials = int(n_trials)
        
        # Print for each file
        if n_trials is not None:
            print(f"File {filename}: n_trials = {n_trials}")
        else:
            print(f"File {filename}: n_trials not found")
        
        # Get sigma_inf: the estimated long-time steady-state value
        # It's stored as sigma_avg (which is actually sigma_mean_avg)
        sigma_inf = data.get('sigma_avg', None)
        
        # Also try roughness_params/mean as alternative
        if sigma_inf is None:
            sigma_inf = data.get('roughness_params_mean', None)
        
        if sigma_inf is None:
            print(f"No sigma_avg or roughness_params/mean found in {filename}, skipping")
            continue
        
        # Convert to float if needed
        if isinstance(sigma_inf, np.ndarray):
            if sigma_inf.size == 1:
                sigma_inf = float(sigma_inf[()])
            else:
                print(f"sigma_avg is not a scalar in {filename}, skipping")
                continue
        else:
            sigma_inf = float(sigma_inf)
        
        # Get Lx and Ly
        Lx = data.get('Lx', None)
        Ly = data.get('Ly', None)
        if isinstance(Lx, np.ndarray):
            Lx = int(Lx[()]) if Lx.size == 1 else None
        if isinstance(Ly, np.ndarray):
            Ly = int(Ly[()]) if Ly.size == 1 else None
        
        if Lx is None or Ly is None:
            print(f"Could not get Lx/Ly from {filename}, skipping")
            continue
        
        # Compute L = min(Lx, Ly)
        L = min(Lx, Ly)
        
        # Get orientation for labeling
        r = data.get('r', False)
        rotate = data.get('rotate', False)
        if isinstance(r, np.ndarray):
            r = bool(r[()]) if r.size == 1 else False
        if isinstance(rotate, np.ndarray):
            rotate = bool(rotate[()]) if rotate.size == 1 else False
        
        orientation = get_orientation(r, rotate)
        
        # Get p from kwargs (defaults to 0.0)
        kwargs = data.get('kwargs', {})
        if isinstance(kwargs, dict):
            p = kwargs.get('p', 0.0)
        else:
            p = 0.0
        # Convert to float if needed
        if isinstance(p, np.ndarray):
            p = float(p[()]) if p.size == 1 else 0.0
        else:
            p = float(p) if p is not None else 0.0
        
        # Fallback: try to extract p from filename if not found in kwargs
        # Filename format: ..._dw_roughnessp0.02_... or ..._dw_roughness_...
        if abs(p) < 1e-10:  # Check if p is effectively zero
            import re
            # Look for pattern: dw_roughnessp0.02 or dw_growthp0.02
            match = re.search(r'_dw_(?:roughness|growth)p([\d.]+)', filename)
            if match:
                try:
                    p = float(match.group(1))
                except:
                    p = 0.0
        
        # Get rule - first try reading directly, then fall back to identifying from R_or/R_and
        file_rule = data.get('rule', None)
        if file_rule is None:
            # Fall back to identifying from R_or and R_and
            R_or = data.get('R_or', None)
            R_and = data.get('R_and', None)
            file_rule = identify_rule(R_or, R_and)
        else:
            # Convert to string if it's an array
            if isinstance(file_rule, np.ndarray):
                if file_rule.size == 1:
                    file_rule = str(file_rule[()])
                else:
                    # Try to decode as string if it's a byte array
                    try:
                        file_rule = file_rule.tobytes().decode('utf-8').strip('\x00')
                    except:
                        file_rule = str(file_rule[0]) if len(file_rule) > 0 else "unknown"
            else:
                file_rule = str(file_rule)
            
            # Clean up: extract just the rule letter (R, F, or M) if it's embedded in other text
            import re
            # First, try to extract just R, F, or M from the string
            match = re.search(r'([RFM])', str(file_rule).upper())
            if match:
                file_rule = match.group(1)
            # If it's already just R, F, or M, keep it
            elif str(file_rule).upper().strip() in ['R', 'F', 'M']:
                file_rule = str(file_rule).upper().strip()
            else:
                file_rule = "unknown"
        
        if rule_name is None:
            rule_name = file_rule
        elif rule_name != file_rule:
            print(f"Warning: Files have different rules ({rule_name} vs {file_rule})")
        
        sigma_inf_list.append(sigma_inf)
        L_list.append(L)
        orientations.append(orientation)
        p_list.append(p)
    
    if len(sigma_inf_list) == 0:
        print("No valid data to plot")
        return
    
    # Convert to numpy arrays
    sigma_inf_arr = np.array(sigma_inf_list)
    L_arr = np.array(L_list)
    
    # Filter out zero or negative values for log scale
    valid = (sigma_inf_arr > 0) & (L_arr > 0)
    sigma_inf_plot = sigma_inf_arr[valid]
    L_plot = L_arr[valid]
    orientations_plot = [orientations[i] for i in range(len(orientations)) if valid[i]]
    p_plot = [p_list[i] for i in range(len(p_list)) if valid[i]]
    
    if len(sigma_inf_plot) < 1:
        print("Not enough valid data points to plot")
        return
    
    # Check if any p values are nonzero
    has_nonzero_p = any(abs(p) > 1e-10 for p in p_plot)  # Use tolerance for float comparison
    marker_shapes = ['o', 's', '^', 'P']  # circle, square, triangle, plus
    p_values = sorted(set(p_plot))
    p_to_marker = {}
    if has_nonzero_p:
        nonzero_p = [p for p in p_values if p != 0.0]
        for i, p_val in enumerate(nonzero_p):
            if i < len(marker_shapes):
                p_to_marker[p_val] = marker_shapes[i]
        if 0.0 in p_values:
            p_to_marker[0.0] = 'o'
    else:
        p_to_marker = {0.0: 'o'}
    
    # Group data by orientation and p
    orientation_data = {}
    for sigma_inf, L_val, orientation, p_val in zip(sigma_inf_plot, L_plot, orientations_plot, p_plot):
        key = (orientation, p_val)
        if key not in orientation_data:
            orientation_data[key] = {'sigma': [], 'L': []}
        orientation_data[key]['sigma'].append(sigma_inf)
        orientation_data[key]['L'].append(L_val)
    
    # Color mapping for orientations
    orientation_colors = {
        r'$\rightarrow$': 'blue',
        r'$\uparrow$': 'red',
        r'$\leftarrow$': 'green',
        r'$\downarrow$': 'purple'
    }
    
    # Create plot (log-log scale)
    fig, ax = plt.subplots(figsize=(4., 4.), dpi=100)
    
    # Fit power law for each orientation and p combination and plot
    legend_handles = []
    legend_labels = []
    
    for (orientation, p_val) in sorted(orientation_data.keys()):
        sigma_orient = np.array(orientation_data[(orientation, p_val)]['sigma'])
        L_orient = np.array(orientation_data[(orientation, p_val)]['L'])
        
        # Get color for this orientation
        color = orientation_colors.get(orientation, 'gray')
        
        # Get marker shape for this p value
        marker = p_to_marker.get(p_val, 'o')
        
        # Fit power law if we have at least 2 points
        alpha = None
        a = None
        if len(sigma_orient) >= 2:
            # On log-log scale: log(sigma_inf) = alpha * log(L) + log(a)
            log_L = np.log(L_orient)
            log_sigma = np.log(sigma_orient)
            
            # Linear regression: [log(L) ones] * [alpha; log(a)] = log(sigma_inf)
            X = np.vstack([log_L, np.ones(len(log_L))]).T
            coeffs = np.linalg.lstsq(X, log_sigma, rcond=None)[0]
            alpha = coeffs[0]  # Power law exponent
            log_a = coeffs[1]  # log of prefactor
            a = np.exp(log_a)  # Prefactor
        
        # Plot data points with orientation-specific color and p-specific marker
        ax.loglog(L_orient, sigma_orient, marker, markerfacecolor=color, 
               markeredgecolor='k', markersize=8, alpha=0.7)
        
        # Plot power law fit as dashed line if we have enough points
        if alpha is not None and a is not None and len(sigma_orient) >= 2:
            # Generate fit curve
            L_fit = np.logspace(np.log10(L_orient.min()), np.log10(L_orient.max()), 100)
            sigma_fit = a * (L_fit ** alpha)
            ax.loglog(L_fit, sigma_fit, '--', color=color, linewidth=linewidth, 
                     alpha=0.6)
        
        # Add to legend
        # Extract orientation symbol from LaTeX string (e.g., r'$\rightarrow$' -> '\rightarrow')
        orientation_symbol = orientation.replace('$', '').strip()
        if alpha is not None:
            # Format alpha to 2 decimal places
            alpha_str = f"{alpha:.2f}"
            if has_nonzero_p:
                p_str = f"{p_val:.2f}"
                label = rf'$\alpha_{{{orientation_symbol}}} = {alpha_str}$, $p = {p_str}$'
            else:
                label = rf'$\alpha_{{{orientation_symbol}}} = {alpha_str}$'
        else:
            if has_nonzero_p:
                p_str = f"{p_val:.2f}"
                label = rf'${orientation_symbol}$, $p = {p_str}$'
            else:
                label = rf'${orientation_symbol}$'
        
        # Create a legend handle with the appropriate color and marker
        legend_handles.append(plt.Line2D([0], [0], marker=marker, color='w', 
                                        markerfacecolor=color, markeredgecolor='k',
                                        markersize=8, linestyle='None'))
        legend_labels.append(label)
    
    # Plot comparison line with L^0.5 scaling only if there's a single orientation
    if len(orientation_data) == 1 and len(sigma_inf_plot) >= 1:
        # Scale the L^0.5 curve to match the data at the first point
        L_ref = np.logspace(np.log10(L_plot.min()), np.log10(L_plot.max()), 100)
        # Scale to match first data point: sigma_ref = sigma_inf_plot[0] * (L_ref / L_plot[0])^0.5
        scale_factor = sigma_inf_plot[0] / (L_plot[0] ** 0.5)
        sigma_ref = scale_factor * (L_ref ** 0.5)
        ax.loglog(L_ref, sigma_ref, '--', color='grey', linewidth=linewidth, 
                 alpha=0.4, label=r'$L^{0.5}$')
        legend_handles.append(plt.Line2D([0], [0], color='grey', linestyle='--', linewidth=linewidth))
        legend_labels.append(r'$L^{0.5}$')
    
    # Labels
    ax.set_xlabel(r'$L$')
    # Check if "_maxheight" is in the first filename
    if len(filenames) > 0 and "_maxheight" in filenames[0]:
        ax.set_ylabel(r'$\widetilde\sigma_\infty(t)$')
    else:
        ax.set_ylabel(r'$\sigma_\infty$')
    
    # Title with rule in mathsf font - just the rule symbol
    if rule_name and rule_name != "unknown":
        title = r'$\mathsf{' + rule_name + '}$'
    else:
        # Try to extract rule name from filename as fallback
        # Filename format: R_dw_roughness_... or data/R_dw_roughness_...
        extracted_rule = None
        if len(filenames) > 0:
            import re
            filename = filenames[0]
            # Extract rule from start of filename (before first underscore)
            match = re.search(r'^([RFM])_|/([RFM])_', filename)
            if match:
                extracted_rule = match.group(1) or match.group(2)
        
        if extracted_rule:
            title = r'$\mathsf{' + extracted_rule + '}$'
        else:
            title = ""
    ax.set_title(title)
    
    # Legend with orientation-specific entries
    if len(legend_handles) > 0:
        ax.legend(legend_handles, legend_labels, loc='best', frameon=True, 
                 fancybox=False, edgecolor='black', framealpha=0.9)
    
    # Grid
    ax.grid(True, alpha=0.3, linestyle='--', linewidth=0.5)
    ax.set_axisbelow(True)
    
    plt.tight_layout()
    
    # Save or show
    if output_filename:
        plt.savefig(output_filename, dpi=300, bbox_inches='tight')
        print(f"Plot saved to {output_filename}")
    else:
        plt.show()
    
    plt.close()

def plot_sigma_vs_L(filenames, output_filename=None):
    """
    Plot sigma_inf vs L, where sigma_inf is the average of sigma(t) over the last half
    of the time interval, and L = max(Lx, Ly) for each file.
    """
    # Convert single filename to list
    if isinstance(filenames, str):
        filenames = [filenames]
    
    # Load data and compute sigma_inf and L for each file
    sigma_inf_list = []
    L_list = []
    orientations = []
    
    for filename in filenames:
        # Try multiple methods to read JLD2 file
        data = None
        
        # Method 1: Try jld2 Python package (if available)
        try:
            import jld2
            with jld2.open(filename, 'r') as f:
                data = dict(f)
        except ImportError:
            pass
        except Exception as e:
            print(f"jld2 package failed: {e}")
        
        # Method 2: Try h5py with our custom loader
        if data is None:
            try:
                data = load_jld2_simple(filename)
            except Exception as e:
                print(f"h5py method failed: {e}")
        
        # Method 3: Try direct h5py access
        if data is None or len(data) == 0:
            try:
                data = {}
                with h5py.File(filename, 'r') as f:
                    for key in ['mode', 'Lx', 'Ly', 'sigma_avg', 'r', 'rotate']:
                        if key in f:
                            try:
                                obj = f[key]
                                if isinstance(obj, h5py.Dataset):
                                    if obj.shape == ():
                                        data[key] = obj[()]
                                    else:
                                        data[key] = np.array(obj[:])
                            except:
                                pass
            except Exception as e:
                print(f"Direct h5py access failed: {e}")
        
        if data is None or len(data) == 0:
            print(f"Could not read {filename}, skipping")
            continue
        
        # Determine mode from filename
        if '_dw_growth' in filename:
            mode = 'dw_growth'
        elif '_dw_roughness' in filename:
            mode = 'dw_roughness'
        else:
            mode = 'unknown'
        
        if mode != 'dw_growth':
            print(f"File {filename} is not in dw_growth mode (filename doesn't contain '_dw_growth'), skipping")
            continue
        
        # Get sigma_avg
        sigma_avg = data.get('sigma_avg', None)
        if sigma_avg is None:
            print(f"No sigma_avg data found in {filename}, skipping")
            continue
        
        if not isinstance(sigma_avg, np.ndarray):
            sigma_avg = np.array(sigma_avg)
        if sigma_avg.ndim > 1:
            sigma_avg = sigma_avg.flatten()
        
        # Compute sigma_inf: average over last half of time interval
        # Skip first point at t=0
        sigma_data = sigma_avg[1:]
        if len(sigma_data) < 2:
            print(f"Not enough data in {filename}, skipping")
            continue
        
        # Average over last half
        n_points = len(sigma_data)
        last_half_start = n_points // 2
        sigma_inf = np.mean(sigma_data[last_half_start:])
        
        # Get Lx and Ly
        Lx = data.get('Lx', None)
        Ly = data.get('Ly', None)
        if isinstance(Lx, np.ndarray):
            Lx = int(Lx[()]) if Lx.size == 1 else None
        if isinstance(Ly, np.ndarray):
            Ly = int(Ly[()]) if Ly.size == 1 else None
        
        if Lx is None or Ly is None:
            print(f"Could not get Lx/Ly from {filename}, skipping")
            continue
        
        # Compute L = max(Lx, Ly)
        L = max(Lx, Ly)
        L = min(Lx, Ly) # poop 
        
        # Get orientation for labeling
        r = data.get('r', False)
        rotate = data.get('rotate', False)
        if isinstance(r, np.ndarray):
            r = bool(r[()]) if r.size == 1 else False
        if isinstance(rotate, np.ndarray):
            rotate = bool(rotate[()]) if rotate.size == 1 else False
        
        orientation = get_orientation(r, rotate)
        
        sigma_inf_list.append(sigma_inf)
        L_list.append(L)
        orientations.append(orientation)
    
    if len(sigma_inf_list) == 0:
        print("No valid data to plot")
        return
    
    # Convert to numpy arrays
    sigma_inf_arr = np.array(sigma_inf_list)
    L_arr = np.array(L_list)
    
    # Filter out zero or negative values for log scale
    valid = (sigma_inf_arr > 0) & (L_arr > 0)
    sigma_inf_plot = sigma_inf_arr[valid]
    L_plot = L_arr[valid]
    orientations_plot = [orientations[i] for i in range(len(orientations)) if valid[i]]
    
    if len(sigma_inf_plot) < 2:
        print("Not enough valid data points for log-log plot")
        return
    
    # Fit power law: sigma_inf ~ a * L^alpha
    # On log-log scale: log(sigma_inf) = alpha * log(L) + log(a)
    log_L = np.log(L_plot)
    log_sigma = np.log(sigma_inf_plot)
    
    # Linear regression: [log(L) ones] * [alpha; log(a)] = log(sigma_inf)
    X = np.vstack([log_L, np.ones(len(log_L))]).T
    coeffs = np.linalg.lstsq(X, log_sigma, rcond=None)[0]
    alpha = coeffs[0]  # Power law exponent
    log_a = coeffs[1]  # log of prefactor
    a = np.exp(log_a)  # Prefactor
    
    # Generate fit curve
    L_fit = np.logspace(np.log10(L_plot.min()), np.log10(L_plot.max()), 100)
    sigma_fit = a * (L_fit ** alpha)
    
    # Format alpha to 2 significant figures
    alpha_str = f"{alpha:.2g}"
    
    # Create plot (log-log scale)
    fig, ax = plt.subplots(figsize=(4., 4.), dpi=100)
    
    # Colors from colormap, ordered by input file order
    global cmap
    colors = cmap(np.linspace(0, 1, len(sigma_inf_plot)))
    
    # Plot each point
    for idx, (sigma_inf, L_val, orientation) in enumerate(zip(sigma_inf_plot, L_plot, orientations_plot)):
        ax.loglog(L_val, sigma_inf, 'o', markerfacecolor=colors[idx], 
               markeredgecolor='k', markersize=8, label=orientation, alpha=0.7)
    
    # Plot power law fit as dashed line
    ax.loglog(L_fit, sigma_fit, '--', color='gray', linewidth=linewidth, 
             alpha=0.6, label=rf'$\alpha={alpha_str}$')
    
    # Labels
    ax.set_xlabel(r'$L$')
    ax.set_ylabel(r'$\sigma_\infty$')
    
    # Legend
    ax.legend(loc='best', frameon=True, fancybox=False, edgecolor='black', framealpha=0.9)
    
    # Grid
    ax.grid(True, alpha=0.3, linestyle='--', linewidth=0.5)
    ax.set_axisbelow(True)
    
    plt.tight_layout()
    
    # Save or show
    if output_filename:
        plt.savefig(output_filename, dpi=300, bbox_inches='tight')
        print(f"Plot saved to {output_filename}")
    else:
        plt.show()
    
    plt.close()

def main():
    """
    Main function to plot results.
    """
    # Check for -vsL flag
    plot_vs_L = '-vsL' in sys.argv
    if plot_vs_L:
        # Remove -vsL from arguments
        sys.argv = [arg for arg in sys.argv if arg != '-vsL']
    
    # Check for -instexp flag
    global plot_instantaneous_exponent
    if '-instexp' in sys.argv:
        plot_instantaneous_exponent = True
        # Remove -instexp from arguments
        sys.argv = [arg for arg in sys.argv if arg != '-instexp']
    else:
        plot_instantaneous_exponent = False
    
    # Check for -allhists flag
    allhists = '-allhists' in sys.argv
    if allhists:
        # Remove -allhists from arguments
        sys.argv = [arg for arg in sys.argv if arg != '-allhists']
    
    if len(sys.argv) < 2:
        # If no arguments, look for jld2 files in data/ directory
        data_dir = Path('data')
        if data_dir.exists():
            # Try to find dw_roughness files first
            jld2_files_roughness = list(data_dir.glob('sqz*_dw_roughness*.jld2'))
            jld2_files_growth = list(data_dir.glob('sqz*_dw_growth_*.jld2'))
            
            if len(jld2_files_roughness) > 0:
                # Plot dw_roughness mode
                print(f"Found {len(jld2_files_roughness)} dw_roughness file(s). Plotting...")
                filenames = [str(f) for f in jld2_files_roughness]
                output = jld2_files_roughness[0].with_suffix('.png')
                plot_dw_roughness(filenames, str(output))
            elif len(jld2_files_growth) > 0:
                # Plot dw_growth mode
                print(f"Found {len(jld2_files_growth)} dw_growth file(s). Plotting...")
                filenames = [str(f) for f in jld2_files_growth]
                output = jld2_files_growth[0].with_suffix('.png')
                if plot_vs_L:
                    plot_sigma_vs_L(filenames, str(output))
                else:
                    plot_dw_growth(filenames, str(output), allhists=allhists)
            else:
                print("No dw_roughness or dw_growth jld2 files found in data/ directory")
                print("Usage: python squeezing_plotter.py <filename1.jld2> [filename2.jld2 ...] [output.png]")
                return
        else:
            print("data/ directory not found")
            print("Usage: python squeezing_plotter.py <filename1.jld2> [filename2.jld2 ...] [output.png]")
            return
    else:
        # Parse arguments - last one might be output filename
        args = sys.argv[1:]
        
        # Check if last argument is an output filename (ends with .png, .pdf, etc.)
        output_filename = None
        if len(args) > 1 and args[-1].endswith(('.png', '.pdf', '.eps', '.svg')):
            output_filename = args[-1]
            input_args = args[:-1]
        else:
            input_args = args
        
        # Expand prefixes to full filenames
        # If an argument doesn't end with .jld2, treat it as a prefix and find all matching files
        filenames = []
        for arg in input_args:
            if arg.endswith('.jld2'):
                # It's a full filename
                filenames.append(arg)
            else:
                # It's a prefix - find all files matching this prefix
                prefix = arg
                # Add .jld2 extension if not present
                if not prefix.endswith('.jld2'):
                    # Use glob to find all files with this prefix
                    pattern = prefix + '*.jld2'
                    matching_files = glob.glob(pattern)
                    if len(matching_files) == 0:
                        print(f"No files found matching prefix: {prefix}")
                    else:
                        matching_files.sort()  # Sort for consistent ordering
                        filenames.extend(matching_files)
                        print(f"Found {len(matching_files)} file(s) matching prefix '{prefix}':")
                        for f in matching_files:
                            print(f"  {f}")
        
        if len(filenames) == 0:
            print("No files to plot")
            return
        
        # Check all files exist
        for filename in filenames:
            if not os.path.exists(filename):
                print(f"File not found: {filename}")
                return
        
        # Determine mode from first filename
        first_filename = filenames[0]
        if '_dw_roughness' in first_filename:
            first_file_mode = 'dw_roughness'
        elif '_dw_growth' in first_filename:
            first_file_mode = 'dw_growth'
        else:
            # Default to dw_growth if can't determine
            first_file_mode = 'dw_growth'
        
        # Plot based on determined mode
        if first_file_mode == 'dw_roughness':
            plot_dw_roughness(filenames, output_filename)
        elif first_file_mode == 'dw_growth':
            if plot_vs_L:
                plot_sigma_vs_L(filenames, output_filename)
            else:
                plot_dw_growth(filenames, output_filename, allhists=allhists)
        else:
            # Default to dw_growth behavior
            if plot_vs_L:
                plot_sigma_vs_L(filenames, output_filename)
            else:
                plot_dw_growth(filenames, output_filename, allhists=allhists)

if __name__ == '__main__':
    main()