pickling-mes/backend/app/services/prediction.py

"""
工艺预测模型 — 灰箱物理模型 + ONNX 神经网络双栈

推理优先级：
  1. ONNX 模型（onnxruntime，若 pt_models/ 目录存在则加载）
  2. 物理灰箱模型（Arrhenius 解析解，始终可用）

训练：运行 backend/train_models.py 重新生成 pt_models/*.onnx
校准：K_cal 系数持久化在 cal_coeffs.json，两个栈都使用同一套 K_cal
"""
import math
import json
import os
from pathlib import Path
from typing import List, Dict, Any, Optional, Tuple
from loguru import logger

# ── 校准系数持久化 ────────────────────────────────────────────────────────────
_CAL_FILE = Path(__file__).parent / "cal_coeffs.json"

def _load_cal() -> Dict[str, float]:
    try:
        with open(_CAL_FILE) as f:
            return json.load(f)
    except Exception:
        return {}

def _save_cal(d: Dict[str, float]):
    with open(_CAL_FILE, "w") as f:
        json.dump(d, f, indent=2)


# ── ONNX 推理层 ───────────────────────────────────────────────────────────────
_PT_DIR   = Path(__file__).parent / "pt_models"
_scalers: Optional[Dict] = None
_sess:    Dict[str, Any] = {}

try:
    import onnxruntime as ort
    import numpy as _np

    _sp = _PT_DIR / "scalers.json"
    if _sp.exists():
        with open(_sp) as f:
            _scalers = json.load(f)

    for _name in ("acid_speed", "tension", "quality"):
        _p = _PT_DIR / f"{_name}.onnx"
        if _p.exists():
            _sess[_name] = ort.InferenceSession(
                str(_p), providers=["CPUExecutionProvider"]
            )
    if _sess:
        logger.info(f"PT models loaded: {list(_sess.keys())}")
except ImportError:
    logger.warning("onnxruntime not installed — using physics fallback")


def _pt_infer(name: str, x_raw: List[float]) -> Optional[List[float]]:
    """标准化 → ONNX 推理 → 反标准化，返回输出向量；失败返回 None。"""
    if name not in _sess or _scalers is None:
        return None
    try:
        sc   = _scalers[name]
        xm   = _np.array(sc["X_mean"], dtype=_np.float32)
        xs   = _np.array(sc["X_std"],  dtype=_np.float32)
        ym   = _np.array(sc["y_mean"], dtype=_np.float32)
        ys   = _np.array(sc["y_std"],  dtype=_np.float32)
        x    = (_np.array(x_raw, dtype=_np.float32) - xm) / xs
        raw  = _sess[name].run(None, {"input": x.reshape(1, -1)})[0][0]
        return (raw * ys + ym).tolist()
    except Exception as e:
        logger.warning(f"PT infer {name} failed: {e}")
        return None


