from datetime import datetime, timedelta
import pandas as pd
import streamlit as st
import isoweek
import altair as alt

from DataAPIRequester import DataAPIRequester
import jtbd
from jtbd import job_group_dict, job_short_name_dict, job_description_dict
from InfoGraphic import draw_concentric_circle, draw_arrow


st.set_page_config(page_title="PLG Dashboard",
                   layout="wide",
                   initial_sidebar_state="collapsed")

st.title("PLG Dashboard")
st.markdown(
    '''
    <div style="text-align: left; font-size: 12px; color: gray;">Dashboard built on data since January 1, 2023</div>
    '''
    , unsafe_allow_html=True
)
week_zero = isoweek.Week(2023, 1)

CLASS5_NAME = "Wandering Phase"
CLASS4_NAME = "Onboarding Phase"
CLASS3_NAME = "Dev & Testing Phase"
CLASS2_NAME = "Production Phase"
CLASS1_NAME = "Paid Production Phase"

CLASS5 = 5
CLASS4 = 4
CLASS3 = 3
CLASS2 = 2
CLASS1 = 1

# Global Metrics
st.header("Global Metrics")


@st.cache_data(ttl=86400)
def get_global_metrics(start_week=0):
    requester = DataAPIRequester('global_metrics',
                                 {'start_week': start_week})
    print(requester.private_key)
    rows = requester.get_data()[0]
    return rows['total_signup_users'], rows['total_active_users'], rows['total_engaged_users'], rows['total_paid_users']


@st.cache_data(ttl=86400)
def get_signup_chart_data(week=isoweek.Week.thisweek() - 12):
    requester = DataAPIRequester('signup',
                                 {'week': week})

    return requester.get_data()


@st.cache_data(ttl=86400)
def get_active_chart_data(week=0):
    requester = DataAPIRequester('active',
                                 {'week': week})

    return requester.get_data()


@st.cache_data(ttl=86400)
def get_engaged_chart_data(week=0):
    requester = DataAPIRequester('engaged',
                                 {'week': week})

    return requester.get_data()


@st.cache_data(ttl=86400)
def get_paid_chart_data(week=0):
    requester = DataAPIRequester('paid',
                                 {'week': week})

    return requester.get_data()


signup, active, engaged, paid = st.columns(4)
total_signup_user_count, total_active_user_count, total_engaged_user_count, total_paid_user_count = get_global_metrics()
signup.metric("Signup Users", total_signup_user_count, help='A user who signs up TiDB Cloud')
active.metric("Active Users", total_active_user_count, help='A user who accesses serverless cluster at least once a '
                                                            'week')
engaged.metric("Engaged Users", total_engaged_user_count, help='A user whose RU metering is > 1M in a week on '
                                                               'serverless cluster')
paid.metric("Paid Users", total_paid_user_count, help='A user who set spend limit beyond free quota on serverless '
                                                      'cluster')

signup_chart, active_chart, engaged_chart, paid_chart = st.columns(4)

signup_data = get_signup_chart_data()
signup_df = pd.DataFrame(signup_data)
signup_df['signup_week'] = signup_df['signup_week'].astype(int)
signup_df = signup_df.sort_values(by=['signup_week'])
signup_df['weekn'] = signup_df['signup_week']
signup_df['user count'] = pd.to_numeric(signup_df['user_count'])
signup_df['date'] = signup_df['end_date_of_signup_week']

with signup_chart:
    st.subheader("Signup Users")
    st.markdown("##### ***Weekly New Signup User Count(#)***")
    _, _, selectbox_col = st.columns(3)
    with selectbox_col:
        signup_option = st.selectbox(
            'Week After Signup',
            (),
            key='signup',
            label_visibility='hidden',
            disabled=True
        )

    chart = alt.Chart(signup_df.tail(12)).mark_line().encode(
        x='date:T',
        y='user count'
    ).interactive()
    signup_chart.altair_chart(chart, use_container_width=True)

with active_chart:
    st.subheader("Active Users")
    st.markdown("##### ***Weekly New User Activation Rate(%)***")
    _, _, selectbox_col = st.columns(3)
    with selectbox_col:
        active_option = st.selectbox(
            'Week After Signup',
            (0, 1, 2),
            key='active'
        )

    data = get_active_chart_data(active_option)
    df = pd.DataFrame(data)
    df['weekn'] = df['weekn'].astype(int)
    df['date'] = df['end_date_of_signup_week']
    df['user count'] = pd.to_numeric(df['data_point'])
    df = df.merge(signup_df[['weekn', 'user count', 'date']], on='weekn', how='outer')
    df = df.sort_values(by=['weekn'])
    df['user count_x'] = df['user count_x'].fillna(0)
    df['conversion rate'] = df['user count_x'] / df['user count_y']
    df['date'] = df['date_y']

    chart = alt.Chart(df.tail(12)).mark_line().encode(
        x='date:T',
        y='conversion rate'
    ).interactive()
    active_chart.altair_chart(chart, use_container_width=True)

