#!/usr/bin/env python3
"""
Generate correlated-k opacity files from petitRADTRANS line-by-line files.

This script converts high-resolution line-by-line opacity files to correlated-k
format using the official petitRADTRANS correlated-k algorithm. This method is
more accurate than simple Gaussian convolution because it preserves the statistical
distribution of opacities within each spectral bin.

Usage:
    python generate_correlated_k_from_lbl.py --input-file <path> --resolution <R>

    Arguments:
        --input-file    Path to line-by-line opacity HDF5 file
        --resolution    Target resolving power (default: 20000)
        --no-plots      Skip diagnostic plot generation
        --rewrite       Overwrite existing output files

Author: Francesco / GUIBRUSHR Team
Date: 2026-01-14
"""

import h5py
import numpy as np
from pathlib import Path
from datetime import datetime
from tqdm import tqdm
import warnings
import argparse
import matplotlib
matplotlib.rcParams['agg.path.chunksize'] = 10000
import os
if not os.environ.get('SPHINX_BUILD'):
    matplotlib.use("TkAgg")
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from matplotlib.colors import LogNorm

# Try to import petitRADTRANS for version info
try:
    import petitRADTRANS
    PRT_VERSION = petitRADTRANS.__version__
except ImportError:
    PRT_VERSION = "unknown"
    warnings.warn("petitRADTRANS not imported, version will be set to 'unknown'")


# =============================================================================
# CONFIGURATION
# =============================================================================
def write_correlated_k(file, doi, wavenumbers, wavenumbers_bins_edges, cross_sections, mol_mass, species,
                       opacities_pressures, opacities_temperatures, g_gauss, weights_gauss,
                       wavelengths=None, n_g=None,
                       contributor=None, description=None):
    if wavelengths is None:
        wavelengths = np.array([1 / wavenumbers[0], 1 / wavenumbers[-1]])

    if n_g is None:
        n_g = g_gauss.size

    with h5py.File(file, "w") as fh5:
        dataset = fh5.create_dataset(
            name='DOI',
            shape=(1,),
            data=doi
        )
        dataset.attrs['long_name'] = 'Data object identifier linked to the data'
        dataset.attrs['contributor'] = str(contributor)
        dataset.attrs['additional_description'] = str(description)

        dataset = fh5.create_dataset(
            name='Date_ID',
            shape=(1,),
            data=f'petitRADTRANS-v{petitRADTRANS.__version__}'
                 f'_{datetime.datetime.now(datetime.timezone.utc).isoformat()}'
        )
        dataset.attrs['long_name'] = 'ISO 8601 UTC time (https://docs.python.org/3/library/datetime.html) ' \
                                     'at which the table has been created, ' \
                                     'along with the version of petitRADTRANS'

        dataset = fh5.create_dataset(
            name='bin_centers',
            data=wavenumbers
        )
        dataset.attrs['long_name'] = 'Centers of the wavenumber bins'
        dataset.attrs['units'] = 'cm^-1'

        dataset = fh5.create_dataset(
            name='bin_edges',
            data=wavenumbers_bins_edges
        )
        dataset.attrs['long_name'] = 'Separations between the wavenumber bins'
        dataset.attrs['units'] = 'cm^-1'

        dataset = fh5.create_dataset(
            name='kcoeff',
            data=cross_sections
        )
        dataset.attrs['long_name'] = ('Table of the k-coefficients with axes '
                                      '(pressure, temperature, wavenumber, g space)')
        dataset.attrs['units'] = 'cm^2/molecule'

        dataset = fh5.create_dataset(
            name='method',
            shape=(1,),
            data='petit_samples'
        )
        dataset.attrs['long_name'] = 'Name of the method used to sample g-space'

        dataset = fh5.create_dataset(
            name='mol_mass',
            shape=(1,),
            data=float(mol_mass)
        )
        dataset.attrs['long_name'] = 'Mass of the species'
        dataset.attrs['units'] = 'AMU'

        dataset = fh5.create_dataset(
            name='mol_name',
            shape=(1,),
            data=species.split('_', 1)[0]
        )
        dataset.attrs['long_name'] = 'Name of the species described'

        dataset = fh5.create_dataset(
            name='ngauss',
            data=n_g
        )
        dataset.attrs['long_name'] = 'Number of points used to sample the g-space'

        dataset = fh5.create_dataset(
            name='p',
            data=opacities_pressures
        )
        dataset.attrs['long_name'] = 'Pressure grid'
        dataset.attrs['units'] = 'bar'

        dataset = fh5.create_dataset(
            name='samples',
            data=g_gauss
        )
        dataset.attrs['long_name'] = 'Abscissas used to sample the k-coefficients in g-space'

        dataset = fh5.create_dataset(
            name='t',
            data=opacities_temperatures
        )
        dataset.attrs['long_name'] = 'Temperature grid'
        dataset.attrs['units'] = 'K'

        dataset = fh5.create_dataset(
            name='temperature_grid_type',
            shape=(1,),
            data='regular'
        )
        dataset.attrs['long_name'] = 'Whether the temperature grid is "regular" ' \
                                     '(same temperatures for all pressures) or "pressure-dependent"'

        dataset = fh5.create_dataset(
            name='weights',
            data=weights_gauss
        )
        dataset.attrs['long_name'] = 'Weights used in the g-space quadrature'

        dataset = fh5.create_dataset(
            name='wlrange',
            data=np.array([wavelengths.min(), wavelengths.max()]) * 1e4  # cm to um
        )
        dataset.attrs['long_name'] = 'Wavelength range covered'
        dataset.attrs['units'] = 'µm'

        dataset = fh5.create_dataset(
            name='wnrange',
            data=np.array([wavenumbers.min(), wavenumbers.max()])
        )
        dataset.attrs['long_name'] = 'Wavenumber range covered'
        dataset.attrs['units'] = 'cm^-1'

class Config:
    """Configuration parameters for correlated-k generation."""

    # Input file (line-by-line high-resolution)
    INPUT_FILE = "/home/francesco/petitRADTRANS/input_data/opacities/lines/line_by_line/H2O/1H2-16O/1H2-16O__POKAZATEL.R1e6_0.3-28mu.xsec.petitRADTRANS.h5"

    # Target resolving power for correlated-k
    R_TARGET = 20000

    # Correlated-k g-space samples (Gauss-Legendre quadrature)
    # Default: 16-point quadrature from petitRADTRANS
    G_SAMPLES = np.array([
        0.0178695646,
        0.0915000852,
        0.2135104155,
        0.3674544109,
        0.5325455891,
        0.6864895845,
        0.8084999148,
        0.8821304354,
        0.9019855072,
        0.9101666761,
        0.9237233795,
        0.9408282679,
        0.9591717321,
        0.9762766205,
        0.9898333239,
        0.9980144928
    ])

    # Corresponding weights for g-space integration
    G_WEIGHTS = np.array([
        0.0455528413,
        0.1000714655,
        0.1411679906,
        0.1632077025,
        0.1632077025,
        0.1411679906,
        0.1000714655,
        0.0455528413,
        0.0050614268,
        0.0111190517,
        0.0156853323,
        0.0181341892,
        0.0181341892,
        0.0156853323,
        0.0111190517,
        0.0050614268
    ])

    # Metadata for output file
    DOI = "Converted from line-by-line using correlated-k method"
    CONTRIBUTOR = "GUIBRUSHR Team"
    DESCRIPTION = f"Correlated-k opacities at R={R_TARGET} generated from high-resolution line-by-line data"

    # Use legacy mode for wavenumber sampling (pRT v2 compatibility)
    USE_LEGACY_MODE = False

    # Rewrite existing files
    REWRITE = True

    # Generate diagnostic plots
    MAKE_PLOTS = True


