Lodge It

import matplotlib
matplotlib.use('agg')
import matplotlib.pyplot as plt
import matplotlib.colors as colors
from matplotlib.lines import Line2D
from matplotlib.patches import Patch
import numpy as np
import h5py
import os
import sys
import re
from collections import defaultdict
from astropy.io import fits
from astropy import units as u
import astropy.constants as const
import multiprocessing

try:
    import seaborn as sns
except ImportError:
    print("Error: seaborn package not found.")
    sys.exit()

sys.path.insert(0, '/home/desika.narayanan/torreylabtools_A1/')
try:
    import simread.readsubfHDF5 as read_subf
    from hyperion.model import ModelOutput
    import PAHFIT_wrapper
    from live_dust_util import SnapshotContainer as Snap
    from live_dust_util import GrainSizeDistribution as GSD
except ImportError as e:
    print(f"Error: Could not import tools. Details: {e}")
    sys.exit()

# ==============================================================================
#                          ⭐ CONFIGURATION ⭐
# ==============================================================================
SIM_LIST_FILE = 'sim_list.txt'
BASE_PD_PATH = '/blue/narayanan/desika.narayanan/pd_runs/smuggle_sfh_ML/cosmic_sands/m100/z0/helena'
OUTPUT_FILENAME_TEMPLATE = 'LPAH_MPAH_LMratio_vs_SFR_{group_key}_metallicity_colorcode.png'

START_SNAP, END_SNAP = 0, 600
PLOT_STYLE = 'seaborn-v0_8-talk'
MAX_REDSHIFT_TO_PLOT = 6.0
MDUST_MIN = 1.

# PAH size threshold for q_PAH (microns)
PAH_SIZE_THRESHOLD = 3.e-3  # a < 30 Angstrom carbonaceous grains

# --- PARALLELISM ---
NUM_WORKERS = 32

# --- CONSTANTS ---
Lsun_cgs = 3.828e33  # erg/s

# --- SHIPLEY+2016 CALIBRATION ---
# SFR [Msun/yr] = PAH_CONST * L_PAH(6.2+7.7+11.3) [erg/s]
# Scatter: 0.14 dex
PAH_CONST = 2.78e-43
SHIPLEY_SCATTER_DEX = 0.14

# --- KENNICUTT & EVANS (2012) ---
KENNICUTT_HALPHA_NORM = 5.5e-42
KENNICUTT_LIR_NORM = 1.49e-10     # SFR = KENNICUTT_LIR_NORM * L(IR)/Lsun

# --- ROBINSON+2026 (arxiv:2601.09810) ---
# Eq 9:  log(L_6.2/Lsun)  = 0.95 * log(SFR) + 6.81
# Eq 10: log(L_11.2/Lsun) = 1.08 * log(SFR) + 6.32
ROBINSON_62 = {'slope': 0.95, 'intercept': 6.81,
                'slope_err': 0.12, 'intercept_err': 0.31,
                'label': r'Robinson+26 ($L_{6.2}$)'}
ROBINSON_112 = {'slope': 1.08, 'intercept': 6.32,
                 'slope_err': 0.16, 'intercept_err': 0.42,
                 'label': r'Robinson+26 ($L_{11.2}$)'}

# --- SHIVAEI & BOOGAARD (2024, A&A 691, L2) ---
# Table 1: log(L(PAH7.7)/Lsun) = 0.84 * log(L(IR)/Lsun) + 0.09
#   sigma_int = 0.20 dex
# Also: alpha_PAH7.7 = (3.08 +/- 1.08) * (4.3/alpha_CO) Msun/Lsun
#   => M_mol = alpha_PAH7.7 * L(PAH7.7)
SHIVAEI_SLOPE = 0.84
SHIVAEI_INTERCEPT = 0.09
SHIVAEI_LIR_TO_SFR_OFFSET = np.log10(1.0 / KENNICUTT_LIR_NORM)  # = 9.827
SHIVAEI_SCATTER_DEX = 0.20


# ==============================================================================
#                            HELPER FUNCTIONS
# ==============================================================================
def find_nearest(array, value):
    return (np.abs(np.asarray(array) - value)).argmin()

