import torch
import numpy as np


class TMPIncreaseModel(torch.nn.Module):
    """
    跨膜压差上升模型
    
    计算过滤阶段的TMP增长量
    """
    def __init__(self):
        super().__init__()
    
    def forward(self, p, L_h):
        """
        计算TMP增长量
        
        参数:
            p: 系统参数对象
            L_h: 过滤时长（小时）
            
        返回:
            TMP增长量
        """
        return float(p.alpha * (p.q_UF ** p.belta) * L_h)


class TMPDecreaseModel(torch.nn.Module):
    """
    跨膜压差恢复模型
    
    计算反洗阶段的TMP恢复比例
    """
    def __init__(self):
        super().__init__()
    
    def forward(self, p, L_s, t_bw_s):
        """
        计算反洗恢复比例
        
        参数:
            p: 系统参数对象
            L_s: 过滤时长（秒）
            t_bw_s: 反洗时长（秒）
            
        返回:
            TMP恢复比例（0到1之间）
        """
        L = max(float(L_s), 1.0)
        t = max(float(t_bw_s), 1e-6)
        
        # 计算恢复比例上限（随过滤时长增加而降低）
        upper_L = p.phi_bw_min + (p.phi_bw_max - p.phi_bw_min) * np.exp(-L / p.L_ref_s)
        
        # 计算时间增益因子（反洗时间越长，恢复越好）
        time_gain = 1.0 - np.exp(-(t / p.tau_bw_s) ** p.gamma_t)
        
        # 综合计算恢复比例
        phi = upper_L * time_gain
        
        return float(np.clip(phi, 0.0, 0.999))


if __name__ == "__main__":
    # 创建模型实例
    model_fp = TMPIncreaseModel()
    model_bw = TMPDecreaseModel()

    # 保存模型参数
    torch.save(model_fp.state_dict(), "uf_fp.pth")
    torch.save(model_bw.state_dict(), "uf_bw.pth")

    print("模型已保存 uf_fp.pth uf_bw.pth")