# ─────────────────────────────────────────────────────────────────────────────
# 1. 酸洗速度模型
# ─────────────────────────────────────────────────────────────────────────────
class AcidSpeedModel:
    """
    灰箱: Arrhenius 动力学 + 二分搜索
    PT栈: 14 维输入 → 最大速度 (m/min)
    输入: [thickness, scale_weight, conc×6, temp×6]
    """
    TANK_LENGTH        = 18.0
    NUM_TANKS          = 6
    K0                 = 0.075
    EA_R               = 5413.0
    T_REF              = 348.15
    C_REF              = 180.0
    N_CONC             = 1.2
    V_MIN              = 20.0
    V_MAX              = 180.0
    SCALE_RATE_FACTOR  = 0.70 * 1.0 + 0.20 * 0.25 + 0.10 * 0.15
    CAL_KEY            = "acid_speed_kcal"

    def __init__(self, thickness, width, steel_grade,
                 acid_conc_list, acid_temp_list,
                 scale_weight=8.5, target_pi=95.0):
        if len(acid_conc_list) != self.NUM_TANKS:
            raise ValueError(f"acid_conc_list 需要 {self.NUM_TANKS} 个元素")
        if len(acid_temp_list) != self.NUM_TANKS:
            raise ValueError(f"acid_temp_list 需要 {self.NUM_TANKS} 个元素")
        self.thickness      = thickness
        self.width          = width
        self.steel_grade    = steel_grade
        self.acid_conc_list = acid_conc_list
        self.acid_temp_list = acid_temp_list
        self.scale_weight   = scale_weight
        self.target_pi      = target_pi
        self.K_cal          = _load_cal().get(self.CAL_KEY, 1.0)

    def _k_i(self, conc, temp_c):
        T_k        = temp_c + 273.15
        arrhenius  = math.exp(-self.EA_R * (1.0/T_k - 1.0/self.T_REF))
        c_factor   = max(conc/self.C_REF, 0.01) ** self.N_CONC
        scale_corr = (8.5 / max(self.scale_weight, 1.0)) ** 0.3
        return self.K0 * arrhenius * c_factor * self.SCALE_RATE_FACTOR * scale_corr * self.K_cal

    def _compute_pi(self, v_mpm):
        v_mps  = v_mpm / 60.0
        pi     = 0.0
        pp, rt = [], []
        for i in range(self.NUM_TANKS):
            t_i = self.TANK_LENGTH / v_mps
            k_i = self._k_i(self.acid_conc_list[i], self.acid_temp_list[i])
            pi  = 100.0 - (100.0 - pi) * math.exp(-k_i * t_i)
            pp.append(round(pi, 2)); rt.append(round(t_i, 1))
        return pi, pp, rt

    def _risk_level(self, speed, pi):
        avg_conc = sum(self.acid_conc_list)/len(self.acid_conc_list)
        avg_temp = sum(self.acid_temp_list)/len(self.acid_temp_list)
        s = 0
        if pi < 85: s += 3
        elif pi < 92: s += 1
        if speed > 140: s += 2
        if avg_conc < 120: s += 2
        if avg_temp < 68: s += 2
        if self.thickness > 4.0: s += 1
        return "HIGH" if s >= 5 else "MEDIUM" if s >= 2 else "LOW"

    def _physics_result(self):
        pi_min, _, _ = self._compute_pi(self.V_MIN)
        if pi_min < self.target_pi:
            pi, pp, rt = self._compute_pi(self.V_MIN)
            return {
                "max_speed": self.V_MIN, "pi_per_tank": pp,
                "residence_time_per_tank": rt, "total_pi": round(pi, 2),
                "under_pickling_risk": self._risk_level(self.V_MIN, pi),
                "warning": "酸液条件不足，建议检查酸浓度和温度",
                "K_cal": self.K_cal, "source": "physics",
            }
        lo, hi, best = self.V_MIN, self.V_MAX, self.V_MIN
        while hi - lo >= 0.5:
            mid = (lo + hi) / 2.0
            if self._compute_pi(mid)[0] >= self.target_pi:
                best = mid; lo = mid + 0.5
            else:
                hi = mid - 0.5
        best = math.floor(best)
        pi, pp, rt = self._compute_pi(best)
        return {
            "max_speed": best, "pi_per_tank": pp,
            "residence_time_per_tank": rt, "total_pi": round(pi, 2),
            "under_pickling_risk": self._risk_level(best, pi),
            "warning": None, "K_cal": self.K_cal, "source": "physics",
        }

    def calculate(self) -> Dict[str, Any]:
        x = [self.thickness, self.scale_weight] + self.acid_conc_list + self.acid_temp_list
        pt = _pt_infer("acid_speed", x)
        if pt is not None:
            raw_speed = pt[0] * self.K_cal
            best = int(max(self.V_MIN, min(self.V_MAX, round(raw_speed))))
            pi, pp, rt = self._compute_pi(best)
            return {
                "max_speed": best, "pi_per_tank": pp,
                "residence_time_per_tank": rt, "total_pi": round(pi, 2),
                "under_pickling_risk": self._risk_level(best, pi),
                "warning": None, "K_cal": self.K_cal, "source": "pt",
            }
        return self._physics_result()

    def calibrate(self, actual_max_speed, actual_quality_ok):
        predicted = self.calculate()["max_speed"]
        if not actual_quality_ok:
            adj = 0.95
        else:
            ratio = actual_max_speed / max(predicted, 1.0)
            adj   = max(0.7, min(1.3, 1.0 + 0.3*(ratio - 1.0)))
        self.K_cal = round(self.K_cal * adj, 4)
        cal = _load_cal(); cal[self.CAL_KEY] = self.K_cal; _save_cal(cal)
        return self.K_cal


