import matplotlib.pyplot as plt
import numpy as np
import h5py
import os
from matplotlib.colors import LogNorm
from matplotlib import cm
import matplotlib as mpl
import matplotlib.ticker as ticker

from .files import load_initial_data
from .lineshapes import get_lineshape_function

def print_calculation_parameters(args):
    """Prints a formatted and context-aware summary of parameters provided to the plotting command.

    This function reads the argparse.Namespace object and prints a neat summary
    of the key parameters that will be used for the current operation.

    Args:
        args (argparse.Namespace): The parsed command-line arguments.
    """
    print("\n--- Calculation Parameters ---")
    p = vars(args)

    def _format_list(items):
        if not isinstance(items, list) or len(items) <= 6:
            return items
        return f"[{', '.join(map(str, items[:3]))}, ..., {', '.join(map(str, items[-3:]))}] ({len(items)} total)"

    # Define a priority list of keys to display if they exist in args
    keys_to_display = [
        # File Inputs/Outputs
        'input_file', 'input_files', 'raman', 'orbach', 
        'output', 'save_csv', 'no_plot', 'show',
        
        # General Plot Settings
        'plot_type', 'transition', 'process', 'mechanisms', 'labels', 'show_both',
        
        # Dimensions & Selections
        'temperatures', 'orientation', 'average',
        'fields', 'field_index',
        
        # Hysteresis / Simulation Specific
        'init_field', 'sweep_rate', 'time_step', 
        'population_propagation', 'tau_dynamics', 'state_limit',
        
        # Axes & Styling
        'scale', 'plot_resolution', 
        'xlim', 'ylim', 'xaxis_units', 'xlim_zoom', 'ylim_zoom',
        
        # Misc
        'n_cores', 'kramers', 'max_modes', 'max_state',
        'lineshape', 'fwhm', 'integration_width'
    ]

    for key in keys_to_display:
        if key in p and p[key] is not None:
            # Skip False booleans (flags not set) to keep output clean
            if isinstance(p[key], bool) and not p[key]:
                continue
                
            display_name = key.replace('_', ' ').capitalize()
            val = p[key]
            
            if isinstance(val, list):
                val_str = _format_list(val)
            else:
                val_str = str(val)
                
            print(f"  {display_name:<28}: {val_str}")
    print("")

def _apply_plot_style():
    """Applies the standard Taupy plot style."""
    mpl.rcParams['mathtext.rm'] = 'Arial'
    mpl.rcParams['mathtext.it'] = 'Arial:italic'
    mpl.rcParams['mathtext.bf'] = 'Arial:italic:bold'
    mpl.rc('font', family='Arial')
    mpl.rcParams['axes.unicode_minus'] = False
    plt.rcParams['pdf.fonttype'] = 42
    cm = 1/2.54
    plt.rcParams.update({'font.size': 6})
    plt.rcParams.update({'font.family':'Arial'})
    plt.rcParams.update({'figure.figsize': [9*cm, 8*cm]})
    plt.rcParams.update({'axes.linewidth': 0.5})
    # Remove grid
    plt.rcParams.update({'axes.grid': False})
    # Standardise tick widths (Major thicker than Minor)
    plt.rcParams.update({'xtick.major.width': 0.5})
    plt.rcParams.update({'xtick.minor.width': 0.5})
    plt.rcParams.update({'ytick.major.width': 0.5})
    plt.rcParams.update({'ytick.minor.width': 0.5})

def _select_and_report_fields(all_mags, args, file_label=""):
    """
    Robustly selects field indices based on args.field_index OR args.fields,
    defaults to index 0 if neither are set, and prints a confirmation table.

    Args:
        all_mags (np.array): Available field magnitudes in the file.
        args (Namespace): Parsed CLI arguments.
        file_label (str): Name of file for logging context.
    
    Returns:
        list of int: Selected indices.
    """
    selected_indices = []
    
    print(f"--- Field Selection for {file_label} ---")
    print(f"{'Index':<6}   {'Requested (T)':<15}   {'Actual (T)':<15}   {'Diff (T)':<15}")
    print("-" * 60)

    # 1. Use Explicit Indices
    if args.field_index is not None:
        # Handle single int vs list
        indices = args.field_index if isinstance(args.field_index, list) else [args.field_index]
        
        for idx in indices:
            if 0 <= idx < len(all_mags):
                selected_indices.append(idx)
                print(f"{idx:<6}   {'Index ' + str(idx):<15}   {all_mags[idx]:<15.6f}   {'-':<15}")
            else:
                print(f"{idx:<6}   {'Index ' + str(idx):<15}   {'INVALID':<15}   {'-':<15}")
                
    # 2. Use Field Magnitudes (if indices not provided)
    elif args.fields is not None:
        vals = args.fields if isinstance(args.fields, list) else [args.fields]
        
        for val in vals:
            idx = np.argmin(np.abs(all_mags - val))
            diff = abs(all_mags[idx] - val)
            selected_indices.append(idx)
            print(f"{idx:<6}   {val:<15.6f}   {all_mags[idx]:<15.6f}   {diff:<15.6e}")

    # 3. Default to 0
    else:
        print(f"{0:<6}   {'Default':<15}   {all_mags[0]:<15.6f}   {'-':<15}")
        selected_indices = [0]

    print()    
    # Remove duplicates but preserve order
    return list(dict.fromkeys(selected_indices))

def plot_1d_from_hdf5(output_prefix, args):
    """Generates and saves plots of the 1D spectral density from HDF5 files.

    Args:
        output_prefix (str): Prefix for the output plot files.
        args (argparse.Namespace): Parsed command-line arguments.
    """
    print_calculation_parameters(args)
    _apply_plot_style()
    
    filenames = args.input_files
    transition = tuple(args.transition)
    processes_to_plot = args.process if args.process else ['pm']

    print(f"Generating 1D spectral density plots for transition {transition}...")

    # We iterate over files here to select fields per-file
    plot_specific_processes(filenames, output_prefix, transition, processes_to_plot, args)

