import torch
import numpy as np

# ===== 膜阻力上升模型 =====
class ResistanceIncreaseModel(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, p, L_s):
        """
        计算膜阻力上升量 ΔR
        """
        A = 128 * 40.0
        J = p.q_UF / A / 3600
        # 膜阻力上升模型(已缩放)
        dR = p.nuK * J * L_s
        return float(dR)


# ===== 膜阻力下降模型 =====
class ResistanceDecreaseModel(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, p, R0, R_end, L_h_start, L_h_next_start, t_bw_s):
        """
        计算物理反冲洗污染去除比例（受反洗时间影响），最大可去除的可逆膜阻力（受过滤时间影响）
        """

        # 计算单次不可逆膜阻力（线性依赖于进水时间）
        # 周期起点和下次起点的理论阻力
        R_start = R0 + p.slope * (L_h_start ** p.power)
        R_next_start = R0 + p.slope * (L_h_next_start ** p.power)

        # 不可逆污染（反洗后残余增加量）
        irreversible_R = max(R_next_start - R_start, 0.0)

        # 本周期的总污染增长量
        total_increase = max(R_end - R_start, 0.0)

        # 可逆污染量 = 本周期总增长 - 不可逆残留
        reversible_R = max(total_increase - irreversible_R, 0.0)

        # 时间因子：反洗时间越长，效果越充分
        time_gain = 1.0 - np.exp(- (t_bw_s / p.tau_bw_s))

        # 实际去除的膜阻力（随机在可去除区间内，乘以时间因子）
        dR_bw = reversible_R * time_gain

        return float(np.clip(dR_bw, 0.0, reversible_R))


# ===== 主程序 =====
if __name__ == "__main__":
    model_fp = ResistanceIncreaseModel()
    model_bw = ResistanceDecreaseModel()

    torch.save(model_fp.state_dict(), "resistance_model_fp.pth")
    torch.save(model_bw.state_dict(), "resistance_model_bw.pth")

    print("模型已安全保存为 resistance_model_fp.pth、resistance_model_bw.pth")