# ─────────────────────────────────────────────────────────────────────────────
# 2. 张力设定模型
# ─────────────────────────────────────────────────────────────────────────────
class TensionModel:
    """
    灰箱: T_base = coef × σ_yield × A，各区段比例系数
    PT栈: 4 维输入 → 10 区段张力 kN
    输入: [thickness, width, yield_strength, tension_coef]
    """
    ZONE_RATIOS = {
        "inlet": 1.00, "s1_roller": 0.85, "acid_entry": 0.78,
        "acid1": 0.72, "acid2": 0.68,    "acid3": 0.68,
        "rinse": 0.70, "leveler": 0.76,  "s2_roller": 0.88, "outlet": 1.00,
    }
    ZONE_NAMES_CN = {
        "inlet": "入口张力辊", "s1_roller": "S1夹送辊",
        "acid_entry": "酸洗入口辊", "acid1": "1#酸槽",
        "acid2": "2#酸槽", "acid3": "3#酸槽",
        "rinse": "漂洗段辊", "leveler": "拉矫机",
        "s2_roller": "S2夹送辊", "outlet": "出口张力辊",
    }

    @staticmethod
    def _zone_cal_key(zone): return f"tension_zone_{zone}"

    def __init__(self, thickness, width, yield_strength, tension_coef=0.25):
        self.thickness      = thickness
        self.width          = width
        self.yield_strength = yield_strength
        self.tension_coef   = tension_coef
        cal = _load_cal()
        self.zone_kcal = {z: cal.get(self._zone_cal_key(z), 1.0) for z in self.ZONE_RATIOS}

    def _physics_zones(self, t_base_kn):
        zones = {}
        for zone, ratio in self.ZONE_RATIOS.items():
            k = self.zone_kcal[zone]
            zones[zone] = {
                "tension_kN": round(t_base_kn * ratio * k, 2),
                "ratio": ratio, "k_cal": k,
                "name_cn": self.ZONE_NAMES_CN[zone],
            }
        return zones

    def calculate(self) -> Dict[str, Any]:
        cross = self.thickness * self.width
        t_base = self.tension_coef * self.yield_strength * cross / 1000.0

        pt = _pt_infer("tension", [self.thickness, self.width, self.yield_strength, self.tension_coef])
        if pt is not None and _scalers and "tension" in _scalers:
            zone_names = _scalers["tension"].get("zone_names", list(self.ZONE_RATIOS.keys()))
            zones = {}
            for i, zone in enumerate(zone_names):
                k = self.zone_kcal.get(zone, 1.0)
                kn = round(max(0.1, pt[i]) * k, 2)
                zones[zone] = {
                    "tension_kN": kn,
                    "ratio": self.ZONE_RATIOS.get(zone, 1.0),
                    "k_cal": k,
                    "name_cn": self.ZONE_NAMES_CN.get(zone, zone),
                }
            source = "pt"
        else:
            zones  = self._physics_zones(t_base)
            source = "physics"

        density    = 7850.0
        mass_per_m = density * (self.thickness/1000.0) * (self.width/1000.0)
        accel_kn   = round(mass_per_m * (30.0/60.0) / 1000.0, 3)
        t_max      = round(t_base * self.zone_kcal.get("inlet", 1.0), 2)

        return {
            "T_max": t_max, "T_base": round(t_base, 2),
            "cross_section_mm2": round(cross, 1),
            "zones": zones,
            "weld_speed_limit": 60.0,
            "weld_tension_kN": round(t_max * 0.60, 2),
            "accel_tension": accel_kn,
            "zone_kcal": self.zone_kcal,
            "source": source,
        }

    def calibrate(self, zone, measured_kn):
        if zone not in self.ZONE_RATIOS:
            raise ValueError(f"未知区段: {zone}")
        t_base = self.tension_coef * self.yield_strength * self.thickness * self.width / 1000.0
        pred   = t_base * self.ZONE_RATIOS[zone] * self.zone_kcal[zone]
        adj    = max(0.5, min(2.0, 1.0 + 0.4*(measured_kn/max(pred,0.1) - 1.0)))
        self.zone_kcal[zone] = round(self.zone_kcal[zone] * adj, 4)
        cal = _load_cal(); cal[self._zone_cal_key(zone)] = self.zone_kcal[zone]; _save_cal(cal)
        return self.zone_kcal