def plot_specific_processes(filenames, output_prefix, transition, processes, args):
    """Plots specific Raman processes for selected fields.

    Args:
        filenames (list): List of input HDF5 files.
        output_prefix (str): Output filename prefix.
        transition (tuple): Transition indices.
        processes (list): List of process strings (e.g., 'pm').
        args (Namespace): CLI args for field selection.
    """
    all_plot_data = {}
    
    for filename in filenames:
        with h5py.File(filename, 'r') as hf:
            if 'orientation_0' not in hf: continue
            results_group = hf['orientation_0/results']
            params_group = hf['orientation_0/parameters']
            all_field_mags = params_group['field_magnitudes'][...]
            field_keys = sorted(results_group.keys(), key=lambda x: int(x.split('_')[1]))

            # --- New Field Selection Logic ---
            selected_indices = _select_and_report_fields(all_field_mags, args, file_label=os.path.basename(filename))

            for idx in selected_indices:
                if idx >= len(field_keys): continue
                field_key = field_keys[idx]
                field_group = results_group[field_key]
                field_mag = all_field_mags[idx]

                if 'raman' not in field_group: continue
                raman_group = field_group['raman']
                omega_grid = raman_group['omega_grid'][...]
                transition_key = f'transition_{transition[0]}_{transition[1]}'
                if transition_key not in raman_group: continue

                transition_group = raman_group[transition_key]
                if len(filenames) > 1:
                    file_label = f"{filename.rsplit('.', 1)[0]} @ {field_mag:.4f} T"
                else:
                    file_label = f"{field_mag:.4f} T"
                all_plot_data[file_label] = {'omega': omega_grid, 'integrands': {}}

                for process in processes:
                    if f'J_{process}' in transition_group:
                        all_plot_data[file_label]['integrands'][process] = transition_group[f'J_{process}'][...]

    if not all_plot_data: 
        print("No matching data found to plot.")
        return

    for scale in ['linear', 'log']:
        fig, ax = plt.subplots(figsize=(10, 6))
        for file_label, data in all_plot_data.items():
            for process, integrand in data['integrands'].items():
                ax.plot(data['omega'], integrand, label=f"{file_label} ({process})")

        ax.set_xlabel(r'$\hbar\omega_1$ (cm$^{-1}$)', fontsize=10)
        ax.set_ylabel(r'Spectral Density Integrand', fontsize=10)
        ax.set_title(rf'Raman Spectral Density for ${transition[0]} \rightarrow {transition[1]}$ ({scale.capitalize()} Scale)', fontsize=10)
        ax.legend()
        ax.set_yscale(scale)
        out_name = f"{output_prefix}_1d_{scale}_processes.pdf"
        fig.savefig(out_name, bbox_inches='tight', dpi=300)
        print(f"Saved 1D plot to {out_name}")
        plt.close(fig)

def plot_2d_from_hdf5(output_prefix, args):
    """Generates 2D heatmap and stacked plots of the spectral density vs. field.

    Args:
        output_prefix (str): Prefix for output files.
        args (argparse.Namespace): Parsed command-line arguments.
    """
    print_calculation_parameters(args)
    _apply_plot_style()
    
    filename = args.input_file
    transition = tuple(args.transition)
    plot_resolution = args.plot_resolution
    processes_to_plot = args.process if args.process else ['total']

    for process in processes_to_plot:
        print(f"Generating 2D plots for transition {transition}, process '{process}' from {filename}...")
        
        all_field_data = []
        with h5py.File(filename, 'r') as hf:
            if 'orientation_0' not in hf:
                print("Error: No 'orientation_0' group found in HDF5 file.")
                return
                
            results_group = hf['orientation_0/results']
            params_group = hf['orientation_0/parameters']
            all_field_mags = params_group['field_magnitudes'][...]
            field_keys = sorted(results_group.keys(), key=lambda x: int(x.split('_')[1]))

            for i, field_key in enumerate(field_keys):
                field_group = results_group[field_key]
                if 'raman' not in field_group:
                    continue
                    
                raman_group = field_group['raman']
                field_mag = all_field_mags[i]
                omega_grid = raman_group['omega_grid'][...]
                transition_key = f'transition_{transition[0]}_{transition[1]}'
                
                if transition_key not in raman_group:
                    continue

                transition_group = raman_group[transition_key]
                
                dataset_name = f'J_{process}'
                if dataset_name in transition_group:
                    integrand = transition_group[dataset_name][...]
                else:
                    continue # Skip if this process doesn't exist for this field
                
                all_field_data.append({'field': field_mag, 'omega': omega_grid, 'integrand': integrand})

        if not all_field_data:
            print(f"Warning: No Raman data found for transition {transition}, process '{process}' in {filename}.")
            continue

        # --- Prepare data for 2D plotting ---
        fields = np.array([d['field'] for d in all_field_data])
        min_omega = min(d['omega'][0] for d in all_field_data)
        max_omega = max(d['omega'][-1] for d in all_field_data)
        num_points = int((max_omega - min_omega) / plot_resolution) + 1
        unified_omega_grid = np.linspace(min_omega, max_omega, num_points)
        
        integrand_grid_2d = np.zeros((len(fields), len(unified_omega_grid)))
        for i, data in enumerate(all_field_data):
            integrand_grid_2d[i, :] = np.interp(unified_omega_grid, data['omega'], data['integrand'])
            
        # --- Generate Plots for the current process ---
        process_suffix = f"_{process}"

        # 1. 2D Heatmap (Linear Scale)
        fig, ax = plt.subplots(figsize=(10, 7))
        c = ax.pcolormesh(fields, unified_omega_grid, integrand_grid_2d.T, shading='gouraud', cmap='viridis')
        fig.colorbar(c, ax=ax, label='Spectral Density (arb. units)')
        ax.set_xlabel('Magnetic Field (T)', fontsize=10)
        ax.set_ylabel(r'$\hbar\omega_1$ (cm$^{-1}$)', fontsize=10)
        ax.set_title(rf'Raman Spectral Density vs. Field (Linear Scale) - Process: {process.upper()}\nTransition ${transition[0]} \rightarrow {transition[1]}$', fontsize=10)
        out_name = f"{output_prefix}_2d_heatmap_linear{process_suffix}.pdf"
        fig.savefig(out_name, bbox_inches='tight', dpi=300)
        print(f"Saved 2D linear heatmap to {out_name}")
        plt.close(fig)

        # 2. 2D Heatmap (Log Scale)
        fig, ax = plt.subplots(figsize=(10, 7))
        integrand_log = integrand_grid_2d + 1e-20
        c = ax.pcolormesh(fields, unified_omega_grid, integrand_log.T, shading='gouraud', cmap='viridis', norm=LogNorm())
        fig.colorbar(c, ax=ax, label='Spectral Density (arb. units)')
        ax.set_xlabel('Magnetic Field (T)', fontsize=10)
        ax.set_ylabel(r'$\hbar\omega_1$ (cm$^{-1}$)', fontsize=10)
        ax.set_title(rf'Raman Spectral Density vs. Field (Log Scale) - Process: {process.upper()}\nTransition ${transition[0]} \rightarrow {transition[1]}$', fontsize=10)
        out_name = f"{output_prefix}_2d_heatmap_log{process_suffix}.pdf"
        fig.savefig(out_name, bbox_inches='tight', dpi=300)
        print(f"Saved 2D log heatmap to {out_name}")
        plt.close(fig)

