"""
超滤膜阻力模型模块
====================
本模块定义了超滤膜阻力的动态变化模型，包括：
1. ResistanceIncreaseModel: 过滤阶段膜阻力上升模型
2. ResistanceDecreaseModel: 反洗阶段膜阻力下降模型

这些模型用于模拟超滤膜在运行过程中的阻力变化，是强化学习环境的核心组件。
"""

import torch
import numpy as np
from uf_train.env.env_params import UFState, UFPhysicsParams


# ==================== 膜阻力上升模型 ====================
class ResistanceIncreaseModel(torch.nn.Module):
    """
    过滤阶段膜阻力上升模型
    
    功能说明：
    - 计算在过滤阶段膜阻力的增长量 ΔR
    - 膜阻力上升主要由污染物在膜表面的累积引起
    - 阻力增长速率与膜通量（J）和过滤时长（L_s）相关
    
    模型公式：
        ΔR = nuK × J × L_s
        其中：
        - nuK: 膜阻力增长系数（反映水质污染特性）
        - J: 膜通量 = q_UF / A / 3600 [m/s]
        - L_s: 过滤时长 [秒]
    """

    def __init__(self, phys: UFPhysicsParams):
        """初始化膜阻力上升模型"""
        super().__init__()
        self.phys = phys

    def forward(self, state: UFState, L_s: float) -> float:
        """
        前向传播：计算膜阻力上升量
        
        参数：
            state (UFState): 超滤运行状态变量对象，包含：
                - q_UF: 过滤进水流量 [m³/h]
                - nuK: 膜阻力增长系数 [m⁻¹/s]
            L_s (float): 过滤时长 [秒]
        
        返回：
            float: 膜阻力上升量 ΔR（已缩放1e10）
        
        注意：
            - 实际膜阻力量级为1e12，为便于数值计算已缩放至1e2量级
            - 膜面积 A = 128组 × 40 m²/组 = 5120 m²，现已优化为 UFPhysicsParams类配置
        """
        # 加载膜有效面积（锡山水厂配置：128组膜，每组40m²）
        A = self.phys.A
        
        # 计算膜通量 J = 流量 / 面积 / 时间单位转换
        # q_UF [m³/h] → J [m³/(m²·s)]
        J = state.q_UF / A / 3600.0  # [m/s]
        
        # 膜阻力上升模型（线性模型，已缩放）
        # nuK: 阻力增长速率，反映水质污染特性
        # J: 膜通量，通量越大污染速率越快
        # L_s: 过滤时间，时间越长累积污染越多
        dR = state.nuK * J * L_s  # [缩放后的阻力单位]
        
        return float(dR)


# ==================== 膜阻力下降模型 ====================
class ResistanceDecreaseModel(torch.nn.Module):
    """
    反洗阶段膜阻力下降模型
    
    功能说明：
    - 计算物理反冲洗能够去除的膜阻力量
    - 区分可逆污染和不可逆污染
    - 反洗时长影响去除效率
    
    模型原理：
    1. 膜污染分为两类：
       - 可逆污染：可通过物理反洗去除（如表面颗粒物）
       - 不可逆污染：无法通过物理反洗去除（如孔内吸附污染）
    
    2. 不可逆污染累积模型：
       R_irr = R0 + slope × t^power
       其中 t 为累积运行时间
    
    3. 反洗效率模型：
       time_gain = 1 - exp(-t_bw / τ)
       反洗时间越长，去除效率越高，但存在上限
    """

    def __init__(self, phys: UFPhysicsParams):
        """初始化膜阻力下降模型（无需训练参数）"""
        super().__init__()
        self.phys = phys

    def forward(self,state: UFState, R0, R_end, L_h_next_start, t_bw_s):
        """
        前向传播：计算物理反洗能够去除的膜阻力
        
        参数：
            state (UFState): 超滤运行状态变量对象，包含：
                - slope: 不可逆污染增长斜率
                - power: 不可逆污染增长幂次
            phys  (UFPhysicsParams): 超滤运行固定参数对象，包含：
                - tau_bw_s: 反洗时长影响的时间尺度

            R0 (float): 本超级周期初始膜阻力
            R_end (float): 过滤结束时的膜阻力（峰值）
            L_h_next_start (float): 下一小周期起始时的累积运行时间 [小时]
            t_bw_s (float): 物理反洗时长 [秒]
        
        返回：
            float: 物理反洗实际去除的膜阻力量
        
        计算步骤：
            1. 基于长期污染模型计算本周期的不可逆污染增量
            2. 计算可逆污染量 = 当前总污染 - 不可逆污染
            3. 应用时间因子（反洗时长的影响）
            4. 返回实际去除的阻力（不超过可逆污染量）
        """
        # ========== 步骤1：计算不可逆污染累积 ==========
        # 使用幂律模型描述长期不可逆污染的累积
        # R_irr(t) = R0 + slope × t^power

        # 下一小周期开始时的理论膜阻力
        delta_R = state.slope * (L_h_next_start ** state.power)
        R_next_start = R0 + delta_R

        # ========== 步骤2：计算可逆污染 ==========
        # 可逆污染 = 当前污染 - 理论长期污染
        reversible_R = max(R_end - R_next_start, 0.0)
        
        # ========== 步骤3：计算反洗时间效率因子 ==========
        # 使用指数衰减模型：time_gain = 1 - exp(-t_bw / τ)
        # τ (tau_bw_s): 时间尺度参数
        # - 反洗时间 t_bw = 0 时，time_gain = 0（无去除效果）
        # - 反洗时间 t_bw → ∞ 时，time_gain → 1（达到最大效率）
        # - 反洗时间 t_bw = τ 时，time_gain ≈ 0.632（去除63.2%）
        # τ参考计算： 以60s去除效率高于95%的要求推算得τ取20s
        # 该取值下：
        # - 反洗时间 t_bw = 40s 时，time_gain ≈ 0.865（去除86.5%）
        # - 反洗时间 t_bw = 60s 时，time_gain ≈ 0.950（去除95.0%）
        tau = self.phys.tau_bw_s
        time_gain = 1.0 - np.exp(-t_bw_s / tau)
        
        # ========== 步骤4：计算实际去除的膜阻力 ==========
        # 实际去除量 = 可逆污染量 × 时间效率因子
        dR_bw = reversible_R * time_gain
        
        # 确保去除量不超过可逆污染总量（物理约束）
        return float(np.clip(dR_bw, 0.0, reversible_R))


# ===== 主程序 =====
if __name__ == "__main__":
    phys = UFPhysicsParams()

    model_fp = ResistanceIncreaseModel(phys)
    model_bw = ResistanceDecreaseModel(phys)

    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")