#!/usr/bin/env python3
"""
generate_data.py — 成品油油库库存-运输协同优化 数据生成器
================================================================
生成合成数据集，用于测试库存-运输协同优化模型。
输出: inventory_transport_data.xlsx (多Sheet)
"""

import numpy as np
from openpyxl import Workbook
from openpyxl.styles import Font, PatternFill, Alignment, Border, Side
from openpyxl.utils import get_column_letter
import os, datetime

np.random.seed(42)

# ═══════════════════════════════════════════════════════════
# 1. 配置参数
# ═══════════════════════════════════════════════════════════
M = 2                # 上游补给源数量
T_short = 30         # 短期规划周期（月度，天）
T_long = 90          # 长期规划周期（季度，天）
S = 100              # 需求场景总数
K_cluster = 30       # 场景聚合后数量

V = 12000.0          # 油库最大库容（吨）
I0 = 3500.0          # 初始库存（吨）
theta = 0.003        # 运输损耗率
alpha = 0.95         # 需求满足率阈值
Ch = 12.0            # 单位库存持有成本（元/吨·天）
Cs = 850.0           # 单位缺货损失（元/吨）
Cl = 7200.0          # 单位损耗成本（元/吨，约汽油均价）

# 补给源参数
supply_sources = {
    0: {'name': '炼厂A', 'Q_max': 800.0, 'L_km': 185.0, 'Ct': 0.45},
    1: {'name': '中转库B', 'Q_max': 600.0, 'L_km': 220.0, 'Ct': 0.52},
}

# ═══════════════════════════════════════════════════════════
# 2. 需求数据生成 (ARIMA(2,1,2) 模拟)
# ═══════════════════════════════════════════════════════════
def generate_demand_arima(T, base=385.0, sigma=82.0, phi1=0.6, phi2=0.25, theta1=-0.3, theta2=-0.15):
    """生成类ARIMA(2,1,2)需求序列"""
    eps = np.random.normal(0, sigma, T + 10)
    y = np.zeros(T + 10)
    y[0], y[1] = base, base
    
    # AR(2) on differences (I(1))
    diff = np.zeros(T + 10)
    diff[0], diff[1] = 0, 0
    for t in range(2, T + 10):
        # MA(2) error
        ma_err = eps[t] + theta1 * eps[t-1] + theta2 * eps[t-2]
        # AR(2) + MA
        diff[t] = phi1 * diff[t-1] + phi2 * diff[t-2] + ma_err
        y[t] = y[t-1] + diff[t]
    
    demand = y[10:]  # 丢弃预热期
    demand = np.maximum(demand, base * 0.4)  # 下限
    demand = np.minimum(demand, base * 2.0)  # 上限
    return demand

# 生成T=30天需求
demand_30 = generate_demand_arima(T_short, base=385.0, sigma=82.0)

# 生成T=90天需求
np.random.seed(123)
demand_90 = generate_demand_arima(T_long, base=385.0, sigma=82.0)

# ═══════════════════════════════════════════════════════════
# 3. 场景生成（S个场景，每个T期）
# ═══════════════════════════════════════════════════════════
def generate_scenarios(T, S, base=385.0, sigma=82.0):
    """生成S个需求场景"""
    mu_t = np.full(T, base)
    # 周期性波动
    seasonal = 40 * np.sin(2 * np.pi * np.arange(T) / 7)  # 周波动
    mu_t += seasonal
    scenarios = np.random.normal(mu_t[:, None], sigma, (T, S))
    scenarios = np.maximum(scenarios, base * 0.3)
    return scenarios

np.random.seed(77)
scenarios_30 = generate_scenarios(T_short, S)
scenarios_90 = generate_scenarios(T_long, S)

# 场景概率（等概率）
p_s = np.ones(S) / S

# ═══════════════════════════════════════════════════════════
# 4. 场景聚合 (K-means)
# ═══════════════════════════════════════════════════════════
def kmeans_cluster_scenarios(scenarios, K):
    """简单K-means场景聚合"""
    T, S = scenarios.shape
    # 初始化聚类中心
    rng = np.random.RandomState(123)
    centers_idx = rng.choice(S, K, replace=False)
    centers = scenarios[:, centers_idx].copy()
    
    for _ in range(50):
        # 分配
        dists = np.array([[np.sum((scenarios[:, s] - centers[:, k])**2) for k in range(K)] for s in range(S)])
        labels = np.argmin(dists, axis=1)
        # 更新中心
        new_centers = np.zeros((T, K))
        for k in range(K):
            members = scenarios[:, labels == k]
            if members.shape[1] > 0:
                new_centers[:, k] = members.mean(axis=1)
            else:
                new_centers[:, k] = centers[:, k]
        if np.allclose(centers, new_centers):
            break
        centers = new_centers
    
    # 聚合概率
    cluster_probs = np.array([np.sum(labels == k) / S for k in range(K)])
    return centers, cluster_probs, labels

