from pathlib import Path
import csv, json, os, sys
os.environ.setdefault('MPLCONFIGDIR', '/tmp/abilene-matplotlib')
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib import font_manager
from matplotlib.patches import Patch
from matplotlib.lines import Line2D
import numpy as np

ROOT = Path(__file__).resolve().parents[1]
ASSETS, DATA = ROOT / 'assets', ROOT / 'data'
font_path = Path('/System/Library/Fonts/Supplemental/Arial Unicode.ttf')
if font_path.exists():
    font_manager.fontManager.addfont(str(font_path))
available_fonts = {font.name for font in font_manager.fontManager.ttflist}
font_families = [name for name in ('Arial Unicode MS', 'Noto Sans CJK SC', 'Microsoft YaHei') if name in available_fonts]
plt.rcParams.update({'font.family': font_families + ['sans-serif'], 'font.size': 12,
    'axes.spines.top': False, 'axes.spines.right': False,
    'axes.titleweight': 'normal', 'axes.unicode_minus': False,
    'svg.fonttype': 'none', 'figure.facecolor': 'white', 'axes.facecolor': 'white'})

def save(fig, name):
    if '--data-only' not in sys.argv:
        fig.savefig(ASSETS / f'{name}.png', dpi=200, facecolor='white')
        fig.savefig(ASSETS / f'{name}.svg', facecolor='white')
    plt.close(fig)

G = 100_000
power_inputs = [
    ('情景 A', 120, 1.10, 1.20),
    ('情景 B', 132, 1.10, 1.30),
    ('Epoch 模型', 132, 1.14, 1.40),
]
power = []
for label, kw, factor, pue in power_inputs:
    rack = G / 72 * kw / 1000
    it = rack * factor
    facility = it * pue
    power.append(dict(scenario=label, gpu_count=G, rack_kw=kw,
        extra_it_factor=factor, pue=pue, rack_mw=rack,
        extra_it_mw=it-rack, facility_overhead_mw=facility-it,
        it_mw=it, facility_mw=facility))
with (DATA / 'power-scenarios.csv').open('w', newline='', encoding='utf-8-sig') as f:
    w=csv.DictWriter(f, fieldnames=list(power[0]));w.writeheader();w.writerows(power)

fig, ax = plt.subplots(figsize=(11.5, 5.7))
fig.subplots_adjust(left=.21, right=.94, top=.71, bottom=.25)
fig.text(.05,.92,'10 万张 GB200：三组设施功率估算',fontsize=21,color='#12324D')
fig.text(.05,.85,'相同 GPU 数量，机柜额定功率、柜外 IT 和 PUE 假设不同',fontsize=12,color='#52616D')
ys=np.arange(3)
left=np.zeros(3)
for col,label,color in [('rack_mw','GPU 机柜（含柜内其他设备）','#185A83'),('extra_it_mw','柜外 IT','#7798AF'),('facility_overhead_mw','设施附加用电','#D4DEE5')]:
    vals=np.array([r[col] for r in power])
    ax.barh(ys, vals, left=left, height=.49, label=label, color=color)
    left+=vals
for y,r in enumerate(power):
    ax.text(r['facility_mw']+4,y,f"{r['facility_mw']:.1f} MW",va='center',fontsize=12,color='#12324D')
ax.set_yticks(ys,[f"{label}\n{kw} kW / IT × {factor:.2f} / PUE {pue:.1f}" for label,kw,factor,pue in power_inputs],fontsize=10)
ax.invert_yaxis();ax.set_xlim(0,345);ax.set_xticks([0,100,200,300]);ax.set_xlabel('设施功率（MW）',labelpad=8)
ax.spines['left'].set_visible(False);ax.spines['bottom'].set_color('#ABB8C1');ax.tick_params(axis='y',length=0)
ax.grid(axis='x',color='#E7ECF0',zorder=0);ax.set_axisbelow(True)
fig.legend(loc='lower left',bbox_to_anchor=(.05,.095),ncol=3,frameon=False,fontsize=10)
fig.text(.05,.035,'容量估算：100,000 ÷ 72 × 机柜 kW × IT 系数 × PUE。',fontsize=10,color='#52616D')
save(fig,'power-scenarios')

# The current article infers active size from serving economics first.
# Training then constrains cumulative tokens and, with recipe assumptions,
# the complete expert-pool size. Read the two reproducible source datasets.
import math
serving=json.loads((DATA/'serving-calibration.json').read_text())
training=json.loads((DATA/'training-constraints.json').read_text())
inputs=training['main_envelope']
rlo,rhi=inputs['total_active_ratio']
taulo,tauhi=inputs['assumed_tokens_per_total_parameter']
cpre=training['compute_budget_ledger'][-1]['value']
active_b=set(float(a) for a in np.linspace(300,400,401))
for ratio in (rlo,rhi):
    for tau in (taulo,tauhi):
        a_t=math.sqrt(cpre/(6e24*ratio*tau))
        if .3 <= a_t <= .4: active_b.add(a_t*1000)
rows=[]
for a in sorted(active_b):
    d=cpre/(6*a*1e9)/1e12
    rmin,rmax=a*rlo/1000,a*rhi/1000
    tmin,tmax=d/tauhi,d/taulo
    lo,hi=max(rmin,tmin),min(rmax,tmax)
    rows.append(dict(active_parameters_b=a, cumulative_pretraining_tokens_t=d,
        architecture_min_total_t=rmin, architecture_max_total_t=rmax,
        recipe_min_total_t=tmin, recipe_max_total_t=tmax,
        joint_min_total_t=lo, joint_max_total_t=hi,
        joint_feasible=lo<=hi+1e-10,
        pretraining_flops=cpre, assumed_r_min=rlo, assumed_r_max=rhi,
        assumed_tau_min=taulo, assumed_tau_max=tauhi))
