import pandas as pd
import numpy as np
from pulp import LpProblem, LpVariable, LpMinimize, LpInteger, lpSum, LpBinary, PULP_CBC_CMD
import os
from datetime import datetime, timedelta  # 新增：用于生成日期

# ==================== 参数配置 ====================
COST_FORMAL = 200  # 正式工每人每天成本（元）
COST_TEMP = 250    # 临时工每人每天成本（元）

shift_names = ["00:00-08:00", "05:00-13:00", "08:00-16:00",
               "12:00-20:00", "14:00-22:00", "16:00-24:00"]
shift_hours = [[0,1,2,3,4,5,6,7], [5,6,7,8,9,10,11,12], [8,9,10,11,12,13,14,15],
               [12,13,14,15,16,17,18,19], [14,15,16,17,18,19,20,21], [16,17,18,19,20,21,22,23]]
CAP_F = 25 * 8
CAP_P = 20 * 8
K = 80
MAX_DAYS_PER_WORKER = int(0.85 * 30)  # 25天
L = 7

# -------------------- 1. 读取历史数据，计算权重 --------------------
df_hist = pd.read_excel("SC4小时货量（last 30days）.xlsx")
df_hist = df_hist.sort_values(["日期", "小时"]).reset_index(drop=True)

dates = df_hist["日期"].unique()
full_data = []
for d in dates:
    day_df = df_hist[df_hist["日期"] == d]
    day_hours = set(day_df["小时"].values)
    for h in range(24):
        if h in day_hours:
            vol = day_df[day_df["小时"] == h]["货量"].values[0]
        else:
            vol = df_hist[df_hist["小时"] == h]["货量"].mean()
        full_data.append([d, h, vol])
full_df = pd.DataFrame(full_data, columns=["日期", "小时", "货量"])
hourly_mean = full_df.groupby("小时")["货量"].mean()
shift_volume = [sum(hourly_mean[h] for h in hours) for hours in shift_hours]
total_day_vol = sum(shift_volume)
weights = [sv / total_day_vol for sv in shift_volume]

print("=== 班次权重 ===")
for i in range(6):
    print(f"{shift_names[i]}: {weights[i]:.2%}")

# -------------------- 2. 读取未来30天每日总货量 --------------------
df_future = pd.read_excel("未来30天预测货量.xlsx")
future_daily_totals = df_future["预测货量"].values
T = len(future_daily_totals)

# 生成未来30天的日期列表（假设第一天为2023-12-01）
start_date = datetime(2023, 12, 1)
date_list = [start_date + timedelta(days=t) for t in range(T)]

# -------------------- 3. 生成两种需求矩阵 --------------------
# 方案A：权重分配
demand_weighted = np.zeros((T, 6), dtype=int)
for t in range(T):
    for i in range(6):
        demand_weighted[t, i] = int(round(future_daily_totals[t] * weights[i]))

# 方案B：均匀分配
demand_uniform = np.zeros((T, 6), dtype=int)
for t in range(T):
    base = future_daily_totals[t] / 6
    alloc = [int(round(base)) for _ in range(6)]
    diff = int(future_daily_totals[t] - sum(alloc))
    for j in range(abs(diff)):
        alloc[j % 6] += 1 if diff > 0 else -1
    demand_uniform[t, :] = alloc

# -------------------- 4. 求解函数（返回模型变量以便提取排班详情） --------------------
def solve_scheduling(demand):
    model = LpProblem("Scheduling", LpMinimize)
    F = {(t, i): LpVariable(f"F_{t}_{i}", lowBound=0, cat=LpInteger)
         for t in range(T) for i in range(6)}
    X = {(k, t, i): LpVariable(f"X_{k}_{t}_{i}", cat=LpBinary)
         for k in range(K) for t in range(T) for i in range(6)}

    # 目标：最小化临时工总人天数
    model += lpSum(F[t, i] for t in range(T) for i in range(6))

    # 产能约束
    for t in range(T):
        for i in range(6):
            model += CAP_F * lpSum(X[k, t, i] for k in range(K)) + CAP_P * F[t, i] >= demand[t, i]
    # 每人每天最多一班
    for k in range(K):
        for t in range(T):
            model += lpSum(X[k, t, i] for i in range(6)) <= 1
    # 出勤率上限
    for k in range(K):
        model += lpSum(X[k, t, i] for t in range(T) for i in range(6)) <= MAX_DAYS_PER_WORKER
    # 连续出勤 ≤7
    for k in range(K):
        for s in range(T - L):
            model += lpSum(X[k, t, i] for t in range(s, s + L + 1) for i in range(6)) <= L

    model.solve(PULP_CBC_CMD(msg=False))
    total_temp = sum(F[t, i].varValue for t in range(T) for i in range(6) if F[t, i].varValue is not None)
    total_formal = sum(X[k, t, i].varValue for k in range(K) for t in range(T) for i in range(6) if X[k, t, i].varValue is not None)
    return total_temp, total_formal, F, X  # 返回变量字典，以便输出详细排班