clustered_30, cprob_30, _ = kmeans_cluster_scenarios(scenarios_30, K_cluster)
clustered_90, cprob_90, _ = kmeans_cluster_scenarios(scenarios_90, K_cluster)

# ═══════════════════════════════════════════════════════════
# 5. 写入Excel
# ═══════════════════════════════════════════════════════════
wb = Workbook()
header_font = Font(bold=True, size=10, color="FFFFFF")
header_fill = PatternFill(start_color="2F5496", end_color="2F5496", fill_type="solid")
thin_border = Border(
    left=Side(style='thin'), right=Side(style='thin'),
    top=Side(style='thin'), bottom=Side(style='thin')
)

def style_header(ws, row, ncols):
    for c in range(1, ncols + 1):
        cell = ws.cell(row=row, column=c)
        cell.font = header_font
        cell.fill = header_fill
        cell.alignment = Alignment(horizontal='center')
        cell.border = thin_border

def style_data(ws, start_row, end_row, ncols):
    for r in range(start_row, end_row + 1):
        for c in range(1, ncols + 1):
            cell = ws.cell(row=r, column=c)
            cell.border = thin_border
            cell.alignment = Alignment(horizontal='center')
            if isinstance(cell.value, float):
                cell.number_format = '0.00'

# ── Sheet 1: 参数配置 ──
s1 = wb.active
s1.title = "参数配置"
params = [
    ["参数", "符号", "值", "单位", "说明"],
    ["补给源数量", "M", M, "个", "炼厂/中转库"],
    ["短期周期数", "T_short", T_short, "天", "月度规划"],
    ["长期周期数", "T_long", T_long, "天", "季度规划"],
    ["总场景数", "S", S, "个", "蒙特卡洛采样"],
    ["聚合场景数", "K", K_cluster, "个", "K-means聚合"],
    ["油库最大库容", "V", V, "吨", "设计库容"],
    ["初始库存", "I0", I0, "吨", "期初库存"],
    ["运输损耗率", "θ", theta, "", "0.3%"],
    ["需求满足率", "α", alpha, "", "95%"],
    ["库存持有成本", "Ch", Ch, "元/吨·天", ""],
    ["缺货损失成本", "Cs", Cs, "元/吨", ""],
    ["损耗单价", "Cl", Cl, "元/吨", "按油品均价"],
]
for i, row in enumerate(params, 1):
    for j, val in enumerate(row, 1):
        s1.cell(row=i, column=j, value=val)
style_header(s1, 1, 5)
style_data(s1, 2, len(params), 5)
for c in range(1, 6):
    s1.column_dimensions[get_column_letter(c)].width = 16

# ── Sheet 2: 补给源参数 ──
s2 = wb.create_sheet("补给源参数")
supp_headers = ["补给源ID", "名称", "Q_max(吨)", "距离(km)", "Ct(元/吨·km)"]
for j, h in enumerate(supp_headers, 1):
    s2.cell(row=1, column=j, value=h)
for i, (sid, info) in enumerate(supply_sources.items(), 2):
    s2.cell(row=i, column=1, value=sid)
    s2.cell(row=i, column=2, value=info['name'])
    s2.cell(row=i, column=3, value=info['Q_max'])
    s2.cell(row=i, column=4, value=info['L_km'])
    s2.cell(row=i, column=5, value=info['Ct'])
style_header(s2, 1, 5)
style_data(s2, 2, len(supply_sources) + 1, 5)
for c in range(1, 6):
    s2.column_dimensions[get_column_letter(c)].width = 16

# ── Sheet 3: 需求序列(T=30) ──
s3 = wb.create_sheet("需求序列_T30")
s3.cell(row=1, column=1, value="周期t")
s3.cell(row=1, column=2, value="需求量D_t(吨)")
s3.cell(row=1, column=3, value="累积需求(吨)")
cum = 0
for t in range(T_short):
    s3.cell(row=t+2, column=1, value=t+1)
    s3.cell(row=t+2, column=2, value=round(demand_30[t], 1))
    cum += demand_30[t]
    s3.cell(row=t+3, column=3, value=round(cum, 1))
