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, figsize=[12, 4])

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]]

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

#1
def arc(df, activity, start=None, end=None, length=3.0, theta=0, 
        hlength=0.5, hwidth=0.5, dummy=False):
    if dummy:
        arc_label = ''
        linestyle = '--'
    else:
        df_slice = df[df['Activity'] == activity]
        activity = df_slice['Activity letter'].to_list()[0]
        duration = df_slice['Duration'].to_list()[0]
        arc_label = f'{activity},{duration:d}'
        linestyle = '-'
    retval = flow.Arrow(headwidth=hwidth, headlength=hlength).\
        label(arc_label).linestyle(linestyle)
    if start:
        retval.at(start)
    if end:
        retval.to(end)
    else:
        retval.length(length).theta(theta)
    return retval

#2
def node(event_name, radius=1.5):
    global df_events
    event_id = df_events.shape[0]
    new_record = {'Event_id':[event_id], 'Event_name':[event_name]}
    df_events = pd.concat([df_events, pd.DataFrame.from_dict(new_record)])
    return flow.State(r=radius).label(f'{event_id:d}')

#3
df_events = pd.DataFrame(columns=['Event_id', 'Event_name'])

#4
dw += (start := node('Start'))
dw += (start_to_mktana := arc(df=df, activity='Market analysis'))
dw += (idea_done := node('Market report available'))
dw += (idea := arc(df=df, activity='Ideation'))
dw += (idea_done := node('First concept'))

#5
dw.push()

#6
dw += (refn := arc(df=df, activity='Refinement', theta=-45))
dw += (refn_done := node('Concept refined'))

#7
dw.pop()
dw += (impl := arc(df, 'Implementation').theta(45))
dw += (impl_done := node('Implementation completed'))

#8
dw += (fict_1 := arc(df, activity=None, start=impl_done.SE, dummy=True).\
       theta(-45))
dw += (design_finalised := node('Design finalised').anchor('W'))
dw += (fict_2 := arc(df, activity=None, start=refn_done.NE, 
                     end=design_finalised.W, dummy=True))

#9
dw += (proto := arc(df, 'Prototyping', start=design_finalised.E).\
       theta(45))
dw += (proto_done := node('Prototyping completed'))
dw += (proto := arc(df, 'Testing & Validation', 
                    start=design_finalised.E).theta(-45))
dw += (tested_and_validated := node('Tested and validated'))
dw += (fict_3 := arc(df, activity=None, start=proto_done.SE, dummy=True).\
       theta(-45))
dw += (fict_4 := arc(df, activity=None, start=tested_and_validated.NE,
                     dummy=True, theta=45))

#10
dw += (ready_for_production := node('Ready for production').anchor('W'))
dw += (prod := arc(df, 'Production', start=ready_for_production.E))
dw += (prod_started := node('Production started'))
dw += (launch := arc(df, 'Launch'))
dw += (prod_started := node('Product launched'))
dw += (post_launch := arc(df, 'Post-launch'))
dw += (prod_started := node('End'))

dw.draw(ax=ax, show=False)
ax.axis('off')
ax.set_aspect(aspect='equal')

#11
legend1_handles, legend1_labels = list(), list()
for _, row in df.iterrows():
    legend1_handles.append(Line2D([0], [0], 
                                  marker=f'${row["Activity letter"]}$', 
                                  linestyle='none'))
    legend1_labels.append(row['Activity'])
legend1 = ax.legend(legend1_handles, legend1_labels, ncols=3, 
                     title='Activities', loc='center', 
                     bbox_to_anchor=(0.5, -0.45))

#12
legend2_handles, legend2_labels = list(), list()
for _, row in df_events.iterrows():
    legend2_handles.append(Line2D([0], [0], marker=f'${row["Event_id"]}$',
                                  linestyle='none'))
    legend2_labels.append(row['Event_name'])
leg_evt = ax.legend(legend2_handles, legend2_labels, ncols=3, 
                    title='Events', loc='center', 
                    bbox_to_anchor=(0.5, -1.2))
ax.add_artist(legend1)

title = 'Product design PERT diagram'
subtitle = 'Activity-on-arc'
set_title_and_subtitle(title=title, subtitle=subtitle, fig=fig, 
                       alignment='left', h_offset=140, v_offset=20)

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