"""Simple analytic Milky Way star-count model used to count stars inside survey beams.

Number density of stars (pc^-3) = thin disk + thick disk + halo (Juric et al. 2008, bias-corrected
parameters) + axisymmetric bulge (McMillan 2017 / Bissantz & Gerhard 2002 form).
The local normalisation is MEASURED from the Gaia Catalogue of Nearby Stars (see calibrate_local()).
The bulge is converted from mass to star number with the same stars-per-solar-mass ratio as the
solar neighbourhood.  No nuclear star cluster / nuclear stellar disk is included (conservative for
Galactic-centre pointings).  Dust is irrelevant at radio wavelengths.
"""
import os
import numpy as np
import pandas as pd

from common import RAW

R0 = 8200.0          # pc, Sun-Galactic centre distance (GRAVITY 2019: 8178 pc)
# Juric et al. (2008) Table 10 bias-corrected values
HZ1, HR1 = 300.0, 2600.0
HZ2, HR2 = 900.0, 3600.0
F_THICK = 0.12
F_HALO, Q_HALO, N_HALO = 0.005, 0.64, 2.77
# McMillan (2017) bulge: rho = rho0 / (1+r'/r0)^alpha * exp(-(r'/rcut)^2), r'^2 = R^2 + (z/q)^2
B_RHO0 = 98.4        # Msun / pc^3   (9.84e10 Msun/kpc^3)
B_R0, B_RCUT, B_ALPHA, B_Q = 75.0, 2100.0, 1.8, 0.5
RHO_STAR_LOCAL = 0.043   # Msun/pc^3 local stellar mass density incl. remnants and BDs (McKee et al. 2015)
RMAX_DISK = 25000.0      # pc, truncation radius of the disks

_N_LOCAL = None


def calibrate_local(rmax_pc=50.0, verbose=False):
    """Local number density of stars (pc^-3) from the Gaia Catalogue of Nearby Stars (GCNS).

    Counts GCNS entries with median distance < rmax_pc (GCNSprob>0.5) and divides by the volume.
    Unresolved companions are not counted (so this is closer to a density of resolved objects/systems);
    the GCNS is ~92% complete for M9 at 100 pc and essentially complete within 50 pc for stars.
    """
    global _N_LOCAL
    fn = os.path.join(RAW, "vizier", "gcns_table1c.csv")
    g = pd.read_csv(fn, usecols=["Dist50", "GCNSprob"])
    d = g["Dist50"].values * 1000.0
    ok = g["GCNSprob"].values > 0.5
    out = {}
    for r in (15.0, 25.0, 50.0, 100.0):
        n = int(np.sum(ok & (d < r)))
        out[r] = (n, n / (4.0 / 3.0 * np.pi * r ** 3))
    _N_LOCAL = out[rmax_pc][1]
    if verbose:
        for r, (n, dens) in out.items():
            print(f"   GCNS: N(<{r:.0f} pc) = {n:7d}  ->  n = {dens:.4f} pc^-3")
    return _N_LOCAL, out


def n_local():
    if _N_LOCAL is None:
        calibrate_local()
    return _N_LOCAL


def density(R, z, components=False):
    """Star number density (pc^-3) at Galactocentric cylindrical radius R (pc) and height z (pc)."""
    nl = n_local()
    n_thin0 = nl / (1.0 + F_THICK + F_HALO)
    az = np.abs(z)
    indisk = (R < RMAX_DISK)
    thin = n_thin0 * np.exp(-(R - R0) / HR1 - az / HZ1) * indisk
    thick = n_thin0 * F_THICK * np.exp(-(R - R0) / HR2 - az / HZ2) * indisk
    rh = np.sqrt(R ** 2 + (z / Q_HALO) ** 2)
    halo = n_thin0 * F_HALO * (R0 / np.maximum(rh, 500.0)) ** N_HALO * (rh < 60000.0)
    rp = np.sqrt(R ** 2 + (z / B_Q) ** 2)
    stars_per_msun = nl / RHO_STAR_LOCAL
    bulge = stars_per_msun * B_RHO0 / (1.0 + rp / B_R0) ** B_ALPHA * np.exp(-(rp / B_RCUT) ** 2)
    if components:
        return thin, thick, halo, bulge
    return thin + thick + halo + bulge


def los_density(l_deg, b_deg, d_pc):
    """Density along lines of sight. l_deg, b_deg arrays (n,), d_pc array (m,) -> (n, m)."""
    l = np.radians(np.atleast_1d(l_deg))[:, None]
    b = np.radians(np.atleast_1d(b_deg))[:, None]
    d = np.atleast_1d(d_pc)[None, :]
    x = R0 - d * np.cos(b) * np.cos(l)
    y = d * np.cos(b) * np.sin(l)
    z = d * np.sin(b)
    return density(np.sqrt(x * x + y * y), z)


# distance grid for all line-of-sight integrals (pc); bin edges and geometric centres
D_EDGES = np.logspace(-0.5, np.log10(60000.0), 321)
D_MID = np.sqrt(D_EDGES[:-1] * D_EDGES[1:])
D_VOL = (D_EDGES[1:] ** 3 - D_EDGES[:-1] ** 3) / 3.0     # int d^2 dd over the bin (pc^3 per sr)


def los_counts_per_sr(l_deg, b_deg):
    """dN per steradian in each distance bin for each line of sight: array (n, len(D_MID))."""
    return los_density(l_deg, b_deg, D_MID) * D_VOL[None, :]