def plot_hysteresis(field_list, magnetisation_dict, temperatures, output_filename='hysteresis', lw=0.75, 
                    xlim=None, ylim=None, xlim_zoom=None, ylim_zoom=None, xaxis_units='T', show_plots=False):
    """Plot hysteresis curves for multiple temperatures with customisable axes.

    Args:
        field_list (array-like): List of magnetic field values.
        magnetisation_dict (dict): Dictionary with temperatures as keys and magnetisation lists as values.
        temperatures (list): List of temperatures to plot.
        output_filename (str): Base name for output files without extension.
        lw (float): Line width for the plots.
        xlim (tuple, optional): X-axis limits for the main plot.
        ylim (tuple, optional): Y-axis limits for the main plot.
        xlim_zoom (tuple, optional): X-axis limits for the zoomed plot in Oe.
        ylim_zoom (tuple, optional): Y-axis limits for the zoomed plot in mu_B x 10^-3.
        xaxis_units (str, optional): Units for the main plot x-axis ('T' or 'Oe').
        show_plots (bool, optional): If True, display plots interactively.
    """
    _apply_plot_style()
    plt.rcParams.update({'font.size': 7})
    plt.rcParams.update({'axes.linewidth': 1})
    
    # Main Plot
    fig_main, ax_main = plt.subplots(figsize=(7.5/2.54, 5.5/2.54))
    colours = mpl.colormaps['rainbow'](np.linspace(0, 1, len(temperatures)))

    plot_fields = np.array(field_list)
    xlabel = r'$H$ / T'
    if xaxis_units == 'Oe':
        plot_fields *= 10000
        xlabel = r'$H$ / Oe'

    for temp, colour in zip(temperatures, colours):
        magnetisation_list = magnetisation_dict[temp]
        ax_main.plot(plot_fields, magnetisation_list, color=colour, lw=lw, label=f'{temp:g}')
        ax_main.plot(-plot_fields, -np.array(magnetisation_list), color=colour, lw=lw)

    ax_main.set_xlabel(xlabel)
    ax_main.set_ylabel(r'$M$ / $\mu_B$')
    ax_main.legend(loc='upper left', bbox_to_anchor=(1, 1), ncol=1, title='T / K', frameon=False)
    
    # Tick customisation
    ax_main.tick_params('both', length=2, width=0.5, direction='in', which='major')
    ax_main.tick_params('both', length=1.5, width=0.5, direction='in', which='minor')
    
    if xlim: ax_main.set_xlim(xlim)
    if ylim: ax_main.set_ylim(ylim)

    fig_main.tight_layout()
    fig_main.subplots_adjust(left=0.14, bottom=0.16, right=0.77, top=0.94)
    if not show_plots:
        fig_main.savefig(output_filename + '.pdf')
        print(f"Saved hysteresis plot to {output_filename}.pdf")

    # Zoomed Plot
    fig_zoom, ax_zoom = plt.subplots(figsize=(7.5/2.54, 5.5/2.54))

    for temp, colour in zip(temperatures, colours):
        magnetisation_list = magnetisation_dict[temp]
        ax_zoom.plot(np.array(field_list), magnetisation_list, color=colour, lw=lw, label=f'{temp:g}')
        ax_zoom.plot(-np.array(field_list), -np.array(magnetisation_list), color=colour, lw=lw)
    
    scale_Oe = 1e-4
    scale_mmu_B = 1e-3
    ticks_x = ticker.FuncFormatter(lambda x, pos: f'{x / scale_Oe:g}')
    ticks_y = ticker.FuncFormatter(lambda y, pos: f'{y / scale_mmu_B:g}')
    ax_zoom.xaxis.set_major_formatter(ticks_x)
    ax_zoom.yaxis.set_major_formatter(ticks_y)

    ax_zoom.set_xlabel(r'$H$ / Oe')
    ax_zoom.set_ylabel(r'$M$ / $\mu_B \times 10^{-3}$ ')
    
    if xlim_zoom:
        ax_zoom.set_xlim(xlim_zoom[0] * scale_Oe, xlim_zoom[1] * scale_Oe)
    else:
        ax_zoom.set_xlim(-120 * scale_Oe, 120 * scale_Oe)
        
    if ylim_zoom:
        ax_zoom.set_ylim(ylim_zoom[0] * scale_mmu_B, ylim_zoom[1] * scale_mmu_B)
    else:
        ax_zoom.set_ylim(-5 * scale_mmu_B, 5 * scale_mmu_B)
    
    ax_zoom.legend(loc='upper left', bbox_to_anchor=(1, 1), ncol=1, title='T / K', frameon=False)
    ax_zoom.tick_params('both', length=2, width=0.5, direction='in', which='major')
    ax_zoom.tick_params('both', length=1.5, width=0.5, direction='in', which='minor')

    fig_zoom.tight_layout()
    fig_zoom.subplots_adjust(left=0.14, bottom=0.16, right=0.77, top=0.94)
    if not show_plots:
        fig_zoom.savefig(output_filename + '_zoom.pdf')
        print(f"Saved zoomed hysteresis plot to {output_filename}_zoom.pdf")
    
    if show_plots:
        plt.show()

    plt.close(fig_main)
    plt.close(fig_zoom)

