"""从开环扫点 CSV 拟合 A_eff(x) 表，输出 valve_model.json。

用法：
    python identify_valve.py --csv open_loop_data/open_loop_xxx.csv --out valve_model.json
    python identify_valve.py --csv "open_loop_data/*.csv" --out valve_model.json --plot
    python identify_valve.py --plot      # 省略 --csv 时默认合并 open_loop_data/ 下全部 CSV

流程：
1. --csv 支持多个路径与 glob 模式，全部 CSV 合并使用（文件名含时间戳，按名称
   排序即按实验先后合并）；省略时默认 open_loop_data/*.csv；
2. 每个 CSV 按 step_index 分组，每组取末尾 --tail 个采样点平均，得到稳态
   (x, Q, P1, P2)，并按开度访问顺序标记逼近方向（上升/下降/该文件起始步）；
3. 对每个稳态点反解 A_eff = Q / (P1_abs * F(r))，同一行程 x 的多个点
   （来自不同次实验）先平均再构表；
4. 按行程升序构表，做单调性检查；
5. 用 ValveModel 存成 JSON，并打印摘要（可选画诊断图，实测点按逼近方向着色）。
"""

import argparse
import csv
import glob
import math
import os
from pathlib import Path

import config
from valve_model import CRITICAL_RATIO, ValveModel, abs_pressure, f_ratio

DEFAULT_TAIL = 20

# 逼近方向：按开度访问顺序，当前开度大于/小于上一开度。
DIRECTION_UP = "up"        # 从更小开度升到当前开度
DIRECTION_DOWN = "down"    # 从更大开度降到当前开度
DIRECTION_ORDER = [DIRECTION_UP, DIRECTION_DOWN, None]
DIRECTION_LABELS = {
    DIRECTION_UP: "ascending",
    DIRECTION_DOWN: "descending",
    None: "start (unknown)",
}
DIRECTION_COLORS = {
    DIRECTION_UP: "#1565C0",
    DIRECTION_DOWN: "#D84315",
    None: "#757575",
}


def _to_float(text):
    """把 CSV 单元格解析成有限浮点数；空/非法/非有限返回 None。"""
    if text is None:
        return None
    text = str(text).strip()
    if text == "":
        return None
    try:
        value = float(text)
    except (TypeError, ValueError):
        return None
    return value if math.isfinite(value) else None


def _average(values):
    """返回非 None 数值的平均；没有有效值时返回 None。"""
    valid = [value for value in values if value is not None]
    return None if not valid else sum(valid) / len(valid)


def load_steady_points(csv_path, tail=DEFAULT_TAIL):
    """提取单个 CSV 每个 step_index 的稳态均值与逼近方向。

    返回按 step（访问顺序）排列的 ``(x, q, p1, p2, direction)`` 列表。
    direction 为 ``"up"``（开度上升逼近）/ ``"down"``（下降逼近）/
    ``None``（该文件第一步，无上一开度可比较）。
    """
    rows_by_step = {}
    opening_by_step = {}
    with open(csv_path, "r", newline="", encoding="utf-8-sig") as file:
        for row in csv.DictReader(file):
            step_text = (row.get("step_index") or "").strip()
            if step_text == "":
                continue
            try:
                step = int(float(step_text))
            except (TypeError, ValueError):
                continue
            rows_by_step.setdefault(step, []).append(row)
            if step not in opening_by_step:
                opening_by_step[step] = _to_float(row.get("opening_pct"))

    points = []
    prev_opening = None
    for step in sorted(rows_by_step):
        opening = opening_by_step.get(step)
        if prev_opening is not None and opening is not None:
            if opening > prev_opening:
                direction = DIRECTION_UP
            elif opening < prev_opening:
                direction = DIRECTION_DOWN
            else:
                direction = None
        else:
            direction = None
        if opening is not None:
            prev_opening = opening

        rows = rows_by_step[step]
        rows.sort(key=lambda r: _to_float(r.get("time_s")) or 0.0)
        tail_rows = rows[-tail:] if tail > 0 else rows

        x = _average(_to_float(r.get("motor_position")) for r in tail_rows)
        q = _average(_to_float(r.get("flow_after_slm")) for r in tail_rows)
        p1 = _average(_to_float(r.get("pressure_before_kpa")) for r in tail_rows)
        p2 = _average(_to_float(r.get("pressure_after_kpa")) for r in tail_rows)

        if None in (x, q, p1, p2):
            continue
        points.append((x, q, p1, p2, direction))
    return points