# =============================================================================
# I/O FUNCTIONS
# =============================================================================

def load_line_by_line_file(filepath: str) -> dict:
    """
    Load line-by-line opacity data from petitRADTRANS HDF5 file.

    Parameters
    ----------
    filepath : str
        Path to the input HDF5 file

    Returns
    -------
    dict
        Dictionary containing:
        - pressures: Pressure grid (bar)
        - temperatures: Temperature grid (K)
        - wavenumbers: Wavenumber bin edges (cm^-1)
        - cross_sections: Cross-section array (n_p, n_t, n_wn) in cm^2/molecule
        - mol_mass: Molecular mass (AMU)
        - mol_name: Molecule name
        - wlrange: Wavelength range (μm)
        - wnrange: Wavenumber range (cm^-1)
    """
    print(f"Loading line-by-line file: {filepath}")

    data = {}
    with h5py.File(filepath, 'r') as f:
        # Load required datasets
        data['pressures'] = f['p'][:]
        data['temperatures'] = f['t'][:]
        data['wavenumbers'] = f['bin_edges'][:]  # in cm^-1
        data['cross_sections'] = f['xsecarr'][:]  # (n_p, n_t, n_wn)

        # Load metadata
        data['mol_mass'] = float(f['mol_mass'][()])
        mol_name_raw = f['mol_name'][()]
        if isinstance(mol_name_raw, bytes):
            data['mol_name'] = mol_name_raw.decode('utf-8')
        elif isinstance(mol_name_raw, np.ndarray):
            data['mol_name'] = str(mol_name_raw[0])
            if isinstance(data['mol_name'], bytes):
                data['mol_name'] = data['mol_name'].decode('utf-8')
        else:
            data['mol_name'] = str(mol_name_raw)

        # Wavelength and wavenumber ranges
        if 'wlrange' in f:
            data['wlrange'] = f['wlrange'][:]
        else:
            # Compute from wavenumbers
            data['wlrange'] = np.array([1e4/data['wavenumbers'][-1], 1e4/data['wavenumbers'][0]])

        if 'wnrange' in f:
            data['wnrange'] = f['wnrange'][:]
        else:
            data['wnrange'] = np.array([data['wavenumbers'][0], data['wavenumbers'][-1]])

    print(f"  Molecule: {data['mol_name']}")
    print(f"  Molecular mass: {data['mol_mass']:.4f} AMU")
    print(f"  Wavenumber range: {data['wnrange'][0]:.2f} - {data['wnrange'][1]:.2f} cm^-1")
    print(f"  Wavelength range: {data['wlrange'][0]:.4f} - {data['wlrange'][1]:.4f} μm")
    print(f"  Pressure grid: {len(data['pressures'])} points")
    print(f"  Temperature grid: {len(data['temperatures'])} points")
    print(f"  Wavenumber grid: {len(data['wavenumbers'])} points")
    print(f"  Cross-section shape: {data['cross_sections'].shape}")

    return data


def write_correlated_k_file(
    filepath: str,
    wavenumbers_centers: np.ndarray,
    wavenumbers_edges: np.ndarray,
    cross_sections: np.ndarray,
    pressures: np.ndarray,
    temperatures: np.ndarray,
    mol_mass: float,
    mol_name: str,
    g_samples: np.ndarray,
    g_weights: np.ndarray,
    doi: str,
    contributor: str = None,
    description: str = None
):
    """
    Write correlated-k opacity file in petitRADTRANS format.

    This follows the exact structure of petitRADTRANS write_correlated_k function.

    Parameters
    ----------
    filepath : str
        Output file path
    wavenumbers_centers : np.ndarray
        Center wavenumbers of bins (cm^-1)
    wavenumbers_edges : np.ndarray
        Edge wavenumbers of bins (cm^-1)
    cross_sections : np.ndarray
        Correlated-k coefficients (n_p, n_t, n_wn, n_g) in cm^2/molecule
    pressures : np.ndarray
        Pressure grid (bar)
    temperatures : np.ndarray
        Temperature grid (K)
    mol_mass : float
        Molecular mass (AMU)
    mol_name : str
        Molecule name
    g_samples : np.ndarray
        G-space sample points
    g_weights : np.ndarray
        G-space quadrature weights
    doi : str
        Data identifier or DOI
    contributor : str, optional
        Contributor name
    description : str, optional
        Additional description
    """
    print(f"\nWriting correlated-k file: {filepath}")

    n_g = g_samples.size
    wavelengths = np.array([1e4 / wavenumbers_centers[-1], 1e4 / wavenumbers_centers[0]])

    with h5py.File(filepath, 'w') as fh5:
        # DOI
        dataset = fh5.create_dataset(
            name='DOI',
            shape=(1,),
            data=doi
        )
        dataset.attrs['long_name'] = 'Data object identifier linked to the data'
        dataset.attrs['contributor'] = str(contributor)
        dataset.attrs['additional_description'] = str(description)

        # Date and version
        dataset = fh5.create_dataset(
            name='Date_ID',
            shape=(1,),
            data=f'petitRADTRANS-v{PRT_VERSION}_{datetime.now().isoformat()}'
        )
        dataset.attrs['long_name'] = 'ISO 8601 UTC time at which the table has been created, along with the version of petitRADTRANS'

        # Bin centers
        dataset = fh5.create_dataset(
            name='bin_centers',
            data=wavenumbers_centers
        )
        dataset.attrs['long_name'] = 'Centers of the wavenumber bins'
        dataset.attrs['units'] = 'cm^-1'

        # Bin edges
        dataset = fh5.create_dataset(
            name='bin_edges',
            data=wavenumbers_edges
        )
        dataset.attrs['long_name'] = 'Separations between the wavenumber bins'
        dataset.attrs['units'] = 'cm^-1'

        # K-coefficients
        dataset = fh5.create_dataset(
            name='kcoeff',
            data=cross_sections
        )
        dataset.attrs['long_name'] = 'Table of the k-coefficients with axes (pressure, temperature, wavenumber, g space)'
        dataset.attrs['units'] = 'cm^2/molecule'

        # Method
        dataset = fh5.create_dataset(
            name='method',
            shape=(1,),
            data='petit_samples'
        )
        dataset.attrs['long_name'] = 'Name of the method used to sample g-space'

        # Molecular mass
        dataset = fh5.create_dataset(
            name='mol_mass',
            shape=(1,),
            data=float(mol_mass)
        )
        dataset.attrs['long_name'] = 'Mass of the species'
        dataset.attrs['units'] = 'AMU'

        # Molecular name
        dataset = fh5.create_dataset(
            name='mol_name',
            shape=(1,),
            data=mol_name.split('_', 1)[0]
        )
        dataset.attrs['long_name'] = 'Name of the species described'

        # Number of g points
        dataset = fh5.create_dataset(
            name='ngauss',
            data=n_g
        )
        dataset.attrs['long_name'] = 'Number of points used to sample the g-space'

        # Pressure grid
        dataset = fh5.create_dataset(
            name='p',
            data=pressures
        )
        dataset.attrs['long_name'] = 'Pressure grid'
        dataset.attrs['units'] = 'bar'

        # G-space samples
        dataset = fh5.create_dataset(
            name='samples',
            data=g_samples
        )
        dataset.attrs['long_name'] = 'Abscissas used to sample the k-coefficients in g-space'

        # Temperature grid
        dataset = fh5.create_dataset(
            name='t',
            data=temperatures
        )
        dataset.attrs['long_name'] = 'Temperature grid'
        dataset.attrs['units'] = 'K'

        # Temperature grid type
        dataset = fh5.create_dataset(
            name='temperature_grid_type',
            shape=(1,),
            data='regular'
        )
        dataset.attrs['long_name'] = 'Whether the temperature grid is "regular" (same temperatures for all pressures) or "pressure-dependent"'

        # G-space weights
        dataset = fh5.create_dataset(
            name='weights',
            data=g_weights
        )
        dataset.attrs['long_name'] = 'Weights used in the g-space quadrature'

        # Wavelength range
        dataset = fh5.create_dataset(
            name='wlrange',
            data=np.array([wavelengths.min(), wavelengths.max()]) * 1e4  # cm to um
        )
        dataset.attrs['long_name'] = 'Wavelength range covered'
        dataset.attrs['units'] = 'µm'

        # Wavenumber range
        dataset = fh5.create_dataset(
            name='wnrange',
            data=np.array([wavenumbers_centers.min(), wavenumbers_centers.max()])
        )
        dataset.attrs['long_name'] = 'Wavenumber range covered'
        dataset.attrs['units'] = 'cm^-1'

    print(f"  Successfully written correlated-k file")
    print(f"  Output shape: {cross_sections.shape}")


