"""
Task 3 - Step 3: 第一站点最优分配计算
=====================================

输入: 02_pairing.xlsx (选中的配对列表)
输出: 03_allocation.xlsx (含最优分配量q*的配对列表)

最优分配公式:
q* = (σ_j * μ_i + σ_i * (Q - μ_j)) / (σ_i + σ_j)

鲁棒性约束:
q_final = clip(q*, μ_i - σ_i, Q - μ_j + σ_j)

物理意义:
- 波动大的站点需要更多"缓冲"
- 当σ_j >> σ_i时，q* → μ_i（精确分配给站点i）
- 当σ_i >> σ_j时，q* → Q - μ_j（为站点j预留更多）
"""

import pandas as pd
import numpy as np
from scipy import stats

# ============================================
# 参数设置
# ============================================
INPUT_FILE = '02_pairing.xlsx'
OUTPUT_FILE = '03_allocation.xlsx'

Q = 400             # 卡车容量
K_ROBUST = 1        # 鲁棒性水平 (84%保护)

# ============================================
# 读取数据
# ============================================
print("=" * 60)
print("Task 3 - Step 3: 第一站点最优分配计算")
print("=" * 60)

# 读取选中的配对
df_pairs = pd.read_excel(INPUT_FILE, sheet_name='selected_pairs')
print(f"\n读取配对数据: {len(df_pairs)} 对")

# ============================================
# 计算最优分配
# ============================================
print(f"\n" + "-" * 40)
print("计算最优分配 q*")
print("-" * 40)

def optimal_allocation(mu_i, sigma_i, mu_j, sigma_j, Q=400):
    """
    计算双站点访问时第一站点的最优分配量

    公式: q* = (σ_j * μ_i + σ_i * (Q - μ_j)) / (σ_i + σ_j)
    """
    if sigma_i + sigma_j == 0:
        # 边界情况：两站点都无波动
        return mu_i

    q_star = (sigma_j * mu_i + sigma_i * (Q - mu_j)) / (sigma_i + sigma_j)
    return q_star


def robust_allocation(q_star, mu_i, sigma_i, mu_j, sigma_j, Q=400, k=1):
    """
    应用鲁棒性约束

    约束: μ_i - k*σ_i ≤ q ≤ Q - μ_j + k*σ_j
    """
    q_lower = max(0, mu_i - k * sigma_i)
    q_upper = min(Q, Q - mu_j + k * sigma_j)

    # 确保上下界合理
    if q_lower > q_upper:
        # 如果约束冲突，取中点
        q_final = (q_lower + q_upper) / 2
    else:
        q_final = np.clip(q_star, q_lower, q_upper)

    return q_final, q_lower, q_upper


def expected_service(q, mu, sigma):
    """
    计算截断期望 E[min(D, q)]

    E[min(D, q)] = μ * Φ(z) - σ * φ(z) + q * (1 - Φ(z))
    where z = (q - μ) / σ

    当q > μ时，大部分时候D < q，期望接近μ
    当q < μ时，经常D > q，期望接近q但略低
    """
    if sigma == 0:
        return min(mu, q)

    z = (q - mu) / sigma
    phi_z = stats.norm.pdf(z)
    Phi_z = stats.norm.cdf(z)

    return mu * Phi_z - sigma * phi_z + q * (1 - Phi_z)


# 计算每个配对的分配
results = []

for idx, row in df_pairs.iterrows():
    mu_i = row['mu_i']
    sigma_i = row['sigma_i']
    mu_j = row['mu_j']
    sigma_j = row['sigma_j']

    # 计算最优分配
    q_star = optimal_allocation(mu_i, sigma_i, mu_j, sigma_j, Q)

    # 应用鲁棒性约束
    q_final, q_lower, q_upper = robust_allocation(
        q_star, mu_i, sigma_i, mu_j, sigma_j, Q, K_ROBUST
    )

    # 计算期望服务量
    E_Si = expected_service(q_final, mu_i, sigma_i)
    E_Sj = expected_service(Q - q_final, mu_j, sigma_j)
    E_total = E_Si + E_Sj

    # 计算分配比例
    alloc_ratio = q_final / Q if Q > 0 else 0

    results.append({
        **row.to_dict(),
        'q_star': q_star,
        'q_lower': q_lower,
        'q_upper': q_upper,
        'q_final': q_final,
        'q_ratio': alloc_ratio,
        'E_Si': E_Si,
        'E_Sj': E_Sj,
        'E_total': E_total,
        'efficiency': E_total / (mu_i + mu_j) if (mu_i + mu_j) > 0 else 0
    })