def get_lir(sedfile, min_wav_lir=100., max_wav_lir=1000.):
    m = ModelOutput(sedfile)
    sed = m.get_sed(inclination='all', aperture=-1)
    lum = sed.val*u.erg/u.s
    nu = sed.nu*u.Hz
    L_nu = lum[0,:]/nu
    min_nu = (const.c/(max_wav_lir*u.micron)).to(u.Hz)
    max_nu = (const.c/(min_wav_lir*u.micron)).to(u.Hz)
    idx1 = find_nearest(nu.value, min_nu.value)
    idx2 = find_nearest(nu.value, max_nu.value)
    start_idx, end_idx = min(idx1, idx2), max(idx1, idx2)
    if start_idx == end_idx: return 0.0 * u.erg / u.s
    integrated_lir = np.trapz(L_nu[start_idx : end_idx+1], nu[start_idx : end_idx+1])
    return integrated_lir.to(u.erg/u.s)

def parse_and_group_sims(filename):
    print(f"Parsing simulation list from: {filename}")
    grouped_sims = defaultdict(list)
    try:
        with open(filename, 'r') as f:
            arepo_paths = [line.strip() for line in f if line.strip()]
    except FileNotFoundError:
        print(f"❌ CRITICAL: Input file '{filename}' not found.")
        sys.exit()

    for path in arepo_paths:
        clean_path = path.strip().rstrip(os.sep)
        if os.path.basename(clean_path) == 'output':
            arepo_output_path = clean_path
            model_name = os.path.basename(os.path.dirname(clean_path))
        else:
            arepo_output_path = os.path.join(clean_path, 'output')
            model_name = os.path.basename(clean_path)

        match = re.search(r'(sobol_\d+)', model_name)
        if not match:
            print(f"Warning: Could not parse group from model: {model_name} (path: {path})")
            continue
        group_key = match.group(1)

        sim_info = {'label': model_name, 'arepo_path': arepo_output_path,
                     'pd_path': os.path.join(BASE_PD_PATH, model_name) + '/'}
        grouped_sims[group_key].append(sim_info)

    print(f"Found {len(arepo_paths)} simulations across {len(grouped_sims)} groups.")
    return grouped_sims

# ==============================================================================
#                CALIBRATION LINE HELPER FUNCTIONS
# ==============================================================================
def shipley16_lpah_vs_sfr(sfr_arr):
    """Shipley+16: L_PAH(6.2+7.7+11.3) = SFR / PAH_CONST  [erg/s -> Lsun]"""
    lpah_ergs = sfr_arr / PAH_CONST
    return lpah_ergs / Lsun_cgs

def robinson26_lpah_vs_sfr(sfr_arr, band='6.2'):
    """Robinson+26: log(L_band/Lsun) = slope * log(SFR) + intercept"""
    params = ROBINSON_62 if band == '6.2' else ROBINSON_112
    log_sfr = np.log10(sfr_arr)
    log_lpah = params['slope'] * log_sfr + params['intercept']
    return 10.**log_lpah

def shivaei24_lpah_vs_sfr(sfr_arr):
    """
    Shivaei & Boogaard (2024):
      log(L(PAH7.7)/Lsun) = 0.84 * log(L(IR)/Lsun) + 0.09
    with log(L(IR)/Lsun) = log(SFR) + 9.827  (Kennicutt & Evans 2012)
    """
    log_sfr = np.log10(sfr_arr)
    log_lir = log_sfr + SHIVAEI_LIR_TO_SFR_OFFSET
    log_lpah = SHIVAEI_SLOPE * log_lir + SHIVAEI_INTERCEPT
    return 10.**log_lpah


