import pandas as pd
import numpy as np
import os
from sklearn.linear_model import LinearRegression

# ============================================================
BASE_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))

FEATURE_PATH = f"{BASE_DIR}/use_data/test_cycle_sample.csv"


# ============================================================
# viscosity
# ============================================================
def xishan_viscosity(T):
    x = (T + 273.15) / 300
    factor = 890 / (
        280.68 * x ** -1.9 +
        511.45 * x ** -7.7 +
        61.131 * x ** -19.6 +
        0.45903 * x ** -40
    )
    return 0.00089 / factor


# ============================================================
# unit mapping
# ============================================================
def get_unit_cols(unit):
    return {
        "flux": f"{unit}_FluxF",
        "dpt1": f"C.M.{unit}_DB@DPT_1",
        "dpt2": f"C.M.{unit}_DB@DPT_2"
    }


# ============================================================
# resistance
# ============================================================
def compute_resistance(df, unit):

    cols = get_unit_cols(unit)

    flux = df[cols["flux"]]
    dpt1 = df[cols["dpt1"]]
    dpt2 = df[cols["dpt2"]]

    mu = xishan_viscosity(df["C.M.RO_TT_ZJS@out"])

    df["R1"] = dpt1 * 3.6e12 / (mu * flux)
    return df


# ============================================================
# V accumulation
# ============================================================
def compute_V(df, unit):

    cols = get_unit_cols(unit)
    flux = df[cols["flux"]].values

    dt = 300
    J = flux * (1e-3 / 3600)

    V = np.zeros(len(df))
    for i in range(1, len(df)):
        V[i] = V[i-1] + J[i-1] * dt

    df["V"] = V
    return df


# ============================================================
# fit model (POWER LAW)
# R = a * V^b
# log R = log a + b log V
# ============================================================
def fit_model(df, R_col):

    df = df.dropna(subset=[R_col, "V"])
    df = df[(df[R_col] > 0) & (df["V"] > 0)]

    if len(df) < 10:
        return None

    V = df["V"].values
    R = df[R_col].values / 1e12

    X = np.log(V + 1e-9).reshape(-1, 1)
    y = np.log(R + 1e-9)

    model = LinearRegression()
    model.fit(X, y)

    return model


# ============================================================
# solve V limit (analytical)
# ============================================================
def solve_v_limit(model, R_limit):

    if model is None:
        return None

    b = model.coef_[0]
    log_a = model.intercept_

    target = np.log(R_limit / 1e12)

    if abs(b) < 1e-12:
        return np.inf

    logV = (target - log_a) / b
    V_limit = np.exp(logV)

    return V_limit


# ============================================================
# main per unit
# ============================================================
def process_unit(unit, df_new, TMP_limit=0.23):

    df = df_new.copy()

    df = compute_resistance(df, unit)
    df = compute_V(df, unit)

    # =========================
    # 当前状态
    # =========================
    mu_avg = xishan_viscosity(df["C.M.RO_TT_ZJS@out"]).mean()

    cols = get_unit_cols(unit)
    flux_avg = df[cols["flux"]].mean()

    V_now = df["V"].iloc[-1]
    R_now = df["R1"].iloc[-1]

    days = max((df["time"].max() - df["time"].min()).total_seconds() / 86400, 1e-6)
    V_daily = V_now / days

    # =========================
    # model
    # =========================
    model = fit_model(df, "R1")

    if model is None:
        return None

    # =========================
    # R limit from TMP
    # =========================
    R_limit = TMP_limit * 3.6e12 / (mu_avg * flux_avg)

    # =========================
    # solve
    # =========================
    V_limit = solve_v_limit(model, R_limit)

    remaining_days = (V_limit - V_now) / max(V_daily, 1e-9)
    remaining_days = min(remaining_days, 90 - days)

    return {
        "unit": unit,
        "R_now": R_now,
        "V_now": V_now,
        "R_limit": R_limit,
        "V_limit": V_limit,
        "remaining_days": remaining_days
    }


# ============================================================
# main
# ============================================================
def main(units, TMP_limit=0.23):

    df_new = pd.read_csv(FEATURE_PATH)
    df_new["time"] = pd.to_datetime(df_new["time"])

    results = []

    for u in units:
        res = process_unit(u, df_new, TMP_limit)
        if res is not None:
            results.append(res)

    return pd.DataFrame(results)


# ============================================================
# run
# ============================================================
if __name__ == "__main__":

    units_to_run = ["RO4"]

    out = main(units_to_run, TMP_limit=0.20)

    print(f"该机组预计下次CIP为：{out['remaining_days'].iloc[0]:.2f} 天后")