"""Stellar population ingredients: present-day mass function, mass-luminosity relation, and the
number of main-sequence stars per unit stellar mass.  Used to turn "a fraction f of stars is
enclosed at temperature T" into a distribution of enclosure radii and lens masses."""
import os

import numpy as np

from common import RAW, radius_au

# Kroupa (2001) IMF segments: (m_lo, m_hi, alpha) with dN/dm ~ m^-alpha
KROUPA = [(0.01, 0.08, 0.3), (0.08, 0.5, 1.3), (0.5, 100.0, 2.3)]
M_MS_LO = 0.08
M_MS_HI = 1.0          # old population (bulge, most of the disc by number): turn-off near 1 Msun


def kroupa_xi(m):
    """Un-normalised continuous Kroupa IMF."""
    m = np.asarray(m, float)
    out = np.zeros_like(m)
    norm = 1.0
    prev_hi = None
    for lo, hi, a in KROUPA:
        if prev_hi is not None:
            norm *= prev_hi ** (-prev_a) / prev_hi ** (-a)
        sel = (m >= lo) & (m < hi)
        out[sel] = norm * m[sel] ** (-a)
        prev_hi, prev_a = hi, a
    return out


def remnant_mass(m):
    """Crude initial-final mass relation: white dwarfs 0.6 Msun (1-8), neutron stars 1.35 (8-20),
    black holes 7 Msun (>20).  Only used for the stellar-mass normalisation."""
    m = np.asarray(m, float)
    return np.where(m < 8, 0.6, np.where(m < 20, 1.35, 7.0))


def mass_grid(n=400, lo=M_MS_LO, hi=M_MS_HI):
    """Mid-points and number weights (sum to 1) of the main-sequence PDMF on a log grid."""
    edges = np.logspace(np.log10(lo), np.log10(hi), n + 1)
    mid = np.sqrt(edges[1:] * edges[:-1])
    w = kroupa_xi(mid) * np.diff(edges)
    return mid, w / w.sum()


def ms_stars_per_msun():
    """Main-sequence stars (0.08-1 Msun) per solar mass of stars + brown dwarfs + remnants."""
    m = np.logspace(-2, 2, 40001)
    xi = kroupa_xi(m)
    dm = np.gradient(m)
    ms = (m >= M_MS_LO) & (m < M_MS_HI)
    bd = m < M_MS_LO
    rem = m >= M_MS_HI
    n_ms = np.sum(xi[ms] * dm[ms])
    mass = np.sum(xi[ms] * m[ms] * dm[ms]) + np.sum(xi[bd] * m[bd] * dm[bd]) + np.sum(xi[rem] * remnant_mass(m[rem]) * dm[rem])
    return {"n_ms_per_msun_total": n_ms / mass,
            "mean_ms_mass": np.sum(xi[ms] * m[ms] * dm[ms]) / n_ms,
            "ms_mass_fraction": np.sum(xi[ms] * m[ms] * dm[ms]) / mass,
            "n_remnant_per_ms": np.sum(xi[rem] * dm[rem]) / n_ms,
            "n_bd_per_ms": np.sum(xi[bd] * dm[bd]) / n_ms}


def lens_mass_function():
    """All compact lenses (brown dwarfs, main sequence, remnants) for the microlensing validation:
    arrays of lens mass and number weight per unit total stellar mass (number per Msun)."""
    m = np.logspace(-2, 2, 4001)
    xi = kroupa_xi(m)
    dm = np.gradient(m)
    mlens = np.where(m < M_MS_HI, m, remnant_mass(m))
    mass = np.sum(xi * mlens * dm)
    return mlens, xi * dm / mass


_MAM = None


def _load_mamajek():
    global _MAM
    if _MAM is None:
        path = os.path.join(RAW, "mamajek_dwarf.txt")
        hdr, rows = None, []
        with open(path) as f:
            for line in f:
                if line.startswith("#SpT"):
                    hdr = line[1:].split()
                    continue
                if line.startswith("#") or not line.strip():
                    continue
                p = line.split()
                if hdr is not None and len(p) >= len(hdr):
                    rows.append(p[:len(hdr)])
        im, il = hdr.index("Msun"), hdr.index("logL")
        mm, ll = [], []
        for r in rows:
            try:
                a, b = float(r[im].replace(":", "")), float(r[il].replace(":", ""))
            except ValueError:
                continue
            mm.append(a)
            ll.append(b)
        mm, ll = np.array(mm), np.array(ll)
        o = np.argsort(mm)
        mm, ll = mm[o], ll[o]
        keep = np.concatenate([[True], np.diff(mm) > 0])
        _MAM = (np.log10(mm[keep]), ll[keep])
    return _MAM


def lum_of_mass(m):
    """Main-sequence luminosity (Lsun) from the Mamajek mean dwarf sequence (Pecaut & Mamajek 2013)."""
    lm, ll = _load_mamajek()
    return 10 ** np.interp(np.log10(m), lm, ll)


def population_radii(T, n=400, lo=M_MS_LO, hi=M_MS_HI):
    """Masses, number weights and enclosure radii (au) for the PDMF at temperature T."""
    m, w = mass_grid(n, lo, hi)
    return m, w, radius_au(lum_of_mass(m), T)