# =============================================================================
# CORRELATED-K ALGORITHM
# =============================================================================

def compute_resolving_power(wavenumbers: np.ndarray) -> float:
    """
    Compute the mean resolving power of a wavenumber grid.

    Parameters
    ----------
    wavenumbers : np.ndarray
        Wavenumber array (cm^-1)

    Returns
    -------
    float
        Mean resolving power R = ν/Δν
    """
    R = wavenumbers[:-1] / np.diff(wavenumbers) + 0.5
    return int(np.mean(R[np.isfinite(R)]))


def create_correlated_k_grid(
    wavenumbers_input: np.ndarray,
    R_input: float,
    R_target: float,
    use_legacy_mode: bool = False
) -> tuple[np.ndarray, np.ndarray]:
    """
    Create wavenumber grid for correlated-k at target resolution.

    This follows the exact algorithm from petitRADTRANS format2petitradtrans.

    Parameters
    ----------
    wavenumbers_input : np.ndarray
        Input wavenumber grid (cm^-1), must be in increasing order
    R_input : float
        Input resolving power
    R_target : float
        Target resolving power
    use_legacy_mode : bool
        Use legacy (pRT v2) wavenumber sampling

    Returns
    -------
    bin_edges : np.ndarray
        Wavenumber bin edges (cm^-1)
    bin_centers : np.ndarray
        Wavenumber bin centers (cm^-1)
    """
    print(f"\nCreating correlated-k wavenumber grid:")
    print(f"  Input resolution: R = {R_input:.0e}")
    print(f"  Target resolution: R = {R_target:,}")

    # Calculate downsampling factor
    downsampling = int(R_input / R_target)
    print(f"  Downsampling factor: {downsampling}")

    if downsampling > len(wavenumbers_input):
        raise ValueError(
            f"Unable to make correlated-k resolving power {R_target} "
            f"from source with resolving power {R_input:.0e}"
        )

    # Create bin edges by downsampling
    if use_legacy_mode:
        print("  Using legacy mode (pRT v2 compatibility)")
        starting_index = 999
    else:
        starting_index = 0

    # Downsample (reverse, downsample, reverse back)
    bin_edges = wavenumbers_input[::-1][starting_index::downsampling][::-1]

    if use_legacy_mode:
        # Add extra edges until we reach the maximum wavenumber
        max_extra_edges = 2
        for i in range(max_extra_edges):
            R_at_edge = bin_edges[-1] / np.diff(bin_edges[-2:]) - 1
            bin_edges = np.append(bin_edges, bin_edges[-1] * (1 + 1 / R_at_edge))

            if bin_edges[-1] > wavenumbers_input[-1]:
                warnings.warn(
                    f"Legacy maximum bin edges ({bin_edges[-1]} cm^-1) "
                    f"is greater than the opacity source maximum wavenumber ({wavenumbers_input[-1]} cm^-1)"
                )
                break
            elif bin_edges[-1] == wavenumbers_input[-1]:
                break

        if bin_edges[-1] < wavenumbers_input[-1]:
            raise ValueError(
                f"Could not reach the source maximum wavenumber ({wavenumbers_input[-1]} cm^-1) "
                f"after adding {max_extra_edges} extra edges ({bin_edges[-1]} cm^-1)"
            )

    # Calculate bin centers
    bin_centers = bin_edges[:-1] + np.diff(bin_edges) / 2

    print(f"  Number of bins: {len(bin_centers)}")
    print(f"  Wavenumber range: {bin_edges[0]:.2f} - {bin_edges[-1]:.2f} cm^-1")
    print(f"  Wavelength range: {1e4/bin_edges[-1]:.4f} - {1e4/bin_edges[0]:.4f} μm")

    return bin_edges, bin_centers