def compute_area_table(steady_points, check_monotonic=True):
    """把稳态点换算成 (x, A_eff) 表，并返回逐点测量值供诊断图着色。

    返回 ``(table, dropped, stats, area_points)``：

    - ``area_points``：每个实测点的 ``(x, A_eff, direction)``，按 x 升序，未平均；
    - ``table``：同一行程 x 的多个 A_eff 先平均；启用检查时再取最长单调连续段；
    - ``dropped``：因非单调被丢弃的 ``(x, A_eff)``；
    - ``stats``：阻塞/亚声速点统计（按原始测量点计数）。
    """
    area_points = []
    choked = subsonic = 0
    for x, q, p1, p2, direction in steady_points:
        p1_abs = abs_pressure(p1)
        p2_abs = abs_pressure(p2)
        if p1_abs <= 0.0:
            continue
        r = p2_abs / p1_abs
        f = f_ratio(r)
        denom = p1_abs * f
        if denom <= 0.0 or q < 0.0:
            continue
        if r <= CRITICAL_RATIO:
            choked += 1
        else:
            subsonic += 1
        area_points.append((x, q / denom, direction))

    area_points.sort(key=lambda item: item[0])

    # 合并多份 CSV 后同一行程会有多个点：先按 x 分组平均（CSV 以 6 位小数
    # 记录行程，round 用于吸收极小浮点差异）。
    by_x = {}
    for x, a, _ in area_points:
        by_x.setdefault(round(x, 6), []).append(a)
    merged = [(float(x), sum(vals) / len(vals)) for x, vals in sorted(by_x.items())]

    ys = [a for _, a in merged]
    if not check_monotonic or _is_monotonic(ys):
        table = merged
        dropped = []
    else:
        start, end = _longest_monotonic_run(ys)
        table = merged[start:end + 1]
        dropped = merged[:start] + merged[end + 1:]

    stats = {"choked": choked, "subsonic": subsonic, "total": choked + subsonic}
    return table, dropped, stats, area_points


def _is_monotonic(ys):
    """序列是否单调（允许相等，但不允许中途反向）。"""
    direction = 0
    for i in range(1, len(ys)):
        delta = ys[i] - ys[i - 1]
        if delta == 0:
            continue
        sign = 1 if delta > 0 else -1
        if direction == 0:
            direction = sign
        elif direction != sign:
            return False
    return True


def _longest_monotonic_run(ys):
    """返回最长单调连续段的闭区间下标 ``(start, end)``。"""
    best_start = best_end = 0
    for start in range(len(ys)):
        direction = 0
        end = start
        for j in range(start + 1, len(ys)):
            delta = ys[j] - ys[j - 1]
            if delta == 0:
                end = j
                continue
            sign = 1 if delta > 0 else -1
            if direction == 0:
                direction = sign
                end = j
            elif direction == sign:
                end = j
            else:
                break
        if end - start > best_end - best_start:
            best_start, best_end = start, end
    return best_start, best_end


