"""Recreate NASA Glenn's stagnation-temperature/altitude chart from its equations.

This is an equation-based reconstruction, not a pixel trace or experimental dataset.
Run with Python's standard library only.
"""
import csv
import math
from pathlib import Path

ROOT = Path(__file__).resolve().parent
GAMMA = 1.4
THETA_R = 5500.0
ALTITUDE_MAX_FT = 100_000
TARGET_TOTAL_TEMPERATURE_R = (500, 1000, 1500, 2000, 3000, 4000, 5000)


def standard_day_static_temperature_r(altitude_ft):
    """NASA Glenn English-unit atmosphere fit, converted from °F to °R."""
    if altitude_ft <= 36_152:
        t_f = 59.0 - 0.00356 * altitude_ft
    elif altitude_ft <= 82_345:
        t_f = -70.0
    else:
        t_f = -205.05 + 0.00164 * altitude_ft
    return t_f + 459.7


def mach_perfect_gas(static_r, total_r):
    if total_r <= static_r:
        return None
    return math.sqrt(2 / (GAMMA - 1) * (total_r / static_r - 1))


def mach_calorically_imperfect(static_r, total_r):
    """NASA Glenn harmonic-vibrator correction for molecular vibration."""
    if total_r <= static_r:
        return None
    x = THETA_R / static_r
    exp_x = math.exp(x)
    gamma_vibrational = 1 + (GAMMA - 1) / (
        1 + (GAMMA - 1) * x * x * exp_x / (exp_x - 1) ** 2
    )
    y = THETA_R / total_r
    exp_y = math.exp(y)
    energy_term = (
        GAMMA / (GAMMA - 1) * (1 - static_r / total_r)
        + THETA_R / total_r * (1 / (exp_y - 1) - 1 / (exp_x - 1))
    )
    mach_squared = 2 * total_r / static_r / gamma_vibrational * energy_term
    return math.sqrt(max(0.0, mach_squared))


def build_data():
    rows = []
    for altitude_ft in range(0, ALTITUDE_MAX_FT + 1, 1000):
        static_r = standard_day_static_temperature_r(altitude_ft)
        for total_r in TARGET_TOTAL_TEMPERATURE_R:
            perfect = mach_perfect_gas(static_r, total_r)
            imperfect = mach_calorically_imperfect(static_r, total_r)
            rows.append((altitude_ft, static_r, total_r, perfect, imperfect))
    return rows


