import plotly.graph_objects as go

def job_stack(jobs, use_unit='cpu', fig_out='', plot_title='', query_bounds=True):
    """Create job stack figure based on a given DataFrame and specified use unit.

    Each job is a rectangle, where the height of the rectangle is based on
    the amount of resources that are being examined. Further, all three job
    states are displayed: queued, running, end time.

    Parameters
    -------
    jobs: DataFrame
        Job DataFrame typically generated by the ccmnt package.
    use_unit: str, optional
        Usage unit to examine. One of: {'cpu', 'cpu-eqv', 'gpu', 'gpu-eqv'}.
        Defaults to 'cpu'.
    fig_out: str, optional
        Writes the generated figure to file as specified.
        If empty, skips writing. Defaults to empty.
    plot_title: str, optional
        Title information passed to the figure object.
    query_bounds: bool, optional
        Draws red lines on the figure to represent where query is valid.
        Defaults to true.
    """

    if use_unit == 'cpu':
        jobs['use_unit'] = jobs['reqcpus']
    elif use_unit == 'cpu-eqv':
        jobs['mem_scale'] = jobs['mem'] / 4000.0
        jobs['use_unit'] = jobs[['mem_scale', 'reqcpus']].max(axis=1)
    elif use_unit == 'gpu':
        jobs['ngpus'] = jobs['r_tres'].str.extract(r'(\d+)(?!.*\d)').astype('float')
        jobs['use_unit'] = jobs['ngpus']
    else:
        print('gpu equiv to be supported soon')
        return

    cumu_sum_units = jobs['use_unit'].cumsum()

    res_count = 0
    x_queue = []
    x_run = []
    x_req = []

    # TODO: Fix this below section into frame calls
    y_cumu = []
    for _, row in jobs.iterrows():
        c_queue = [row['submit'], row['start'], row['start'], row['submit'], row['submit'], None]
        x_queue = x_queue + c_queue

        c_run = [row['start'], row['end'], row['end'], row['start'], row['start'], None]
        x_run = x_run + c_run

        c_req = [row['end'], row['start']+row['timelimit'],
                 row['start']+row['timelimit'], row['end'], row['end'], None]
        x_req = x_req + c_req

        c_cumres = [res_count, res_count, res_count+row['use_unit'],
                    res_count+row['use_unit'], res_count, None]
        y_cumu = y_cumu + c_cumres
        res_count = res_count+row['use_unit']

    fig = go.Figure()
    fig.add_trace(go.Scatter(
        x=x_queue,
        y=y_cumu,
        fill='toself',
        fillcolor='rgba(200,200,200,.5)',
        line_color='rgba(200,200,200,.3)',
        name='queued',
        hoveron='points+fills',
        text=jobs['jobid']
    ))
    fig.add_trace(go.Scatter(
        x=x_run,
        y=y_cumu,
        fill='toself',
        fillcolor='rgba(140,180,140,.9)',
        line_color='rgba(140,180,140,.1)',
        name='running'
    ))
    fig.add_trace(go.Scatter(
        x=x_req,
        y=y_cumu,
        fill='toself',
        fillcolor='rgba(120,120,180,.2)',
        line_color='rgba(120,120,180,.1)',
        name='requested'
    ))
    fig.add_trace(go.Scatter(
        x=jobs['submit'],
        y=cumu_sum_units-jobs['use_unit'],
        mode='markers',
        name='submit',
        marker_color='rgba(100,100,100,.3)',
        hovertext=jobs['jobid']
    ))
    fig.add_trace(go.Scatter(
        x=jobs['start'],
        y=cumu_sum_units-jobs['use_unit'],
        mode='markers',
        name='start',
        marker_color='rgba(20,120,20,.3)',
        hovertext=jobs['jobid']
    ))
    fig.add_trace(go.Scatter(
        x=jobs['end'],
        y=cumu_sum_units-jobs['use_unit'],
        mode='markers',
        name='end',
        marker_color='rgba(120,20,20,.3)',
        hovertext=jobs['jobid']
    ))
    fig.add_trace(go.Scatter(
        x=jobs['start']+jobs['timelimit'],
        y=cumu_sum_units-jobs['use_unit'],
        mode='markers',
        name='end',
        marker_color='rgba(100,100,140,.3)',
        hovertext=jobs['jobid']
    ))

    if query_bounds:
        Warning('Query bounds not yet implemented.')

    fig.update_layout(
        title_text="Job stack: "+plot_title,
        yaxis_title='Cumulative resources requested ('+str(use_unit)+')',
        xaxis_title='Date Time',
        showlegend=True)

    if fig_out != '':
        fig.write_html(fig_out)