def identify(csv_paths, tail=DEFAULT_TAIL, out_path=None):
    """从一份或多份开环扫点 CSV 端到端辨识。

    返回 ``(model, steady_points, table, dropped, stats, area_points)``。
    ``csv_paths`` 为路径列表，其稳态点合并后统一拟合。
    """
    steady = []
    for path in csv_paths:
        steady.extend(load_steady_points(path, tail))
    check_monotonic = config.IDENTIFY_VALVE_MONOTONIC_CHECK_ENABLED
    table, dropped, stats, area_points = compute_area_table(
        steady,
        check_monotonic=check_monotonic,
    )
    if len(table) < 2:
        raise ValueError(
            f"有效稳态点不足（{len(table)} 个），无法构表。请确认 CSV 为 "
            "open_loop.py 产出（需含 step_index 列），且 "
            "pressure_before_kpa / pressure_after_kpa / flow_after_slm "
            "列读数合理。"
        )
    model = ValveModel(
        table,
        config.MOTOR_OPEN_POSITION,
        config.MOTOR_CLOSED_POSITION,
        enforce_monotonic=check_monotonic,
    )
    if out_path:
        model.save(out_path)
    return model, steady, table, dropped, stats, area_points


def resolve_csv_paths(items=None):
    """把 CLI 传入的路径/glob 展开为去重、按名称排序的 CSV 路径列表。

    ``items`` 为空时默认使用工程目录下 ``open_loop_data/*.csv``（文件名含
    时间戳，排序后即按实验先后合并）。
    """
    if not items:
        default_dir = Path(__file__).resolve().parent / "open_loop_data"
        items = [str(default_dir / "*.csv")]
    paths = []
    for item in items:
        pattern = str(item).replace("\\", "/")
        if glob.has_magic(pattern):
            paths.extend(sorted(glob.glob(pattern)))
        else:
            paths.append(pattern)
    # 去重：glob 结果与显式路径的斜杠方向可能不同，用 normcase 归一化比较
    # （Windows 下同时统一大小写）。
    unique = []
    seen = set()
    for path in paths:
        key = os.path.normcase(path)
        if key not in seen:
            seen.add(key)
            unique.append(path)
    return unique


def _count_directions(steady_points):
    counts = {DIRECTION_UP: 0, DIRECTION_DOWN: 0, None: 0}
    for point in steady_points:
        counts[point[4]] += 1
    return counts


def _scatter_by_direction(axis, points):
    """按逼近方向分色画散点；``points`` 为 ``(x, y, direction)`` 列表。"""
    by_direction = {d: [] for d in DIRECTION_ORDER}
    for x, y, direction in points:
        by_direction.setdefault(direction, []).append((x, y))
    for direction in DIRECTION_ORDER:
        pts = by_direction.get(direction, [])
        if not pts:
            continue
        axis.scatter(
            [p[0] for p in pts],
            [p[1] for p in pts],
            color=DIRECTION_COLORS[direction],
            label=DIRECTION_LABELS[direction],
            s=18,
            alpha=0.85,
        )


def create_diagnostic_plot(steady_points, table, dropped, area_points, image_path):
    """画 x vs A_eff 与 x vs Q_ss 两张子图，实测点按逼近方向着色。

    同一行程下上升/下降两团点明显分离，说明阀门迟滞显著，前馈表应考虑
    分方向处理；完全重叠则迟滞可忽略。
    """
    import matplotlib

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

    xs = [x for x, _ in table]
    areas = [a for _, a in table]

    figure, axes = plt.subplots(2, 1, figsize=(10, 8), constrained_layout=True)
    figure.suptitle("Identified valve characteristic")

    axes[0].plot(xs, areas, color="#263238", linewidth=1.2, label="fitted A_eff (per-x mean)")
    _scatter_by_direction(axes[0], area_points)
    if dropped:
        axes[0].scatter(
            [x for x, _ in dropped],
            [a for _, a in dropped],
            color="#B71C1C",
            marker="x",
            label="dropped (non-monotonic)",
        )
    axes[0].set_ylabel("A_eff (slm/kPa)")
    axes[0].set_title("Effective area vs stroke")
    axes[0].legend(loc="best")
    axes[0].grid(True, alpha=0.3, linestyle="--")

    _scatter_by_direction(
        axes[1], [(p[0], p[1], p[4]) for p in steady_points]
    )
    axes[1].set_ylabel("Flow (SLM)")
    axes[1].set_xlabel("Motor stroke x")
    axes[1].set_title("Steady flow vs stroke")
    axes[1].legend(loc="best")
    axes[1].grid(True, alpha=0.3, linestyle="--")

    figure.savefig(image_path, dpi=config.PLOT_DPI, bbox_inches="tight")
    plt.close(figure)


