#!/usr/bin/env python
"""Extract extinction-density profiles A'(d) along each <3kpc HMSFR sightline from 3D dust maps.

Maps:
  - Edenhofer+2024 (Gaia BP/RP; local lbd cube; 69-1250 pc, 1 pc steps, 13.7' pix)
  - Bayestar19 (PS1+2MASS+Gaia; dec>-30; via dustmaps, once downloaded)
  - DECaPS (Zhang/Green/Saydjari southern plane; via dustmaps, once downloaded)

For each target: cumulative extinction profile + dA/dd density; find jumps (cloud candidates)
via peak-finding on the density; save profiles npz + summary csv.
"""
import os
import numpy as np
from astropy.io import fits
from astropy.coordinates import SkyCoord
import astropy.units as u
import csv

EDEN = '/orange/adamginsburg/galactic_plane_surveys/dustmaps/edenhofer_mean_and_std_lbd.fits'
OUT = 'extinction_profiles'
os.makedirs(OUT, exist_ok=True)


def load_targets():
    targs = []
    with open('master_distances.csv') as fh:
        for r in csv.DictReader(fh):
            if r['in_3kpc'] != 'True' and r['in_2kpc'] != 'True':
                continue
            targs.append(dict(name=r['Name'], glon=float(r['GLON']), glat=float(r['GLAT']),
                              d_best=float(r['d_best']), method=r['method'], ref=r['ref'],
                              conf=r['conf'], sample='RMS'))
    with open('curated_nearby_sample.csv') as fh:
        for r in csv.DictReader(fh):
            targs.append(dict(name=r['Name'], glon=float(r['GLON']), glat=float(r['GLAT']),
                              d_best=float(r['Dist_kpc']), method=r['Dist_method'],
                              ref=r['Dist_ref'], conf='curated', sample='curated'))
    return targs


class EdenhoferCube:
    """Direct reader for the lbd-grid Edenhofer mean+std cube."""

    def __init__(self, path=EDEN):
        # load fully into RAM: strided memmap access over GPFS is pathologically slow
        self.hdul = fits.open(path, memmap=False)
        hd = self.hdul['MEAN'].header
        # axes: NAXIS1=glon(1572), NAXIS2=glat(786), NAXIS3=distance(1182)
        self.nl, self.nb, self.nd = hd['NAXIS1'], hd['NAXIS2'], hd['NAXIS3']
        self.crval1, self.cdelt1, self.crpix1 = hd['CRVAL1'], hd['CDELT1'], hd['CRPIX1']
        self.crval2, self.cdelt2, self.crpix2 = hd['CRVAL2'], hd['CDELT2'], hd['CRPIX2']
        self.dist = hd['CRVAL3'] + hd['CDELT3'] * np.arange(self.nd)  # pc

    def pix(self, glon, glat):
        gl = ((glon + 180) % 360) - 180
        i = int(round(self.crpix1 - 1 + (gl - self.crval1) / self.cdelt1))
        j = int(round(self.crpix2 - 1 + (glat - self.crval2) / self.cdelt2))
        return i, j

    def profile(self, glon, glat, box=1):
        """Return (dist_pc, density, density_std) averaged over (2*box+1)^2 pixels."""
        i, j = self.pix(glon, glat)
        i0, i1 = max(0, i - box), min(self.nl, i + box + 1)
        j0, j1 = max(0, j - box), min(self.nb, j + box + 1)
        # data layout [dist, glat, glon]
        dens = np.nanmean(self.hdul['MEAN'].data[:, j0:j1, i0:i1], axis=(1, 2))
        std = np.nanmean(self.hdul['STD.'].data[:, j0:j1, i0:i1], axis=(1, 2))
        return self.dist, dens.astype(float), std.astype(float)


def find_jumps(dist, dens, smooth_pc=15, min_prom_frac=0.05):
    """Peaks in extinction density = cloud distance candidates.
    Returns list of (d_peak_pc, dA_integrated, prominence_frac)."""
    from scipy.ndimage import gaussian_filter1d
    from scipy.signal import find_peaks
    dd = np.median(np.diff(dist))
    sm = gaussian_filter1d(dens, smooth_pc / dd)
    total = np.trapezoid(sm, dist)
    if total <= 0:
        return []
    pk, props = find_peaks(sm, prominence=min_prom_frac * np.nanmax(sm))
    out = []
    for p, prom in zip(pk, props['prominences']):
        # integrate peak between surrounding minima
        lo = props['left_bases'][list(pk).index(p)]
        hi = props['right_bases'][list(pk).index(p)]
        dA = np.trapezoid(sm[lo:hi + 1], dist[lo:hi + 1])
        out.append((float(dist[p]), float(dA), float(prom / np.nanmax(sm))))
    out.sort(key=lambda t: -t[1])
    return out


def main():
    targs = load_targets()
    cube = EdenhoferCube()
    rows = []
    prof_store = {}
    for t in targs:
        d, dens, std = cube.profile(t['glon'], t['glat'])
        jumps = find_jumps(d, dens)
        prof_store[t['name']] = np.vstack([d, dens, std])
        # nearest significant jump to adopted distance, if within cube
        d_adopt = t['d_best'] * 1000
        best_jump = None
        if jumps and d_adopt < 1400:
            cands = [j for j in jumps if j[1] > 0.1 * sum(x[1] for x in jumps)]
            if cands:
                best_jump = min(cands, key=lambda j: abs(j[0] - d_adopt))
        rows.append(dict(
            name=t['name'], sample=t['sample'], glon=round(t['glon'], 4), glat=round(t['glat'], 4),
            d_best_kpc=t['d_best'], method=t['method'], conf=t['conf'],
            n_jumps=len(jumps),
            jumps_pc=';'.join(f'{j[0]:.0f}' for j in jumps[:5]),
            jumps_dA=';'.join(f'{j[1]:.2f}' for j in jumps[:5]),
            eden_match_pc=(f'{best_jump[0]:.0f}' if best_jump else ''),
            eden_offset_pct=(f'{100 * (best_jump[0] - d_adopt) / d_adopt:.0f}' if best_jump else ''),
        ))
        print(f"{t['name']:24s} d={t['d_best']:.2f} {t['method']:16s} "
              f"jumps={rows[-1]['jumps_pc']} match={rows[-1]['eden_match_pc']}")
    np.savez_compressed(os.path.join(OUT, 'edenhofer_profiles.npz'), **prof_store)
    with open(os.path.join(OUT, 'edenhofer_jumps.csv'), 'w', newline='') as fh:
        w = csv.DictWriter(fh, fieldnames=list(rows[0].keys()))
        w.writeheader()
        w.writerows(rows)
    print(f'\n{len(rows)} targets -> {OUT}/edenhofer_jumps.csv')


if __name__ == '__main__':
    main()