# ─────────────────────────────────────────────────────────────────────────────
# 3. 质量预测模型
# ─────────────────────────────────────────────────────────────────────────────
class QualityPredictionModel:
    """
    灰箱: Arrhenius PI 计算 + 速度惩罚
    PT栈: 6 维输入 → [pi_score, surface_score]
    输入: [thickness, avg_speed, acid_conc_avg, acid_temp_avg, scale_weight, fe_conc_avg]
    """
    EA_R    = 5413.0
    T_REF   = 348.15
    C_REF   = 180.0
    N_CONC  = 1.2
    CAL_KEY = "quality_kcal"

    def __init__(self, thickness, avg_speed, acid_conc_avg, acid_temp_avg,
                 scale_weight=8.5, fe_conc_avg=60.0):
        self.thickness     = thickness
        self.avg_speed     = avg_speed
        self.acid_conc_avg = acid_conc_avg
        self.acid_temp_avg = acid_temp_avg
        self.scale_weight  = scale_weight
        self.fe_conc_avg   = fe_conc_avg
        self.K_cal         = _load_cal().get(self.CAL_KEY, 1.0)

    def _pi(self):
        T_k        = self.acid_temp_avg + 273.15
        arr        = math.exp(-self.EA_R*(1.0/T_k - 1.0/self.T_REF))
        c_factor   = max(self.acid_conc_avg/self.C_REF, 0.01)**self.N_CONC
        fe_inh     = 1.0 - max(0.0, (self.fe_conc_avg-80.0)/200.0)*0.35
        scale_corr = (8.5/max(self.scale_weight,1.0))**0.3
        exposure   = (1.2*arr*c_factor*fe_inh*scale_corr*18.0*6) / (self.avg_speed/60.0)
        return min(max(100.0*(1.0-math.exp(-exposure/10.0))*self.K_cal, 0), 100)

    def _surface(self, pi):
        if self.avg_speed < 60:
            ss = 80.0
        elif self.avg_speed <= 140:
            ss = 80.0 + 15.0*(self.avg_speed-60)/80.0
        else:
            ss = 95.0 - 30.0*((self.avg_speed-140)/40.0)
        return min(max(pi*0.65 + ss*0.35, 0), 100)

    def _grade(self, pi, suf):
        c = (pi+suf)/2.0
        if c >= 90: return "A1"
        if c >= 80: return "A2"
        if c >= 70: return "B1"
        if c >= 60: return "B2"
        return "C"

    def _recommendations(self, pi, suf):
        recs = []
        if self.fe_conc_avg > 80:
            recs.append(f"铁离子浓度偏高（{self.fe_conc_avg:.0f} g/L），建议加速换酸")
        if pi < 80:
            recs.append("酸洗指数偏低，建议提高酸液浓度至 180 g/L 以上，或升温至 80°C")
        if pi < 65:
            recs.append(f"欠酸洗风险高，建议将线速降至 {max(self.avg_speed*0.75, 20):.0f} m/min")
        if self.acid_temp_avg < 70:
            recs.append(f"酸液温度偏低（{self.acid_temp_avg:.1f}°C），建议升温至 75~85°C")
        if self.acid_conc_avg < 120:
            recs.append(f"游离酸浓度偏低（{self.acid_conc_avg:.0f} g/L），建议补充新酸")
        if self.avg_speed > 150:
            recs.append(f"线速过高（{self.avg_speed:.0f} m/min），欠酸洗风险")
        if self.scale_weight > 12.0:
            recs.append(f"氧化铁皮偏重（{self.scale_weight:.1f} g/m²），建议检查加热炉气氛")
        if not recs:
            recs.append("工艺参数在正常范围内，当前设定可继续保持")
        return recs

    def calculate(self) -> Dict[str, Any]:
        x  = [self.thickness, self.avg_speed, self.acid_conc_avg,
               self.acid_temp_avg, self.scale_weight, self.fe_conc_avg]
        pt = _pt_infer("quality", x)

        if pt is not None:
            pi  = round(float(min(max(pt[0]*self.K_cal, 0), 100)), 1)
            suf = round(float(min(max(pt[1]*self.K_cal, 0), 100)), 1)
            src = "pt"
        else:
            pi  = round(self._pi(), 1)
            suf = round(self._surface(pi), 1)
            src = "physics"

        return {
            "pi_score": pi, "surface_score": suf,
            "overall_grade": self._grade(pi, suf),
            "recommendations": self._recommendations(pi, suf),
            "K_cal": self.K_cal, "source": src,
        }

    def calibrate(self, actual_grade):
        grade_map = {"A1": 95, "A2": 85, "B1": 75, "B2": 65, "C": 50}
        target = grade_map.get(actual_grade, 75)
        res    = self.calculate()
        pred   = (res["pi_score"] + res["surface_score"]) / 2.0
        adj    = max(0.7, min(1.3, 1.0 + 0.3*(target/max(pred,1.0) - 1.0)))
        self.K_cal = round(self.K_cal * adj, 4)
        cal = _load_cal(); cal[self.CAL_KEY] = self.K_cal; _save_cal(cal)
        return self.K_cal


