"""Equal-area wings and a fixed-lift induced-drag model; write PNG to cwd."""

from pathlib import Path

import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
import numpy as np


def induced_drag(span, lift=4800.0, pressure=500.0, efficiency=0.8):
    return lift**2 / (np.pi * efficiency * pressure * np.asarray(span)**2)


def main():
    area = 16.0
    spans = np.array([8.0, 12.0])
    chords = area / spans
    aspect_ratios = spans**2 / area
    drag = induced_drag(spans)
    np.testing.assert_allclose(spans * chords, [area, area])
    np.testing.assert_allclose(aspect_ratios, [4.0, 9.0])
    np.testing.assert_allclose(drag, [286.478897565, 127.323954474], rtol=1e-10)
    np.testing.assert_allclose(drag[1] / drag[0], 4.0 / 9.0)
    # Compare both equivalent versions, including the varying coefficient.
    for changed_area in [12.0, 16.0, 24.0]:
        lift_coefficient = 4800.0 / (500.0 * changed_area)
        ar = 8.0**2 / changed_area
        from_coefficients = 500.0 * changed_area * lift_coefficient**2 / (np.pi * 0.8 * ar)
        np.testing.assert_allclose(from_coefficients, drag[0])

    plt.rcParams.update({"font.size": 10, "figure.facecolor": "white", "axes.facecolor": "white"})
    fig = plt.figure(figsize=(10, 4), dpi=100)
    grid = fig.add_gridspec(2, 2, width_ratios=[1.05, 1.0], hspace=0.72, wspace=0.33)
    colors = ["#276FBF", "#D15D21"]
    for i, (span, chord, ar, color) in enumerate(zip(spans, chords, aspect_ratios, colors)):
        ax = fig.add_subplot(grid[i, 0])
        ax.add_patch(Rectangle((-span / 2, 0), span, chord, facecolor=color, alpha=0.28,
                               edgecolor=color, linewidth=2))
        ax.axvline(0, color="#777777", linewidth=0.7, linestyle=":")
        ax.set_xlim(-6.7, 6.7)
        ax.set_ylim(-0.35, 2.65)
        ax.set_aspect("equal", adjustable="box")
        ax.set_xticks([-6, -4, -2, 0, 2, 4, 6])
        ax.set_yticks([0, 1, 2])
        ax.set_ylabel("Chord (m)")
        ax.set_title(f"Span {span:g} m; area 16 m²; AR = {ar:g}", fontsize=10, pad=5)
        if i == 1:
            ax.set_xlabel("Spanwise coordinate (m)")
        ax.spines[["top", "right"]].set_visible(False)

    ax = fig.add_subplot(grid[:, 1])
    span_grid = np.linspace(6.5, 13, 200)
    ax.plot(span_grid, induced_drag(span_grid), color="#343434", linewidth=2)
    for span, value, color in zip(spans, drag, colors):
        ax.scatter(span, value, color=color, s=45, zorder=3)
        ax.annotate(f"{value:.2f} N", (span, value), xytext=((8, 12) if span == 8 else (-8, 12)),
                    ha=("left" if span == 8 else "right"),
                    textcoords="offset points", color=color, fontsize=11)
    ax.set_xlim(6.5, 13.2)
    ax.set_ylim(0, 470)
    ax.set_xlabel("Wingspan b (m)")
    ax.set_ylabel("Lift-induced drag (N)")
    ax.set_title("Fixed L = 4.8 kN, q = 500 Pa, e = 0.8", fontsize=10)
    ax.text(0.98, 0.96, "Dᵢ = L² / (π e q b²)", transform=ax.transAxes,
            ha="right", va="top", fontsize=11)
    ax.grid(alpha=0.18)
    ax.spines[["top", "right"]].set_visible(False)
    fig.suptitle("Longer span at equal wing area: a model comparison", fontsize=13, y=0.98)
    fig.subplots_adjust(left=0.075, right=0.98, top=0.86, bottom=0.15)
    output = Path(__file__).with_suffix(".png").name
    fig.savefig(output, dpi=100, facecolor="white")
    plt.close(fig)
    print(f"{output}: 1000×400; AR={aspect_ratios.tolist()}; drag={drag.tolist()}")


if __name__ == "__main__":
    main()