def compute_correlated_k_coefficients(
    cross_sections: np.ndarray,
    wavenumbers: np.ndarray,
    bin_edges: np.ndarray,
    g_samples: np.ndarray
) -> np.ndarray:
    """
    Compute correlated-k coefficients from line-by-line cross-sections.

    This implements the exact algorithm from petitRADTRANS format2petitradtrans.
    For each spectral bin:
    1. Select cross-sections in that bin
    2. Sort them in increasing order
    3. Divide into g-space intervals
    4. Average within each g interval

    Parameters
    ----------
    cross_sections : np.ndarray
        Line-by-line cross-sections (n_p, n_t, n_wn) in cm^2/molecule
    wavenumbers : np.ndarray
        Wavenumber grid (cm^-1)
    bin_edges : np.ndarray
        Correlated-k bin edges (cm^-1)
    g_samples : np.ndarray
        G-space sample points

    Returns
    -------
    np.ndarray
        Correlated-k coefficients (n_p, n_t, n_bins, n_g) in cm^2/molecule
    """
    n_p, n_t, n_wn = cross_sections.shape
    n_bins = len(bin_edges) - 1
    n_g = len(g_samples)

    print(f"\nComputing correlated-k coefficients:")
    print(f"  Input shape: {cross_sections.shape}")
    print(f"  Output shape: ({n_p}, {n_t}, {n_bins}, {n_g})")
    print(f"  G-space points: {n_g}")

    # Initialize output array
    correlated_k = np.zeros((n_p, n_t, n_bins, n_g), dtype=np.float64)

    # Define g-space bin edges
    samples_edges = np.zeros(n_g + 1)
    samples_edges[-1] = 1.0
    for i in range(n_g - 1):
        samples_edges[i + 1] = (g_samples[i + 1] + g_samples[i]) / 2

    # Minimum selection size to ensure enough points for g-space sampling
    min_selection_size = int(np.ceil(1 / np.min(np.diff(g_samples)))) * 2

    print(f"  Minimum selection size for g-space sampling: {min_selection_size}")
    print(f"  Processing {n_bins} spectral bins...")

    # Process each spectral bin
    for i in tqdm(range(n_bins), desc="Computing c-k"):
        # Select wavenumbers in this bin
        selection = np.nonzero(
            np.logical_and(
                np.greater_equal(wavenumbers, bin_edges[i]),
                np.less(wavenumbers, bin_edges[i + 1])
            )
        )[0]

        # Handle edge cases where selection is too small
        if selection.size < 2:
            selection_low = np.nonzero(np.less(wavenumbers, bin_edges[i + 1]))[0]
            selection_high = np.nonzero(np.greater_equal(wavenumbers, bin_edges[i]))[0]

            if selection_low.size == 0:
                selection_low = 0
            else:
                selection_low = selection_low[-1]

            if selection_high.size == 0:
                selection_high = n_wn - 1
            else:
                selection_high = selection_high[0]

            if selection_low == selection_high:
                selection_low = max((selection_low - 1, 0))
                selection_high = min((selection_high + 1, n_wn - 1))

            if selection_low > n_wn - 2:
                selection_low = n_wn - 2

            if selection_high < 1:
                selection_high = 1
            elif selection_high == n_wn - 1:
                selection_high = None

            selection = [selection_low, selection_high]

        # Extract cross-sections for this bin
        _cross_sections = cross_sections[:, :, selection[0]:selection[-1]]

        # If selection is too small, interpolate to get enough points
        if np.size(selection) < min_selection_size:
            from scipy.interpolate import interp1d

            _wavenumbers = wavenumbers[selection[0]:selection[-1]]
            _wavenumbers_interp = np.linspace(_wavenumbers[0], _wavenumbers[-1], min_selection_size)

            interpolator = interp1d(_wavenumbers, _cross_sections, axis=-1)
            _cross_sections = interpolator(_wavenumbers_interp)

            selection_size = min_selection_size
        else:
            selection_size = selection.size

        # Sort cross-sections in increasing order
        _cross_sections = np.sort(_cross_sections, axis=-1)

        # Create g-space grid (cumulative distribution function)
        g_sort = np.linspace(0, 1, selection_size)

        # Average within each g interval
        for j in range(n_g):
            g_selection = np.nonzero(
                np.logical_and(
                    np.greater_equal(g_sort, samples_edges[j]),
                    np.less_equal(g_sort, samples_edges[j + 1])
                )
            )[0]

            # Mean over the g interval
            correlated_k[:, :, i, j] = np.mean(
                _cross_sections[:, :, g_selection[0]:g_selection[-1]],
                axis=-1
            )

    print(f"  Correlated-k computation complete")

    return correlated_k


# =============================================================================
# DIAGNOSTIC PLOTTING FUNCTIONS
# =============================================================================

def plot_resolution_check(bin_centers: np.ndarray, R_target: float, output_dir: Path, mol_name: str):
    """
    Plot resolving power vs wavenumber to verify target resolution.

    Parameters
    ----------
    bin_centers : np.ndarray
        Wavenumber bin centers (cm^-1)
    R_target : float
        Target resolving power
    output_dir : Path
        Directory to save plot
    mol_name : str
        Molecule name for title
    """
    print("\n  Generating resolution check plot...")

    # Convert to wavelength
    wavelengths = 1e4 / bin_centers  # μm

    # Compute local resolving power
    R_local = bin_centers[:-1] / np.diff(bin_centers) + 0.5

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))

    # Plot vs wavenumber
    ax1.plot(bin_centers[:-1], R_local, 'b-', lw=0.5, alpha=0.7, label='Measured')
    ax1.axhline(R_target, color='r', ls='--', lw=2, label=f'Target R = {R_target:,}')
    ax1.set_xlabel('Wavenumber (cm$^{-1}$)', fontsize=12)
    ax1.set_ylabel('Resolving Power R = ν/Δν', fontsize=12)
    ax1.set_title(f'Correlated-k Resolution Check - {mol_name}', fontsize=13, fontweight='bold')
    ax1.legend()
    ax1.grid(alpha=0.3)

    # Check deviation from target
    R_mean = np.mean(R_local[np.isfinite(R_local)])
    R_std = np.std(R_local[np.isfinite(R_local)])
    ax1.text(0.05, 0.95, f'Mean R = {R_mean:.0f}\nStd R = {R_std:.0f}',
             transform=ax1.transAxes, fontsize=10, verticalalignment='top',
             bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))

    # Plot vs wavelength
    ax2.plot(wavelengths[:-1], R_local, 'b-', lw=0.5, alpha=0.7, label='Measured')
    ax2.axhline(R_target, color='r', ls='--', lw=2, label=f'Target R = {R_target:,}')
    ax2.set_xlabel('Wavelength (μm)', fontsize=12)
    ax2.set_ylabel('Resolving Power R = λ/Δλ', fontsize=12)
    ax2.set_title('Resolution vs Wavelength', fontsize=13)
    ax2.legend()
    ax2.grid(alpha=0.3)

    plt.tight_layout()
    plt.savefig(output_dir / f'diagnostic_1_resolution_{mol_name}.png', dpi=150)
    plt.close()

    print("    Resolution check plot saved")