with (DATA/'joint-parameter-constraints.csv').open('w',newline='',encoding='utf-8-sig') as f:
    w=csv.DictWriter(f,fieldnames=list(rows[0]));w.writeheader();w.writerows(rows)
x=np.array([r['active_parameters_b'] for r in rows])
def series(name): return np.array([r[name] for r in rows])
fig,ax=plt.subplots(figsize=(11.5,7.2))
fig.subplots_adjust(left=.10,right=.95,top=.72,bottom=.24)
fig.text(.05,.94,'活跃参数与总参数的联合条件',fontsize=21,color='#12324D')
fig.text(.05,.885,'主预训练预算 6.2208 × 10²⁶ FLOPs；数据配方固定为 τ = 30',fontsize=12,color='#52616D')
ax.fill_between(x,series('architecture_min_total_t'),series('architecture_max_total_t'),
    facecolor='#DCEAF3',edgecolor='none')
ax.plot(x,series('architecture_min_total_t'),color='#185A83',linewidth=1.4,label='架构比例 r ≈ 26.9–32.7')
ax.plot(x,series('architecture_max_total_t'),color='#185A83',linewidth=1.4)
assert taulo == tauhi, 'This chart depicts one fixed recipe ratio'
ax.plot(x,series('recipe_min_total_t'),color='#52616D',linestyle='--',linewidth=1.4)
central=dict(serving['central_model_scenario'])
central_a,central_t=central['active_parameters_b'],central['total_parameters_t']
ax.scatter([central_a],[central_t],s=65,color='#12324D',edgecolors='white',linewidths=1.0,zorder=6)
ax.annotate('350B 活跃 / 约 10T 总量',xy=(central_a,central_t),xytext=(352,7.7),
    fontsize=12,color='#12324D',arrowprops=dict(arrowstyle='-',color='#12324D',linewidth=1.0))
ax.set_xlim(298,402);ax.set_ylim(7,14)
ax.set_xticks([300,325,350,375,400]);ax.set_yticks([8,10,12,14])
ax.set_xlabel('每 token 活跃参数（B，十亿）',labelpad=10)
ax.set_ylabel('完整模型总参数（T，万亿）',labelpad=10)
ax.spines['left'].set_color('#ABB8C1');ax.spines['bottom'].set_color('#ABB8C1')
ax.grid(axis='y',color='#E7ECF0',linewidth=.7);ax.set_axisbelow(True)
handles=[Patch(facecolor='#DCEAF3',edgecolor='#185A83',label='架构参照 r ≈ 27–33'),
    Line2D([0],[0],color='#52616D',linestyle='--',linewidth=1.4,label='数据配方 τ = 30')]
fig.legend(handles=handles,loc='upper left',bbox_to_anchor=(.07,.83),ncol=2,frameon=False,fontsize=10)
fig.text(.05,.115,'代表情景：350B 活跃 → 296T 累计 token → 9.87T 总参数；对应 r ≈ 28.2。',fontsize=11,color='#52616D')
fig.text(.05,.06,'本文假设：10 万 GPU、120 天、40% 有效比例、60% 主预训练；r 参照 Kimi K3 与 V4 Pro。',fontsize=10,color='#52616D')
save(fig,'joint-parameter-constraints')
central=dict(serving['central_model_scenario'])
central['tokens_per_total_parameter']=central['cumulative_training_tokens_t']/central['total_parameters_t']
summary=dict(as_of='2026-09-10',power_scenarios=power,
    epoch_per_building_gpu=50400,epoch_per_building_racks=700,
    epoch_per_building_it_mw=700*.132*1.14,
    epoch_per_building_facility_mw=700*.132*1.14*1.4,
    central_model_scenario=central,
    preferred_active_parameters_b=[300,400],
    joint_total_parameter_envelope_t=inputs['total_parameters_t'],
    central_scenario_basis='Conditional author scenario: infer active size from matched agent-serving economics using Kimi K3 and GPT-5.5/GPT-5.6 Sol prices; use hardware compute, open-MoE ratios and token recipes to constrain total size. Exact representative: 350B active / 9.8742857T total / 296.228571T processed tokens at tau=30; rounded article headline: 350B / 10T / 300T.',
    serving_calibration_file='serving-calibration.json',
    training_constraint_file='training-constraints.json',
    figure_sources={'power-scenarios':'power-scenarios.csv','joint-parameter-constraints':'joint-parameter-constraints.csv'})
(DATA/'calculations.json').write_text(json.dumps(summary,ensure_ascii=False,indent=2)+'\n')
assert math.isclose(central['cumulative_training_tokens_t'],296.2285714285714)
assert math.isclose(central['total_parameters_t'],central['active_parameters_b']*central['total_active_ratio']/1000)
assert math.isclose(central['total_parameters_t'],training['midpoint_example']['total_parameters_t'])
assert math.isclose(central['tokens_per_total_parameter'],taulo)
assert math.isclose(min(r['joint_min_total_t'] for r in rows if r['joint_feasible']),inputs['total_parameters_t'][0])
assert math.isclose(max(r['joint_max_total_t'] for r in rows if r['joint_feasible']),inputs['total_parameters_t'][1])
assert abs(power[2]['facility_mw']-292.6)<1e-9
print(json.dumps({'chart_points':len(rows),'central':central,'joint_envelope':inputs['total_parameters_t']},ensure_ascii=False))