#!/usr/bin/env python
"""Bayestar19 (dec>-30) + DECaPS (southern plane) cumulative-extinction profiles for all
<3kpc HMSFR sightlines; jump-finding on d(cum)/dd; writes profiles npz + jumps csv.
Run after extinction_profiles.py (reuses its load_targets/find_jumps)."""
import os
import csv
import numpy as np
from astropy.coordinates import SkyCoord
import astropy.units as u
from extinction_profiles import load_targets, find_jumps, OUT


def profile_from_cumulative(query, coord_grid):
    """coord_grid: SkyCoord array along one sightline at many distances -> cumulative, density."""
    raise NotImplementedError


def run_map(name, query, dists_pc, coords_lb, dec_min=None, dec_max=None):
    rows, store = [], {}
    for t in coords_lb:
        c0 = SkyCoord(t['glon'] * u.deg, t['glat'] * u.deg, frame='galactic')
        dec = c0.icrs.dec.deg
        if dec_min is not None and dec < dec_min:
            continue
        if dec_max is not None and dec > dec_max:
            continue
        coords = SkyCoord([t['glon']] * len(dists_pc) * u.deg,
                          [t['glat']] * len(dists_pc) * u.deg,
                          distance=dists_pc * u.pc, frame='galactic')
        try:
            cum = query(coords, mode='mean')
        except (ValueError, TypeError) as e:
            print(f"{t['name']}: query fail {e}")
            continue
        if cum is None or not np.any(np.isfinite(cum)):
            print(f"{t['name']}: no coverage")
            continue
        cum = np.asarray(cum, float)
        dens = np.gradient(cum, dists_pc)
        dens[~np.isfinite(dens)] = 0
        dens[dens < 0] = 0
        jumps = find_jumps(dists_pc, dens, smooth_pc=30)
        store[t['name']] = np.vstack([dists_pc, cum, dens])
        d_adopt = t['d_best'] * 1000
        best = None
        if jumps:
            sig = [j for j in jumps if j[1] > 0.1 * sum(x[1] for x in jumps)]
            if sig:
                best = min(sig, key=lambda j: abs(j[0] - d_adopt))
        rows.append(dict(name=t['name'], sample=t['sample'], 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[:6]),
                         jumps_dA=';'.join(f'{j[1]:.3f}' for j in jumps[:6]),
                         match_pc=(f'{best[0]:.0f}' if best else ''),
                         offset_pct=(f'{100 * (best[0] - d_adopt) / d_adopt:.0f}' if best else '')))
        print(f"{t['name']:24s} d={t['d_best']:.2f} jumps={rows[-1]['jumps_pc']} "
              f"match={rows[-1]['match_pc']}")
    np.savez_compressed(os.path.join(OUT, f'{name}_profiles.npz'), **store)
    if rows:
        with open(os.path.join(OUT, f'{name}_jumps.csv'), 'w', newline='') as fh:
            w = csv.DictWriter(fh, fieldnames=list(rows[0].keys()))
            w.writeheader()
            w.writerows(rows)
    print(f'{name}: {len(rows)} targets written')


def main(which):
    targs = load_targets()
    dists = np.arange(60, 6001, 10.0)  # pc
    if which in ('bayestar', 'both'):
        from dustmaps.bayestar import BayestarQuery
        bq = BayestarQuery(version='bayestar2019')
        run_map('bayestar19', bq, dists, targs, dec_min=-28.0)
    if which in ('decaps', 'both'):
        from dustmaps.decaps import DECaPSQueryLite
        dq = DECaPSQueryLite(mean_only=True)
        run_map('decaps', dq, dists, targs, dec_max=0.0)


if __name__ == '__main__':
    import sys
    main(sys.argv[1] if len(sys.argv) > 1 else 'both')