from github import Github
import os 
import streamlit as st
import datetime
import plotly.figure_factory as ff
import plotly.graph_objects as go
import pandas as pd
import math
import copy

st.set_page_config(layout="wide")
name2repo = [("Training", "bigcode-project/Megatron-LM"),
             ("Dataset", "bigcode-project/bigcode-dataset"),
             ("Evaluation", "bigcode-project/bigcode-evaluation-harness"),
             ("Inference", "bigcode-project/bigcode-inference-benchmark"),
             ("Legal", "bigcode-project/admin"),
             ("Demo", "bigcode-project/bigcode-demo")
            ]

name2num_milestones = dict()
g = Github(os.environ.get('github'))
df = list()
all_status = list()
bad_milestones = list()


for name, repo_name in name2repo:
    repo = g.get_repo(repo_name)
    num_milestones = 0
    for milestone in repo.get_milestones():
        try:
            num_milestones += 1
            desc = dict()
            for line in milestone.description.split('\n'):
                tmp = line.split(":")
                if len(tmp) > 1:
                    key = tmp[0].lower()
                    value = tmp[1].strip()
                    if key == 'status':
                        value = value.lower()
                    desc[key] = value
            task_name = f"""<a href="https://www.github.com/{repo_name}/milestone/{milestone.number}", target="_black">{milestone.title}</a>"""
            if desc['status'] not in all_status:
                all_status.append(desc['status'])
            df.append(dict(Task=task_name, 
                           Start=desc['start date'], 
                           Finish=milestone.due_on.strftime('%Y-%m-%d'), 
                           Resource=desc['status'], 
                           Description=desc['leader']))
        except:
            num_milestones -= 1
            task_name = f"""<a href="https://www.github.com/{repo_name}/milestone/{milestone.number}", target="_black">{milestone.title}</a>"""
            bad_milestones.append(task_name)
    name2num_milestones[name] = num_milestones

copy_df = copy.deepcopy(df)
colors = {'not started': 'rgb(217, 217, 217)',
          'in progress': 'rgb(147, 196, 125)',
          'high priority - on track': 'rgb(234, 153, 153)',
          'high priority - help needed': 'rgb(255, 0, 0)',
          'completed': 'rgb(56, 118, 29)'}

if len(all_status) == 0:
    task_name = "None"

for key in colors.keys():
    if key not in all_status:
        copy_df.append(dict(Task=task_name, 
                   Start='2023-04-09', 
                   Finish='2023-04-09', 
                   Resource=key))

fig = ff.create_gantt(copy_df, colors=colors, 
                      index_col='Resource', 
                      show_colorbar=True, 
                      show_hover_fill=True,
                      group_tasks=True,
                      title="BigCode planning")

fig.update_xaxes(ticks= "outside",
                 ticklabelmode= "period", 
                 tickformat="%b",
                 tickcolor= "black", 
                 ticklen=10, 
                 range=[datetime.datetime(2022, 12, 30),
                        datetime.datetime(2023, 4, 16)],
                 minor=dict(
                     ticklen=4,  
                     dtick=7*24*60*60*1000,  
                     tick0="2023-01-01", 
                     griddash='dot', 
                     gridcolor='white')
                )

fig.update_layout(margin=go.layout.Margin(l=250))
fig.layout.xaxis.rangeselector = None # remove range selector on top

# Add today line
fig.add_vline(x=datetime.datetime.now().strftime('%Y-%m-%d'), line_width=3, line_dash="dash", line_color="black")
fig.add_annotation({
            "x": datetime.datetime.now().strftime('%Y-%m-%d'),
            "y": fig.layout.yaxis['range'][1],
            "yshift": 10,
            "text": "Today",
            "showarrow": False,
        })

# Add The Stack 1.2
fig.add_vline(x='2023-03-05', line_width=3, line_dash="dash", line_color="red")
fig.add_annotation({
            "x": '2023-03-05',
            "y": fig.layout.yaxis['range'][1],
            "yshift": 10,
            "text": "The Stack 1.2",
            "showarrow": False,
        })

# Add PII
fig.add_vline(x='2023-03-19', line_width=3, line_dash="dash", line_color="red")
fig.add_annotation({
            "x": '2023-03-19',
            "y": fig.layout.yaxis['range'][1],
            "yshift": 10,
            "text": "+ PII redacted",
            "showarrow": False,
        })

# Add release line
fig.add_vline(x='2023-04-09', line_width=3, line_dash="dash", line_color="red")
fig.add_annotation({
            "x": '2023-04-09',
            "y": fig.layout.yaxis['range'][1],
            "yshift": 10,
            "text": "Model release",
            "showarrow": False,
        })

# Add point of contacts
fig.add_annotation({
            "x": "2023-01-01",
            "y": fig.layout.yaxis['range'][1],
            "yshift": 10,
            "xanchor": "left",
            "text": "Contact",
            "showarrow": False})
for i, entry in enumerate(df[::-1]):
    fig.add_annotation(x='2023-01-01', y=i,
                text=entry['Description'],
                showarrow=False,
                xanchor="left",
                xref="x")

# Add working group annotations
fig.add_hline(y=-0.5, line_width=1, line_color="grey")
height = -0.5
for name, _ in name2repo[::-1]:
    if name2num_milestones[name] > 0:
        fig.add_annotation(x='2023-04-13', y=height + name2num_milestones[name]/2,
                text=name,
                showarrow=False,
                align="center",
                textangle=-90)
        height += name2num_milestones[name]
        fig.add_hline(y=height, line_width=1, line_color="grey")

st.plotly_chart(fig, use_container_width=True)

if len(bad_milestones):
    with st.expander("Bad Milestones"):
        for bms in bad_milestones:
            st.markdown(bms + "\n\n", unsafe_allow_html=True)

if st.button("Refresh"):
    st.experimental_rerun()