"""逐周期保存控制数据，并在运行结束时生成四联曲线图。"""

import csv
from datetime import datetime
import math
from pathlib import Path


ATMOSPHERIC_PRESSURE_KPA = 101.325
CRITICAL_PRESSURE_RATIO = 0.528


CSV_FIELDS = (
    "time_s",
    "timestamp",
    "target_flow_slm",
    "measured_flow_slm",
    "opening_pct",
    "pressure_kpa",
    "pressure_before_kpa",
    "P_abs_ratio",
    "error_slm",
    "motor_position",
    "feedforward_pct",
    "correction_pct",
    "gain_scale",
    "actual_dt_s",
    "status",
)


class FlowRunRecorder:
    """把每个控制周期写入 CSV，并在结束时输出一张四联图。"""

    def __init__(self, output_directory, *, plot_dpi=160):
        output_dir = Path(output_directory)
        output_dir.mkdir(parents=True, exist_ok=True)

        run_id = datetime.now().strftime("%Y%m%d_%H%M%S")
        self.csv_path = output_dir / f"flow_run_{run_id}.csv"
        self.image_path = output_dir / f"flow_run_{run_id}.png"
        self.plot_dpi = int(plot_dpi)

        self._file = self.csv_path.open("w", newline="", encoding="utf-8-sig")
        self._writer = csv.DictWriter(self._file, fieldnames=CSV_FIELDS)
        self._writer.writeheader()
        self._file.flush()
        self._first_timestamp = None
        self.sample_count = 0
        self._closed = False

    def record(self, result):
        """保存一个 ControlStepResult；每次写入后立即刷新到磁盘。"""
        if self._closed:
            raise RuntimeError("记录器已经关闭")

        if self._first_timestamp is None:
            self._first_timestamp = float(result.timestamp)
        elapsed_s = max(0.0, float(result.timestamp) - self._first_timestamp)

        p_abs_ratio = self._absolute_pressure_ratio(
            result.pressure_before_kpa,
            result.pressure_kpa,
        )
        self._writer.writerow(
            {
                "time_s": f"{elapsed_s:.6f}",
                "timestamp": f"{float(result.timestamp):.6f}",
                "target_flow_slm": self._format_value(result.target_flow_slm),
                "measured_flow_slm": self._format_value(
                    result.measured_flow_slm
                ),
                "opening_pct": self._format_value(result.opening_pct),
                "pressure_kpa": self._format_value(result.pressure_kpa),
                "pressure_before_kpa": self._format_value(
                    result.pressure_before_kpa
                ),
                "P_abs_ratio": self._format_value(p_abs_ratio),
                "error_slm": self._format_value(result.error_slm),
                "motor_position": self._format_value(result.motor_position),
                "feedforward_pct": self._format_value(result.feedforward_pct),
                "correction_pct": self._format_value(result.correction_pct),
                "gain_scale": self._format_value(result.gain_scale),
                "actual_dt_s": self._format_value(result.actual_dt_s),
                "status": result.status,
            }
        )
        self._file.flush()
        self.sample_count += 1

    def finalize(self):
        """关闭 CSV 并生成 PNG；无采样点时只保留 CSV。"""
        self.close()
        if self.sample_count == 0:
            return None
        self._create_plot()
        return self.image_path

    def close(self):
        if not self._closed:
            self._file.flush()
            self._file.close()
            self._closed = True

    def _create_plot(self):
        # 使用无界面后端，保证从终端运行和无显示器环境都能保存图片。
        import matplotlib

        matplotlib.use("Agg")
        import matplotlib.pyplot as plt

        time_s = []
        target_flow = []
        measured_flow = []
        opening = []
        feedforward = []
        pressure_before = []
        pressure_after = []
        p_abs_ratio = []

        with self.csv_path.open("r", newline="", encoding="utf-8-sig") as file:
            for row in csv.DictReader(file):
                time_s.append(float(row["time_s"]))
                target_flow.append(self._float_or_nan(row["target_flow_slm"]))
                measured_flow.append(
                    self._float_or_nan(row["measured_flow_slm"])
                )
                opening.append(self._float_or_nan(row["opening_pct"]))
                feedforward.append(
                    self._float_or_nan(row.get("feedforward_pct"))
                )
                before = self._float_or_nan(row["pressure_before_kpa"])
                after = self._float_or_nan(row["pressure_kpa"])
                pressure_before.append(before)
                pressure_after.append(after)
                ratio = self._float_or_nan(row.get("P_abs_ratio"))
                if not math.isfinite(ratio):
                    ratio_value = self._absolute_pressure_ratio(before, after)
                    ratio = math.nan if ratio_value is None else ratio_value
                p_abs_ratio.append(ratio)

        figure, axes = plt.subplots(
            4,
            1,
            figsize=(12, 14),
            sharex=True,
            constrained_layout=True,
        )
        figure.suptitle("Gas Flow Control Response", fontsize=15)

        axes[0].plot(
            time_s,
            measured_flow,
            color="#1565C0",
            linewidth=1.4,
            label="Measured flow",
        )
        axes[0].plot(
            time_s,
            target_flow,
            color="#D84315",
            linewidth=1.5,
            linestyle="--",
            label="Target flow",
        )
        axes[0].set_ylabel("Flow (SLM)")
        axes[0].set_title("Flow response")
        axes[0].legend(loc="best")

        axes[1].plot(
            time_s,
            opening,
            color="#2E7D32",
            linewidth=1.4,
            label="Opening",
        )
        if any(math.isfinite(value) for value in feedforward):
            axes[1].plot(
                time_s,
                feedforward,
                color="#EF6C00",
                linewidth=1.2,
                linestyle="--",
                label="Feedforward",
            )
            axes[1].legend(loc="best")
        axes[1].set_ylabel("Opening (%)")
        axes[1].set_title("Valve opening")
        axes[1].set_ylim(-2, 102)

        self._plot_series(
            axes[2],
            time_s,
            pressure_before,
            "Pressure before valve",
            "#1565C0",
        )
        self._plot_series(
            axes[2],
            time_s,
            pressure_after,
            "Pressure after valve",
            "#D84315",
        )
        axes[2].set_ylabel("Pressure (kPa)")
        axes[2].set_title("Pressure across control valve")
        axes[2].legend(loc="best")

        self._plot_series(
            axes[3],
            time_s,
            p_abs_ratio,
            "P_abs_ratio (after/before)",
            "#EF6C00",
        )
        axes[3].axhline(
            CRITICAL_PRESSURE_RATIO,
            color="#B71C1C",
            linestyle="--",
            linewidth=1.2,
            label=f"Critical ratio {CRITICAL_PRESSURE_RATIO:.3f}",
        )
        axes[3].set_ylabel("P_abs_ratio")
        axes[3].set_xlabel("Time (s)")
        axes[3].set_title("Absolute pressure ratio across control valve")
        axes[3].legend(loc="best")

        for axis in axes:
            axis.grid(True, alpha=0.3, linestyle="--")
            axis.margins(x=0)

        figure.savefig(self.image_path, dpi=self.plot_dpi, bbox_inches="tight")
        plt.close(figure)

    @staticmethod
    def _plot_series(axis, time_s, series, label, color):
        if any(math.isfinite(value) for value in series):
            axis.plot(
                time_s,
                series,
                color=color,
                linewidth=1.4,
                label=label,
            )
        else:
            axis.text(
                0.5,
                0.5,
                f"No {label.lower()} data",
                transform=axis.transAxes,
                ha="center",
                va="center",
                color="gray",
            )

    @staticmethod
    def _absolute_pressure_ratio(pressure_before_kpa, pressure_after_kpa):
        """返回阀后绝对压力 / 阀前绝对压力；无效读数返回 None。"""
        if pressure_before_kpa is None or pressure_after_kpa is None:
            return None
        try:
            before_abs = float(pressure_before_kpa) + ATMOSPHERIC_PRESSURE_KPA
            after_abs = float(pressure_after_kpa) + ATMOSPHERIC_PRESSURE_KPA
        except (TypeError, ValueError):
            return None
        if before_abs <= 0.0:
            return None
        ratio = after_abs / before_abs
        return ratio if math.isfinite(ratio) else None

    @staticmethod
    def _format_value(value):
        if value is None:
            return ""
        number = float(value)
        return "" if not math.isfinite(number) else f"{number:.6f}"

    @staticmethod
    def _float_or_nan(value):
        if value in (None, ""):
            return math.nan
        try:
            return float(value)
        except (TypeError, ValueError):
            return math.nan