# ─────────────────────────────────────────────────────────────────────────────
# 4. 消耗预测模型（无 PT 版本，定额+修正公式足够）
# ─────────────────────────────────────────────────────────────────────────────
class AcidConsumptionModel:
    ACID_WITH_REGEN    = 2.0
    ACID_WITHOUT_REGEN = 35.0
    STEAM_UNIT         = 39.8
    POWER_UNIT         = 14.0
    COOLING_UNIT       = 1.21

    def __init__(self, thickness, width, coil_weight_kg,
                 has_regen_station=True, fe_conc_avg=60.0):
        self.thickness         = thickness
        self.width             = width
        self.coil_weight_kg    = coil_weight_kg
        self.has_regen_station = has_regen_station
        self.fe_conc_avg       = fe_conc_avg

    def calculate(self) -> Dict[str, Any]:
        wt        = self.coil_weight_kg / 1000.0
        acid_base = self.ACID_WITH_REGEN if self.has_regen_station else self.ACID_WITHOUT_REGEN
        fe_factor = 1.0 + max(0.0, (self.fe_conc_avg-100.0)/100.0)*0.4
        acid_unit = round(acid_base * fe_factor, 3)
        return {
            "coil_weight_t":              round(wt, 3),
            "acid_consumption_kg":        round(acid_unit * wt, 2),
            "acid_unit_kg_per_t":         acid_unit,
            "steam_consumption_kg":       round(self.STEAM_UNIT * wt, 2),
            "steam_unit_kg_per_t":        self.STEAM_UNIT,
            "power_consumption_kwh":      round(self.POWER_UNIT * wt, 2),
            "power_unit_kwh_per_t":       self.POWER_UNIT,
            "cooling_water_m3":           round(self.COOLING_UNIT * wt, 3),
            "cooling_water_unit_m3_per_t": self.COOLING_UNIT,
            "fe_conc_factor":             round(fe_factor, 3),
        }