def plot_g_space_distribution(
    correlated_k: np.ndarray,
    bin_centers: np.ndarray,
    g_samples: np.ndarray,
    g_weights: np.ndarray,
    pressures: np.ndarray,
    temperatures: np.ndarray,
    output_dir: Path,
    mol_name: str,
    n_bins_to_plot: int = 6
):
    """
    Plot g-space distribution for selected spectral bins.

    Shows how opacities are distributed across g-space for different bins.

    Parameters
    ----------
    correlated_k : np.ndarray
        Correlated-k coefficients (n_p, n_t, n_bins, n_g)
    bin_centers : np.ndarray
        Wavenumber bin centers (cm^-1)
    g_samples : np.ndarray
        G-space sample points
    g_weights : np.ndarray
        G-space weights
    pressures : np.ndarray
        Pressure grid (bar)
    temperatures : np.ndarray
        Temperature grid (K)
    output_dir : Path
        Directory to save plot
    mol_name : str
        Molecule name
    n_bins_to_plot : int
        Number of spectral bins to plot
    """
    print("\n  Generating g-space distribution plots...")

    # Select representative P-T point (mid-range)
    p_idx = len(pressures) // 2
    t_idx = len(temperatures) // 2
    P = pressures[p_idx]
    T = temperatures[t_idx]

    # Select bins uniformly distributed across spectrum
    n_bins = len(bin_centers)
    bin_indices = np.linspace(0, n_bins - 1, n_bins_to_plot, dtype=int)

    fig, axes = plt.subplots(2, 3, figsize=(15, 10))
    axes = axes.flatten()

    for idx, bin_idx in enumerate(bin_indices):
        ax = axes[idx]

        wn = bin_centers[bin_idx]
        wl = 1e4 / wn

        # Get k-coefficients for this bin
        k_coeffs = correlated_k[p_idx, t_idx, bin_idx, :]

        # Plot k vs g
        ax.plot(g_samples, k_coeffs, 'o-', color='steelblue', ms=5, lw=2, label='k-coefficient')
        ax.fill_between(g_samples, 0, k_coeffs, alpha=0.3, color='steelblue')

        # Mark g-space sample points with weights
        for i, (g, k, w) in enumerate(zip(g_samples, k_coeffs, g_weights)):
            ax.plot([g, g], [0, k], 'k--', alpha=0.3, lw=0.5)
            if i % 2 == 0:  # Label every other point to avoid crowding
                ax.text(g, k * 1.1, f'w={w:.3f}', fontsize=7, ha='center', rotation=45)

        ax.set_xlabel('g (cumulative probability)', fontsize=10)
        ax.set_ylabel('k-coefficient (cm²/molecule)', fontsize=10)
        ax.set_title(f'Bin {bin_idx}: {wn:.1f} cm⁻¹ ({wl:.3f} μm)', fontsize=10)
        ax.set_yscale('log')
        ax.set_xlim(-0.05, 1.05)
        ax.grid(alpha=0.3)

    plt.suptitle(f'G-space Distribution for {mol_name}\nP = {P:.2e} bar, T = {T:.0f} K',
                 fontsize=14, fontweight='bold')
    plt.tight_layout()
    plt.savefig(output_dir / f'diagnostic_2_g_space_{mol_name}.png', dpi=150)
    plt.close()

    print("    G-space distribution plots saved")


def plot_lbl_vs_ck_comparison(
    cross_sections_lbl: np.ndarray,
    correlated_k: np.ndarray,
    wavenumbers_lbl: np.ndarray,
    bin_centers_ck: np.ndarray,
    bin_edges_ck: np.ndarray,
    g_weights: np.ndarray,
    pressures: np.ndarray,
    temperatures: np.ndarray,
    output_dir: Path,
    mol_name: str,
    R_input: float,
    R_target: float
):
    """
    Compare line-by-line and correlated-k opacities.

    Reconstructs spectrum from correlated-k and compares with original line-by-line.

    Parameters
    ----------
    cross_sections_lbl : np.ndarray
        Line-by-line cross-sections (n_p, n_t, n_wn)
    correlated_k : np.ndarray
        Correlated-k coefficients (n_p, n_t, n_bins, n_g)
    wavenumbers_lbl : np.ndarray
        Line-by-line wavenumbers (cm^-1)
    bin_centers_ck : np.ndarray
        Correlated-k bin centers (cm^-1)
    bin_edges_ck : np.ndarray
        Correlated-k bin edges (cm^-1)
    g_weights : np.ndarray
        G-space weights
    pressures : np.ndarray
        Pressure grid (bar)
    temperatures : np.ndarray
        Temperature grid (K)
    output_dir : Path
        Directory to save plot
    mol_name : str
        Molecule name
    R_input : float
        Input resolution
    R_target : float
        Target resolution
    """
    print("\n  Generating line-by-line vs correlated-k comparison...")

    # Select representative P-T point
    p_idx = len(pressures) // 2
    t_idx = len(temperatures) // 2
    P = pressures[p_idx]
    T = temperatures[t_idx]

    # Reconstruct spectrum from correlated-k (weighted mean over g)
    ck_reconstructed = np.sum(correlated_k[p_idx, t_idx, :, :] * g_weights[np.newaxis, :], axis=1)

    # Convert to wavelength
    wavelengths_lbl = 1e4 / wavenumbers_lbl
    wavelengths_ck = 1e4 / bin_centers_ck

    # Create figure with 3 panels
    fig = plt.figure(figsize=(16, 12))
    gs = gridspec.GridSpec(3, 2, figure=fig, hspace=0.3, wspace=0.3)

    # Panel A: Full spectrum comparison
    ax1 = fig.add_subplot(gs[0, :])
    ax1.plot(wavelengths_lbl, cross_sections_lbl[p_idx, t_idx, :], 'b-', lw=0.3, alpha=0.5,
             label=f'Line-by-line (R={R_input:.0e})')
    ax1.plot(wavelengths_ck, ck_reconstructed, 'r-', lw=1.0, alpha=0.9,
             label=f'Correlated-k reconstructed (R={R_target:,})')
    ax1.set_xlabel('Wavelength (μm)', fontsize=12)
    ax1.set_ylabel('Cross-section (cm²/molecule)', fontsize=12)
    ax1.set_title(f'A) Full Spectrum Comparison - {mol_name}\nP = {P:.2e} bar, T = {T:.0f} K',
                  fontsize=13, fontweight='bold')
    ax1.set_yscale('log')
    ax1.legend(fontsize=10)
    ax1.grid(alpha=0.3)

    # Panel B: Zoom on H2O band (1.35-1.45 μm)
    ax2 = fig.add_subplot(gs[1, 0])
    wl_min, wl_max = 1.35, 1.45
    mask_lbl = (wavelengths_lbl > wl_min) & (wavelengths_lbl < wl_max)
    mask_ck = (wavelengths_ck > wl_min) & (wavelengths_ck < wl_max)
    ax2.plot(wavelengths_lbl[mask_lbl], cross_sections_lbl[p_idx, t_idx, mask_lbl], 'b-', lw=0.5, alpha=0.7, label='LBL')
    ax2.plot(wavelengths_ck[mask_ck], ck_reconstructed[mask_ck], 'r-', lw=1.0, alpha=0.9, label='C-K')
    ax2.set_xlabel('Wavelength (μm)', fontsize=11)
    ax2.set_ylabel('Cross-section (cm²/molecule)', fontsize=11)
    ax2.set_title(f'B) Zoom: {wl_min}-{wl_max} μm', fontsize=12)
    ax2.legend(fontsize=9)
    ax2.grid(alpha=0.3)

    # Panel C: Another zoom (2.0-2.1 μm)
    ax3 = fig.add_subplot(gs[1, 1])
    wl_min, wl_max = 2.0, 2.1
    mask_lbl = (wavelengths_lbl > wl_min) & (wavelengths_lbl < wl_max)
    mask_ck = (wavelengths_ck > wl_min) & (wavelengths_ck < wl_max)
    ax3.plot(wavelengths_lbl[mask_lbl], cross_sections_lbl[p_idx, t_idx, mask_lbl], 'b-', lw=0.5, alpha=0.7, label='LBL')
    ax3.plot(wavelengths_ck[mask_ck], ck_reconstructed[mask_ck], 'r-', lw=1.0, alpha=0.9, label='C-K')
    ax3.set_xlabel('Wavelength (μm)', fontsize=11)
    ax3.set_ylabel('Cross-section (cm²/molecule)', fontsize=11)
    ax3.set_title(f'C) Zoom: {wl_min}-{wl_max} μm', fontsize=12)
    ax3.legend(fontsize=9)
    ax3.grid(alpha=0.3)

    # Panel D: Ratio (interpolate c-k onto lbl grid)
    ax4 = fig.add_subplot(gs[2, 0])
    ck_interp = np.interp(wavenumbers_lbl, bin_centers_ck, ck_reconstructed)
    valid = cross_sections_lbl[p_idx, t_idx, :] > 0
    ratio = ck_interp[valid] / cross_sections_lbl[p_idx, t_idx, valid]
    ratio_filtered = ratio[(ratio > 0.1) & (ratio < 10)]
    ax4.hist(ratio_filtered, bins=100, alpha=0.7, edgecolor='black', color='steelblue')
    ax4.axvline(1.0, color='r', ls='--', lw=2, label='Ratio = 1')
    ax4.axvline(np.median(ratio_filtered), color='g', ls='--', lw=2,
                label=f'Median = {np.median(ratio_filtered):.3f}')
    ax4.set_xlabel('C-K / LBL Ratio', fontsize=11)
    ax4.set_ylabel('Count', fontsize=11)
    ax4.set_title('D) Ratio Distribution', fontsize=12)
    ax4.legend(fontsize=9)
    ax4.grid(alpha=0.3)

    # Panel E: Bin width distribution
    ax5 = fig.add_subplot(gs[2, 1])
    bin_widths = np.diff(bin_edges_ck)
    ax5.hist(bin_widths, bins=50, alpha=0.7, edgecolor='black', color='coral')
    ax5.axvline(np.mean(bin_widths), color='r', ls='--', lw=2,
                label=f'Mean = {np.mean(bin_widths):.2f} cm⁻¹')
    ax5.set_xlabel('Bin Width (cm⁻¹)', fontsize=11)
    ax5.set_ylabel('Count', fontsize=11)
    ax5.set_title('E) Correlated-k Bin Width Distribution', fontsize=12)
    ax5.legend(fontsize=9)
    ax5.grid(alpha=0.3)

    plt.savefig(output_dir / f'diagnostic_3_lbl_vs_ck_{mol_name}.png', dpi=150)
    plt.close()

    print("    Line-by-line vs correlated-k comparison saved")


