from pathlib import Path
import csv,json,itertools
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
P=Path(__file__).resolve().parent

def schedule(n,durations):
    ends=[0.,0.,0.]; rows=[]
    for j in range(n):
        previous=0.
        for s,d in enumerate(durations):
            start=max(previous,ends[s]); finish=start+d
            rows.append(dict(job=j+1,stage=s+1,duration_min=d,start_min=start,finish_min=finish))
            ends[s]=finish;previous=finish
    return rows,ends[-1] if n else 0.

def mix(ea,eb):
    a=30*(1+ea);b=70*(1+eb);total=a+b
    return dict(error_A_fraction=ea,error_B_fraction=eb,mass_A_g=a,mass_B_g=b,total_g=total,fraction_A=a/total,composition_error_pp=100*(a/total-.3))

def save(name,rows):
    with (P/name).open('w',newline='',encoding='utf-8') as f:
        w=csv.DictWriter(f,fieldnames=list(rows[0]));w.writeheader();w.writerows(rows)

summary=[]
for n in range(21):
    _,t=schedule(n,[2,3,1]);summary.append(dict(jobs=n,makespan_min=t,serial_min=6*n))
ops,_=schedule(5,[2,3,1]);grid=[]
for ds in itertools.product(range(1,6),repeat=3):
    _,t=schedule(5,ds);grid.append(dict(dispense_min=ds[0],mix_min=ds[1],inspect_min=ds[2],jobs=5,makespan_min=t,serial_min=5*sum(ds)))
mixes=[mix(i/100,j/100) for i in range(-10,11) for j in range(-10,11)]
save('job-counts.csv',summary);save('five-job-operations.csv',ops);save('duration-grid.csv',grid);save('mixing.csv',mixes)
checks=0
for n in range(21):
    for ds in itertools.product(range(1,6),repeat=3):
        rr,t=schedule(n,ds)
        assert t==(sum(ds)+(n-1)*max(ds) if n else 0);checks+=1
        for r in rr:
            assert r['finish_min']-r['start_min']==ds[r['stage']-1]
            if r['stage']>1:assert r['start_min']>=rr[(r['job']-1)*3+r['stage']-2]['finish_min']
            if r['job']>1:assert r['start_min']>=rr[(r['job']-2)*3+r['stage']-1]['finish_min']
for r in mixes:
    assert abs(r['fraction_A']+(r['mass_B_g']/r['total_g'])-1)<1e-14
    if r['error_A_fraction']==r['error_B_fraction']:assert abs(r['fraction_A']-.3)<1e-14
assert abs(mix(.1,-.1)['fraction_A']-33/96)<1e-14
assert abs(mix(-.1,.1)['fraction_A']-27/104)<1e-14
for i in range(20):
    for j in range(21):assert mixes[(i+1)*21+j]['fraction_A']>mixes[i*21+j]['fraction_A']
for i in range(21):
    for j in range(20):assert mixes[i*21+j+1]['fraction_A']<mixes[i*21+j]['fraction_A']
plt.rcParams.update({'font.size':10,'axes.spines.top':False,'axes.spines.right':False,'savefig.dpi':180})
fig,ax=plt.subplots(1,2,figsize=(11,4.5),layout='constrained')
colors=['#3769a2','#7957a8','#278c89']
for r in ops:
    ax[0].barh(r['job'],r['duration_min'],left=r['start_min'],color=colors[r['stage']-1],height=.65)
from matplotlib.patches import Patch
ax[0].legend(handles=[Patch(color=colors[k],label=label) for k,label in enumerate(['Dispense','Mix','Inspect'])],fontsize=8)
ax[0].set(xlabel='Time (hypothetical min)',ylabel='Job',yticks=range(1,6),title='A  Five jobs, three exclusive resources');ax[0].invert_yaxis()
ax[1].plot([r['jobs'] for r in summary],[r['serial_min'] for r in summary],label='One complete job at a time',color='#999999')
ax[1].plot([r['jobs'] for r in summary],[r['makespan_min'] for r in summary],label='Overlapping stages',color=colors[1],marker='o',markersize=3)
ax[1].set(xlabel='Number of jobs',ylabel='Time until all jobs finish (min)',title='B  First completion and continuing throughput');ax[1].legend(fontsize=8);ax[1].grid(alpha=.2)
fig.suptitle('Synthetic workflow: 2 min dispensing, 3 min mixing, 1 min inspection')
for ext in ['png','svg','pdf']:fig.savefig(P/f'workflow-schedule.{ext}')
plt.close(fig)
fig,ax=plt.subplots(figsize=(7.5,5.5),layout='constrained')
z=np.array([r['composition_error_pp'] for r in mixes]).reshape(21,21)
m=ax.pcolormesh(np.arange(-10,11),np.arange(-10,11),z,shading='nearest',cmap='RdBu_r',vmin=-4.4,vmax=4.4)
ax.plot([-10,10],[-10,10],color='black',ls='--',lw=1,label='Equal fractional errors: composition unchanged')
ax.set(xlabel='Dose error in B (%)',ylabel='Dose error in A (%)',title='Synthetic 30 g A + 70 g B recipe\nComposition error, not a sensory prediction');ax.legend(loc='upper left',fontsize=8)
fig.colorbar(m,ax=ax,label='Change in A mass fraction (percentage points)')
for ext in ['png','svg','pdf']:fig.savefig(P/f'mixture-errors.{ext}')
plt.close(fig)
validation=dict(all_pass=True,closed_form_schedule_cases=checks,resource_and_precedence_checks=True,equal_error_invariance=True,mass_fraction_closure=True,monotonicity=True,corner_fraction_min=27/104,corner_fraction_max=33/96,csv_rows=602)
(P/'validation.json').write_text(json.dumps(validation,indent=2)+'\n')
(P/'study.json').write_text(json.dumps(dict(title='Coordinating work and measuring mixtures',synthetic=True,schedule=dict(durations_min=[2,3,1],jobs=list(range(21)),duration_grid=[1,2,3,4,5],resources_per_stage=1,buffers='unlimited',failures=False,preemptive=False),mixture=dict(nominal_mass_g=[30,70],fractional_error_min=-.1,fractional_error_max=.1,step=.01),counts=dict(job_counts=21,operations=15,duration_grid=125,mixing=441),limitations=['No cleaning, setup, transport time, finite buffers, variability or failed checks','Inspection duration does not imply inspection acceptance','Mixture fractions are mass accounting only; no material compatibility or sensory model']),indent=2)+'\n')
print(json.dumps(validation))