# ==============================================================================
#                   WORKER FUNCTION
# ==============================================================================
def process_single_snapshot(args):
    """Worker function to process a single snapshot."""
    sim_details, snap_num = args

    snap_num_str = str(snap_num).zfill(3)
    arepo_snap_file = os.path.join(sim_details['arepo_path'], f'snapdir_{snap_num_str}', f'snapshot_{snap_num_str}.0.hdf5')
    pd_sed_file = os.path.join(sim_details['pd_path'], f'snap{snap_num_str}', f'snap{snap_num_str}.galaxy0.rtout.sed')

    if not os.path.exists(arepo_snap_file) or not os.path.exists(pd_sed_file):
        return None

    try:
        with h5py.File(arepo_snap_file, 'r') as f:
             a = f['Header'].attrs['Time']
        redshift = (1./a) - 1 if a > 0 else -1

        if redshift < 0 or redshift > MAX_REDSHIFT_TO_PLOT:
            return None

        cat = read_subf.subfind_catalog(sim_details['arepo_path'], snap_num,
                                        keysel=['SubhaloGasMetallicitySfrWeighted', 'SubhaloGrNr',
                                               'GroupFirstSub', 'SubhaloSFR', 'SubhaloMassType'])
        if cat.nsubs == 0:
            return None

        central_bool = np.arange(cat.nsubs) == cat.GroupFirstSub[cat.SubhaloGrNr]
        sfr_weighted_metallicity = np.array(cat.SubhaloGasMetallicitySfrWeighted[central_bool])[0]
        if sfr_weighted_metallicity <= 0:
            return None

        sfr = np.array(cat.SubhaloSFR[central_bool])[0]
        if sfr <= 0:
            return None

        dust_mass_code_units = np.array(cat.SubhaloMassType[central_bool, 3])[0]
        dust_mass_msun = dust_mass_code_units * 1.e10 / 0.7

        if dust_mass_msun < MDUST_MIN:
            return None

        oh12_snap = np.log10(sfr_weighted_metallicity * 0.5 / (0.70 * 16.0)) + 12.0

        # --- Compute total L_PAH from PAHFIT (all bands) ---
        temp_file = f'dummy_{os.getpid()}.ecsv'
        PAHFIT_wrapper.PAHFIT_wrapper(pd_sed_file, temp_file)
        lpah = PAHFIT_wrapper.get_LPAH(temp_file) * u.erg / u.s

        lfir = get_lir(pd_sed_file)

        if os.path.exists(temp_file):
            os.remove(temp_file)

        if lfir.value <= 0 or lpah.value <= 0:
            return None

        # --- Compute M_PAH from GSD (physical grain sizes) ---
        mpah = np.nan
        filtered_snap_dir = os.path.join(sim_details['arepo_path'], f'snapdir_{snap_num_str}', 'filtered_snaps/')
        filtered_hdf5 = os.path.join(filtered_snap_dir, f'snapshot_{snap_num_str}.hdf5')
        if os.path.exists(filtered_hdf5):
            try:
                gsd_snapshot = Snap(snap_num, filtered_snap_dir)
                gsd = GSD(gsd_snapshot, a=10.**np.linspace(-3, 0, 16))
                qpah, _, _ = gsd.compute_q_pah(size=PAH_SIZE_THRESHOLD)
                mpah = qpah * dust_mass_msun  # M_PAH in Msun
            except Exception as e_gsd:
                # GSD computation failed; leave mpah as NaN
                pass

        return {
            'model': sim_details['label'],
            'snap': snap_num,
            'redshift': redshift,
            'oh12': oh12_snap,
            'lpah': lpah.value,        # total L_PAH in erg/s (all PAHFIT bands)
            'sfr': sfr,                # SFR in Msun/yr
            'mpah': mpah,              # M_PAH in Msun (NaN if GSD unavailable)
            'mdust': dust_mass_msun,   # total dust mass in Msun
        }

    except Exception as e:
        return None