def plot_pt_grid_heatmap(
    correlated_k: np.ndarray,
    bin_centers: np.ndarray,
    pressures: np.ndarray,
    temperatures: np.ndarray,
    g_weights: np.ndarray,
    output_dir: Path,
    mol_name: str
):
    """
    Plot heatmap of mean opacity across P-T grid.

    Parameters
    ----------
    correlated_k : np.ndarray
        Correlated-k coefficients (n_p, n_t, n_bins, n_g)
    bin_centers : np.ndarray
        Wavenumber bin centers (cm^-1)
    pressures : np.ndarray
        Pressure grid (bar)
    temperatures : np.ndarray
        Temperature grid (K)
    g_weights : np.ndarray
        G-space weights
    output_dir : Path
        Directory to save plot
    mol_name : str
        Molecule name
    """
    print("\n  Generating P-T grid heatmap...")

    # Compute mean opacity (weighted over g, averaged over wavelength)
    ck_weighted = np.sum(correlated_k * g_weights[np.newaxis, np.newaxis, np.newaxis, :], axis=3)
    mean_opacity = np.mean(ck_weighted, axis=2)  # (n_p, n_t)

    # Select a few spectral bins to show
    n_bins = len(bin_centers)
    bin_indices = [0, n_bins // 4, n_bins // 2, 3 * n_bins // 4, n_bins - 1]

    fig, axes = plt.subplots(2, 3, figsize=(16, 10))
    axes = axes.flatten()

    # Plot overall mean opacity
    # Handle potential zeros by setting a minimum value
    mean_opacity_plot = np.where(mean_opacity > 0, mean_opacity, np.nan)
    vmin = np.nanmin(mean_opacity_plot)
    vmax = np.nanmax(mean_opacity_plot)

    im = axes[0].imshow(mean_opacity_plot.T, aspect='auto', origin='lower',
                        cmap='viridis', norm=LogNorm(vmin=vmin, vmax=vmax),
                        extent=[0, len(pressures) - 1, 0, len(temperatures) - 1])
    axes[0].set_xlabel('Pressure index', fontsize=10)
    axes[0].set_ylabel('Temperature index', fontsize=10)
    axes[0].set_title('Mean Opacity (all λ)', fontsize=11, fontweight='bold')
    cbar = plt.colorbar(im, ax=axes[0])
    cbar.set_label('σ (cm²/molecule)', fontsize=9)

    # Add pressure and temperature labels
    axes[0].set_xticks(range(len(pressures)))
    axes[0].set_xticklabels([f'{p:.1e}' for p in pressures], rotation=45, ha='right', fontsize=7)
    axes[0].set_yticks(range(len(temperatures)))
    axes[0].set_yticklabels([f'{t:.0f}' for t in temperatures], fontsize=7)

    # Plot opacity at selected wavelengths
    for idx, bin_idx in enumerate(bin_indices):
        ax = axes[idx + 1]
        wn = bin_centers[bin_idx]
        wl = 1e4 / wn

        opacity_at_bin = np.sum(correlated_k[:, :, bin_idx, :] * g_weights[np.newaxis, np.newaxis, :], axis=2)

        # Handle potential zeros
        opacity_plot = np.where(opacity_at_bin > 0, opacity_at_bin, np.nan)
        vmin_bin = np.nanmin(opacity_plot)
        vmax_bin = np.nanmax(opacity_plot)

        im = ax.imshow(opacity_plot.T, aspect='auto', origin='lower',
                       cmap='viridis', norm=LogNorm(vmin=vmin_bin, vmax=vmax_bin),
                       extent=[0, len(pressures) - 1, 0, len(temperatures) - 1])
        ax.set_xlabel('Pressure index', fontsize=10)
        ax.set_ylabel('Temperature index', fontsize=10)
        ax.set_title(f'{wn:.1f} cm⁻¹ ({wl:.3f} μm)', fontsize=10)
        cbar = plt.colorbar(im, ax=ax)
        cbar.set_label('σ (cm²/molecule)', fontsize=8)

        ax.set_xticks(range(len(pressures)))
        ax.set_xticklabels([f'{p:.1e}' for p in pressures], rotation=45, ha='right', fontsize=7)
        ax.set_yticks(range(len(temperatures)))
        ax.set_yticklabels([f'{t:.0f}' for t in temperatures], fontsize=7)

    plt.suptitle(f'Opacity Heatmap across P-T Grid - {mol_name}', fontsize=14, fontweight='bold')
    plt.tight_layout()
    plt.savefig(output_dir / f'diagnostic_4_pt_heatmap_{mol_name}.png', dpi=150)
    plt.close()

    print("    P-T grid heatmap saved")


def plot_statistics_summary(
    correlated_k: np.ndarray,
    bin_centers: np.ndarray,
    g_samples: np.ndarray,
    g_weights: np.ndarray,
    output_dir: Path,
    mol_name: str,
    R_target: float
):
    """
    Plot statistical summary of correlated-k coefficients.

    Parameters
    ----------
    correlated_k : np.ndarray
        Correlated-k coefficients (n_p, n_t, n_bins, n_g)
    bin_centers : np.ndarray
        Wavenumber bin centers (cm^-1)
    g_samples : np.ndarray
        G-space samples
    g_weights : np.ndarray
        G-space weights
    output_dir : Path
        Directory to save plot
    mol_name : str
        Molecule name
    R_target : float
        Target resolution
    """
    print("\n  Generating statistics summary...")

    wavelengths = 1e4 / bin_centers

    fig, axes = plt.subplots(2, 2, figsize=(14, 10))

    # Panel A: Mean opacity vs wavelength (averaged over P-T and g)
    ax = axes[0, 0]
    mean_opacity = np.mean(np.mean(np.sum(correlated_k * g_weights[np.newaxis, np.newaxis, np.newaxis, :], axis=3),
                                     axis=0), axis=0)
    ax.plot(wavelengths, mean_opacity, 'b-', lw=1.5)
    ax.set_xlabel('Wavelength (μm)', fontsize=11)
    ax.set_ylabel('Mean Opacity (cm²/molecule)', fontsize=11)
    ax.set_title('A) Mean Opacity vs Wavelength\n(averaged over P-T and g)', fontsize=12)
    ax.set_yscale('log')
    ax.grid(alpha=0.3)

    # Panel B: Opacity range (min-max) vs wavelength
    ax = axes[0, 1]
    min_opacity = np.min(correlated_k, axis=(0, 1, 3))
    max_opacity = np.max(correlated_k, axis=(0, 1, 3))
    ax.fill_between(wavelengths, min_opacity, max_opacity, alpha=0.3, color='steelblue', label='Range')
    ax.plot(wavelengths, mean_opacity, 'r-', lw=1.5, label='Mean')
    ax.set_xlabel('Wavelength (μm)', fontsize=11)
    ax.set_ylabel('Opacity (cm²/molecule)', fontsize=11)
    ax.set_title('B) Opacity Range vs Wavelength\n(across all P-T and g)', fontsize=12)
    ax.set_yscale('log')
    ax.legend(fontsize=9)
    ax.grid(alpha=0.3)

    # Panel C: G-space weight contribution
    ax = axes[1, 0]
    ax.bar(g_samples, g_weights, width=0.03, alpha=0.7, edgecolor='black', color='coral')
    ax.set_xlabel('g (cumulative probability)', fontsize=11)
    ax.set_ylabel('Weight', fontsize=11)
    ax.set_title('C) G-space Quadrature Weights', fontsize=12)
    ax.grid(alpha=0.3, axis='y')
    ax.text(0.5, max(g_weights) * 0.9, f'Total: {np.sum(g_weights):.4f}',
            fontsize=10, ha='center',
            bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))

    # Panel D: Statistics table
    ax = axes[1, 1]
    ax.axis('off')

    # Calculate statistics
    ck_weighted = np.sum(correlated_k * g_weights[np.newaxis, np.newaxis, np.newaxis, :], axis=3)
    stats_text = f"""
    CORRELATED-K STATISTICS
    {'=' * 40}

    Molecule: {mol_name}
    Target Resolution: R = {R_target:,}

    Number of bins: {len(bin_centers):,}
    G-space points: {len(g_samples)}

    Wavelength range: {wavelengths.min():.3f} - {wavelengths.max():.3f} μm
    Wavenumber range: {bin_centers.min():.1f} - {bin_centers.max():.1f} cm⁻¹

    Opacity Statistics (weighted mean):
    - Min:    {np.min(ck_weighted):.3e} cm²/molecule
    - Max:    {np.max(ck_weighted):.3e} cm²/molecule
    - Mean:   {np.mean(ck_weighted):.3e} cm²/molecule
    - Median: {np.median(ck_weighted):.3e} cm²/molecule

    G-space weights sum: {np.sum(g_weights):.6f}
    """

    ax.text(0.1, 0.9, stats_text, transform=ax.transAxes,
            fontsize=10, verticalalignment='top', family='monospace',
            bbox=dict(boxstyle='round', facecolor='lightgray', alpha=0.5))

    plt.suptitle(f'Statistical Summary - {mol_name}', fontsize=14, fontweight='bold')
    plt.tight_layout()
    plt.savefig(output_dir / f'diagnostic_5_statistics_{mol_name}.png', dpi=150)
    plt.close()

    print("    Statistics summary saved")


