"""Main-sequence relations (Mamajek / Pecaut & Mamajek 2013 mean dwarf sequence) and the
astrometric mass-ratio function (AMRF) of Shahaf et al. (2019, 2023)."""
import os

import numpy as np
import pandas as pd

from common import RAW


def load_mamajek(path=None):
    path = path or 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)])
    df = pd.DataFrame(rows, columns=hdr)
    for c in df.columns:
        if c not in ("SpT", "#SpT"):
            df[c] = pd.to_numeric(df[c].astype(str).str.replace(":", "", regex=False), errors="coerce")
    return df


class MainSequence:
    """Interpolators in stellar mass for the mean dwarf sequence, 0.075 - 2.5 Msun."""

    def __init__(self):
        t = load_mamajek()
        t = t[(t.Msun >= 0.075) & (t.Msun <= 3.5)].dropna(subset=["Msun", "M_G", "logL"])
        t = t.sort_values("Msun").drop_duplicates("Msun")
        # enforce monotonic M_G(M) (the table has small non-monotonic wiggles around A6-A9)
        mg = np.minimum.accumulate(t.M_G.values)
        self.m = t.Msun.values
        self.logm = np.log10(self.m)
        self.MG = mg
        self.logL = t.logL.values
        tk = t.dropna(subset=["M_Ks", "Ks-W1"])
        self._mk_logm = np.log10(tk.Msun.values)
        self._MW1 = (tk.M_Ks - tk["Ks-W1"]).values
        tc = t.dropna(subset=["Bp-Rp"])
        self._c_logm = np.log10(tc.Msun.values)
        self._bprp = tc["Bp-Rp"].values
        self._c_MG = np.minimum.accumulate(tc.M_G.values)

    def MG_of_mass(self, m):
        return np.interp(np.log10(m), self.logm, self.MG)

    def mass_of_MG(self, MG):
        return 10 ** np.interp(-np.asarray(MG), -self.MG, self.logm)

    def logL_of_mass(self, m):
        return np.interp(np.log10(m), self.logm, self.logL)

    def MW1_of_mass(self, m):
        return np.interp(np.log10(m), self._mk_logm, self._MW1)

    def bprp_of_MG(self, MG):
        return np.interp(-np.asarray(MG), -self._c_MG, self._bprp)

    def flux_ratio_G(self, m2, m1):
        """G-band flux ratio of a main-sequence star of mass m2 to one of mass m1."""
        m2 = np.asarray(m2, float)
        s = 10 ** (-0.4 * (self.MG_of_mass(np.clip(m2, 0.075, 3.5)) - self.MG_of_mass(m1)))
        return np.where(m2 < 0.075, 0.0, s)


def amrf_dark(q):
    """AMRF of a non-luminous companion: A = q (1+q)^(-2/3)."""
    q = np.asarray(q, float)
    return q / (1 + q) ** (2.0 / 3.0)


_QGRID = np.logspace(-3, 2.5, 4000)
_AGRID = amrf_dark(_QGRID)


def q_from_amrf_dark(A):
    """Invert A = q (1+q)^(-2/3) (monotonic)."""
    return np.interp(A, _AGRID, _QGRID)


def amrf_luminous(q, S):
    """AMRF when the companion contributes a flux ratio S = F2/F1 in the G band."""
    q = np.asarray(q, float)
    return np.abs(q / (1 + q) - S / (1 + S)) * (1 + q) ** (1.0 / 3.0)


def amrf_limits(ms, m1, nq=400):
    """Largest AMRF reachable with (i) a single main-sequence companion (q <= 1) and
    (ii) a companion that is itself a close pair of equal main-sequence stars (q <= 2).
    These are the class-I/II and class-II/III boundaries of Shahaf et al. (2019)."""
    m1 = np.atleast_1d(np.asarray(m1, float))
    q1 = np.linspace(0.01, 1.0, nq)
    q2 = np.linspace(0.02, 2.0, nq)
    a_ms = np.empty(len(m1))
    a_tr = np.empty(len(m1))
    for i, m in enumerate(m1):
        S1 = ms.flux_ratio_G(q1 * m, m)
        a_ms[i] = np.max(amrf_luminous(q1, S1))
        S2 = 2 * ms.flux_ratio_G(q2 * m / 2, m)
        a_tr[i] = np.max(amrf_luminous(q2, S2))
    return a_ms, a_tr


def a0_from_thiele_innes(A, B, F, G):
    u = 0.5 * (A * A + B * B + F * F + G * G)
    v = A * G - B * F
    return np.sqrt(u + np.sqrt(np.maximum(u * u - v * v, 0.0)))