from string import ascii_uppercase
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
import pandas as pd
import schemdraw
import schemdraw.flow as flow
from mpl_ornaments.titles import set_title_and_subtitle

fig, ax = plt.subplots(nrows=1, ncols=1)

#1
def node_label(df, activity):
    df_slice = df[df['Activity'] == activity]
    activity = df_slice['Activity letter'].to_list()[0]
    duration = df_slice['Duration'].to_list()[0]
    return f'{activity},{duration:d}'

#2
df = pd.read_csv('data/product-design.csv', comment='#')
df['Duration'] = df['End'] - df['Start']
df['Activity letter'] = list(ascii_uppercase)[0:df.shape[0]]

#3
dw = schemdraw.Drawing()
dw.config(font='Times New Roman')

#4
arrow_props = {'headwidth': 0.25, 'headlength': 0.35}

#5
dw += (mktana := flow.State().label(node_label(df, 'Market analysis')))

#6
dw += (mktana_to_ideation := flow.Arrow(**arrow_props).\
       at(mktana.E).right(length=1))

#7
dw += (idea := flow.State().label(node_label(df, 'Ideation')))
dw.push()

#8
dw += (idea_to_refinement := flow.Arrow(**arrow_props).\
       theta(-45).length(1))

#9
dw += (refn := flow.State().label(node_label(df, 'Refinement')))
dw.pop()

#10
dw += (ideation_to_implementation := flow.Arrow(**arrow_props).\
       theta(45).length(1))
dw += (impl := flow.State().label(node_label(df, 'Implementation')))

#11
dw += (impl_to_proto := flow.Arrow(**arrow_props).\
       at(impl.E).right(length=1))
dw += (proto := flow.State().label(node_label(df, 'Prototyping')))

#12
dw += (refn_to_test_and_val := flow.Arrow(**arrow_props).\
       at(refn.E).right(length=1))
dw += (test_and_val := flow.State().\
       label(node_label(df, 'Testing & Validation')))

#13
dw += (test_and_val_to_prod := flow.Arrow(**arrow_props).
       at(test_and_val.NE).theta(45).length(1))
dw += (prod := flow.State().anchor('W').\
       label(node_label(df, 'Production')))

#14
dw += (prod_to_launch := flow.Arrow(**arrow_props).\
       at(prod.E).right(length=1))
dw += (launch := flow.State().label(node_label(df, 'Launch')))
dw += (launch_to_post_launch := flow.Arrow(**arrow_props).\
       at(launch.E).right(length=1))
dw += (post_launch := flow.State().label(node_label(df, 'Post-launch')))

#15
dw += (proto_to_prod := flow.Arrow(**arrow_props).at(proto.SE).to(prod.W))
dw += (refn_to_proto := flow.Arrow(**arrow_props).\
       at(refn.NE).to(proto.SW))
dw += (impl_to_test_and_val := flow.Arrow(**arrow_props).at(impl.SE).\
       to(test_and_val.NW))

#16
dw.draw(ax=ax, show=False)

#17
ax.axis('off')
ax.set_aspect(aspect='equal')

#18
handles, labels = (list(), list())
for _, row in df.iterrows():
    labels.append(row['Activity'])
    handles.append(Line2D(xdata=[0], ydata=[0], linestyle='none',
                          marker=f'${row["Activity letter"]}$'))

#19
ax.legend(handles=handles, labels=labels, ncols=3, title='Activities', 
          loc='center', bbox_to_anchor=(0.5, -0.45))

title = 'Product design PERT diagram'
subtitle = 'Activity-on-node'
set_title_and_subtitle(title=title, subtitle=subtitle, fig=fig, 
                       alignment='left', h_offset=75, v_offset=55)

fig.savefig('charts/pert-aon.png', bbox_inches='tight', dpi=600)