def plot_rate(output_prefix, args):
    """Main function for plotting relaxation rates.
    
    Args:
        output_prefix (str): Prefix for output files.
        args (argparse.Namespace): Parsed command-line arguments.
    """
    print_calculation_parameters(args)
    _apply_plot_style()

    plot_type = args.plot_type
    
    # Common Args
    common_args = {
        'output_prefix': output_prefix,
        'mechanisms': args.mechanisms,
        'show_both': args.show_both,
        'raman_files': args.raman,
        'orbach_files': args.orbach,
        'xlim': args.xlim,
        'ylim': args.ylim,
        'labels': args.labels,
        'save_csv': args.save_csv,
        'no_plot': args.no_plot,
        'show': args.show,
        'average': args.average,
        'args': args # Pass full args for field selection
    }

    if plot_type == 'rate-vs-temp':
        plot_rate_vs_temp(scale=args.scale, orientation=args.orientation, **common_args)
    elif plot_type == 'tau-vs-inv-temp':
        plot_tau_vs_inv_temp(orientation=args.orientation, **common_args)
    elif plot_type == 'rate-vs-field':
        # Field plots don't use the field selector, they use temperature
        plot_rate_vs_field(scale=args.scale, orientation=args.orientation, temperatures=args.temperatures, **common_args)

def _load_and_prepare_data(raman_files, orbach_files, orientation, field_indices, labels=None, vs_field=False, target_temp=None, average=False):
    """
    Loads rate data from files and prepares it for plotting, handling multiple systems and mechanisms.
    
    Args:
        ...
        field_indices (list): List of selected indices to load. Can be None if vs_field is True.
    """
    systems = []
    num_systems = 0
    if raman_files:
        num_systems = len(raman_files)
    elif orbach_files:
        num_systems = len(orbach_files)

    for i in range(num_systems):
        if labels and i < len(labels):
            sys_name = labels[i]
        else:
            sys_name = f'System {i+1}'

        # Load data for this system
        # If vs_field is False (Rate vs T), we might have multiple fields selected (field_indices is a list)
        # We need to create a system entry for EACH selected field if there's more than one.
        
        current_indices = field_indices if field_indices is not None else [0] # Default wrapper for vs_field logic
        
        for f_idx in current_indices:
            # Construct a system name that includes the field if multiple are selected
            if len(current_indices) > 1 and not vs_field:
                # We can't know the exact field B value here without peeking inside the file first 
                # or returning it from _load. _load_rate_data now returns (temp, rates, field_val)
                eff_sys_name = f"{sys_name} (Field Idx {f_idx})" 
            else:
                eff_sys_name = sys_name

            system_data = {'name': eff_sys_name, 'raman': None, 'orbach': None, 'total': None}
            
            # Helper to unpack and store
            def unpack_load(file_path, mech):
                if not file_path: return None
                if vs_field:
                    return _load_rate_vs_field_data(file_path, orientation, target_temp, mech, average)
                else:
                    return _load_rate_data(file_path, orientation, f_idx, mech, average)

            # Load Raman
            if raman_files and i < len(raman_files):
                res = unpack_load(raman_files[i], 'raman')
                if res:
                    x, y, f_val = res
                    system_data['raman'] = (x, y)
                    if len(current_indices) > 1 and not vs_field:
                        system_data['name'] = f"{sys_name} ({f_val:.4f} T)"

            # Load Orbach
            if orbach_files and i < len(orbach_files):
                res = unpack_load(orbach_files[i], 'orbach')
                if res:
                    x, y, f_val = res
                    system_data['orbach'] = (x, y)
                    if len(current_indices) > 1 and not vs_field and system_data['name'] == eff_sys_name:
                         system_data['name'] = f"{sys_name} ({f_val:.4f} T)"

            # Compute Total
            if system_data['raman'] and system_data['orbach']:
                x_r, y_r = system_data['raman']
                x_o, y_o = system_data['orbach']
                
                # Check if x-axes (Temp or Field) match
                if not np.allclose(x_r, x_o, atol=1e-6):
                    # Only warn if lengths are significant
                    if len(x_r) > 1 and len(x_o) > 1:
                         print(f"[{sys_name}] Warning: Raman and Orbach x-axes differ. Interpolating to common set.")
                    # Create a common sorted union of points
                    common_x = np.union1d(x_r, x_o)
                    # Interpolate both to common x
                    y_r_interp = np.interp(common_x, x_r, y_r)
                    y_o_interp = np.interp(common_x, x_o, y_o)
                    system_data['total'] = (common_x, y_r_interp + y_o_interp)
                else:
                    system_data['total'] = (x_r, y_r + y_o)

            elif system_data['raman']:
                system_data['total'] = system_data['raman']
            elif system_data['orbach']:
                system_data['total'] = system_data['orbach']
            
            systems.append(system_data)
            
            if vs_field: break # Only run once per system if vs_field
        
    return systems

