# -*- coding: utf-8 -*-
"""Demo: one-at-a-time design-parameter sensitivity for an RC Brayton cycle."""

from pathlib import Path
import csv
import sys

import matplotlib.pyplot as plt


PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from brayton_cycle import local_rc_design_sensitivity  # noqa: E402


DESIGN_VARIABLES = (
    "T_low",
    "T_high",
    "p_low",
    "p_high",
    "ploss",
    "x",
)


def build_base_case():
    fixed_params = {
        "T_high": 650 + 273.15,
        "T_low": 42 + 273.15,
        "p_high": 20.0e3,
        "p_low": 9.09e3,
        "ploss": 0.01,
    }
    params = {
        "x": 0.279,
        "compressor_eff": 0.9,
        "recompressor_eff": 0.9,
        "turbine_eff": 0.93,
        "recuperator_eff": 0.94,
        "highT_recuperator_eff": 0.96,
    }
    return fixed_params, params


def sort_by_importance(rows):
    return sorted(
        rows,
        key=lambda row: abs(row["normalized_sensitivity"] or 0.0),
        reverse=True,
    )


def print_summary(rows):
    print("RC design-parameter local sensitivity")
    print("Efficiency values are decimals; 0.46 means 46%.")
    print()
    print(
        f"{'variable':<10} {'base':>12} {'eff-':>12} "
        f"{'eff+':>12} {'norm_sens':>14}"
    )
    print("-" * 66)
    for row in sort_by_importance(rows):
        sensitivity = row["normalized_sensitivity"]
        sensitivity_text = "nan" if sensitivity is None else f"{sensitivity: .6f}"
        print(
            f"{row['variable']:<10} "
            f"{row['base_value']:>12.6g} "
            f"{row['minus_efficiency']:>12.6f} "
            f"{row['plus_efficiency']:>12.6f} "
            f"{sensitivity_text:>14}"
        )


def save_csv(rows, output_path):
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with output_path.open("w", newline="", encoding="utf-8") as file:
        writer = csv.DictWriter(file, fieldnames=list(rows[0].keys()))
        writer.writeheader()
        writer.writerows(rows)


def plot_sensitivity(rows, output_path=None, show=False):
    sorted_rows = sort_by_importance(rows)
    variables = [row["variable"] for row in sorted_rows]
    sensitivities = [
        row["normalized_sensitivity"] or 0.0
        for row in sorted_rows
    ]
    colors = ["#1f77b4" if value >= 0 else "#d62728" for value in sensitivities]

    fig, ax = plt.subplots(figsize=(8, 4.8), dpi=130)
    ax.barh(variables, sensitivities, color=colors)
    ax.axvline(0.0, color="black", linewidth=0.8)
    ax.set_xlabel("Normalized sensitivity")
    ax.set_title("RC design-parameter sensitivity")
    ax.invert_yaxis()
    fig.tight_layout()

    if output_path is not None:
        output_path.parent.mkdir(parents=True, exist_ok=True)
        fig.savefig(output_path)
    if show:
        plt.show()

    return fig, ax


def main(save_outputs=True, show_plot=False):
    fixed_params, params = build_base_case()
    rows = local_rc_design_sensitivity(
        fixed_params,
        params,
        variables=DESIGN_VARIABLES,
        relative_step=0.01,
    )

    print_summary(rows)

    if save_outputs:
        output_dir = PROJECT_ROOT / "examples" / "output"
        save_csv(rows, output_dir / "rc_design_sensitivity.csv")
        plot_sensitivity(
            rows,
            output_path=output_dir / "rc_design_sensitivity.png",
            show=show_plot,
        )
    elif show_plot:
        plot_sensitivity(rows, show=True)

    return rows


if __name__ == "__main__":
    main()