# ==============================================================================
#                   MAIN
# ==============================================================================
if __name__ == '__main__':
    multiprocessing.set_start_method('spawn', force=True)

    print("="*70)
    print("L_PAH, M_PAH, & L/M_PAH vs SFR — COLOR-CODED BY METALLICITY")
    print(f"Using {NUM_WORKERS} workers")
    print("="*70)

    # ---- Parse simulations ----
    grouped_simulations = parse_and_group_sims(SIM_LIST_FILE)

    if not grouped_simulations:
        print("❌ CRITICAL: No valid simulations found. Exiting.")
        sys.exit()

    # Build list of all tasks
    all_tasks = []
    task_to_group = {}
    for group_key, sim_list in grouped_simulations.items():
        for sim_details in sim_list:
            for snap_num in range(START_SNAP, END_SNAP + 1):
                task_idx = len(all_tasks)
                all_tasks.append((sim_details, snap_num))
                task_to_group[task_idx] = group_key

    print(f"Total tasks to process: {len(all_tasks)}")
    print("Processing...\n")

    # Process in parallel
    results_by_group = defaultdict(list)
    processed_count = 0
    valid_count = 0

    with multiprocessing.Pool(processes=NUM_WORKERS) as pool:
        for task_idx, result in enumerate(pool.imap(process_single_snapshot, all_tasks, chunksize=50)):
            processed_count += 1
            if processed_count % 1000 == 0:
                print(f"  Processed {processed_count}/{len(all_tasks)} tasks, {valid_count} valid points...", flush=True)

            if result is not None:
                valid_count += 1
                group_key = task_to_group[task_idx]
                results_by_group[group_key].append(result)

    print(f"\nProcessing complete! {valid_count} valid points total.")

    # ==========================================================================
    #       PLOTTING: THREE-PANEL FIGURE
    #       Left:   L_PAH (total) vs SFR, color = metallicity
    #       Center: M_PAH vs SFR, color = metallicity
    #       Right:  L_PAH / M_PAH vs SFR, color = metallicity
    # ==========================================================================
    plots_generated = 0

    for group_key, results in results_by_group.items():
        if not results:
            print(f"⚠️ Warning: No valid data for group '{group_key}'. Skipping plot.")
            continue

        print(f"\nGenerating plot for group: {group_key} ({len(results)} points)...")

        # Organize simulation data
        oh12     = np.array([r['oh12'] for r in results])
        lpah     = np.array([r['lpah'] for r in results])      # erg/s
        sfr      = np.array([r['sfr']  for r in results])      # Msun/yr
        mpah     = np.array([r['mpah'] for r in results])      # Msun (may contain NaN)
        mdust    = np.array([r['mdust'] for r in results])     # Msun

        lpah_lsun = lpah / Lsun_cgs

        # Mask for valid M_PAH points (needed for center and right panels)
        valid_mpah = np.isfinite(mpah) & (mpah > 0)
        n_valid_mpah = np.sum(valid_mpah)
        print(f"  Valid L_PAH points: {len(results)}, Valid M_PAH points: {n_valid_mpah}")

        # L_PAH / M_PAH ratio (Lsun / Msun) — only where M_PAH is valid
        lm_ratio = np.full_like(lpah_lsun, np.nan)
        lm_ratio[valid_mpah] = lpah_lsun[valid_mpah] / mpah[valid_mpah]

        plt.style.use(PLOT_STYLE)
        fig, (ax_left, ax_center, ax_right) = plt.subplots(1, 3, figsize=(42, 12))

        # SFR array for calibration lines
        sfr_line = np.logspace(-2, 4, 300)

        # Shared colormap/norm
        cmap = plt.cm.RdYlBu_r
        norm = colors.Normalize(vmin=7.5, vmax=9.0)

        # ==================================================================
        #  LEFT PANEL: L_PAH (total) vs SFR
        # ==================================================================
        ax = ax_left

        # --- Shipley+16 ---
        lpah_shipley = shipley16_lpah_vs_sfr(sfr_line)
        sigma_s = SHIPLEY_SCATTER_DEX
        ax.fill_between(sfr_line, lpah_shipley * 10**(-sigma_s), lpah_shipley * 10**(+sigma_s),
                         color='grey', alpha=0.20, zorder=3)
        ax.plot(sfr_line, lpah_shipley,
                color='black', ls='--', lw=2.5, zorder=4,
                label=r'Shipley+16 ($L_{6.2+7.7+11.3}$)')

        # --- Robinson+26: 6.2 µm ---
        lpah_r62 = robinson26_lpah_vs_sfr(sfr_line, band='6.2')
        ax.plot(sfr_line, lpah_r62,
                color='tab:blue', ls='-.', lw=2.0, zorder=4,
                label=ROBINSON_62['label'])

        # --- Robinson+26: 11.2 µm ---
        lpah_r112 = robinson26_lpah_vs_sfr(sfr_line, band='11.2')
        ax.plot(sfr_line, lpah_r112,
                color='tab:cyan', ls='-.', lw=2.0, zorder=4,
                label=ROBINSON_112['label'])

        # --- Shivaei & Boogaard 2024: PAH 7.7 µm ---
        lpah_sb24 = shivaei24_lpah_vs_sfr(sfr_line)
        sigma_sb = SHIVAEI_SCATTER_DEX
        ax.fill_between(sfr_line, lpah_sb24 * 10**(-sigma_sb), lpah_sb24 * 10**(+sigma_sb),
                         color='tab:orange', alpha=0.15, zorder=3)
        ax.plot(sfr_line, lpah_sb24,
                color='tab:orange', ls=':', lw=2.5, zorder=4,
                label=r'Shivaei & Boogaard 24 ($L_{7.7}$)')

        # --- Simulation data ---
        sc_left = ax.scatter(sfr, lpah_lsun,
                        marker='o',
                        c=oh12, cmap=cmap, norm=norm,
                        s=180, alpha=0.85,
                        edgecolor='black', linewidth=0.5, zorder=15)

        ax.set_xlabel(r'SFR $[\mathrm{M}_\odot\;\mathrm{yr}^{-1}]$', fontsize=24)
        ax.set_ylabel(r'$L_{\mathrm{PAH}}\;[\mathrm{total}]\;[L_\odot]$', fontsize=24)
        ax.set_xscale('log')
        ax.set_yscale('log')
        ax.grid(True, which='major', linestyle='--', alpha=0.4)
        ax.tick_params(axis='both', which='major', labelsize=16)
        ax.legend(fontsize=11, loc='upper left', framealpha=0.9)

        # ==================================================================
        #  CENTER PANEL: M_PAH vs SFR
        # ==================================================================
        ax = ax_center

        # --- Simulation data (M_PAH) ---
        if n_valid_mpah > 0:
            ax.scatter(sfr[valid_mpah], mpah[valid_mpah],
                       marker='o',
                       c=oh12[valid_mpah], cmap=cmap, norm=norm,
                       s=180, alpha=0.85,
                       edgecolor='black', linewidth=0.5, zorder=15)
        else:
            ax.text(0.5, 0.5, 'No valid M_PAH data\n(filtered_snaps not found?)',
                    transform=ax.transAxes, ha='center', va='center', fontsize=16, color='red')

        ax.set_xlabel(r'SFR $[\mathrm{M}_\odot\;\mathrm{yr}^{-1}]$', fontsize=24)
        ax.set_ylabel(r'$M_{\mathrm{PAH}}\;[M_\odot]$', fontsize=24)
        ax.set_xscale('log')
        ax.set_yscale('log')
        ax.grid(True, which='major', linestyle='--', alpha=0.4)
        ax.tick_params(axis='both', which='major', labelsize=16)

        # ==================================================================
        #  RIGHT PANEL: L_PAH / M_PAH vs SFR
        # ==================================================================
        ax = ax_right

        valid_ratio = np.isfinite(lm_ratio) & (lm_ratio > 0)
        n_valid_ratio = np.sum(valid_ratio)

        if n_valid_ratio > 0:
            ax.scatter(sfr[valid_ratio], lm_ratio[valid_ratio],
                       marker='o',
                       c=oh12[valid_ratio], cmap=cmap, norm=norm,
                       s=180, alpha=0.85,
                       edgecolor='black', linewidth=0.5, zorder=15)
        else:
            ax.text(0.5, 0.5, 'No valid L/M data',
                    transform=ax.transAxes, ha='center', va='center', fontsize=16, color='red')

        ax.set_xlabel(r'SFR $[\mathrm{M}_\odot\;\mathrm{yr}^{-1}]$', fontsize=24)
        ax.set_ylabel(r'$L_{\mathrm{PAH}} / M_{\mathrm{PAH}}\;[L_\odot / M_\odot]$', fontsize=24)
        ax.set_xscale('log')
        ax.set_yscale('log')
        ax.grid(True, which='major', linestyle='--', alpha=0.4)
        ax.tick_params(axis='both', which='major', labelsize=16)

        # ------------------------------------------------------------------
        # Shared colorbar
        # ------------------------------------------------------------------
        fig.subplots_adjust(right=0.92)
        cbar_ax = fig.add_axes([0.935, 0.15, 0.012, 0.7])
        cbar = fig.colorbar(sc_left, cax=cbar_ax)
        cbar.set_label(r'12 + log(O/H)', fontsize=20)
        cbar.ax.tick_params(labelsize=16)

        fig.suptitle(f'{group_key}', fontsize=22, y=0.98)

        output_filename = OUTPUT_FILENAME_TEMPLATE.format(group_key=group_key)
        plt.savefig(output_filename, dpi=300, bbox_inches='tight')
        plt.close('all')
        print(f"✅ Plot saved as {output_filename}")
        plots_generated += 1

    print(f"\n\n🎉 All done! Successfully generated {plots_generated} plot(s).")