df_allocation = pd.DataFrame(results)

# ============================================
# 统计信息
# ============================================
print(f"\n分配统计:")
print(f"  - q* 范围: [{df_allocation['q_star'].min():.1f}, {df_allocation['q_star'].max():.1f}]")
print(f"  - q_final 范围: [{df_allocation['q_final'].min():.1f}, {df_allocation['q_final'].max():.1f}]")
print(f"  - 分配比例范围: [{df_allocation['q_ratio'].min():.2%}, {df_allocation['q_ratio'].max():.2%}]")
print(f"  - 平均分配比例: {df_allocation['q_ratio'].mean():.2%}")

print(f"\n期望服务量统计:")
print(f"  - E[S_i] 范围: [{df_allocation['E_Si'].min():.1f}, {df_allocation['E_Si'].max():.1f}]")
print(f"  - E[S_j] 范围: [{df_allocation['E_Sj'].min():.1f}, {df_allocation['E_Sj'].max():.1f}]")
print(f"  - E[total] 范围: [{df_allocation['E_total'].min():.1f}, {df_allocation['E_total'].max():.1f}]")
print(f"  - 平均效率: {df_allocation['efficiency'].mean():.2%}")

# 检查是否有被约束裁剪的情况
clipped_lower = (df_allocation['q_final'] <= df_allocation['q_lower'] + 0.01).sum()
clipped_upper = (df_allocation['q_final'] >= df_allocation['q_upper'] - 0.01).sum()
print(f"\n鲁棒性约束影响:")
print(f"  - 触及下界: {clipped_lower} 对")
print(f"  - 触及上界: {clipped_upper} 对")
print(f"  - 未被裁剪: {len(df_allocation) - clipped_lower - clipped_upper} 对")

# ============================================
# 显示关键配对的分配
# ============================================
print(f"\n" + "-" * 40)
print("高价值配对的分配方案 (Top 10)")
print("-" * 40)

display_cols = ['site_i_name', 'site_j_name', 'mu_i', 'mu_j', 'q_final', 'E_Si', 'E_Sj', 'E_total']
df_top = df_allocation.nlargest(10, 'value')[display_cols].copy()
df_top['site_i_name'] = df_top['site_i_name'].str[:25]
df_top['site_j_name'] = df_top['site_j_name'].str[:25]
print(df_top.to_string(index=False))

# ============================================
# 保存结果
# ============================================
with pd.ExcelWriter(OUTPUT_FILE, engine='openpyxl') as writer:
    # Sheet 1: 完整分配结果
    df_allocation.to_excel(writer, sheet_name='allocation', index=False)

    # Sheet 2: 简化视图
    simple_cols = ['site_i_id', 'site_j_id', 'site_i_name', 'site_j_name',
                   'mu_i', 'mu_j', 'sigma_i', 'sigma_j', 'k_i', 'k_j',
                   'q_final', 'E_Si', 'E_Sj', 'E_total', 'value']
    df_allocation[simple_cols].to_excel(writer, sheet_name='allocation_simple', index=False)

    # Sheet 3: 参数记录
    params = pd.DataFrame({
        'parameter': ['Q', 'K_ROBUST'],
        'value': [Q, K_ROBUST],
        'description': ['卡车容量', '鲁棒性水平(k*σ)']
    })
    params.to_excel(writer, sheet_name='parameters', index=False)

    # Sheet 4: 汇总统计
    summary = pd.DataFrame({
        'metric': ['total_pairs', 'avg_q_ratio', 'avg_E_total', 'avg_efficiency',
                  'min_E_total', 'max_E_total', 'clipped_lower', 'clipped_upper'],
        'value': [len(df_allocation),
                 df_allocation['q_ratio'].mean(),
                 df_allocation['E_total'].mean(),
                 df_allocation['efficiency'].mean(),
                 df_allocation['E_total'].min(),
                 df_allocation['E_total'].max(),
                 clipped_lower, clipped_upper]
    })
    summary.to_excel(writer, sheet_name='summary', index=False)

print(f"\n结果已保存至: {OUTPUT_FILE}")
print("  - Sheet 'allocation': 完整分配结果")
print("  - Sheet 'allocation_simple': 简化视图")
print("  - Sheet 'parameters': 参数记录")
print("  - Sheet 'summary': 汇总统计")
print("\n" + "=" * 60)
