#!/usr/bin/env python3
"""Reproduce the published 1-D constant-power NEP lower-bound trade points.

This standalone implementation follows Appendix A, Eqs. 34-40 of
Turyshev, arXiv:2602.04198v3. Values are model outputs, not flight performance.
Run from this directory; the script writes nep-trade.csv and nep-trade.svg.
"""
from __future__ import annotations

import csv
import math
from pathlib import Path

HERE = Path(__file__).resolve().parent
AU_M = 149_597_870_700.0
YEAR_S = 365.25 * 24.0 * 60.0 * 60.0
G0 = 9.80665
M0_KG = 20_000.0
PAYLOAD_KG = 800.0
ETA = 0.7
ISP_S = 9_000.0
DISTANCE_M = 650.0 * AU_M
ALPHA_VALUES = (10.0, 12.0, 15.0, 20.0, 30.0, 40.0)


def design(alpha: float, dry_mass: float) -> dict[str, float]:
    exhaust_velocity = G0 * ISP_S
    power_kw = (dry_mass - PAYLOAD_KG) / alpha
    thrust_n = 2.0 * ETA * power_kw * 1000.0 / exhaust_velocity
    mass_flow = thrust_n / exhaust_velocity
    burn_time = (M0_KG - dry_mass) / mass_flow
    final_speed = exhaust_velocity * math.log(M0_KG / dry_mass)
    burn_distance = exhaust_velocity / mass_flow * (
        M0_KG - dry_mass - dry_mass * math.log(M0_KG / dry_mass)
    )
    if burn_distance >= DISTANCE_M:
        lo, hi = 0.0, burn_time
        for _ in range(100):
            mid = (lo + hi) / 2.0
            mass = M0_KG - mass_flow * mid
            distance = exhaust_velocity / mass_flow * (
                M0_KG - mass - mass * math.log(M0_KG / mass)
            )
            if distance < DISTANCE_M:
                lo = mid
            else:
                hi = mid
        arrival_time = (lo + hi) / 2.0
    else:
        arrival_time = burn_time + (DISTANCE_M - burn_distance) / final_speed
    return {
        "power_kw": power_kw,
        "mass_ratio": M0_KG / dry_mass,
        "propellant_t": (M0_KG - dry_mass) / 1000.0,
        "terminal_km_s": final_speed / 1000.0,
        "burn_years": burn_time / YEAR_S,
        "time_years": arrival_time / YEAR_S,
    }


def optimize(alpha: float) -> dict[str, float]:
    # Golden-section minimization over the dry mass, as in the appendix sweep.
    lo, hi = PAYLOAD_KG + 1e-5, M0_KG - 1e-5
    phi = (math.sqrt(5.0) - 1.0) / 2.0
    x1, x2 = hi - phi * (hi - lo), lo + phi * (hi - lo)
    y1, y2 = design(alpha, x1)["time_years"], design(alpha, x2)["time_years"]
    for _ in range(160):
        if y1 < y2:
            hi, x2, y2 = x2, x1, y1
            x1 = hi - phi * (hi - lo)
            y1 = design(alpha, x1)["time_years"]
        else:
            lo, x1, y1 = x1, x2, y2
            x2 = lo + phi * (hi - lo)
            y2 = design(alpha, x2)["time_years"]
    dry_mass = (lo + hi) / 2.0
    result = design(alpha, dry_mass)
    result["alpha_kg_per_kw"] = alpha
    result["dry_mass_kg"] = dry_mass
    return result


def make_svg(rows: list[dict[str, float]]) -> str:
    width, height = 980, 560
    left, right, top, bottom = 100, 45, 60, 95
    x0, x1 = 8.0, 42.0
    y0, y1 = 24.0, 43.0
    sx = lambda x: left + (x - x0) / (x1 - x0) * (width - left - right)
    sy = lambda y: top + (y1 - y) / (y1 - y0) * (height - top - bottom)
    parts = [
        '<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 980 560" role="img" aria-labelledby="title desc">',
        '<title id="title">Modeled NEP outbound time increases with integrated stage specific mass</title>',
        '<desc id="desc">Six points from the published one-dimensional constant-power lower-bound model: 10 to 40 kilograms per kilowatt correspond to about 26.9 to 40.2 years to travel 650 astronomical units. This is not a complete mission timeline.</desc>',
        '<rect width="980" height="560" rx="22" fill="#07182e"/>',
        '<g font-family="Arial,sans-serif" fill="#e8f2ff">',
        '<text x="100" y="35" font-size="21" font-weight="700">ONE-DIMENSIONAL CONSTANT-POWER MODEL · 20 t SPACECRAFT</text>',
    ]
    for y in (25, 30, 35, 40):
        parts += [f'<line x1="{left}" y1="{sy(y):.1f}" x2="{width-right}" y2="{sy(y):.1f}" stroke="#294462"/>',
                  f'<text x="{left-16}" y="{sy(y)+6:.1f}" text-anchor="end" font-size="15">{y}</text>']
    for x in (10, 15, 20, 25, 30, 35, 40):
        parts += [f'<line x1="{sx(x):.1f}" y1="{top}" x2="{sx(x):.1f}" y2="{height-bottom}" stroke="#1b3551"/>',
                  f'<text x="{sx(x):.1f}" y="{height-bottom+25}" text-anchor="middle" font-size="15">{x}</text>']
    pts = " ".join(f'{sx(r["alpha_kg_per_kw"]):.1f},{sy(r["time_years"]):.1f}' for r in rows)
    parts += [f'<polyline points="{pts}" fill="none" stroke="#59d9d0" stroke-width="4"/>']
    for r in rows:
        x, y = sx(r["alpha_kg_per_kw"]), sy(r["time_years"])
        parts += [f'<circle cx="{x:.1f}" cy="{y:.1f}" r="7" fill="#ff72c8" stroke="#fff" stroke-width="2"/>',
                  f'<text x="{x:.1f}" y="{y-15:.1f}" text-anchor="middle" font-size="14">{r["time_years"]:.1f}</text>']
    parts += [
        f'<text x="{left+(width-left-right)/2}" y="{height-22}" text-anchor="middle" font-size="17">Integrated power + propulsion specific mass, α (kg/kWe)</text>',
        f'<text x="26" y="{(top+height-bottom)/2}" text-anchor="middle" font-size="17" transform="rotate(-90 26 {(top+height-bottom)/2})">Reported outbound-leg time to 650 AU (years)</text>',
        '<text x="100" y="526" font-size="14" fill="#b9c9dc">Optimistic lower-bound trade; launch, injection, steering and operational margin are excluded.</text>',
        '</g></svg>'
    ]
    return "\n".join(parts) + "\n"


def main() -> None:
    rows = [optimize(alpha) for alpha in ALPHA_VALUES]
    with (HERE / "nep-trade.csv").open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=("alpha_kg_per_kw", "power_kw", "mass_ratio", "propellant_t", "terminal_km_s", "burn_years", "time_years"))
        writer.writeheader()
        for row in rows:
            writer.writerow({key: f"{row[key]:.8f}" for key in writer.fieldnames})
    (HERE / "nep-trade.svg").write_text(make_svg(rows), encoding="utf-8")
    for row in rows:
        print(f"alpha={row['alpha_kg_per_kw']:g}: power={row['power_kw']:.3f} kWe, t={row['time_years']:.4f} yr")


if __name__ == "__main__":
    main()