def parse_args(argv=None):
    parser = argparse.ArgumentParser(
        description="从开环扫点 CSV 拟合 A_eff(x) 阀特性表（支持多 CSV/glob 合并）"
    )
    parser.add_argument(
        "--csv", nargs="+",
        help="一个或多个 CSV 路径或 glob 模式；省略时默认合并 open_loop_data/*.csv",
    )
    parser.add_argument("--out", default="valve_model.json", help="输出 JSON 路径")
    parser.add_argument(
        "--tail", type=int, default=DEFAULT_TAIL,
        help=f"每组末尾用于平均的采样点数（默认 {DEFAULT_TAIL}）",
    )
    parser.add_argument("--plot", action="store_true", help="生成诊断图 PNG")
    args = parser.parse_args(argv)
    if args.tail < 1:
        parser.error("--tail 必须 >= 1")
    return args


def main(argv=None):
    args = parse_args(argv)
    csv_paths = resolve_csv_paths(args.csv)
    if not csv_paths:
        print("未找到任何 CSV：请确认 open_loop_data/ 下有开环扫点数据，")
        print('或用 --csv 指定路径或 glob 模式（如 --csv "open_loop_data/*.csv"）。')
        return 2
    for path in csv_paths:
        if not Path(path).exists():
            print(f"CSV 不存在：{path}")
            return 2

    try:
        model, steady, table, dropped, stats, area_points = identify(
            csv_paths, tail=args.tail, out_path=args.out
        )
    except (ValueError, OSError) as exc:
        print(f"辨识失败：{exc}")
        return 2

    if len(csv_paths) <= 8:
        print("合并的 CSV：")
        for path in csv_paths:
            print(f"  {path}")
    else:
        print(f"合并的 CSV：{csv_paths[0]} 等共 {len(csv_paths)} 个文件")
    directions = _count_directions(steady)
    print(
        f"读取稳态点 {len(steady)} 个：上升逼近 {directions[DIRECTION_UP]}、"
        f"下降逼近 {directions[DIRECTION_DOWN]}、起始点 {directions[None]}；"
        f"阻塞 {stats['choked']} 个，亚声速 {stats['subsonic']} 个，"
        f"合计 {stats['total']} 个。"
    )
    print(
        f"有效行程范围 {table[0][0]:.1f} ~ {table[-1][0]:.1f}"
        f"（{len(table)} 点，同一行程已平均）。"
    )

    if not config.IDENTIFY_VALVE_MONOTONIC_CHECK_ENABLED:
        print("提示：单调性检查已关闭，所有按行程平均后的点均被保留。")

    if dropped:
        print(
            "警告：A_eff(x) 非单调，以下点被丢弃"
            "（建议缩小扫点范围到单调段重扫）。"
        )
    print("行程 x -> A_eff：")
    for x, a in table:
        print(f"  x={x:8.3f}  A_eff={a:.6f}")
    if dropped:
        print("被丢弃的点：")
        for x, a in dropped:
            print(f"  x={x:8.3f}  A_eff={a:.6f}  [DROPPED]")

    print(f"已保存阀特性模型：{args.out}")

    if args.plot:
        image_path = Path(args.out).with_suffix(".png")
        try:
            create_diagnostic_plot(steady, table, dropped, area_points, image_path)
            print(f"诊断图已保存：{image_path}")
        except Exception as exc:
            print(f"生成诊断图失败（不影响 JSON）：{exc}")

    return 0


if __name__ == "__main__":
    import sys

    sys.exit(main())