def _load_rate_data(input_file, orientation, field_idx, mechanism, average=False):
    """
    Loads rate vs temperature data for a specific field index.
    Returns (temps, rates, field_value).
    """
    if input_file.endswith('.h5') or input_file.endswith('.hdf5'):
        with h5py.File(input_file, 'r') as hf:
            if 'orientation_0' not in hf:
                print(f"Warning: 'orientation_0' not found in {input_file}.")
                return None
            
            ref_group = hf['orientation_0']
            fields = ref_group['parameters/field_magnitudes'][()]
            
            if field_idx >= len(fields):
                print(f"Warning: Field index {field_idx} out of range for {input_file} (max {len(fields)-1}).")
                return None
            
            field_mag = fields[field_idx]

            if average:
                # Powder average logic
                if 'orientation_0' not in hf:
                    print(f"Warning: 'orientation_0' not found in {input_file}, cannot calculate average.")
                    return None
                
                # Assume all orientations share the same field parameters (safe assumption for Taupy)
                temps = ref_group['parameters/temperatures'][()]
                total_rates = np.zeros_like(temps)
                total_weight = 0.0
                
                for key in hf.keys():
                    if key.startswith('orientation_'):
                        orient_group = hf[key]
                        weight = orient_group.attrs.get('weight', 0.0)
                        
                        try:
                            rate_path = f'results/field_{field_idx}/{mechanism}_rates'
                            if rate_path in orient_group:
                                rates = orient_group[rate_path][()]
                                total_rates += rates * weight
                                total_weight += weight
                        except KeyError:
                            pass
                
                if total_weight > 0:
                    total_rates /= total_weight
                    
                return (temps, total_rates, field_mag)
            
            else:
                # Specific orientation logic
                if f'orientation_{orientation}' not in hf:
                    print(f"Warning: orientation {orientation} not found in {input_file}. Skipping.")
                    return None
                
                orient_group = hf[f'orientation_{orientation}']
                temps = orient_group['parameters/temperatures'][()]
                
                if f'results/field_{field_idx}/{mechanism}_rates' not in orient_group:
                     return None
                rates = orient_group[f'results/field_{field_idx}/{mechanism}_rates'][()]
                return (temps, rates, field_mag)

    elif input_file.endswith('.csv'):
        # CSV support is basic; assumes single field or pre-selected
        if average:
            print("Warning: Powder averaging not supported for CSV input.")
        data = np.loadtxt(input_file, delimiter=',', skiprows=1)
        temps = data[:, 0]
        rates = data[:, 1]
        return (temps, rates, 0.0)
    
    return None

def _load_rate_vs_field_data(input_file, orientation, target_temp, mechanism, average=False):
    """Loads rate vs field data for a specific temperature. Returns (fields, rates, temp_val)"""
    if target_temp is None: target_temp = 2.0 

    if input_file.endswith('.h5') or input_file.endswith('.hdf5'):
        with h5py.File(input_file, 'r') as hf:
            if average:
                 if 'orientation_0' not in hf: return None
                 ref_group = hf['orientation_0']
                 
                 # Determine common x-axis (fields) and temp index
                 fields_arr = ref_group['parameters/field_magnitudes'][()]
                 temps_arr = ref_group['parameters/temperatures'][()]
                 temp_idx = np.argmin(np.abs(temps_arr - target_temp))
                 
                 if abs(temps_arr[temp_idx] - target_temp) > 1.0:
                     print(f"Warning: Closest temperature found ({temps_arr[temp_idx]:.1f} K) is far from requested ({target_temp} K).")

                 total_rates_vs_field = np.zeros_like(fields_arr)
                 total_weight = 0.0
                 
                 for key in hf.keys():
                    if key.startswith('orientation_'):
                        orient_group = hf[key]
                        weight = orient_group.attrs.get('weight', 0.0)
                        
                        # Accumulate rates for each field
                        results_group = orient_group['results']
                        rates_for_orient = []
                        
                        for i, _ in enumerate(fields_arr):
                             field_key = f'field_{i}'
                             rate_val = 0.0
                             if field_key in results_group:
                                 mech_key = f'{mechanism}_rates'
                                 if mech_key in results_group[field_key]:
                                     rate_val = results_group[field_key][mech_key][temp_idx]
                             rates_for_orient.append(rate_val)
                        
                        total_rates_vs_field += np.array(rates_for_orient) * weight
                        total_weight += weight
                 
                 if total_weight > 0:
                     total_rates_vs_field /= total_weight
                 
                 return (fields_arr, total_rates_vs_field, temps_arr[temp_idx])
            
            else:
                if f'orientation_{orientation}' not in hf:
                    print(f"Warning: orientation {orientation} not found in {input_file}. Skipping.")
                    return None
                
                orient_group = hf[f'orientation_{orientation}']
                fields_arr = orient_group['parameters/field_magnitudes'][()]
                temps_arr = orient_group['parameters/temperatures'][()]
                
                # Find closest temperature index
                temp_idx = np.argmin(np.abs(temps_arr - target_temp))
                
                # Safety check: ensure the found temp is reasonably close
                if abs(temps_arr[temp_idx] - target_temp) > 1.0:
                     print(f"Warning: Closest temperature found ({temps_arr[temp_idx]:.1f} K) is far from requested ({target_temp} K).")

                rates_vs_field = []
                valid_fields = []

                # Iterate through field keys to ensure alignment
                results_group = orient_group['results']
                # We assume results are stored as field_0, field_1 matching field_magnitudes order
                for i, field_val in enumerate(fields_arr):
                    field_key = f'field_{i}'
                    if field_key in results_group:
                         mech_key = f'{mechanism}_rates'
                         if mech_key in results_group[field_key]:
                             r_val = results_group[field_key][mech_key][temp_idx]
                             rates_vs_field.append(r_val)
                             valid_fields.append(field_val)
                
                return (np.array(valid_fields), np.array(rates_vs_field), temps_arr[temp_idx])
            
    print("Warning: Only HDF5 supported for rate-vs-field currently.")
    return None