# =============================================================================
# MAIN PROCESSING
# =============================================================================

def generate_correlated_k(config: Config):
    """
    Main function to generate correlated-k opacities from line-by-line data.

    Parameters
    ----------
    config : Config
        Configuration object with all parameters
    """
    print("=" * 80)
    print("CORRELATED-K OPACITY GENERATOR")
    print("=" * 80)
    print(f"Started at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print()

    # Validate input file
    input_path = Path(config.INPUT_FILE)
    if not input_path.exists():
        raise FileNotFoundError(f"Input file not found: {config.INPUT_FILE}")

    # Load line-by-line data
    print("=" * 80)
    print("LOADING INPUT DATA")
    print("=" * 80)
    data = load_line_by_line_file(config.INPUT_FILE)

    # Extract data
    wavenumbers = data['wavenumbers']
    cross_sections = data['cross_sections']
    pressures = data['pressures']
    temperatures = data['temperatures']
    mol_mass = data['mol_mass']
    mol_name = data['mol_name']

    # Compute input resolution
    R_input = compute_resolving_power(wavenumbers)

    # Validate target resolution
    if config.R_TARGET >= R_input:
        raise ValueError(
            f"Target resolution ({config.R_TARGET}) must be lower than "
            f"input resolution ({R_input:.0e})"
        )

    # Create correlated-k wavenumber grid
    print("=" * 80)
    print("CREATING CORRELATED-K GRID")
    print("=" * 80)
    bin_edges, bin_centers = create_correlated_k_grid(
        wavenumbers,
        R_input,
        config.R_TARGET,
        config.USE_LEGACY_MODE
    )

    # Compute correlated-k coefficients
    print("=" * 80)
    print("COMPUTING CORRELATED-K COEFFICIENTS")
    print("=" * 80)
    correlated_k = compute_correlated_k_coefficients(
        cross_sections,
        wavenumbers,
        bin_edges,
        config.G_SAMPLES
    )

    # Generate output filename
    print("=" * 80)
    print("SAVING OUTPUT")
    print("=" * 80)

    output_dir = input_path.parent
    output_dir = Path(str(output_dir).replace("line_by_line", "correlated_k"))

    # Create filename following pRT convention
    wl_min = 1e4 / bin_edges[-1]
    wl_max = 1e4 / bin_edges[0]

    # Extract molecule info from input filename
    input_stem = input_path.stem
    parts = input_stem.split('__')
    if len(parts) >= 2:
        molecule_info = parts[0]
        source_info = parts[1].split('.')[0]
        output_filename = f"{molecule_info}__{source_info}.R{int(config.R_TARGET)}_{wl_min:.1f}-{int(np.ceil(wl_max))}mu.ktable.petitRADTRANS.h5"
    else:
        output_filename = f"{mol_name}.R{int(config.R_TARGET)}_{wl_min:.1f}-{int(np.ceil(wl_max))}mu.ktable.petitRADTRANS.h5"

    output_path = output_dir / output_filename

    # Check if file exists
    if output_path.exists() and not config.REWRITE:
        raise FileExistsError(
            f"Output file already exists: {output_path}\n"
            f"Set REWRITE=True to overwrite"
        )

    # Write output file
    write_correlated_k_file(
        str(output_path),
        bin_centers,
        bin_edges,
        correlated_k,
        pressures,
        temperatures,
        mol_mass,
        mol_name,
        config.G_SAMPLES,
        config.G_WEIGHTS,
        config.DOI,
        config.CONTRIBUTOR,
        config.DESCRIPTION
    )

    # Generate diagnostic plots
    if config.MAKE_PLOTS:
        print()
        print("=" * 80)
        print("GENERATING DIAGNOSTIC PLOTS")
        print("=" * 80)

        # Create output directory for plots
        plots_dir = output_dir / f"diagnostic_plots_ck_R{int(config.R_TARGET)}_{mol_name}"
        plots_dir.mkdir(parents=True, exist_ok=True)
        print(f"Saving plots to: {plots_dir}")

        # Plot 1: Resolution check
        plot_resolution_check(bin_centers, config.R_TARGET, plots_dir, mol_name)

        # Plot 2: G-space distribution
        # plot_g_space_distribution(
        #     correlated_k, bin_centers, config.G_SAMPLES, config.G_WEIGHTS,
        #     pressures, temperatures, plots_dir, mol_name
        # )

        # Plot 3: Line-by-line vs correlated-k comparison
        plot_lbl_vs_ck_comparison(
            cross_sections, correlated_k, wavenumbers, bin_centers, bin_edges,
            config.G_WEIGHTS, pressures, temperatures, plots_dir, mol_name,
            R_input, config.R_TARGET
        )

        # Plot 4: P-T grid heatmap
        # plot_pt_grid_heatmap(
        #     correlated_k, bin_centers, pressures, temperatures,
        #     config.G_WEIGHTS, plots_dir, mol_name
        # )

        # Plot 5: Statistics summary
        plot_statistics_summary(
            correlated_k, bin_centers, config.G_SAMPLES, config.G_WEIGHTS,
            plots_dir, mol_name, config.R_TARGET
        )

        print(f"\n  All diagnostic plots saved to: {plots_dir}")

    # Final summary
    print()
    print("=" * 80)
    print("SUMMARY")
    print("=" * 80)
    print(f"Input file: {config.INPUT_FILE}")
    print(f"Output file: {output_path}")
    print(f"Input resolution: R = {R_input:.0e}")
    print(f"Output resolution: R = {config.R_TARGET:,}")
    print(f"Compression factor: {len(wavenumbers) / len(bin_centers):.1f}x")
    print(f"Input grid points: {len(wavenumbers):,}")
    print(f"Output grid points: {len(bin_centers):,}")
    print(f"G-space points: {len(config.G_SAMPLES)}")
    print(f"P-T grid: {len(pressures)} × {len(temperatures)} = {len(pressures) * len(temperatures)}")
    print()
    print(f"Finished at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print("=" * 80)
    print("DONE!")
    print("=" * 80)


# =============================================================================
# ENTRY POINT
# =============================================================================

def main():
    """Main entry point."""
    # Parse command line arguments
    parser = argparse.ArgumentParser(
        description="Convert line-by-line opacities to correlated-k format",
        formatter_class=argparse.RawDescriptionHelpFormatter
    )

    parser.add_argument(
        '--input-file',
        type=str,
        required=False,
        help='Path to line-by-line opacity HDF5 file'
    )

    parser.add_argument(
        '--resolution',
        type=int,
        default=20000,
        help='Target resolving power (default: 20000)'
    )

    parser.add_argument(
        '--no-plots',
        action='store_true',
        help='Skip diagnostic plot generation'
    )

    parser.add_argument(
        '--rewrite',
        action='store_true',
        help='Overwrite existing output files'
    )

    parser.add_argument(
        '--legacy-mode',
        action='store_true',
        help='Use legacy wavenumber sampling (pRT v2 compatibility)'
    )

    args = parser.parse_args()

    try:
        # Create configuration
        config = Config()

        # Override with command line arguments if provided
        if args.input_file:
            config.INPUT_FILE = args.input_file

        config.R_TARGET = args.resolution
        config.MAKE_PLOTS = not args.no_plots
        config.REWRITE = args.rewrite or config.REWRITE  # Use CLI flag or default True
        config.USE_LEGACY_MODE = args.legacy_mode

        # Update description with actual resolution
        config.DESCRIPTION = f"Correlated-k opacities at R={config.R_TARGET} generated from high-resolution line-by-line data"

        # Validate input file
        if not Path(config.INPUT_FILE).exists():
            raise FileNotFoundError(f"Input file not found: {config.INPUT_FILE}")

        # Run conversion
        generate_correlated_k(config)

    except Exception as e:
        print(f"\n{'='*80}")
        print("ERROR")
        print(f"{'='*80}")
        print(f"{type(e).__name__}: {e}")
        print(f"{'='*80}")
        raise


if __name__ == '__main__':
    main()