# matplotlib.use('Agg') # Removed to allow inline notebook plotting
import os
import psutil

# Set OMP_NUM_THREADS to avoid overhead on large machines
if "OMP_NUM_THREADS" not in os.environ:
    n_cores = psutil.cpu_count(logical=False)
    os.environ["OMP_NUM_THREADS"] = str(n_cores)
    print(f"OMP_NUM_THREADS set to {n_cores}")

from pysm3 import Sky
from pysm3 import units as u

import healpy as hp
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
from astropy.table import Table
from pysm3.models.websky import y2uK_CMB

import logging
logging.getLogger("pysm3").setLevel(logging.INFO)

def remove_monopole(m):
    mono = np.mean(m)
    return m - mono, mono

def joint_limits(map_dict, percentile=99.5):
    stacked = np.concatenate([m.value.ravel() for m in map_dict.values()])
    vmax = np.percentile(np.abs(stacked), percentile)
    if vmax <= 0:
        vmax = np.max(np.abs(stacked))
    if vmax == 0:
        vmax = 1.0
    vmax = float(f"{vmax:.2g}")
    return -vmax, vmax

# Kinematic SZ Comparisons
ksz_models = ["ksz1", "ksz2", "ksz3", "ksz5", "ksz6"]
ksz_titles = {
    "ksz1": "WebSky",
    "ksz2": "AGORA (unlensed)",
    "ksz3": "AGORA",
    "ksz5": "FLAMINGO L1_m9",
    "ksz6": r"FLAMINGO fgas-$8\sigma$",
}
ksz_colors = {
    "ksz1": "#ff9900",  # WebSky orange in the reference figure
    "ksz2": "#8b0000",  # Dark red distinguishes the unlensed variant
    "ksz3": "#ff0000",
    "ksz5": "#332288",  # FLAMINGO purple in the reference figure
    "ksz6": "#9467bd",  # Lighter purple distinguishes strong feedback
}
ksz_linestyles = {
    "ksz1": "-",
    "ksz2": "--",
    "ksz3": "-",
    "ksz5": "-",
    "ksz6": "--",
}
ksz_markers = {"ksz2": "o", "ksz6": "s"}
ksz_zorder = {"ksz2": 4, "ksz6": 4}
ksz_freq = 143
os.makedirs("data", exist_ok=True)
ksz_maps = {}
for model in ksz_models:
    cache_file = f"data/cache_ksz_{model}_{ksz_freq}GHz.fits"
    if os.path.exists(cache_file):
        print(f"Loading {model} kSZ from cache: {cache_file}")
        ksz_maps[model] = hp.read_map(cache_file, field=0) * u.uK_CMB
    else:
        print(f"Computing {model} kSZ at {ksz_freq} GHz...")
        sky = Sky(nside=2048, preset_strings=[model], output_unit=u.uK_CMB)
        ksz_maps[model] = sky.get_emission(ksz_freq * u.GHz)[0].copy()
        hp.write_map(cache_file, ksz_maps[model].to_value(u.uK_CMB), overwrite=True)
        del sky

ksz_unit = str(next(iter(ksz_maps.values())).unit)
print(
    f"Loaded kSZ templates at {ksz_freq} GHz (units: {ksz_unit}) -> "
    + ", ".join(ksz_titles[model] for model in ksz_models)
)

Computing ksz1 kSZ at 143 GHz...

setting the output map dtype to [dtype('float64')]

Computing ksz2 kSZ at 143 GHz...

setting the output map dtype to [dtype('float64')]

Computing ksz3 kSZ at 143 GHz...

setting the output map dtype to [dtype('float64')]

Computing ksz5 kSZ at 143 GHz...

setting the output map dtype to [dtype('float64')]

Computing ksz6 kSZ at 143 GHz...

setting the output map dtype to [dtype('float64')]

Loaded kSZ templates at 143 GHz (units: uK_CMB) -> WebSky, AGORA (unlensed), AGORA, FLAMINGO L1_m9, FLAMINGO fgas-$8\sigma$

fig = plt.figure(figsize=(4.0 * len(ksz_models), 4.0))
stacked = np.concatenate([ksz_maps[model].value.ravel() for model in ksz_models])
vmax = np.percentile(np.abs(stacked), 99.5)
if vmax <= 0:
    vmax = np.max(np.abs(stacked))
if vmax == 0:
    vmax = 1.0
vmax = float(f"{vmax:.2g}")
vmin = -vmax
ksz_monopoles = {}
for col_idx, model in enumerate(ksz_models, start=1):
    component = ksz_maps[model]
    centered_map, mono = remove_monopole(component)
    ksz_monopoles[model] = mono
    hp.mollview(
        centered_map,
        sub=(1, len(ksz_models), col_idx),
        fig=fig,
        min=vmin,
        max=vmax,
        unit=str(component.unit),
        title="",
    )
    ax = plt.gca()
    ax.text(
        0.5,
        1.03,
        f"{ksz_titles[model]} @ {ksz_freq} GHz",
        transform=ax.transAxes,
        ha="center",
        va="bottom",
        fontsize=10,
        fontweight="bold",
    )

plt.subplots_adjust(top=0.85, wspace=0.25)

print("kSZ monopoles removed (units in", ksz_unit, "):")
for model in ksz_models:
    print(f"  {ksz_titles[model]}: {ksz_monopoles[model]:.3e}")

kSZ monopoles removed (units in uK_CMB ):
  WebSky: -1.866e-01 uK_CMB
  AGORA (unlensed): 4.191e-02 uK_CMB
  AGORA: 4.148e-02 uK_CMB
  FLAMINGO L1_m9: -7.113e-02 uK_CMB
  FLAMINGO fgas-$8\sigma$: -7.092e-02 uK_CMB

lmax = 6000
for model in ksz_models:
    cls = hp.anafast(ksz_maps[model].value, lmax=lmax)
    ls = np.arange(len(cls))
    lfac = (ls+1)*ls/2/np.pi
    plt.loglog(
        ls,
        lfac * cls,
        color=ksz_colors[model],
        linestyle=ksz_linestyles[model],
        linewidth=3,
        marker=ksz_markers.get(model),
        markevery=300,
        markersize=4,
        markerfacecolor="white",
        markeredgewidth=1.5,
        zorder=ksz_zorder.get(model, 3),
        label=ksz_titles[model],
    )

plt.xlim(50, lmax)
plt.ylim(0.05, None)
plt.legend(ncol=1, fontsize=13)
plt.ylabel(r"$D_\ell^{\rm kSZ}\; [\mu\mathrm{K}^2]$", fontsize=16)
plt.xlabel(r"$\ell$", fontsize=16)

Text(0.5, 0, '$\\ell$')

Reference: FLAMINGO (integrated maps) Figure 7 (kSZ)

Comparison note: The FLAMINGO fiducial model shown here is L1_m9, while the reference figure shows L2p8_m9. Stronger baryonic feedback leads to more diffuse gas and to lower kSZ power on small scales. On large scales, the WebSky and FLAMINGO simulations capture the field component from free electrons outside of halos.