def _save_data_csv(systems, filename_prefix, suffix=""):
    """Saves the plotted data to CSV files."""
    for sys in systems:
        sys_name_safe = sys['name'].replace(' ', '_').replace('(', '').replace(')', '').replace('.', 'p')
        if sys['total'] is not None:
            x, y = sys['total']
            out_name = f"{filename_prefix}_{sys_name_safe}{suffix}_total.csv"
            np.savetxt(out_name, np.column_stack((x, y)), delimiter=',', header='X,Rate', comments='')
            print(f"Saved CSV data to {out_name}")

def _resolve_filename(output_prefix, default_ext='.pdf'):
    """Resolves output filename, respecting provided extension if any."""
    if not output_prefix:
        return 'plot' + default_ext
    _, ext = os.path.splitext(output_prefix)
    if ext in ['pdf','png','ps','eps','svg']:
        return output_prefix
    return output_prefix + default_ext

def plot_rate_vs_temp(output_prefix, scale, orientation, args, mechanisms, show_both, raman_files, orbach_files, 
                      xlim=None, ylim=None, labels=None, save_csv=False, no_plot=False, show=False, average=False):
    """Plots 1/tau vs. T."""
    
    # 1. Determine Field Indices (Per File, technically, but we use the first to report)
    # We need to peek at the first file to get available fields for the report
    first_file = raman_files[0] if raman_files else (orbach_files[0] if orbach_files else None)
    selected_indices = [0]
    
    if first_file and (first_file.endswith('.h5') or first_file.endswith('.hdf5')):
        with h5py.File(first_file, 'r') as hf:
            if 'orientation_0' in hf:
                 all_mags = hf['orientation_0/parameters/field_magnitudes'][()]
                 selected_indices = _select_and_report_fields(all_mags, args, file_label=os.path.basename(first_file))

    systems = _load_and_prepare_data(raman_files, orbach_files, orientation, selected_indices, labels, average=average)
    
    if save_csv:
        _save_data_csv(systems, output_prefix)

    fig, ax = plt.subplots()

    # Simplify labels if only one system
    single_system = len(systems) == 1

    for system in systems:
        sys_prefix = "" if single_system else f"{system['name']} "
        
        for mechanism in mechanisms:
            if system[mechanism]:
                temps, rates = system[mechanism]
                label = f"{sys_prefix}{mechanism.capitalize()}"
                ax.plot(temps, rates, label=label)
        if show_both and system['raman'] and system['orbach']:
            temps, raman_rates = system['raman']
            _, orbach_rates = system['orbach']
            ax.plot(temps, raman_rates, label=f"{sys_prefix}Raman")
            ax.plot(temps, orbach_rates, label=f"{sys_prefix}Orbach")

    ax.set_xlabel("Temperature / K")
    ax.set_ylabel(r"$\tau^{-1}$  / s$^{-1}$")
    
    # Parse scale string (e.g. log-lin, log-log)
    # Assumes format: xscale-yscale
    if '-' in scale:
        x_s, y_s = scale.split('-')
    else:
        # Fallback for simple "log" or "linear"
        x_s = scale
        y_s = scale
        
    # Map 'lin' to 'linear' for matplotlib
    x_s = 'linear' if x_s == 'lin' else x_s
    y_s = 'linear' if y_s == 'lin' else y_s
        
    ax.set_xscale(x_s)
    ax.set_yscale(y_s)
    
    # Defaults
    if xlim is None: ax.set_xlim(1, 120)
    else: ax.set_xlim(xlim)
    
    if ylim is None: ax.set_ylim(1e-5, 1e5)
    else: ax.set_ylim(ylim)

    ax.legend()
    plt.tight_layout()
    
    if not no_plot:
        outfile = _resolve_filename(f"{output_prefix}_rate_vs_temp", ".pdf")
        plt.savefig(outfile)
        print(f"Saved plot to {outfile}")
    
    if show:
        plt.show()