def chart_svg(rows):
    width, height = 1200, 760
    left, right, top, bottom = 110, 1145, 136, 640
    mach_max = 8.5
    colors = {
        500: "#392b77", 1000: "#a02db8", 1500: "#c22c85",
        2000: "#bb4772", 3000: "#4c7a4a", 4000: "#276aa3", 5000: "#8f3654",
    }

    def x(altitude_ft):
        return left + altitude_ft / ALTITUDE_MAX_FT * (right - left)

    def y(mach):
        return bottom - mach / mach_max * (bottom - top)

    elements = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" role="img" aria-labelledby="title desc">',
        '<title id="title">Stagnation-temperature contours: perfect and vibrationally imperfect air</title>',
        '<desc id="desc">Faithful equation-based redraw of NASA Glenn’s chart. Seven constant stagnation-temperature contours show the Mach number versus altitude for a standard-day atmosphere. Solid lines use constant specific heat; dashed lines include the source page’s simple molecular-vibration correction. Curves are model calculations, not measurements.</desc>',
        '<rect width="1200" height="760" rx="24" fill="#f7f9ff"/>',
        '<text x="58" y="46" fill="#ae35d5" font-family="Arial,sans-serif" font-size="17" font-weight="700" letter-spacing="2">SOURCE FIGURE RECONSTRUCTION · HYPERSONIC FLOW</text>',
        '<text x="58" y="92" fill="#172646" font-family="Arial,sans-serif" font-size="31" font-weight="700">Real-gas effects reshape stagnation-temperature contours</text>',
    ]

    # Grid and dual-unit axes: kilometres are prominent; feet retain source context.
    for mach_tick in (0, 2, 4, 6, 8):
        yy = y(mach_tick)
        elements.append(f'<line x1="{left}" y1="{yy:.1f}" x2="{right}" y2="{yy:.1f}" stroke="#d9e1f0" stroke-width="1.4"/>')
        elements.append(f'<text x="{left-18}" y="{yy+6:.1f}" text-anchor="end" fill="#364766" font-family="Arial,sans-serif" font-size="17">{mach_tick}</text>')
    for altitude_ft_tick in range(0, ALTITUDE_MAX_FT + 1, 20_000):
        km_tick = altitude_ft_tick * 0.0003048
        xx = x(altitude_ft_tick)
        elements.append(f'<line x1="{xx:.1f}" y1="{top}" x2="{xx:.1f}" y2="{bottom}" stroke="#e7ecf5" stroke-width="1.2"/>')
        km_label = "0" if altitude_ft_tick == 0 else f"{km_tick:.1f}"
        elements.append(f'<text x="{xx:.1f}" y="{bottom+29}" text-anchor="middle" fill="#364766" font-family="Arial,sans-serif" font-size="16">{km_label}</text>')
        elements.append(f'<text x="{xx:.1f}" y="{bottom+52}" text-anchor="middle" fill="#56677f" font-family="Arial,sans-serif" font-size="12">{altitude_ft_tick:,} ft</text>')

    elements.extend([
        f'<line x1="{left}" y1="{bottom}" x2="{right}" y2="{bottom}" stroke="#172646" stroke-width="2.5"/>',
        f'<line x1="{left}" y1="{top}" x2="{left}" y2="{bottom}" stroke="#172646" stroke-width="2.5"/>',
        f'<text x="{(left+right)/2:.1f}" y="{height-24}" text-anchor="middle" fill="#172646" font-family="Arial,sans-serif" font-size="19" font-weight="700">Altitude (km; original figure in feet)</text>',
        f'<text transform="translate(34 {(top+bottom)/2:.1f}) rotate(-90)" text-anchor="middle" fill="#172646" font-family="Arial,sans-serif" font-size="19" font-weight="700">Mach number</text>',
    ])

    # Source-figure legend: line style, not color, identifies the thermodynamic model.
    legend_y = 122
    elements.extend([
        f'<line x1="{left+12}" y1="{legend_y}" x2="{left+82}" y2="{legend_y}" stroke="#364766" stroke-width="4"/>',
        f'<text x="{left+94}" y="{legend_y+6}" fill="#263653" font-family="Arial,sans-serif" font-size="16">Calorically perfect · solid</text>',
        f'<line x1="{left+365}" y1="{legend_y}" x2="{left+435}" y2="{legend_y}" stroke="#364766" stroke-width="3" stroke-dasharray="11 8"/>',
        f'<text x="{left+447}" y="{legend_y+6}" fill="#263653" font-family="Arial,sans-serif" font-size="16">Vibrational correction · dashed</text>',
    ])

    for total_r in TARGET_TOTAL_TEMPERATURE_R:
        color = colors[total_r]
        selected = [row for row in rows if row[2] == total_r]
        for model_index, dash, stroke_width in ((3, "", 3.2), (4, ' stroke-dasharray="10 7"', 2.5)):
            points = [(x(row[0]), y(row[model_index])) for row in selected if row[model_index] is not None and row[model_index] <= mach_max]
            if len(points) < 2:
                continue
            d = " ".join(("M" if i == 0 else "L") + f"{px:.1f},{py:.1f}" for i, (px, py) in enumerate(points))
            model_label = "constant-cp" if model_index == 3 else "vibrational correction"
            elements.append(f'<path d="{d}" fill="none" stroke="{color}" stroke-width="{stroke_width}"{dash} tabindex="0" data-tooltip="Tt {total_r} °R ({total_r*5/9:.0f} K), {model_label}"/>')

        # Direct contour label on the plateau, with both source and SI values.
        label_altitude = 53_000
        static_r = standard_day_static_temperature_r(label_altitude)
        label_mach = mach_perfect_gas(static_r, total_r)
        if label_mach is not None and label_mach < mach_max:
            lx, ly = x(label_altitude), y(label_mach)
            kelvin = total_r * 5 / 9
            label = f"{kelvin:.0f} K · {total_r:,} °R"
            elements.append(f'<rect x="{lx+9:.1f}" y="{ly-14:.1f}" width="148" height="22" rx="5" fill="#ffffff" fill-opacity="0.88"/>')
            elements.append(f'<text x="{lx+16:.1f}" y="{ly+3:.1f}" fill="{color}" font-family="Arial,sans-serif" font-size="14" font-weight="700">{label}</text>')

    elements.append('</svg>')
    return "\n".join(elements) + "\n"


def main():
    # Sanity checks against the source's standard-atmosphere breakpoints and visible contour spacing.
    assert abs(standard_day_static_temperature_r(0) - 518.7) < 1e-9
    assert abs(standard_day_static_temperature_r(50_000) - 389.7) < 1e-9
    low = mach_calorically_imperfect(standard_day_static_temperature_r(50_000), 1000)
    low_perfect = mach_perfect_gas(standard_day_static_temperature_r(50_000), 1000)
    high = mach_calorically_imperfect(standard_day_static_temperature_r(50_000), 5000)
    high_perfect = mach_perfect_gas(standard_day_static_temperature_r(50_000), 5000)
    assert abs(low - low_perfect) < 0.1
    assert high > high_perfect + 0.4
    rows = build_data()
    with (ROOT / "real-gas-stagnation-contours.csv").open("w", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle)
        writer.writerow(["altitude_ft", "altitude_km", "static_temperature_R", "target_stagnation_temperature_R", "target_stagnation_temperature_K", "mach_perfect_gas", "mach_vibrational_model"])
        for altitude_ft, static_r, total_r, perfect, imperfect in rows:
            writer.writerow([
                altitude_ft,
                f"{altitude_ft * 0.0003048:.4f}",
                f"{static_r:.3f}",
                total_r,
                f"{total_r * 5 / 9:.3f}",
                "" if perfect is None else f"{perfect:.6f}",
                "" if imperfect is None else f"{imperfect:.6f}",
            ])
    (ROOT / "real-gas-stagnation-contours.svg").write_text(chart_svg(rows), encoding="utf-8")
    print(f"Recreated {len(TARGET_TOTAL_TEMPERATURE_R)} NASA stagnation-temperature contour families from the published equations; wrote {len(rows)} rows.")


if __name__ == "__main__":
    main()