# ---------------------------------------------------------------- simple kinematic model (for the HI rest-frame prior)
# A beacon tuned to the 21-cm line in the rest frame of its own star arrives Doppler-shifted by the star's
# line-of-sight velocity.  Mean rotation V_phi and 1-D velocity dispersion per component (km/s); flat
# rotation curve, V0 = 233 km/s at R0 (McMillan 2017); velocities relative to the local standard of rest.
V0 = 233.0
KIN = ((V0, 30.0),          # thin disc
       (V0 - 50.0, 50.0),   # thick disc (asymmetric drift ~50 km/s)
       (0.0, 120.0),        # halo (no net rotation)
       (50.0, 110.0))       # bulge (slow rotator)


def los_counts_vwin(l_deg, b_deg, windows):
    """Like los_counts_per_sr, but also returns, for every velocity window (v_lo, v_hi) in km/s, the counts
    weighted by the probability that a star's LSR line-of-sight velocity lies inside the window.
    Returns (counts (n, m), [counts_in_window (n, m) for each window])."""
    from scipy.special import ndtr
    l = np.radians(np.atleast_1d(l_deg))[:, None]
    b = np.radians(np.atleast_1d(b_deg))[:, None]
    d = D_MID[None, :]
    x = R0 - d * np.cos(b) * np.cos(l)
    y = d * np.cos(b) * np.sin(l)
    z = d * np.sin(b)
    R = np.sqrt(x * x + y * y)
    comps = density(R, z, components=True)
    slcb = np.sin(l) * np.cos(b)
    rr = R0 / np.maximum(R, 1.0)
    tot = np.zeros(R.shape)
    out = [np.zeros(R.shape) for _ in windows]
    for dens, (vphi, sig) in zip(comps, KIN):
        cnt = dens * D_VOL[None, :]
        tot += cnt
        if windows:
            vbar = (vphi * rr - V0) * slcb
            for k, (vlo, vhi) in enumerate(windows):
                out[k] += cnt * (ndtr((vhi - vbar) / sig) - ndtr((vlo - vbar) / sig))
    return tot, out


def total_stars():
    """Total number of stars in the model Galaxy and per component (numerical integral)."""
    R = np.linspace(0.5, 60000.0, 6000)
    z = np.concatenate([np.linspace(0, 3000, 3001), np.linspace(3010, 60000, 2000)])
    RR, ZZ = np.meshgrid(R, z, indexing="ij")
    comps = density(RR, ZZ, components=True)
    out = []
    for c in comps:
        iz = 2.0 * np.trapezoid(c, z, axis=1)
        out.append(np.trapezoid(iz * 2.0 * np.pi * R, R))
    return dict(thin=out[0], thick=out[1], halo=out[2], bulge=out[3], total=float(sum(out)))


def sky_grid():
    """Quasi equal-importance (l, b) grid with fine latitude sampling near the plane.
    Returns l (deg), b (deg), solid angle (sr) per cell."""
    b_edges = np.concatenate([np.arange(-90, -10, 1.0), np.arange(-10, 10, 0.125), np.arange(10, 90.01, 1.0)])
    b_mid = 0.5 * (b_edges[:-1] + b_edges[1:])
    dsinb = np.sin(np.radians(b_edges[1:])) - np.sin(np.radians(b_edges[:-1]))
    ls, bs, om = [], [], []
    for bm, ds in zip(b_mid, dsinb):
        dl = 0.25 if abs(bm) < 10 else 1.0
        l_mid = np.arange(dl / 2, 360, dl)
        ls.append(l_mid); bs.append(np.full_like(l_mid, bm)); om.append(np.full_like(l_mid, ds * np.radians(dl)))
    return np.concatenate(ls), np.concatenate(bs), np.concatenate(om)


def gal_to_equ(l_deg, b_deg):
    """Galactic -> ICRS (RA, Dec) in degrees (standard J2000 rotation)."""
    l = np.radians(l_deg); b = np.radians(b_deg)
    ra_gp, dec_gp, l_ncp = np.radians(192.85948), np.radians(27.12825), np.radians(122.93192)
    sd = np.sin(b) * np.sin(dec_gp) + np.cos(b) * np.cos(dec_gp) * np.cos(l_ncp - l)
    dec = np.arcsin(np.clip(sd, -1, 1))
    y = np.cos(b) * np.sin(l_ncp - l)
    x = np.sin(b) * np.cos(dec_gp) - np.cos(b) * np.sin(dec_gp) * np.cos(l_ncp - l)
    ra = (np.arctan2(y, x) + ra_gp) % (2 * np.pi)
    return np.degrees(ra), np.degrees(dec)


def equ_to_gal(ra_deg, dec_deg):
    """ICRS -> Galactic (l, b) in degrees."""
    ra = np.radians(ra_deg); dec = np.radians(dec_deg)
    ra_gp, dec_gp, l_ncp = np.radians(192.85948), np.radians(27.12825), np.radians(122.93192)
    sb = np.sin(dec) * np.sin(dec_gp) + np.cos(dec) * np.cos(dec_gp) * np.cos(ra - ra_gp)
    b = np.arcsin(np.clip(sb, -1, 1))
    y = np.cos(dec) * np.sin(ra - ra_gp)
    x = np.sin(dec) * np.cos(dec_gp) - np.cos(dec) * np.sin(dec_gp) * np.cos(ra - ra_gp)
    l = (l_ncp - np.arctan2(y, x)) % (2 * np.pi)
    return np.degrees(l), np.degrees(b)


if __name__ == "__main__":
    nl, tab = calibrate_local(verbose=True)
    print("local density adopted:", nl)
    print({k: f"{v:.3e}" for k, v in total_stars().items()})
    print(equ_to_gal(266.405, -28.936), gal_to_equ(0.0, 0.0))