def plot_tau_vs_inv_temp(output_prefix, orientation, args, mechanisms, show_both, raman_files, orbach_files, 
                         xlim=None, ylim=None, labels=None, save_csv=False, no_plot=False, show=False, average=False):
    """Plots tau vs. 1/T."""

    # 1. Determine Field Indices
    first_file = raman_files[0] if raman_files else (orbach_files[0] if orbach_files else None)
    selected_indices = [0]
    
    if first_file and (first_file.endswith('.h5') or first_file.endswith('.hdf5')):
        with h5py.File(first_file, 'r') as hf:
            if 'orientation_0' in hf:
                 all_mags = hf['orientation_0/parameters/field_magnitudes'][()]
                 selected_indices = _select_and_report_fields(all_mags, args, file_label=os.path.basename(first_file))

    systems = _load_and_prepare_data(raman_files, orbach_files, orientation, selected_indices, labels, average=average)
    
    if save_csv:
        _save_data_csv(systems, output_prefix)

    fig, ax = plt.subplots()
    single_system = len(systems) == 1
    data_min_T, data_max_T = np.inf, -np.inf
    has_data = False

    for system in systems:
        sys_prefix = "" if single_system else f"{system['name']} "
        for mechanism in mechanisms:
            if system[mechanism]:
                temps, rates = system[mechanism]
                valid = (temps > 0) & (rates > 0)
                temps = temps[valid]
                rates = rates[valid]
                
                if len(temps) > 0:
                    data_min_T = min(data_min_T, np.min(temps))
                    data_max_T = max(data_max_T, np.max(temps))
                    has_data = True

                tau = 1.0 / rates
                inv_temps = 1.0 / temps
                label = f"{sys_prefix}{mechanism.capitalize()}"
                ax.plot(inv_temps, tau, label=label)
        
        if show_both and system['raman'] and system['orbach']:
            temps, raman_rates = system['raman']
            _, orbach_rates = system['orbach']
            ax.plot(1/temps, 1/raman_rates, label=f"{sys_prefix}Raman")
            ax.plot(1/temps, 1/orbach_rates, label=f"{sys_prefix}Orbach")

    ax.set_xlabel("T / K")
    ax.set_ylabel(r"$\tau$ / s")
    ax.set_yscale('log')
    
    if not has_data:
        data_min_T, data_max_T = 1, 120

    # --- Fix Limits ---
    # Invert T to get 1/T limits. 
    # High T -> Low 1/T (Left)
    # Low T -> High 1/T (Right)
    
    if xlim is None:
        # Default Auto-Scaling if NO xlim provided
        x_min = 1.0 / data_max_T
        x_max = 1.0 / data_min_T
        padding = (x_max - x_min) * 0.05
        ax.set_xlim(x_min - padding, x_max + padding)
    else:
        # Strictly follow user limits if provided
        # User provides T, we plot 1/T.
        # xlim[0] is Min T -> Max 1/T (Right limit)
        # xlim[1] is Max T -> Min 1/T (Left limit)
        # We generally want Min 1/T on left.
        # So left limit = 1/xlim[1], right limit = 1/xlim[0]
        ax.set_xlim(1.0/xlim[1], 1.0/xlim[0])

    # --- Smart Ticks ---
    # Calculate ticks based on the *actual* plot limits we just set
    curr_xlim = ax.get_xlim()
    
    # Avoid division by zero if limits are weird
    safe_min = curr_xlim[0] if abs(curr_xlim[0]) > 1e-9 else 1e-9
    safe_max = curr_xlim[1] if abs(curr_xlim[1]) > 1e-9 else 1e-9
    
    vis_max_T = 1.0 / safe_min
    vis_min_T = 1.0 / safe_max
    
    # Sort T range
    vis_T = sorted([vis_min_T, vis_max_T])
    vis_min_T, vis_max_T = vis_T[0], vis_T[1]
    
    potential_ticks = [1, 2, 5, 10, 20, 30, 40, 50, 100, 200, 300]
    ticks_to_use = [t for t in potential_ticks if vis_min_T <= t <= vis_max_T]
    
    if len(ticks_to_use) < 3:
         ticks_to_use = np.linspace(vis_min_T, vis_max_T, 5)

    tick_locs = 1.0 / np.array(ticks_to_use)
    ax.set_xticks(tick_locs)
    ax.set_xticklabels([f"{t:g}" for t in ticks_to_use])

    if ylim is None: ax.set_ylim(1e-5, 1e5)
    else: ax.set_ylim(ylim)

    ax.legend()
    plt.tight_layout()
    
    if not no_plot:
        outfile = _resolve_filename(f"{output_prefix}_tau_vs_inv_temp", ".pdf")
        plt.savefig(outfile)
        print(f"Saved plot to {outfile}")
    
    if show:
        plt.show()

def plot_rate_vs_field(output_prefix, scale, orientation, args, mechanisms, show_both, raman_files, orbach_files, 
                       xlim=None, ylim=None, labels=None, save_csv=False, no_plot=False, show=False, temperatures=None, average=False):
    """Plots 1/tau vs. H."""
    
    if temperatures is None:
        temperatures = [2.0]
        print("No temperature specified for rate-vs-field plot. Defaulting to 2.0 K.")

    fig, ax = plt.subplots()
    
    # Loop over temperatures to create multiple lines
    for temp in temperatures:
        print(f"Plotting rate vs field at T = {temp} K")
        # Passing None for field_indices as we want all fields (vs Field plot)
        systems = _load_and_prepare_data(raman_files, orbach_files, orientation, None, labels, vs_field=True, target_temp=temp, average=average)
        
        if save_csv:
             _save_data_csv(systems, output_prefix, suffix=f"_{temp}K")

        single_system = len(systems) == 1

        for system in systems:
            sys_name = system['name']
            lbl_suffix = f" ({temp} K)" if len(temperatures) > 1 else ""
            sys_prefix = f"{sys_name} " if not single_system else ""

            for mechanism in mechanisms:
                if system[mechanism]:
                    fields, rates = system[mechanism]
                    label = f"{sys_prefix}{mechanism.capitalize()}{lbl_suffix}"
                    ax.plot(fields, rates, label=label)
            
            if show_both and system['raman'] and system['orbach']:
                fields, raman_rates = system['raman']
                _, orbach_rates = system['orbach']
                ax.plot(fields, raman_rates, label=f"{sys_prefix}Raman{lbl_suffix}")
                ax.plot(fields, orbach_rates, label=f"{sys_prefix}Orbach{lbl_suffix}")

    ax.set_xlabel("Magnetic Field / T")
    ax.set_ylabel(r"$\tau^{-1}$  / s$^{-1}$")
    
    # Parse scale string (e.g. log-lin, log-log)
    if '-' in scale:
        x_s, y_s = scale.split('-')
    else:
        # Fallback for simple "log" or "linear"
        x_s = scale
        y_s = scale
        
    # Map 'lin' to 'linear' for matplotlib
    x_s = 'linear' if x_s == 'lin' else x_s
    y_s = 'linear' if y_s == 'lin' else y_s
        
    ax.set_xscale(x_s)
    ax.set_yscale(y_s)
    
    if xlim is None:
        if 'log' in x_s: ax.set_xlim(0.01, 7)
        else: ax.set_xlim(0, 7)
    else: ax.set_xlim(xlim)
        
    if ylim is None: ax.set_ylim(1e-5, 1e5)
    else: ax.set_ylim(ylim)

    ax.legend()
    plt.tight_layout()
    
    if not no_plot:
        outfile = _resolve_filename(f"{output_prefix}_rate_vs_field", ".pdf")
        plt.savefig(outfile)
        print(f"Saved plot to {outfile}")
    
    if show:
        plt.show()