style_header(s3, 1, 3)
style_data(s3, 2, T_short + 1, 3)
for c in range(1, 4):
    s3.column_dimensions[get_column_letter(c)].width = 16

# ── Sheet 4: 需求序列(T=90) ──
s4 = wb.create_sheet("需求序列_T90")
s4.cell(row=1, column=1, value="周期t")
s4.cell(row=1, column=2, value="需求量D_t(吨)")
cum = 0
for t in range(T_long):
    s4.cell(row=t+2, column=1, value=t+1)
    s4.cell(row=t+2, column=2, value=round(demand_90[t], 1))
style_header(s4, 1, 2)
style_data(s4, 2, T_long + 1, 2)
for c in range(1, 3):
    s4.column_dimensions[get_column_letter(c)].width = 16

# ── Sheet 5: 需求统计 ──
s5 = wb.create_sheet("需求统计")
stats = [
    ["统计量", "T=30", "T=90"],
    ["均值μ_t", round(demand_30.mean(), 1), round(demand_90.mean(), 1)],
    ["标准差σ_t", round(demand_30.std(), 1), round(demand_90.std(), 1)],
    ["最小值", round(demand_30.min(), 1), round(demand_90.min(), 1)],
    ["最大值", round(demand_30.max(), 1), round(demand_90.max(), 1)],
    ["变异系数CV", round(demand_30.std()/demand_30.mean(), 3), round(demand_90.std()/demand_90.mean(), 3)],
    ["总需求(吨)", round(demand_30.sum(), 0), round(demand_90.sum(), 0)],
]
for i, row in enumerate(stats, 1):
    for j, val in enumerate(row, 1):
        s5.cell(row=i, column=j, value=val)
style_header(s5, 1, 3)
style_data(s5, 2, len(stats), 3)
for c in range(1, 4):
    s5.column_dimensions[get_column_letter(c)].width = 18

# ── Sheet 6: 场景数据(T=30前10场景) ──
s6 = wb.create_sheet("场景数据_T30_前10")
s6.cell(row=1, column=1, value="周期\\场景")
for s_idx in range(min(10, S)):
    s6.cell(row=1, column=s_idx+2, value=f"场景{s_idx+1}")
for t in range(T_short):
    s6.cell(row=t+2, column=1, value=t+1)
    for s_idx in range(min(10, S)):
        s6.cell(row=t+2, column=s_idx+2, value=round(scenarios_30[t, s_idx], 1))
style_header(s6, 1, min(10, S) + 1)
style_data(s6, 2, T_short + 1, min(10, S) + 1)
s6.column_dimensions['A'].width = 10
for c in range(2, min(10, S) + 2):
    s6.column_dimensions[get_column_letter(c)].width = 12

# ── Sheet 7: 聚合场景(T=30, K=30) ──
s7 = wb.create_sheet("聚合场景_T30")
s7.cell(row=1, column=1, value="周期\\聚合场景")
for k in range(K_cluster):
    s7.cell(row=1, column=k+2, value=f"簇{k+1}")
    s7.cell(row=1, column=k+2+K_cluster+1, value=f"概率")
s7.cell(row=1, column=K_cluster*2+2, value="—")
for t in range(T_short):
    s7.cell(row=t+2, column=1, value=t+1)
    for k in range(K_cluster):
        s7.cell(row=t+2, column=k+2, value=round(clustered_30[t, k], 1))
for k in range(K_cluster):
    s7.cell(row=2, column=K_cluster+k+2, value=round(cprob_30[k], 4))
style_header(s7, 1, K_cluster*2+2)
for c in range(1, K_cluster*2+3):
    s7.column_dimensions[get_column_letter(c)].width = 8
s7.column_dimensions['A'].width = 10

# ── Save ──
outpath = '/home/ubuntu/chengzi-buluo/orange-family/docs/inventory_transport_data.xlsx'
os.makedirs(os.path.dirname(outpath), exist_ok=True)
wb.save(outpath)
print(f"✅ Data Excel saved: {outpath}")
print(f"   Sheets: {wb.sheetnames}")
print(f"   T=30 demand: μ={demand_30.mean():.1f}, σ={demand_30.std():.1f}")
print(f"   T=90 demand: μ={demand_90.mean():.1f}, σ={demand_90.std():.1f}")
print(f"   Scenarios: {S} raw → {K_cluster} clustered")