import argparse
import argcomplete
import sys
import numpy as np
import h5py

from .lineshapes import get_integration_width

class AdvancedHelpAction(argparse.Action):
    """
    A custom argparse action to allow us to hide advanced/debug options behind a
    secondary help menu, --advanced_help.
    """
    def __init__(self, option_strings, dest, advanced_help_data, **kwargs):
        self.advanced_help_data = advanced_help_data
        super().__init__(option_strings, dest, nargs=0, **kwargs)

    def __call__(self, parser, namespace, values, option_string=None):
        print("Advanced Calculation Options:")
        for options, help_text in self.advanced_help_data.items():
            print(f"  {options:<30} {help_text}")
        sys.exit(0)

def _add_calculation_setup_args(parser):
    """Adds calculation setup arguments to a parser."""
    calc_group = parser.add_argument_group('Calculation Setup')
    calc_group.add_argument(
        'input_file', type=str,
        help="Path to the input tau.hdf5 file produced by 'spin_phonon_suite prep tau_direct'."
    )
    calc_group.add_argument(
        '-s', '--max_state', type=int, default=None,
        help="Number of electronic states to include in the calculation. For Raman, this is the number of intermediate states (c). Default: all states."
    )
    calc_group.add_argument(
        '--states', type=int, nargs='+',
        help="List of electronic states to include in the calculation (0-indexed). For Raman, this is the number of intermediate states (c)."
    )
    calc_group.add_argument(
        '-m', '--max_modes', type=int,
        help="Number of phonon modes to include, starting from the lowest energy."
    )
    calc_group.add_argument(
        '--modes', type=int, nargs='+',
        help="List of phonon modes to include (1 indexed, as in the tau.hdf5)."
    )
    calc_group.add_argument(
        '-c', '--n_cores', type=int, default=1,
        help="Number of CPU cores to use. (Default: 1)"
    )
    calc_group.add_argument(
        '--dashboard', action='store_true', default=False,
        help="Start the Dask dashboard, which is useful for debugging/calculation monitoring."
    )

    kramers_group = calc_group.add_mutually_exclusive_group(required=True)
    kramers_group.add_argument(
        '--kramers', action='store_true', dest='kramers',
        help="Kramers system (odd number of electrons, doubly degenerate states at zero field)."
    )
    kramers_group.add_argument(
        '--non_kramers', action='store_false', dest='kramers',
        help='Non-Kramers system.'
    )

def _add_lineshape_args(parser, default_integration_width):
    """Adds phonon lineshape arguments to a parser."""
    lineshape_group = parser.add_argument_group('Phonon Lineshape')
    lineshape_group.add_argument(
        '-f', '--fwhm', type=float, default=10.0,
        help="Full-width half-maximum for the phonon lineshape in cm-1. (Default: 10.0 cm-1)"
    )
    lineshape_group.add_argument(
        '-l', '--lineshape', type=str, default='antilorentzian',
        choices=['antilorentzian', 'lorentzian', 'gaussian'],
        help="The phonon lineshape function. (Default: antilorentzian)"
    )
    lineshape_group.add_argument(
        '-i', '--integration_width', type=float, default=default_integration_width,
        help="Number of standard devations (equiv) for integration. "
             "0 integrates over the full energy range. (Default: %(default)s)"
    )

