import pandas as pd
import joblib
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import train_test_split

# === 配置路径 ===
csv_path = 'C:\\Users\\Administrator\\Desktop\\defrost\\feedback_data.csv'
model_save_path = "defrost_time_corrector.pkl"

# === 特征列定义 ===
feature_columns = [
    "w", "rho_coal", "rho_ice", "C_coal", "C_ice", "L", "k_coal", "k_ice", "h",
    "T_air", "T_initial", "T_m", "a", "b", "c"
]

# === 1. 读取CSV并预处理 ===
df = pd.read_csv(csv_path, parse_dates=["t_formula", "t_real"], encoding='utf-8')

# 确保字段类型正确
df["material_name"] = df["material_name"].astype(str)
df["manufactured_goods"] = df["manufactured_goods"].astype(str)

# 计算真实解冻时长（小时）
df["t_real_hours"] = (df["t_real"] - df["t_formula"]).dt.total_seconds() / 3600

# === 2. 训练模型（用已有所有历史数据） ===
X = df[feature_columns]
y = df["t_real_hours"]

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

model = RandomForestRegressor(n_estimators=100, random_state=42)
model.fit(X_train, y_train)

# 保存模型
joblib.dump(model, model_save_path)
print(f"✅ 模型训练完成，已保存为 {model_save_path}")

# === 3. 用最新数据预测（比如最后一条或多条） ===

# 假设你要预测最后新增的一条数据（如果多条可以改）
new_sample = df.tail(1)  # 取最后一行，也可以是 tail(n) 最后n行

X_new = new_sample[feature_columns]
predicted_time = model.predict(X_new)[0]

# 把预测值写回DataFrame
df.loc[new_sample.index, "predicted_t_real_hours"] = predicted_time

# === 4. 保存带预测值的CSV ===
df.to_csv(csv_path, encoding='gbk', index=False)
print(f"✅ 最新数据预测完成，已更新到 {csv_path}")

print(f"\n📊 预测最后一条数据真实解冻时间为：{predicted_time:.2f} 小时")