def plot_orbach_spectral_density(output_prefix, args):
    """
    Calculates and plots the Orbach spectral density (Frobenius norm squared of the coupling matrix,
    deconvoluted by the phonon lineshape) for each phonon mode.

    J(hbar*omega) = Sum_q ( Frobenius_Norm_q^2 * Lineshape(hbar*omega - energy_q) )

    Args:
        output_prefix (str): Prefix for the output plot files.
        args (argparse.Namespace): Parsed command-line arguments.
    """
    print_calculation_parameters(args)
    _apply_plot_style()
    
    input_file = args.input_file
    kramers = args.kramers
    max_state = args.max_state
    max_modes = args.max_modes
    
    # Extract lineshape parameters with defaults if not present
    lineshape_name = getattr(args, 'lineshape', 'antilorentzian')
    fwhm = getattr(args, 'fwhm', 10.0)
    
    data = load_initial_data(input_file, kramers, max_state, max_modes)
    
    couplings = data['symmetrised_couplings']
    mode_energies = data['all_mode_energies']
    
    # --- 1. Calculate Coupling Strengths (Squared Frobenius Norm) ---
    coupling_strengths_sq = []
    phonon_energies = []
    
    sorted_indices = sorted(couplings.keys())
    
    if args.dos:
        plot_type="phonon DOS"
        # Phonon DOS is equal to Orbach J, with coupling strength set to 1
        coupling_strengths_sq = [1] * len(sorted_indices)
        phonon_energies = mode_energies[sorted_indices]
    else:
        plot_type="Orbach spectral density"
        for idx in sorted_indices:
            matrix = couplings[idx]
            # Frobenius norm squared: sum(|element|^2)
            norm_sq = np.sum(np.abs(matrix)**2)
            coupling_strengths_sq.append(norm_sq)
            phonon_energies.append(mode_energies[idx])
    
    phonon_energies = np.array(phonon_energies)
    coupling_strengths_sq = np.array(coupling_strengths_sq)
    
    # --- 2. Setup Energy Grid for Spectral Density ---
    # Determine range
    erange = getattr(args, 'erange', None)
    if erange is None:
        # Default range: 0 to max phonon energy + buffer
        max_e = np.max(phonon_energies) if len(phonon_energies) > 0 else 100.0
        # Buffer usually depends on width (FWHM). 
        # Using 5*FWHM as a reasonable tail buffer.
        buffer = 5.0 * fwhm 
        erange = [0.0, max_e + buffer]
        
    grid_resolution = getattr(args, 'grid_resolution', 1.0) # Default 1 cm-1 if not specified
    
    num_points = int((erange[1] - erange[0]) / grid_resolution) + 1
    energy_grid = np.linspace(erange[0], erange[1], num_points)
    
    # --- 3. Convolve with Lineshape ---
    spectral_density = np.zeros_like(energy_grid)
    lineshape_func = get_lineshape_function(lineshape_name, fwhm)
            
    print(f"Calculating {plot_type} on grid [{erange[0]:.1f}, {erange[1]:.1f}] cm-1 "
          f"with resolution {grid_resolution:.2f} cm-1...")
    
    # We sum over all modes q
    for i, e_q in enumerate(phonon_energies):
        strength = coupling_strengths_sq[i]
        # Evaluate lineshape centered at e_q on the whole grid
        # lineshape_func returns the density value at x for a mode at mu
        profile = lineshape_func(energy_grid, e_q)
        spectral_density += strength * profile
        
    # --- 4. Save and Plot ---
    
    # Save CSV
    if args.save_csv:
        csv_filename = _resolve_filename(f"{output_prefix}_{plot_type.replace(' ', '_')}", ".csv")
        if args.dos: 
            header='Energy(cm-1),Phonon DOS'
        else: 
            header='Energy(cm-1),SpectralDensity(cm-1)'
        np.savetxt(csv_filename, np.column_stack((energy_grid, spectral_density)), 
                   delimiter=',', header=header, comments='')
        print(f"Saved {plot_type} data to {csv_filename}")

    # Plot
    fig, ax = plt.subplots(figsize=(10, 6))
    
    # Plot the continuous spectral density
    ax.plot(energy_grid, spectral_density, color='k', lw=1.5, label=plot_type)
        
    ax.set_xlabel(r'Phonon Energy (cm$^{-1}$)', fontsize=10)
    if args.dos:
        ax.set_ylabel(r'Phonon DOS', fontsize=10)
    else:
        ax.set_ylabel(r'Spectral Density $J(\hbar\omega)$ (cm$^{-2}$)', fontsize=10)
    ax.set_title(plot_type, fontsize=10)
    
    # Manual Limits
    xlim = getattr(args, 'xlim', None)
    ylim = getattr(args, 'ylim', None)
    
    if xlim: ax.set_xlim(xlim)
    if ylim: ax.set_ylim(ylim)
    
    no_plot = getattr(args, 'no_plot', False)
    show = getattr(args, 'show', False)
    
    if not no_plot:
        outfile = _resolve_filename(f"{output_prefix}_{plot_type.replace(' ', '_')}", ".pdf")
        fig.savefig(outfile, bbox_inches='tight', dpi=300)
        print(f"Saved plot to {outfile}")
    
    if show:
        plt.show()
    
    plt.close(fig)