def _add_zeeman_args(parser):
    """Adds magnetic field arguments to a parser."""
    zeeman_group = parser.add_argument_group('Magnetic field')
    zeeman_group.add_argument(
        '-d', '--field_direction', type=float, nargs=3, metavar=('X', 'Y', 'Z'),
        help="A vector defining the direction of the applied Zeeman field. If --field magnitudes is given, this vector is normalised."
             "If --field-magnitudes is not given, the magnitude of this vector is used."
    )
    zeeman_group.add_argument(
        '-B', '--field_magnitudes', type=float, nargs='+', default=[0.0002],
        help="A list of field magnitudes (T) over which to calculate rates. "
             "If given, normalises and uses the direction from --zeeman-field (or the default axis). (Default: 2 Oe = 0.0002 T)"
    )
    zeeman_group.add_argument(
        '--no_reorient', dest='reorient',action='store_false',
        help="Apply the field vector in the input coordinate frame,"
             "instead of reorienting to the principal magnetic axis of the ground doublet, which is the default behaviour."
    )
    powder_avg_group = parser.add_argument_group('Powder Average')
    powder_avg_group.add_argument(
        '--orientations', type=int, default=1,
        help="Number of orientations for powder averaging. (Default: 1, for a single orientation)"
    )
    powder_avg_group.add_argument(
        '--quadrature', type=str, default='fibonacci',
        choices=['fibonacci'], # Add more choices here, e.g., 'lebedev'
        help="Quadrature scheme for generating orientation vectors. (Default: fibonacci)"
    )

    powder_avg_group.add_argument(
        '--serial_orientations', action='store_true', dest='serial_orientations', default=None,
        help="Calculate multiple field orientations serially, rather than in parallel."
    )
    powder_avg_group.add_argument(
        '--parallel_orientations', action='store_false', dest='serial_orientations',
        help="Calculate multiple field orientations in parallel. By default, this is set based on the number of field magnitudes"
             " and orientations."
    )

def _add_post_processing_args(parser):
    """Adds post-processing arguments to a parser."""
    post_group = parser.add_argument_group('Post-processing')
    post_group.add_argument(
        '-o', '--output', type=str,
        help="Output file name for the HDF5 archive. (Default: [input_file]_out.hdf5)"
    )
    post_group.add_argument(
        '-T', '--temperatures', type=float, nargs='+',
        help="A list of temperatures (K) at which to calculate relaxation rates."
    )
    post_group.add_argument(
        '--output_csv', type=str, default=None,
        help="Filename for the output CSV file containing rates. (Default: [output h5].csv)"
    )
    post_group.add_argument(
        '--save_gamma', action='store_true',
        help="Save the full Gamma rate matrices to the output HDF5 file."
    )
    post_group.add_argument(
        '--save_detbal', action='store_true',
        help="Save the detailed balance check matrices to the output HDF5 file."
    )