with engaged_chart:
    st.subheader("Engaged Users")
    st.markdown("##### ***Weekly New User Engagement Rate(%)***")
    _, _, selectbox_col = st.columns(3)
    with selectbox_col:
        engaged_option = st.selectbox(
            'Week After Signup',
            (0, 1, 2),
            key='engaged'
        )
    data = get_engaged_chart_data(engaged_option)
    df = pd.DataFrame(data)
    df['weekn'] = df['weekn'].astype(int)
    df['date'] = df['end_date_of_signup_week']
    df['user count'] = pd.to_numeric(df['data_point'])
    df = df.merge(signup_df[['weekn', 'user count', 'date']], on='weekn', how='outer')
    df = df.sort_values(by=['weekn'])
    df['user count_x'] = df['user count_x'].fillna(0)
    df['conversion rate'] = df['user count_x'] / df['user count_y']
    df['date'] = df['date_y']
    chart = alt.Chart(df.tail(12)).mark_line().encode(
        x='date:T',
        y='conversion rate'
    ).interactive()
    engaged_chart.altair_chart(chart, use_container_width=True)

with paid_chart:
    st.subheader("Paid Users")
    st.markdown("##### ***Weekly New Paid User Count(#)***")
    _, _, selectbox_col = st.columns(3)
    with selectbox_col:
        paid_option = st.selectbox(
            'Week After Signup',
            (),
            key='paid',
            label_visibility='hidden',
            disabled=True
        )
    data = get_paid_chart_data(paid_option)
    df = pd.DataFrame(data)
    df['weekn'] = df['weekn'].astype(int)
    df = df.sort_values(by=['weekn'])
    df['date'] = df['end_date_of_signup_week']
    df['user count'] = pd.to_numeric(df['data_point'])
    df = df.merge(signup_df, on='weekn', how='outer')
    df['user count'] = df['user count_x'].fillna(0)
    df = df.sort_values(by=['signup_week'])
    df['date'] = df['date_y']
    chart = alt.Chart(df.tail(12)).mark_line().encode(
        x='date:T',
        y='user count'
    ).interactive()

    paid_chart.altair_chart(chart, use_container_width=True)

# Granular Metrics
st.header("Serverless User Growth Dashboard")


@st.cache_data(ttl=86400)
def get_user_count_of_annotation():
    requester = DataAPIRequester('user_count_of_annotation', {})

    return requester.get_data()


@st.cache_data(ttl=86400)
def get_annotation_user(annotation: str):
    requester = DataAPIRequester('annotation', {'annotation': annotation})

    return requester.get_data()


@st.cache_data(ttl=86400)
def get_annotation_user_count_with_time_range(annotation: str, start_time: str, end_time: str):
    if annotation == jtbd.Wandering:
        requester = DataAPIRequester('wandering_statistic',
                                     {'start_time': start_time,
                                      'end_time': end_time}
                                     )
    else:
        requester = DataAPIRequester('annotation_statistic',
                                     {'annotation': annotation,
                                      'start_time': start_time,
                                      'end_time': end_time}
                                     )
    return requester.get_data()


@st.cache_data(ttl=86400)
def get_new_user_last_week(annotation: str):
    requester = DataAPIRequester('new_user_last_week', {'annotation': annotation})

    return requester.get_data()


@st.cache_data(ttl=86400)
def get_new_user_last_month(annotation: str):
    requester = DataAPIRequester('new_user_last_month', {'annotation': annotation})

    return requester.get_data()


@st.cache_data(ttl=86400)
def get_converted_user_count(origin_class: str, target_class: str, time1: str, time2: str, time3: str):
    requester = DataAPIRequester('converted_user_count', {'origin_class': origin_class,
                                                          'target_class': target_class,
                                                          'time1': time1,
                                                          'time2': time2,
                                                          'time3': time3})

    return requester.get_data()


