"""
06 — CFD vs POD-NN time analysis.

Two figures (paper Fig 4):
  a) Cumulative wall-clock time crossover (CFD vs POD-NN).
  b) Per-component time bar chart (log scale): dataset generation,
     NN training, POD fit, and single inference.

Both are analytical plots driven by measured timing constants.
"""

import numpy as np
import matplotlib.pyplot as plt

from _utils import PLOTS_DIR

CFD_PER_ITER       = 99.0           # s / CFD simulation
PODNN_DATASET      = 99.0 * 100     # 100 CFD sims for training
PODNN_POD_TRAINING = 0.0318         # mean POD fit time, s
PODNN_NN_TRAINING  = 46.7705        # mean NN training time, s
PODNN_TRAINING     = PODNN_POD_TRAINING + PODNN_NN_TRAINING
PODNN_UPFRONT      = PODNN_DATASET + PODNN_TRAINING
PODNN_PER_ITER     = 0.24e-3        # 0.24 ms per inference


def main():
    n_cross = PODNN_UPFRONT / (CFD_PER_ITER - PODNN_PER_ITER)
    t_cross = CFD_PER_ITER * n_cross
    print(f"Crossover: {n_cross:.1f} iters  ({t_cross/3600:.2f} hrs)")

    n_max = int(n_cross * 2.5)
    n = np.arange(0, n_max + 1)
    t_cfd = CFD_PER_ITER * n
    t_pod = PODNN_UPFRONT + PODNN_PER_ITER * n

    FS = 26
    fig, ax = plt.subplots(figsize=(10, 6))
    ax.plot(n, t_cfd / 3600, linewidth=2.5, color="tab:red",  label="CFD")
    ax.plot(n, t_pod / 3600, linewidth=2.5, color="tab:blue", label="POD-NN")
    ax.axvline(n_cross, color="gray", linestyle="--", linewidth=1, alpha=0.7)
    ax.plot(n_cross, t_cross / 3600, "ko", markersize=10, zorder=5)
    ax.annotate(f"Crossover: {n_cross:.0f} iterations\n({t_cross/3600:.1f} hrs)",
                xy=(n_cross, t_cross / 3600),
                xytext=(n_cross + n_max * 0.08, t_cross / 3600 * 0.6),
                fontsize=FS - 8, ha="left",
                arrowprops=dict(arrowstyle="->", color="black", lw=1.5))
    ax.set_xlabel("Number of Design Iterations", fontsize=FS)
    ax.set_ylabel("Cumulative Wall-Clock Time (hrs)", fontsize=FS)
    ax.tick_params(labelsize=FS - 4)
    ax.legend(fontsize=FS - 4)
    ax.grid(True, alpha=0.3)
    ax.set_xlim(0, n_max); ax.set_ylim(0, None)

    plt.tight_layout()
    PLOTS_DIR.mkdir(parents=True, exist_ok=True)
    out = PLOTS_DIR / "time_crossover.png"
    plt.savefig(out, dpi=300, bbox_inches="tight")
    print(f"Saved {out}")
    plt.close()

    # --- Per-component bar chart (paper Fig 4b) ---
    components = [
        ("Dataset\nGeneration", PODNN_DATASET,      "tab:orange"),
        ("NN\nTraining",        PODNN_NN_TRAINING,  "tab:blue"),
        ("POD\nFit",            PODNN_POD_TRAINING, "tab:green"),
        ("Single\nInference",   PODNN_PER_ITER,     "tab:purple"),
    ]
    labels = [c[0] for c in components]
    values = [c[1] for c in components]
    colors = [c[2] for c in components]

    def fmt_time(s):
        if s >= 3600: return f"{s/3600:.1f} hr"
        if s >= 1:    return f"{s:.1f} s"
        if s >= 1e-3: return f"{s*1e3:.1f} ms"
        return f"{s*1e6:.1f} µs"

    fig, ax = plt.subplots(figsize=(10, 6))
    bars = ax.barh(labels, values, color=colors, edgecolor="white", height=0.5)
    ax.set_xscale("log")
    ax.set_xlabel("Wall-Clock Time (s)", fontsize=FS)
    ax.tick_params(labelsize=FS - 2)
    ax.grid(True, which="both", axis="x", alpha=0.3)
    ax.invert_yaxis()
    max_val = max(values)
    for bar, val in zip(bars, values):
        if val == max_val:
            ax.text(val * 0.6, bar.get_y() + bar.get_height() / 2,
                    fmt_time(val), va="center", ha="right", fontsize=FS - 4,
                    color="white", fontweight="bold")
        else:
            ax.text(val * 1.5, bar.get_y() + bar.get_height() / 2,
                    fmt_time(val), va="center", ha="left", fontsize=FS - 4)

    plt.tight_layout()
    out_bar = PLOTS_DIR / "time_components_bar.png"
    plt.savefig(out_bar, dpi=300, bbox_inches="tight")
    print(f"Saved {out_bar}")
    plt.close()


if __name__ == "__main__":
    main()