# -------------------- 5. 求解两种方案 --------------------
print("\n求解中...")
temp_w, formal_w, F_w, X_w = solve_scheduling(demand_weighted)
temp_u, formal_u, F_u, X_u = solve_scheduling(demand_uniform)

# -------------------- 6. 成本计算与对比 --------------------
cost_w = formal_w * COST_FORMAL + temp_w * COST_TEMP
cost_u = formal_u * COST_FORMAL + temp_u * COST_TEMP

print("\n" + "=" * 70)
print("成本对比分析")
print("=" * 70)
print(f"{'方案':<15}{'正式工成本(元)':<18}{'临时工成本(元)':<18}{'总成本(元)':<15}")
print(f"{'基于权重':<15}{formal_w * COST_FORMAL:<18.0f}{temp_w * COST_TEMP:<18.0f}{cost_w:<15.0f}")
print(f"{'均匀分配':<15}{formal_u * COST_FORMAL:<18.0f}{temp_u * COST_TEMP:<18.0f}{cost_u:<15.0f}")
print(f"\n基于权重的方案比均匀分配方案节约成本: {cost_u - cost_w:.0f} 元 ({(cost_u - cost_w) / cost_u * 100:.2f}%)")

# -------------------- 7. 输出未来一周（前7天）的具体排班结果（基于权重方案） --------------------
print("\n" + "=" * 70)
print("未来一周（2023年12月1日 - 12月7日）基于权重方案的各班次人员安排")
print("=" * 70)

week_rows = []
for t in range(min(7, T)):
    date_str = date_list[t].strftime("%Y-%m-%d")
    print(f"\n日期: {date_str}  总货量={future_daily_totals[t]:.0f} 件")
    print(f"{'班次':<18}{'货量需求':<10}{'正式工人数':<10}{'临时工人数':<10}{'总人数':<10}")
    print("-" * 58)
    for i in range(6):
        f_cnt = sum(X_w[k, t, i].varValue for k in range(K) if X_w[k, t, i].varValue is not None)
        t_cnt = F_w[t, i].varValue if F_w[t, i].varValue is not None else 0
        print(f"{shift_names[i]:<18}{demand_weighted[t,i]:<10}{int(f_cnt):<10}{int(t_cnt):<10}{int(f_cnt+t_cnt):<10}")
        week_rows.append({
            "日期": date_str,
            "总货量": future_daily_totals[t],
            "班次": shift_names[i],
            "货量需求": demand_weighted[t, i],
            "正式工人数": int(f_cnt),
            "临时工人数": int(t_cnt),
            "总人数": int(f_cnt + t_cnt)
        })

# 保存一周排班结果到CSV
os.makedirs("outputs", exist_ok=True)
df_week = pd.DataFrame(week_rows)
week_file = "outputs/未来一周排班结果.csv"
df_week.to_csv(week_file, index=False, encoding="utf-8-sig")
print(f"\n未来一周排班结果已保存至: {week_file}")

# 保存成本对比结果
pd.DataFrame({
    "方案": ["基于权重", "均匀分配"],
    "正式工总人天数": [formal_w, formal_u],
    "临时工总人天数": [temp_w, temp_u],
    "总成本(元)": [cost_w, cost_u]
}).to_csv("outputs/成本对比.csv", index=False, encoding="utf-8-sig")
print("成本对比结果已保存至 outputs/成本对比.csv")