def parse_cli_args():
    """
    Parses command-line arguments for the program.

    Returns:
        argparse.Namespace: An object containing the parsed command-line args.
    """
    parser = argparse.ArgumentParser(
        description="Calculate spin-lattice relaxation rates."
    )

    # --- Mode Selection ---
    subparsers = parser.add_subparsers(
        dest='mode', required=True,
        help="The operational mode of Taupy."
    )

    # --- 'raman' Mode Parser ---
    raman_example = """Example:
  tau2 raman tau.hdf5 --kramers --fwhm 1.0 -T {1..120} -o tau2_raman.hdf5
"""
    parser_calc = subparsers.add_parser(
        'raman',
        help="Calculate the raman spectral density, and optionally calculate rates using this.",
        epilog=raman_example,
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser_calc.set_defaults(func=None)

    # Add shared argument groups
    _add_calculation_setup_args(parser_calc)
    _add_lineshape_args(parser_calc, default_integration_width=2)
    _add_zeeman_args(parser_calc)
    _add_post_processing_args(parser_calc)

    # --- Arguments specific to 'raman' mode ---
    raman_group = parser_calc.add_argument_group('Raman Options')
    raman_group.add_argument(
        '--raman_max_state', type=int, default=2,
        help="Number of inital/final states to include in the calculation. (Default: 2, i.e between state 0 and 1) "
    )
    raman_group.add_argument(
        '--raman_states', type=int, nargs='+',
        help="List of inital/final states to include in the calculation (0 indexed, within states included in --states if set)."
    )

    raman_group.add_argument(
        '--n_chunks', type=int, default=None,
        help="Number of chunks to divide phonon modes, per field. By default, this is set automatically based on the size of the calculation (number of fields, number of modes)."
    )

    raman_group.add_argument(
        '--rate_integrator', type=str, default='trapz',
        choices=['trapz', 'simpson'],
        help="The numerical integrator for rate calculation. 'trapz' (default) is stable; "
             "'simpson' is more accurate for smooth data. (Default: trapz)"
    )

    grid_group = parser_calc.add_argument_group('Integration Grid')
    grid_group.add_argument(
        '-e', '--erange', type=float, nargs=2,
        help="Energy range [min, max] for the integration grid. Can be "
             "inferred from --max-modes."
    )

    grid_group.add_argument(
        '-r', '--grid_resolution', type=float, default=None,
        help="Resolution (in cm-1) for the uniform grid the Raman spectral density is calculated on. Default: FWHM / 10"
    )
    grid_group.add_argument(
        '-g', '--grid_points', type=int,
        help="Number of points for the uniform global integration grid. (Default: 200)"
    )
    grid_group.add_argument(
        '--grid_variable', type=str, default='omega1',
        choices=['omega1', 'omega2', 'variable'],
        help="Strategy for choosing the integration variable. 'omega1' is default. "
             "'omega2' swaps the roles of the two phonons. 'variable' chooses based on the process. "
             "(Default: omega1)"
    )
    grid_group.add_argument(
        '--window_type', type=str, default='union',
        choices=['union', 'intersection', 'phononq'],
        help="Strategy for choosing the phonon integration window. 'union' is default. "
             "'union' applies a window around both phonon q and phonon r. The integration window is "
             "the union of these windows (i.e. both windows, combining them if they overlap). "
             "'intersection' uses the intersection of these two windows. 'phononq' checks whether "
             "windows on phonon q and phonon r overlap, and if they do uses only the phonon q window "
             "(emulating the behaviour of Tau)"
    )

    debug_group = parser_calc.add_argument_group('Debugging')
    debug_group.add_argument(
        '--plot_transition', type=int, nargs=2,
        help="Generate 1D plots for the specified transition (e.g., 0 1). "
             "Only available for single field calculations."
    )
    debug_group.add_argument(
        '--debug_summary', action='store_true',
        help="Save a summary CSV and print a console summary of the most "
             "significant phonon pair contributions for each field."
    )
    debug_group.add_argument(
        '--debug_verbose', action='store_true',
        help="Save detailed CSVs showing the integrand for the most significant phonon pairs."
    )
    debug_group.add_argument(
        '--debug_top_n', type=int, default=5,
        help="The number of top contributing pairs to show in the summary or save in verbose debug mode. (Default: 5)"
    )
    debug_group.add_argument(
        '--profile', action='store_true',
        help="Enable profiling and save statistics to 'profile_stats.prof'."
    )

    # --- 'raman_rates' Mode Parser ---
    parser_raman_rates = subparsers.add_parser(
        'raman_rates',
        help="Calculate relaxation rates from a saved Taupy HDF5 file containing the Raman spectral density."
    )
    parser_raman_rates.add_argument(
        'input_file', type=str,
        help="Path to the input tau.hdf5 file to process."
    )
    parser_raman_rates.add_argument(
        '-o', '--output', type=str,
        help="Output file name for the rates CSV file. (Default: [input_file]_rates.csv)"
    )
    parser_raman_rates.add_argument(
        '-T', '--temperatures', type=float, nargs='+', required=True,
        help="A list of temperatures (K) at which to calculate rates."
    )
    parser_raman_rates.add_argument(
        '--rate_integrator', type=str, default='trapz',
        choices=['trapz', 'simpson'],
        help="The numerical integrator for the final rate calculation. (Default: trapz, which is more stable)"
    )
    parser_raman_rates.set_defaults(func=None)


    # --- 'orbach' Mode Parser ---
    orbach_example = """Example:
  tau2 orbach tau.hdf5 --kramers --fwhm 1.0 -T {1..120} -o tau2_orbach.hdf5
"""
    parser_orbach = subparsers.add_parser(
        'orbach',
        help="Calculate Orbach relaxation rates.",
        epilog=orbach_example,
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser_orbach.set_defaults(func=None)

    # Add shared argument groups
    _add_calculation_setup_args(parser_orbach)
    _add_lineshape_args(parser_orbach, default_integration_width=0)
    _add_zeeman_args(parser_orbach)
    _add_post_processing_args(parser_orbach)

    orbach_group = parser_orbach.add_argument_group('Orbach Options')
    orbach_group.add_argument(
        '--gamma_diagonalisation', type=str,
        choices=['arb','mpmath','deflation'], default='arb',
        help="Technique used to diagonalise gamma matrix. By default, 'arb': high precision diagonalisation using 'arb', which does arbitrary-precision ball arithmetic. 'mpmath' is a slower python-only alternative. Alternatively, 'deflation': the gamma matrix is symmetrised and deflated to remove the approx. zero 'slow tau' eigenvalue. This allows for diagonalisation at double precision. 'deflation' is fast, but may fail for well isolated ground states below about 5K (such as non-Kramers systems, or Kramers systems in large magnetic fields)."
    )
    orbach_group.add_argument(
        '--gamma_precision', type=int, default=113,
        help="For 'arb' and 'mpmath' gamma diagonalisation techniques, set the precision, in bits. 113, the default, corresponds to quad precision, 53 is double. Arbitrary precision is supported."
    )

    # --- 'hyst' Mode Parser ---
    parser_hyst = subparsers.add_parser(
        'hyst',
        help="Simulate magnetic hysteresis using pre-calculated rate data."
    )
    parser_hyst.set_defaults(func=None)

    parser_hyst.add_argument(
        'input_files', type=str, nargs='+',
        help='One or more HDF5 files containing rate data from `tau2 raman/orbach`.'
    )
    parser_hyst.add_argument(
        '-o', '--output', type=str, default='hysteresis',
        help='Prefix for output .csv and .pdf files. (Default: hysteresis)'
    )
    parser_hyst.add_argument(
        '-T', '--temperatures', nargs='+', type=float, required=True,
        help='Temperatures to run hysteresis simulation at (K).'
    )
    parser_hyst.add_argument(
        '--init_field', nargs='+', type=float, default=[7],
        help='Initial field to sweep from, to -field (T). Default: 7T.'
    )
    parser_hyst.add_argument(
        '--sweep_rate', nargs='+', type=float, default=[22/10000],
        help='Sweep rate (T/s). Default: 22/10000 T/s = 22 Oe/s. Optionally, a list for variable sweep rates.'
    )
    parser_hyst.add_argument(
        '--time_step', type=float, default=1,
        help='Time step over which rates are propagated (s). Default: 1 s.'
    )
    parser_hyst.add_argument(
        '--population_propagation', type=str, default='implicit-euler',
        choices=['fixed', 'implicit-euler'],
        help="Integration method for population dynamics. (Default: implicit-euler)"
    )
    parser_hyst.add_argument(
        '-c', '--n_cores', type=int, default=1,
        help='Number of threads used to parallelise simulation over field orientations. (Default: 1)'
    )
    parser_hyst.add_argument(
        '--save_csv', action=argparse.BooleanOptionalAction, default=True,
        help='Save M/H csv file.'
    )
    parser_hyst.add_argument(
        '--save_fig', action=argparse.BooleanOptionalAction, default=True,
        help='Save M/H figure.'
    )
    parser_hyst.add_argument(
        '--show_plots', action=argparse.BooleanOptionalAction, default=False,
        help='Display plots interactively instead of only saving to file.'
    )
    parser_hyst.add_argument(
        '--state_limit', type=int, default=4,
        help='Number of states to include in simulation. Only Kramers systems supported. (Default: 4)'
    )
    parser_hyst.add_argument(
        '--tau_dynamics', type=str, nargs='+', default=None,
        choices=['all', 'raman', 'orbach'],
        help='Mechanism(s) of rates to include in simulation. (Default: all)'
    )
    # Plotting arguments
    plot_group = parser_hyst.add_argument_group('Plotting Options')
    plot_group.add_argument(
        '--xlim', type=float, nargs=2, default=None,
        help='X-axis limits for the main plot.'
    )
    plot_group.add_argument(
        '--ylim', type=float, nargs=2, default=None,
        help='Y-axis limits for the main plot.'
    )
    plot_group.add_argument(
        '--xlim-zoom', dest='xlim_zoom', type=float, nargs=2, default=None,
        help='X-axis limits for the zoomed plot (in Oe).'
    )
    plot_group.add_argument(
        '--ylim-zoom', dest='ylim_zoom', type=float, nargs=2, default=None,
        help='Y-axis limits for the zoomed plot (in mu_B x 10^-3).'
    )
    plot_group.add_argument(
        '--xaxis-units', dest='xaxis_units', type=str, default='T',
        choices=['T', 'Oe'],
        help="Units for the main plot x-axis. (Default: T)"
    )

    # --- 'plot_raman_j' Mode Parser ---
    parser_plot = subparsers.add_parser(
        'plot_raman_j',
        help="Plot a 1D spectral density from a Taupy HDF5 file containing the Raman spectral density."
    )
    parser_plot.add_argument(
        'input_files', type=str, nargs='+',
        help="Path(s) to the input hdf5 file(s) to plot, containing the Raman spectral density calculated by 'tau2 raman'"
    )
    parser_plot.add_argument(
        '-o', '--output', type=str,
        help="Output prefix for the plot file(s). (Default: based on input file name(s))"
    )
    parser_plot.add_argument(
        '-t', '--transition', type=int, nargs=2, required=True,
        help="The transition to plot (e.g., 0 1)."
    )
    parser_plot.add_argument(
        '--process', type=str, nargs='+', default=None,
        choices=['pp', 'mm', 'pm', 'mp'],
        help="The Raman process(es) to plot; 'pp', 'mm', 'pm', or 'mp'. The default, 'total', plots the sum of all processes."
    )
    # Field Selection - Mutually Exclusive
    field_group = parser_plot.add_mutually_exclusive_group()
    field_group.add_argument(
        '--field_index', type=int, nargs='+', default=None,
        help="The index/indices of the magnetic field(s) to plot. Incompatible with --fields. (Default: 0 if neither set)"
    )
    field_group.add_argument(
        '--fields', type=float, nargs='+', default=None,
        help="The value(s) of the magnetic field(s) to plot in Tesla. Incompatible with --field_index."
    )
    parser_plot.set_defaults(func=None)

    # --- 'plot_raman_j_2d' Mode Parser ---
    parser_plot2d = subparsers.add_parser(
        'plot_raman_j_2d',
        help="Plot a 2D spectral density heatmap from a Taupy HDF5 file containing the Raman spectral density."
    )
    parser_plot2d.add_argument(
        'input_file', type=str,
        help="Path to the input tau.hdf5 file to plot."
    )
    parser_plot2d.add_argument(
        '-o', '--output', type=str,
        help="Output prefix for the plot file. (Default: based on input file name)"
    )
    parser_plot2d.add_argument(
        '-t', '--transition', type=int, nargs=2, required=True,
        help="The transition to plot (e.g., 0 1)."
    )
    parser_plot2d.add_argument(
        '--plot_resolution', type=float, default=0.5,
        help="Resolution (in cm-1) for the energy axis of the 2D plot. (Default: 0.5)"
    )
    parser_plot2d.add_argument(
        '--process', type=str, nargs='+', default=['total'],
        choices=['pp', 'mm', 'pm', 'mp', 'total'],
        help="The Raman process(es) to plot; 'pp', 'mm', 'pm', or 'mp'. The default, 'total', plots the sum of all processes."
    )
    parser_plot2d.set_defaults(func=None)

    # --- 'plot_rate' Mode Parser ---
    parser_plotrate = subparsers.add_parser(
        'plot_rate',
        help="Plot relaxation rates from an HDF5 or CSV file."
    )
    parser_plotrate.add_argument(
        '--raman', type=str, nargs='+',
        help="Path(s) to the input hdf5 or csv file(s) containing Raman rates."
    )
    parser_plotrate.add_argument(
        '--orbach', type=str, nargs='+',
        help="Path(s) to the input hdf5 or csv file(s) containing Orbach rates."
    )
    parser_plotrate.add_argument(
        '-o', '--output', type=str,
        help="Output prefix for the plot file(s). (Default: based on input file name(s))"
    )
    parser_plotrate.add_argument(
        '--plot_type', type=str, default='rate-vs-temp',
        choices=['rate-vs-temp', 'tau-vs-inv-temp', 'rate-vs-field'],
        help="The type of plot to generate. (Default: rate-vs-temp)"
    )
    parser_plotrate.add_argument(
        '--scale', type=str, default='log-log',
        choices=['log-log', 'log-lin', 'lin-log', 'lin-lin'],
        help="The scale of the plot. (Default: log-log)"
    )
    parser_plotrate.add_argument(
        '--orientation', type=int, default=0,
        help="The index of the orientation to plot. (Default: 0). Ignored if --average is set."
    )
    parser_plotrate.add_argument(
        '--average', action='store_true',
        help="Plot the powder averaged rates (weighted sum over all orientations). Overrides --orientation."
    )
    
    # Field Selection - Mutually Exclusive Group
    field_group_rate = parser_plotrate.add_mutually_exclusive_group()
    field_group_rate.add_argument(
        '--field_index', type=int, nargs='+', default=None,
        help="The index/indices of the magnetic field(s) to plot rates for. (Default: 0 if neither set)"
    )
    field_group_rate.add_argument(
        '--fields', type=float, nargs='+', default=None,
        help="The value(s) of the magnetic field(s) to plot rates for in Tesla."
    )
    parser_plotrate.add_argument(
        '--mechanisms', type=str, nargs='+', default=['total'],
        choices=['raman', 'orbach', 'total'],
        help="The relaxation mechanism(s) to plot. (Default: total)"
    )
    parser_plotrate.add_argument(
        '--show-both', action='store_true',
        help="Show both Raman and Orbach rates on the same plot."
    )
    parser_plotrate.add_argument(
        '--labels', nargs='+', default=None,
        help="Custom labels for the systems being plotted. Defaults to 'System 1', 'System 2', etc."
    )
    parser_plotrate.add_argument(
        '--xlim', type=float, nargs=2, default=None,
        help='X-axis limits for the plot.'
    )
    parser_plotrate.add_argument(
        '--ylim', type=float, nargs=2, default=None,
        help='Y-axis limits for the plot.'
    )
    parser_plotrate.add_argument(
        '--save_csv', action='store_true',
        help="Save the plotted data to a CSV file."
    )
    parser_plotrate.add_argument(
        '--no_plot', action='store_true',
        help="Do not save the plot to a file."
    )
    parser_plotrate.add_argument(
        '--show', action='store_true',
        help="Display the plot interactively."
    )
    parser_plotrate.add_argument(
        '-T', '--temperatures', type=float, nargs='+', default=None,
        help="Temperature(s) (K) for rate-vs-field plots. If not specified, defaults to 2 K."
    )
    parser_plotrate.set_defaults(func=None)

    # --- 'orbach_spectral_density' Mode Parser ---
    parser_orbach_spectral = subparsers.add_parser(
        'plot_orbach_j',
        help="Calculate and plot the Orbach spectral density from a Tau HDF5 file."
    )
    parser_orbach_spectral.add_argument(
        'input_file', type=str,
        help="Path to the input Tau HDF5 file (from 'spin_phonon_suite prep tau_direct')."
    )
    # Reuse setup args for loading correct states/modes
    calc_group_orbach_spec = parser_orbach_spectral.add_argument_group('Calculation Setup')
    calc_group_orbach_spec.add_argument(
        '-s', '--max_state', type=int, default=None,
        help="Number of electronic states to include in the calculation. (Default: all states)"
    )
    calc_group_orbach_spec.add_argument(
        '-m', '--max_modes', type=int,
        help="Number of phonon modes to include."
    )
    kramers_group_orbach_spec = calc_group_orbach_spec.add_mutually_exclusive_group(required=True)
    kramers_group_orbach_spec.add_argument(
        '--kramers', action='store_true', dest='kramers',
        help="Kramers system (odd number of electrons)."
    )
    kramers_group_orbach_spec.add_argument(
        '--non_kramers', action='store_false', dest='kramers',
        help='Non-Kramers system.'
    )
    
    # Reuse Lineshape args for convolution
    _add_lineshape_args(parser_orbach_spectral, default_integration_width=0)
    
    # Add Grid/Plotting Args
    grid_group_orbach = parser_orbach_spectral.add_argument_group('Grid & Plotting')
    grid_group_orbach.add_argument(
        '-e', '--erange', type=float, nargs=2,
        help="Energy range [min, max] for the spectral density grid."
    )
    grid_group_orbach.add_argument(
        '-r', '--grid_resolution', type=float, default=1.0,
        help="Resolution (in cm-1) for the spectral density grid."
    )
    grid_group_orbach.add_argument(
        '-o', '--output', type=str,
        help="Output prefix for the plot/csv file(s)."
    )
    grid_group_orbach.add_argument(
        '--xlim', type=float, nargs=2, default=None,
        help='X-axis limits for the plot.'
    )
    grid_group_orbach.add_argument(
        '--ylim', type=float, nargs=2, default=None,
        help='Y-axis limits for the plot.'
    )
    grid_group_orbach.add_argument(
        '--save_csv', action='store_true',
        help="Save the calculated spectral density to a CSV file."
    )
    grid_group_orbach.add_argument(
        '--no_plot', action='store_true',
        help="Do not save the plot to a file."
    )
    grid_group_orbach.add_argument(
        '--show', action='store_true',
        help="Display the plot interactively."
    )
    calc_group_orbach_spec.add_argument(
        '--dos', action='store_true',
        help="Plot the phonon DOS instead of the Orbach spectral density."
    )

    parser_orbach_spectral.set_defaults(func=None)

    # --- Raman Advanced Help ---
    calc_advanced_option_strings = [
        '--grid_variable',
        '--plot_transition', 
        '--debug_summary', 
        '--debug_verbose',
        '--debug_top_n', 
        '--profile',
        '--n_chunks',
        '--no_reorient',
        '--quadrature',
        '--serial_orientations', 
        '--parallel_orientations',
        '--save_gamma', 
        '--save_detbal',
        '--rate_integrator',
        '--grid_points',
        '--integration_width',
        '--states',
        '--raman_states',
        '--window_type',
        '--dashboard',
    ]

    calc_advanced_help_dict = {}
    for action in parser_calc._actions:
        if any(opt in calc_advanced_option_strings for opt in action.option_strings):
            key = ', '.join(action.option_strings)
            # Manually interpolate %(default)s
            help_text = action.help
            if help_text and '%(default)' in help_text:
                help_text = help_text % {'default': action.default}
            calc_advanced_help_dict[key] = help_text
            action.help = argparse.SUPPRESS
    parser_calc.add_argument(
        '--advanced_help',
        action=AdvancedHelpAction,
        advanced_help_data=calc_advanced_help_dict,
        help='Show debugging options and advanced options which should not normally need to be set.'
    )

    # --- Orbach Advanced Help ---
    orbach_advanced_option_strings = [
        '--plot_transition', 
        '--debug_summary', 
        '--debug_verbose',
        '--debug_top_n', 
        '--profile',
        '--no_reorient',
        '--quadrature',
        '--serial_orientations', 
        '--parallel_orientations',
        '--save_gamma', 
        '--save_detbal',
        '--integration_width',
        '--gamma_diagonalisation', 
        '--gamma_precision',
        '--states',
        '--dashboard',
    ]

    orbach_advanced_help_dict = {}

    for action in parser_orbach._actions:
        if any(opt in orbach_advanced_option_strings for opt in action.option_strings):
            key = ', '.join(action.option_strings)
            # Manually interpolate %(default)s
            help_text = action.help
            if help_text and '%(default)' in help_text:
                help_text = help_text % {'default': action.default}
            orbach_advanced_help_dict[key] = help_text
            action.help = argparse.SUPPRESS

    parser_orbach.add_argument(
        '--advanced_help',
        action=AdvancedHelpAction,
        advanced_help_data=orbach_advanced_help_dict,
        help='Show debugging options and advanced options which should not normally need to be set.'
    )

    argcomplete.autocomplete(parser)

    return parser.parse_args()


def prepare_and_validate_args(args, all_mode_energies, all_state_energies):
    """
    Validates and automatically sets interdependent arguments.
    """
    if args.mode not in ['raman', 'orbach']:
        return args

    # --- 1. State Selection ---
    # args.states is the master list of states to load (basis set).
    # We sort it immediately to match h5py loading order.
    if args.states:
        args.states_to_load = sorted(list(set(args.states)))
        args.states = args.states_to_load
        args.max_state = len(args.states_to_load)
    elif args.max_state is not None:
        args.states_to_load = list(range(args.max_state))
        args.states = args.states_to_load
    else:
        args.states_to_load = list(range(len(all_state_energies)))
        args.states = args.states_to_load
        args.max_state = len(all_state_energies)

    # Raman specific: Interpret --raman_states as INDICES into the loaded basis
    if args.mode == 'raman':
        if not args.raman_states:
            # Default to the first N states of the loaded set, where N is raman_max_state
            args.raman_states = list(range(args.raman_max_state))
        
        args.raman_max_state = len(args.raman_states)

    # --- 2. Mode Selection ---
    if args.modes is not None:
        args.max_modes = len(args.modes)
        # Validation
        idxs = [m - 1 for m in args.modes]
        if any(m < 0 or m >= len(all_mode_energies) for m in idxs):
             raise ValueError(f"Modes {args.modes} out of range.")
             
    elif args.max_modes is None:
        # Auto-detection
        if args.mode == 'orbach':
            args.max_modes = len(all_mode_energies)
            print(f"Orbach: defaulting to including all {args.max_modes} modes.")
        else:
            # Raman: Auto-detect based on electronic structure or erange
            width = get_integration_width(args.lineshape, args.fwhm, args.integration_width)
            if args.erange: 
                limit = args.erange[1]
            else:
                idx = 2 # First excited doublet/state
                limit = (all_state_energies[idx] if idx < len(all_state_energies) else all_state_energies[-1]) - width
                if limit <= 0: limit = 100.0
            
            args.max_modes = np.sum(all_mode_energies <= limit)
            if args.max_modes == 0: 
                raise ValueError("Could not set the number of modes automatically. " \
                        "Check phonon mode energies are correct, and set --max_modes "
                        "and --erange manually.")
            print(f"Auto-set max_modes={args.max_modes} based on: ")
            if args.kramers:
                print(f"  First excited Kramers doublet: {all_state_energies[2]:.2f} cm-1")
            else:
                print(f"  Second excited state: {all_state_energies[2]:.2f} cm-1")

            print(f"  and integration width {width:.2f} cm-1 = {args.integration_width} standard deviation (equivalents).\n")
            if any(field > 2/10000 for field in args.field_magnitudes):
                print(f"Warning: setting modes automatically might not work for finite field rate calculations!\n\n")

    # Orbach requires no further grid setup
    if args.mode == 'orbach':
        return args

    # --- 3. Raman Grid Setup ---
    width = get_integration_width(args.lineshape, args.fwhm, args.integration_width)
    
    # Infer erange if missing
    if args.erange is None:
        if width == np.inf:
            raise ValueError("Cannot use automatic range/mode determination with "
                    "--integration_width 0/-i 0.")
        if args.modes:
            max_e = np.max(all_mode_energies[[m-1 for m in args.modes]])
        else:
            max_e = all_mode_energies[args.max_modes - 1] if args.max_modes > 0 else 0.0
        args.erange = [0.0, max_e + width]
        print(f"Auto-set erange to [0.0, {args.erange[1]:.2f}] based on modes.\n")

    if args.grid_resolution is None:
        args.grid_resolution = args.fwhm / 10.0
        print(f"Setting grid resolution to {args.grid_resolution:.2f} cm-1: (FWHM = {args.fwhm:.2f} cm-1) / 10")
    
    args.grid_points = int((args.erange[1] - args.erange[0]) / args.grid_resolution) + 1 if (args.erange[1] > args.erange[0]) else 2

    return args
