#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Generate publication-ready PDF parameter-space maps for a possible moon of
(612533) 2002 XV93.

The plotted quantity is tidal power dissipated in the primary:

    P_tide = (21/2) (k2/Q) G m_s^2 R^5 n e^2 / a^6

where n = sqrt(G(M+m_s)/a^3).

Four (sometimes overlapping) translucent regions are drawn:

    Green  : viable -> P_tide >= 10 MW and r_peri = a(1-e) > R_TNO
    Color A: orbit crosses inside the fluid Roche limit (a < a_Roche)
    Color B: pericenter inside the primary, r_peri = a(1-e) < R_TNO
    Color C: insufficient tidal power, P_tide < 10 MW

Each region is filled at alpha=0.25 so overlaps blend visibly.
"""

import math
from pathlib import Path

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import to_rgba
from matplotlib.patches import Patch


# ----------------------------
# Physical assumptions
# ----------------------------

G = 6.67430e-11              # m^3 kg^-1 s^-2

R_TNO = 255e3                # m
RHO_TNO = 1000.0             # kg m^-3 
RHO_MOON = 1000.0            # kg m^-3

M_TNO = (4.0 / 3.0) * math.pi * RHO_TNO * R_TNO**3

ROCHE_LIMIT_R_TNO_fluid = 2.44 * (RHO_TNO / RHO_MOON) ** (1.0 / 3.0)  # Fluid Roche limit
ROCHE_LIMIT_R_TNO_rigid = 1.26 * (RHO_TNO / RHO_MOON) ** (1.0 / 3.0)  # Fluid Roche limit

K2_OVER_Q = 1.0e-4
P_THRESHOLD = 10.0e6         # W

MOON_DIAMETERS_KM = [10, 25, 32, 50, 63, 75, 100]

# ----------------------------
# Numerical and plotting setup
# ----------------------------

A_MIN, A_MAX = 0.0, 15.0     # a/R_TNO
E_MIN, E_MAX = 0.0, 1.0

N_A = 1000
N_E = 800

OUTDIR = Path("xv93_tidal_maps_pdf")

# Region colors (each drawn at alpha=0.25).
GREEN        = "#3cb44b"     # viable
COLOR_ROCHE  = "#fffac8"     # a < Roche limit
COLOR_PERI   = "#ffd8b1"     # r_peri < R_TNO
COLOR_POWER  = "#fabed4"     # insufficient power

ALPHA = 1.0

DPI = 128
plt.rcParams['font.family'] = 'serif'
plt.rcParams['font.sans-serif'] = ['Times']
plt.rcParams.update({'font.size': 15})

plt.rcParams['mathtext.fontset'] = 'custom'
plt.rcParams['mathtext.rm'] = 'Times'
plt.rcParams['mathtext.it'] = 'Times:italic'
plt.rcParams['mathtext.bf'] = 'Times:bold'


def tidal_power(d_moon_km, a_over_R, ecc):
    """Return tidal power in W for a moon of diameter d_moon_km."""
    r_moon = 0.5 * d_moon_km * 1e3
    m_moon = (4.0 / 3.0) * math.pi * RHO_MOON * r_moon**3

    a = np.where(a_over_R > 0.0, a_over_R * R_TNO, np.nan)
    n = np.sqrt(G * (M_TNO + m_moon) / a**3)

    return ((21.0 / 2.0) * K2_OVER_Q * G * m_moon**2 * R_TNO**5 * n * ecc**2 / a**6)


def _overlay(ax, condition, color):
    """Fill the grid cells where `condition` is True with a translucent color."""
    rgba = np.zeros(condition.shape + (4,), dtype=float)
    r, g, b, _ = to_rgba(color)
    rgba[..., 0] = r
    rgba[..., 1] = g
    rgba[..., 2] = b
    rgba[..., 3] = np.where(condition, ALPHA, 0.0)

    ax.imshow(
        rgba,
        origin="lower",
        extent=[A_MIN, A_MAX, E_MIN, E_MAX],
        aspect="auto",
        interpolation="nearest",
    )


def make_map(d_moon_km, outdir=OUTDIR):
    """Create one PDF map for a given moon diameter."""
    outdir.mkdir(parents=True, exist_ok=True)

    fig, ax = plt.subplots(figsize=(1.5 * 1280 / 2 / DPI, 1 * 1280 / 2 / DPI), dpi=DPI)

    a_vals = np.linspace(A_MIN, A_MAX, N_A)
    e_vals = np.linspace(E_MIN, E_MAX, N_E)
    A, E = np.meshgrid(a_vals, e_vals)

    p_tide = tidal_power(d_moon_km, A, E)

    r_peri = A * (1.0 - E)

    # Individual conditions.
    # inside_roche      = A < ROCHE_LIMIT_R_TNO        # orbit crosses Roche limit
    inside_roche_f      = r_peri < ROCHE_LIMIT_R_TNO_fluid        # orbit crosses Roche limit
    inside_roche_r      = r_peri < ROCHE_LIMIT_R_TNO_rigid        # orbit crosses Roche limit
    peri_inside_tno   = r_peri < 1.0                 # r_peri < R_TNO
    insufficient_pow  = p_tide < P_THRESHOLD         # P_tide < 10 MW

    viable = (~peri_inside_tno) & (~insufficient_pow)

    # Four translucent overlays (all at alpha=0.25 so overlaps blend).
    _overlay(ax, viable, GREEN)
    _overlay(ax, insufficient_pow, COLOR_PERI)
    # _overlay(ax, inside_roche_f, COLOR_ROCHE)
    # _overlay(ax, inside_roche_r, COLOR_ROCHE)
    _overlay(ax, peri_inside_tno, COLOR_ROCHE)

    # R_TNO boundary: a(1-e)=R_TNO -> e=1-1/a.
    e_curve = np.linspace(0.0, 0.99, 500)
    # a_curve = 1.0 / (1.0 - e_curve)
    # mask1 = a_curve <= A_MAX
    # ax.plot(a_curve[mask1], e_curve[mask1], "--", lw=1.4, c="#f58231", alpha=0.5)

    # # R_Roche boundary: a(1-e)=R_TNO -> e=1-1/a.
    a_curve = ROCHE_LIMIT_R_TNO_fluid / (1.0 - e_curve)
    mask2 = a_curve <= A_MAX
    ax.plot(a_curve[mask2], e_curve[mask2], "--", lw=1.4, c='k')

    a_curve = ROCHE_LIMIT_R_TNO_rigid / (1.0 - e_curve)
    mask2 = a_curve <= A_MAX
    ax.plot(a_curve[mask2], e_curve[mask2], "--", lw=1.4, c="k")

    # ax.axvline(ROCHE_LIMIT_R_TNO_fluid, ls="--", c='k')
    # ax.axvline(ROCHE_LIMIT_R_TNO_rigid, ls="--", c='k')

    # Fluid Roche limit.
    # ax.axvline(ROCHE_LIMIT_R_TNO, color="k", linestyle="--", lw=1.4)
    ax.text(
        # ROCHE_LIMIT_R_TNO_fluid - 0.55,
        10,
        0.92,
        r"$R_{\mathrm{Roche,\ rigid}}$",
        rotation=0,
        # va="top",
        # ha="left",
    )
    ax.text(
        # ROCHE_LIMIT_R_TNO_rigid - 0.55,
        10,
        0.82,
        r"$R_{\mathrm{Roche,\ fluid}}$",
        rotation=0,
        # va="top",
        # ha="left",
    )

    ax.text(
        0.5, 0.9,
        r"$\mathbf{D_{\rm s}=%d~km}$" % d_moon_km, fontsize=22
    )

    ax.set_xlim(A_MIN, A_MAX)
    ax.set_ylim(E_MIN, E_MAX)
    ax.set_xlabel(r"Semi-major axis $(R_{\rm a})$")
    ax.set_ylabel(r"Eccentricity")

    legend_handles = [
        Patch(facecolor=to_rgba(COLOR_ROCHE, ALPHA), label=r"$r_{\rm peri} < R_{\rm a}$"),
        # Patch(facecolor=to_rgba(COLOR_ROCHE, ALPHA), label=r"$a < R_{\rm Roche}$"),
        # Patch(facecolor=to_rgba(COLOR_ROCHE, ALPHA), label=r"$r_{\rm peri} < R_{\rm Roche}$"),
        Patch(facecolor=to_rgba(COLOR_PERI, ALPHA), label=r"$\dot{E}_{\rm tide} < 10$ MW"),
        Patch(facecolor=to_rgba(GREEN, ALPHA), label=r"Viable"),
    ]
    ax.legend(handles=legend_handles, loc="lower right", framealpha=0.9)

    fig.tight_layout()

    output_path = outdir / f"2002_xv93_tidal_{d_moon_km:g}km-1000.pdf"
    fig.savefig(output_path, bbox_inches="tight")
    plt.close(fig)
    return output_path


def main():
    print(f"\n M_TNO = {M_TNO:.6e} kg")
    print(f" Fluid Roche limit = {ROCHE_LIMIT_R_TNO_fluid:.3f} R_TNO")
    print(f" Fluid Roche limit = {ROCHE_LIMIT_R_TNO_rigid:.3f} R_TNO")
    print(f"\n Output directory: {OUTDIR.resolve()}\n")

    for d_moon_km in MOON_DIAMETERS_KM:
        output_path = make_map(d_moon_km)
        print(" " + str(output_path))


if __name__ == "__main__":
    main()
    print("\n")