def display_info_graphic(job_list: list, class_type: int, class_name: str):
    cols = st.columns(8)
    last_last_month_time = (datetime.now() - timedelta(days=90)).strftime('%Y-%m-%d')
    last_month_time = (datetime.now() - timedelta(days=60)).strftime('%Y-%m-%d')
    start_time = (datetime.now() - timedelta(days=30)).strftime('%Y-%m-%d')
    end_time = datetime.now().strftime('%Y-%m-%d')
    num0_list = []
    num1_list = []
    num2_list = []
    num3_list = []
    num4_list = []
    for i in range(len(job_list)):
        num0 = 0
        num1 = 0
        num2 = 0
        num3 = 0
        num4 = 0
        for sub_job in job_group_dict[job_list[i]]:

            sub_num0 = int(get_annotation_user_count_with_time_range(sub_job, last_last_month_time, last_month_time)[0][
                               'user_count'])
            num0 += sub_num0
            sub_num1 = int(
                get_annotation_user_count_with_time_range(sub_job, last_month_time, start_time)[0]['user_count'])
            num1 += sub_num1
            sub_num2 = int(get_annotation_user_count_with_time_range(sub_job, start_time, end_time)[0]['user_count'])
            num2 += sub_num2
            sub_num3 = int(get_new_user_last_week(sub_job)[0]['user_count'])
            num3 += sub_num3
            sub_num4 = int(get_new_user_last_month(sub_job)[0]['user_count'])
            num4 += sub_num4
        num0_list.append(num0)
        num1_list.append(num1)
        num2_list.append(num2)
        num3_list.append(num3)
        num4_list.append(num4)

    with cols[0]:
        df = pd.DataFrame(
            {
                'user count': [sum(num0_list), sum(num1_list), sum(num2_list)],
                'date': [last_last_month_time, last_month_time, start_time],
                'month': ['M-2', 'M-1', 'M'],
            }
        )

        y_limit = df['user count'].max() * 1.2

        job_chart = alt.Chart(df).encode(
            x=alt.X('month:N', sort=None, axis=alt.Axis(title=None, labelAngle=0)),
            y=alt.Y('user count:Q', axis=None, scale=alt.Scale(domain=[0, y_limit])),
            text='user count:Q'
        ).properties(
            height=240
        )

        bar_chart = job_chart.mark_bar(width=30)

        chart_with_text = bar_chart + job_chart.mark_text(align='center', dy=-5)

        cols[0].altair_chart(chart_with_text, use_container_width=True)

    for i in range(len(job_list)):
        with cols[i + 1]:
            num1 = num1_list[i]
            num2 = num2_list[i]
            num3 = num3_list[i]
            num4 = num4_list[i]
            st.markdown(
                f"""
                    <div style="display: flex; justify-content: space-between;">
                        <div style="text-align: left; font-size: 16px; font-weight: bold;">{job_short_name_dict[job_list[i]]}</div>
                    </div>
                """,
                help=job_description_dict[job_list[i]],
                unsafe_allow_html=True
            )
            if num1 == 0 and num2 == 0:
                st.header("0")
            else:
                st.pyplot(
                    draw_concentric_circle(num1, num2, num3, num4, max_num=max(num1_list), min_num=min(num1_list)))

    with cols[6]:
        st.markdown(
            f""" <div style="display: flex; justify-content: space-between;"> <div style="text-align: center; 
            font-size: 16px; font-weight: bold; font-style: italic; line-height: 2.5; background-color:	LemonChiffon; 
            border-radius: 10px;">{class_name}</div>
                </div>
            """,
            unsafe_allow_html=True
        )
        if class_type != CLASS1 and class_type != CLASS5:
            user_count = get_converted_user_count(str(class_type), str(class_type-1), last_month_time, start_time, end_time)[0]['user_count']
            arrow = draw_arrow(user_count, class_name)
            st.pyplot(arrow)

    # with cols[7]:
    #     st.info('informational message', icon="ℹ️")


@st.cache_data(ttl=86400)
def generate_job_annotations():
    job_list = jtbd.class1_job_list + jtbd.class2_job_list + jtbd.class3_job_list + jtbd.class4_job_list
    for job in job_list:
        for sub_job in job_group_dict[job]:
            api_name = 'generate_' + sub_job.lower().replace(' ', '_')
            requester = DataAPIRequester(api_name)
            print(f'{sub_job}: {requester.get_data()}')


with st.expander("Job Dashboard Value"):
    st.markdown("""
    - 增长看板设计理念：
      - 如何实现用户增长，要求我们理解增长的本质。我们认为本质在于用户越来越多地使用产品功能，从而转化为活跃的产品用户。没有这个过程，用户数就不会增长。因此，增长看板旨在展示用户增长的过程，特别是用户对产品功能利用的变化。
      - 我们对 TiDB Serverless 提供的众多产品功能进行了抽象和建模，最终确定了13个关键任务。这些任务代表用户在使用 DBaaS 产品去解决业务问题时所需要用到的产品功能模块划分。增长看板将从不同维度展示用户对这13个关键任务的利用情况，反映用户增长是否真实发生。
    - 增长看板提供以下信息：
      - 各关键任务上月和本月的用户留存数，反映本月各任务的用户增长或流失情况
      - 各关键任务本月新增使用用户数，反映本月各任务的新用户增长是否符合预期
      - 各关键任务近一周新增使用用户数，反映近期各任务的新用户增长是否符合预期
      - 基于关键任务构建的用户增长转化漏斗，提供过去三个月各层用户的留存数，反映漏斗各层的用户增长情况
      - 漏斗各层本月的转化用户数，反映漏斗转化情况
""")

generate_job_annotations()
display_info_graphic(jtbd.class5_job_list, CLASS5, CLASS5_NAME)
display_info_graphic(jtbd.class4_job_list, CLASS4, CLASS4_NAME)
display_info_graphic(jtbd.class3_job_list, CLASS3, CLASS3_NAME)
display_info_graphic(jtbd.class2_job_list, CLASS2, CLASS2_NAME)
display_info_graphic(jtbd.class1_job_list, CLASS1, CLASS1_NAME)