import config
import pandas as pd
import chart_studio
from plotly.subplots import make_subplots
import plotly.graph_objects as go
import chart_studio.plotly as py
from datetime import datetime

from scripts.trends import get_data_filename, get_group_queries
from scripts.trends.predict import arima_predict, mean_confidence_interval
from utils.io import mkdir_if_not_exist

# COVID_START_DATE = "2020-01-11"
COVID_START_DATE = "2020-03-15" # "2020-01-12"
REOPEN_DATE = "2020-06-09"
REOPEN_DATE_MINUS_1 = "2020-06-07"

DATA_START_DATE = "2020-01-01"
DATA_END_DATE = "2020-08-31"
SOCIAL_DISTANCE_ORDER_DATE = "2020-03-16"
PREDICT_FROM_DATE = REOPEN_DATE


def plot_trends(group, country="US", state=None, place=None, predictive_method="ARIMA"):
    """ Plot trends of queries
    @param group: Query group
    @param country: Country code
    @param state: State code
    @param place: Place data in dictionary
    @param predictive_method: ARIMA or GAM
    """
    print(f"* Plotting Google Trends of `{group}` for {country} - {state or 'All'}")
    group_queries = get_group_queries(group, only_root=True)

    n_queries = len(group_queries)
    n_cols = 3
    n_rows = int(n_queries / n_cols) + (1 if n_queries % n_cols else 0)

    # Annotations
    annotations = []

    # Initialize figure with subplots
    subplot_titles = ["%s..." % t[:22] if len(t) >= 22 else t for t in group_queries]
    fig = make_subplots(
        rows=n_rows, cols=n_cols, subplot_titles=subplot_titles,
        shared_yaxes=True,
        print_grid=True
    )

    # Marked Dates
    covid_start_date = COVID_START_DATE
    reopen_date = REOPEN_DATE
    reopen_date_minus_1 = REOPEN_DATE_MINUS_1
    data_start_date = DATA_START_DATE
    data_end_date = DATA_END_DATE

    # Figure variable
    baseline = 0
    value_range = [0, 100]

    # Model params
    model_params = []

    for idx, query in enumerate(group_queries):
        row = int(idx / n_cols) + 1
        col = idx % n_cols + 1
        showlegend = idx == 0

        query_file_path = get_data_filename(group, query, country=country, state=state, full=True)
        df = pd.read_csv(query_file_path, parse_dates=True)
        count = df["date"].count()

        # ARIMA Model
        if query in df.columns:
            print("Query: ", query)
            # get_arima_params(df[query])
            df, model = arima_predict(df, from_date=PREDICT_FROM_DATE, value_col=query)
            params = model.get_params()
            model_params.append([query, str(params["order"])])
            # return False
        
        # No data
        if count == 0:
            continue

        # Process
        stayhome_order_date = place.get("ClosedFrom") if place else SOCIAL_DISTANCE_ORDER_DATE

        df = df[(df["date"] >= data_start_date) & (df["date"] <= data_end_date)]
        df_before = df[(df["date"] <= reopen_date)]
        df_after = df[(df["date"] >= reopen_date_minus_1)]
        df_prediction = df[df["is_predicted"] == 1]

        # Normalize
        if config.TRENDS_APPLY_NORMALIZATION:
            max_value = df[query].max()
            baseline = df_before[query].median()
            df["value"] = df[query].apply(lambda x: (x - baseline) / max_value)
            df_before["value"] = df_before[query].apply(lambda x: (x - baseline) / max_value)
            df_after["value"] = df_after[query].apply(lambda x: (x - baseline) / max_value)
            baseline = 0
            value_range = [-1, 1]
        else:
            max_value = df[query].max()
            baseline = df_before[query].median()
            df["value"] = df[query]
            df_before["value"] = df_before[query]
            df_after["value"] = df_after[query]

        # Compute difference
        query_text = query.split("+")[0].strip() + " + ..." if "+" in query else query
        actual_mean, actual_meanCI95min, actual_meanCI95max = mean_confidence_interval(df_prediction[query])
        predict_mean = df_prediction["prediction"].mean()
        diff = round(100 * (actual_mean - predict_mean) / predict_mean, 1)
        diffCI95min = round(100 * (actual_meanCI95min - predict_mean) / predict_mean, 1)
        diffCI95max = round(100 * (actual_meanCI95max - predict_mean) / predict_mean, 1)
        x_date = list(df['date'])[int(df["date"].count()/2)]
        diff_annot = go.layout.Annotation(
            text=f'<b>{query_text}</b><br><sub><b style="color:{config.COLOR_UPTREND if diff >= 0 else config.COLOR_DOWNTREND}">{diff}%</b>; 95%CI, [{diffCI95min}%, {diffCI95max}%]</sub>',
            showarrow=False, xanchor="center", yanchor="top", 
            x=x_date,
            y=0.0,
            xshift=0,
            yshift=-5,
            xref=f"x{'' if idx == 0 else idx + 1}",
            yref=f"y{'' if idx == 0 else idx + 1}"
        )
        annotations.append(diff_annot)

        # Lockdown period
        max_y = max(df[query].max(), abs(df[query].min()))
        min_y = -max_y
        shape_lockdown = go.layout.Shape(**{"type": "rect","y0":100,"y1": -100,"x0":COVID_START_DATE, 
                    "x1":REOPEN_DATE,"xref":"x1","yref":"y1","layer":"below",
                    "fillcolor":"#eeeeee", "line":dict(width=0), "line_width": 0})
        fig.add_shape(shape_lockdown, row=row, col=col)

        # Horizontal line 
        shape = go.layout.Shape(**{"type": "line","y0":baseline,"y1": baseline,"x0":str(df["date"].values[0]), 
                    "x1":str(df["date"].values[-1]),"xref":"x1","yref":"y1","layer":"below",
                    "line": {"color": "rgb(200, 200, 200)","width": 1.5}})
        fig.add_shape(shape, row=row, col=col)

        # Stay home order
        if stayhome_order_date:
            shape_stayhome_order = go.layout.Shape(**{"type": "line","y0":-0.25,"y1": 0.25,"x0":stayhome_order_date, 
                    "x1":stayhome_order_date,"xref":"x1","yref":"y1",
                    "line": {"color": "blue","width": 1.5, "dash": "dot"}})
            fig.add_shape(shape_stayhome_order, row=row, col=col)

        # Plot
        subplot_before = go.Scatter(x=df_before["date"], y=df_before["value"], 
                            mode="lines", name="Before Lockdown",
                            line=dict(width=1, color=config.LINE_COLOR_BEFORE), 
                            line_shape="linear", showlegend=False) # linear or spline 
        subplot_after = go.Scatter(x=df_after["date"], y=df_after["value"], 
                            mode="lines", name="Actual Queries",
                            line=dict(width=1.5, color=config.LINE_COLOR_AFTER), 
                            line_shape="linear", showlegend=showlegend) # linear or spline 
        subplot_prediction = go.Scatter(x=df_prediction["date"], y=df_prediction["prediction"], 
                            mode="lines", name="Expected Queries",
                            line=dict(width=2, color=config.LINE_COLOR_BEFORE, dash="dot"), 
                            line_shape="linear", showlegend=showlegend) # linear or spline 
        subplot_lockdown_legend = go.Bar(x=[reopen_date,], y=[0,], 
                            name="Early Lockdown Phase", 
                            showlegend=showlegend,
                            marker_color="#eeeeee")
        fig.add_trace(subplot_before, row=row, col=col)
        fig.add_trace(subplot_after, row=row, col=col)
        fig.add_trace(subplot_prediction, row=row, col=col)
        if idx == 0:
            fig.add_trace(subplot_lockdown_legend, row=row, col=col)

        # break

    # Caption
    # caption = go.layout.Annotation(
    #     showarrow=False,
    #     text="",
    #     xanchor="center",
    #     x=0.5,
    #     yanchor="top",
    #     y=0.0,
    #     yshift=0,
    # )

    # Layout
    # location = f"{country}.{state}" if state else country
    # fig_title = f"""Term: {group}. Location: {location}<br>
    # <span style="font-size: 14px;line-height:1">Period: {data_start_date} - {data_end_date}
    # <br>Lockdown Period: {covid_start_date} - {PREDICT_FROM_DATE}</span>"""
    fig_title = ""
    fig.update_layout(title={"text": fig_title, "x":0.5, "xanchor": "center"}, 
                    title_font=dict(size=12),
                    height=50 + n_rows * 175, width=250 * n_cols, coloraxis=dict(colorscale="Bluered_r"), 
                    showlegend=True, plot_bgcolor="rgb(255,255,255)", titlefont={"size": 30},
                    margin={"t": 50},
                    annotations=annotations,
                    legend=dict(
                        orientation="v",
                        yanchor="bottom",
                        y=0,
                        xanchor="right",
                        x=1,
                        bgcolor="white",
                        bordercolor="#333",
                        borderwidth=1
                    )
                )
    fig.update_xaxes(showgrid=False, showticklabels=False, showline=False)
    fig.update_yaxes(showgrid=False, showticklabels=False, showline=True, range=value_range)

    # Store model parameters
    mkdir_if_not_exist(config.TRENDS_OUTPUT_DIR)
    df_params = pd.DataFrame(model_params, columns=["Query", "Order"])
    df_params.to_csv("%s/ARIMA_orders_%s.csv" % (config.TRENDS_OUTPUT_DIR, group), index=False)

    # Create online URL
    url = py.iplot(fig, filename=group, file_id=group)
    print("URL:", url.src)

    if config.TRENDS_EXPORT_FIGURES:
        # Save
        mkdir_if_not_exist(config.TRENDS_FIGURES_DIR)
        fig.write_image("%s/%s_%s_%s.jpg" % (config.TRENDS_FIGURES_DIR, country, state or "All", group))
        # fig.show()
    else:
        # Show
        fig.show()
9f7626262e3c10a32ba747ccd88b4df5cc2d884e
<mask token>

<mask token>


def plot_trends(group, country='US', state=None, place=None,
    predictive_method='ARIMA'):
    """ Plot trends of queries
    @param group: Query group
    @param country: Country code
    @param state: State code
    @param place: Place data in dictionary
    @param predictive_method: ARIMA or GAM
    """
    print(
        f"* Plotting Google Trends of `{group}` for {country} - {state or 'All'}"
        )
    group_queries = get_group_queries(group, only_root=True)
    n_queries = len(group_queries)
    n_cols = 3
    n_rows = int(n_queries / n_cols) + (1 if n_queries % n_cols else 0)
    annotations = []
    subplot_titles = [('%s...' % t[:22] if len(t) >= 22 else t) for t in
        group_queries]
    fig = make_subplots(rows=n_rows, cols=n_cols, subplot_titles=
        subplot_titles, shared_yaxes=True, print_grid=True)
    covid_start_date = COVID_START_DATE
    reopen_date = REOPEN_DATE
    reopen_date_minus_1 = REOPEN_DATE_MINUS_1
    data_start_date = DATA_START_DATE
    data_end_date = DATA_END_DATE
    baseline = 0
    value_range = [0, 100]
    model_params = []
    for idx, query in enumerate(group_queries):
        row = int(idx / n_cols) + 1
        col = idx % n_cols + 1
        showlegend = idx == 0
        query_file_path = get_data_filename(group, query, country=country,
            state=state, full=True)
        df = pd.read_csv(query_file_path, parse_dates=True)
        count = df['date'].count()
        if query in df.columns:
            print('Query: ', query)
            df, model = arima_predict(df, from_date=PREDICT_FROM_DATE,
                value_col=query)
            params = model.get_params()
            model_params.append([query, str(params['order'])])
        if count == 0:
            continue
        stayhome_order_date = place.get('ClosedFrom'
            ) if place else SOCIAL_DISTANCE_ORDER_DATE
        df = df[(df['date'] >= data_start_date) & (df['date'] <= data_end_date)
            ]
        df_before = df[df['date'] <= reopen_date]
        df_after = df[df['date'] >= reopen_date_minus_1]
        df_prediction = df[df['is_predicted'] == 1]
        if config.TRENDS_APPLY_NORMALIZATION:
            max_value = df[query].max()
            baseline = df_before[query].median()
            df['value'] = df[query].apply(lambda x: (x - baseline) / max_value)
            df_before['value'] = df_before[query].apply(lambda x: (x -
                baseline) / max_value)
            df_after['value'] = df_after[query].apply(lambda x: (x -
                baseline) / max_value)
            baseline = 0
            value_range = [-1, 1]
        else:
            max_value = df[query].max()
            baseline = df_before[query].median()
            df['value'] = df[query]
            df_before['value'] = df_before[query]
            df_after['value'] = df_after[query]
        query_text = query.split('+')[0].strip(
            ) + ' + ...' if '+' in query else query
        actual_mean, actual_meanCI95min, actual_meanCI95max = (
            mean_confidence_interval(df_prediction[query]))
        predict_mean = df_prediction['prediction'].mean()
        diff = round(100 * (actual_mean - predict_mean) / predict_mean, 1)
        diffCI95min = round(100 * (actual_meanCI95min - predict_mean) /
            predict_mean, 1)
        diffCI95max = round(100 * (actual_meanCI95max - predict_mean) /
            predict_mean, 1)
        x_date = list(df['date'])[int(df['date'].count() / 2)]
        diff_annot = go.layout.Annotation(text=
            f'<b>{query_text}</b><br><sub><b style="color:{config.COLOR_UPTREND if diff >= 0 else config.COLOR_DOWNTREND}">{diff}%</b>; 95%CI, [{diffCI95min}%, {diffCI95max}%]</sub>'
            , showarrow=False, xanchor='center', yanchor='top', x=x_date, y
            =0.0, xshift=0, yshift=-5, xref=
            f"x{'' if idx == 0 else idx + 1}", yref=
            f"y{'' if idx == 0 else idx + 1}")
        annotations.append(diff_annot)
        max_y = max(df[query].max(), abs(df[query].min()))
        min_y = -max_y
        shape_lockdown = go.layout.Shape(**{'type': 'rect', 'y0': 100, 'y1':
            -100, 'x0': COVID_START_DATE, 'x1': REOPEN_DATE, 'xref': 'x1',
            'yref': 'y1', 'layer': 'below', 'fillcolor': '#eeeeee', 'line':
            dict(width=0), 'line_width': 0})
        fig.add_shape(shape_lockdown, row=row, col=col)
        shape = go.layout.Shape(**{'type': 'line', 'y0': baseline, 'y1':
            baseline, 'x0': str(df['date'].values[0]), 'x1': str(df['date']
            .values[-1]), 'xref': 'x1', 'yref': 'y1', 'layer': 'below',
            'line': {'color': 'rgb(200, 200, 200)', 'width': 1.5}})
        fig.add_shape(shape, row=row, col=col)
        if stayhome_order_date:
            shape_stayhome_order = go.layout.Shape(**{'type': 'line', 'y0':
                -0.25, 'y1': 0.25, 'x0': stayhome_order_date, 'x1':
                stayhome_order_date, 'xref': 'x1', 'yref': 'y1', 'line': {
                'color': 'blue', 'width': 1.5, 'dash': 'dot'}})
            fig.add_shape(shape_stayhome_order, row=row, col=col)
        subplot_before = go.Scatter(x=df_before['date'], y=df_before[
            'value'], mode='lines', name='Before Lockdown', line=dict(width
            =1, color=config.LINE_COLOR_BEFORE), line_shape='linear',
            showlegend=False)
        subplot_after = go.Scatter(x=df_after['date'], y=df_after['value'],
            mode='lines', name='Actual Queries', line=dict(width=1.5, color
            =config.LINE_COLOR_AFTER), line_shape='linear', showlegend=
            showlegend)
        subplot_prediction = go.Scatter(x=df_prediction['date'], y=
            df_prediction['prediction'], mode='lines', name=
            'Expected Queries', line=dict(width=2, color=config.
            LINE_COLOR_BEFORE, dash='dot'), line_shape='linear', showlegend
            =showlegend)
        subplot_lockdown_legend = go.Bar(x=[reopen_date], y=[0], name=
            'Early Lockdown Phase', showlegend=showlegend, marker_color=
            '#eeeeee')
        fig.add_trace(subplot_before, row=row, col=col)
        fig.add_trace(subplot_after, row=row, col=col)
        fig.add_trace(subplot_prediction, row=row, col=col)
        if idx == 0:
            fig.add_trace(subplot_lockdown_legend, row=row, col=col)
    fig_title = ''
    fig.update_layout(title={'text': fig_title, 'x': 0.5, 'xanchor':
        'center'}, title_font=dict(size=12), height=50 + n_rows * 175,
        width=250 * n_cols, coloraxis=dict(colorscale='Bluered_r'),
        showlegend=True, plot_bgcolor='rgb(255,255,255)', titlefont={'size':
        30}, margin={'t': 50}, annotations=annotations, legend=dict(
        orientation='v', yanchor='bottom', y=0, xanchor='right', x=1,
        bgcolor='white', bordercolor='#333', borderwidth=1))
    fig.update_xaxes(showgrid=False, showticklabels=False, showline=False)
    fig.update_yaxes(showgrid=False, showticklabels=False, showline=True,
        range=value_range)
    mkdir_if_not_exist(config.TRENDS_OUTPUT_DIR)
    df_params = pd.DataFrame(model_params, columns=['Query', 'Order'])
    df_params.to_csv('%s/ARIMA_orders_%s.csv' % (config.TRENDS_OUTPUT_DIR,
        group), index=False)
    url = py.iplot(fig, filename=group, file_id=group)
    print('URL:', url.src)
    if config.TRENDS_EXPORT_FIGURES:
        mkdir_if_not_exist(config.TRENDS_FIGURES_DIR)
        fig.write_image('%s/%s_%s_%s.jpg' % (config.TRENDS_FIGURES_DIR,
            country, state or 'All', group))
    else:
        fig.show()

<mask token>
COVID_START_DATE = '2020-03-15'
REOPEN_DATE = '2020-06-09'
REOPEN_DATE_MINUS_1 = '2020-06-07'
DATA_START_DATE = '2020-01-01'
DATA_END_DATE = '2020-08-31'
SOCIAL_DISTANCE_ORDER_DATE = '2020-03-16'
PREDICT_FROM_DATE = REOPEN_DATE


def plot_trends(group, country='US', state=None, place=None,
    predictive_method='ARIMA'):
    """ Plot trends of queries
    @param group: Query group
    @param country: Country code
    @param state: State code
    @param place: Place data in dictionary
    @param predictive_method: ARIMA or GAM
    """
    print(
        f"* Plotting Google Trends of `{group}` for {country} - {state or 'All'}"
        )
    group_queries = get_group_queries(group, only_root=True)
    n_queries = len(group_queries)
    n_cols = 3
    n_rows = int(n_queries / n_cols) + (1 if n_queries % n_cols else 0)
    annotations = []
    subplot_titles = [('%s...' % t[:22] if len(t) >= 22 else t) for t in
        group_queries]
    fig = make_subplots(rows=n_rows, cols=n_cols, subplot_titles=
        subplot_titles, shared_yaxes=True, print_grid=True)
    covid_start_date = COVID_START_DATE
    reopen_date = REOPEN_DATE
    reopen_date_minus_1 = REOPEN_DATE_MINUS_1
    data_start_date = DATA_START_DATE
    data_end_date = DATA_END_DATE
    baseline = 0
    value_range = [0, 100]
    model_params = []
    for idx, query in enumerate(group_queries):
        row = int(idx / n_cols) + 1
        col = idx % n_cols + 1
        showlegend = idx == 0
        query_file_path = get_data_filename(group, query, country=country,
            state=state, full=True)
        df = pd.read_csv(query_file_path, parse_dates=True)
        count = df['date'].count()
        if query in df.columns:
            print('Query: ', query)
            df, model = arima_predict(df, from_date=PREDICT_FROM_DATE,
                value_col=query)
            params = model.get_params()
            model_params.append([query, str(params['order'])])
        if count == 0:
            continue
        stayhome_order_date = place.get('ClosedFrom'
            ) if place else SOCIAL_DISTANCE_ORDER_DATE
        df = df[(df['date'] >= data_start_date) & (df['date'] <= data_end_date)
            ]
        df_before = df[df['date'] <= reopen_date]
        df_after = df[df['date'] >= reopen_date_minus_1]
        df_prediction = df[df['is_predicted'] == 1]
        if config.TRENDS_APPLY_NORMALIZATION:
            max_value = df[query].max()
            baseline = df_before[query].median()
            df['value'] = df[query].apply(lambda x: (x - baseline) / max_value)
            df_before['value'] = df_before[query].apply(lambda x: (x -
                baseline) / max_value)
            df_after['value'] = df_after[query].apply(lambda x: (x -
                baseline) / max_value)
            baseline = 0
            value_range = [-1, 1]
        else:
            max_value = df[query].max()
            baseline = df_before[query].median()
            df['value'] = df[query]
            df_before['value'] = df_before[query]
            df_after['value'] = df_after[query]
        query_text = query.split('+')[0].strip(
            ) + ' + ...' if '+' in query else query
        actual_mean, actual_meanCI95min, actual_meanCI95max = (
            mean_confidence_interval(df_prediction[query]))
        predict_mean = df_prediction['prediction'].mean()
        diff = round(100 * (actual_mean - predict_mean) / predict_mean, 1)
        diffCI95min = round(100 * (actual_meanCI95min - predict_mean) /
            predict_mean, 1)
        diffCI95max = round(100 * (actual_meanCI95max - predict_mean) /
            predict_mean, 1)
        x_date = list(df['date'])[int(df['date'].count() / 2)]
        diff_annot = go.layout.Annotation(text=
            f'<b>{query_text}</b><br><sub><b style="color:{config.COLOR_UPTREND if diff >= 0 else config.COLOR_DOWNTREND}">{diff}%</b>; 95%CI, [{diffCI95min}%, {diffCI95max}%]</sub>'
            , showarrow=False, xanchor='center', yanchor='top', x=x_date, y
            =0.0, xshift=0, yshift=-5, xref=
            f"x{'' if idx == 0 else idx + 1}", yref=
            f"y{'' if idx == 0 else idx + 1}")
        annotations.append(diff_annot)
        max_y = max(df[query].max(), abs(df[query].min()))
        min_y = -max_y
        shape_lockdown = go.layout.Shape(**{'type': 'rect', 'y0': 100, 'y1':
            -100, 'x0': COVID_START_DATE, 'x1': REOPEN_DATE, 'xref': 'x1',
            'yref': 'y1', 'layer': 'below', 'fillcolor': '#eeeeee', 'line':
            dict(width=0), 'line_width': 0})
        fig.add_shape(shape_lockdown, row=row, col=col)
        shape = go.layout.Shape(**{'type': 'line', 'y0': baseline, 'y1':
            baseline, 'x0': str(df['date'].values[0]), 'x1': str(df['date']
            .values[-1]), 'xref': 'x1', 'yref': 'y1', 'layer': 'below',
            'line': {'color': 'rgb(200, 200, 200)', 'width': 1.5}})
        fig.add_shape(shape, row=row, col=col)
        if stayhome_order_date:
            shape_stayhome_order = go.layout.Shape(**{'type': 'line', 'y0':
                -0.25, 'y1': 0.25, 'x0': stayhome_order_date, 'x1':
                stayhome_order_date, 'xref': 'x1', 'yref': 'y1', 'line': {
                'color': 'blue', 'width': 1.5, 'dash': 'dot'}})
            fig.add_shape(shape_stayhome_order, row=row, col=col)
        subplot_before = go.Scatter(x=df_before['date'], y=df_before[
            'value'], mode='lines', name='Before Lockdown', line=dict(width
            =1, color=config.LINE_COLOR_BEFORE), line_shape='linear',
            showlegend=False)
        subplot_after = go.Scatter(x=df_after['date'], y=df_after['value'],
            mode='lines', name='Actual Queries', line=dict(width=1.5, color
            =config.LINE_COLOR_AFTER), line_shape='linear', showlegend=
            showlegend)
        subplot_prediction = go.Scatter(x=df_prediction['date'], y=
            df_prediction['prediction'], mode='lines', name=
            'Expected Queries', line=dict(width=2, color=config.
            LINE_COLOR_BEFORE, dash='dot'), line_shape='linear', showlegend
            =showlegend)
        subplot_lockdown_legend = go.Bar(x=[reopen_date], y=[0], name=
            'Early Lockdown Phase', showlegend=showlegend, marker_color=
            '#eeeeee')
        fig.add_trace(subplot_before, row=row, col=col)
        fig.add_trace(subplot_after, row=row, col=col)
        fig.add_trace(subplot_prediction, row=row, col=col)
        if idx == 0:
            fig.add_trace(subplot_lockdown_legend, row=row, col=col)
    fig_title = ''
    fig.update_layout(title={'text': fig_title, 'x': 0.5, 'xanchor':
        'center'}, title_font=dict(size=12), height=50 + n_rows * 175,
        width=250 * n_cols, coloraxis=dict(colorscale='Bluered_r'),
        showlegend=True, plot_bgcolor='rgb(255,255,255)', titlefont={'size':
        30}, margin={'t': 50}, annotations=annotations, legend=dict(
        orientation='v', yanchor='bottom', y=0, xanchor='right', x=1,
        bgcolor='white', bordercolor='#333', borderwidth=1))
    fig.update_xaxes(showgrid=False, showticklabels=False, showline=False)
    fig.update_yaxes(showgrid=False, showticklabels=False, showline=True,
        range=value_range)
    mkdir_if_not_exist(config.TRENDS_OUTPUT_DIR)
    df_params = pd.DataFrame(model_params, columns=['Query', 'Order'])
    df_params.to_csv('%s/ARIMA_orders_%s.csv' % (config.TRENDS_OUTPUT_DIR,
        group), index=False)
    url = py.iplot(fig, filename=group, file_id=group)
    print('URL:', url.src)
    if config.TRENDS_EXPORT_FIGURES:
        mkdir_if_not_exist(config.TRENDS_FIGURES_DIR)
        fig.write_image('%s/%s_%s_%s.jpg' % (config.TRENDS_FIGURES_DIR,
            country, state or 'All', group))
    else:
        fig.show()

import config
import pandas as pd
import chart_studio
from plotly.subplots import make_subplots
import plotly.graph_objects as go
import chart_studio.plotly as py
from datetime import datetime
from scripts.trends import get_data_filename, get_group_queries
from scripts.trends.predict import arima_predict, mean_confidence_interval
from utils.io import mkdir_if_not_exist
COVID_START_DATE = '2020-03-15'
REOPEN_DATE = '2020-06-09'
REOPEN_DATE_MINUS_1 = '2020-06-07'
DATA_START_DATE = '2020-01-01'
DATA_END_DATE = '2020-08-31'
SOCIAL_DISTANCE_ORDER_DATE = '2020-03-16'
PREDICT_FROM_DATE = REOPEN_DATE


def plot_trends(group, country='US', state=None, place=None,
    predictive_method='ARIMA'):
    """ Plot trends of queries
    @param group: Query group
    @param country: Country code
    @param state: State code
    @param place: Place data in dictionary
    @param predictive_method: ARIMA or GAM
    """
    print(
        f"* Plotting Google Trends of `{group}` for {country} - {state or 'All'}"
        )
    group_queries = get_group_queries(group, only_root=True)
    n_queries = len(group_queries)
    n_cols = 3
    n_rows = int(n_queries / n_cols) + (1 if n_queries % n_cols else 0)
    annotations = []
    subplot_titles = [('%s...' % t[:22] if len(t) >= 22 else t) for t in
        group_queries]
    fig = make_subplots(rows=n_rows, cols=n_cols, subplot_titles=
        subplot_titles, shared_yaxes=True, print_grid=True)
    covid_start_date = COVID_START_DATE
    reopen_date = REOPEN_DATE
    reopen_date_minus_1 = REOPEN_DATE_MINUS_1
    data_start_date = DATA_START_DATE
    data_end_date = DATA_END_DATE
    baseline = 0
    value_range = [0, 100]
    model_params = []
    for idx, query in enumerate(group_queries):
        row = int(idx / n_cols) + 1
        col = idx % n_cols + 1
        showlegend = idx == 0
        query_file_path = get_data_filename(group, query, country=country,
            state=state, full=True)
        df = pd.read_csv(query_file_path, parse_dates=True)
        count = df['date'].count()
        if query in df.columns:
            print('Query: ', query)
            df, model = arima_predict(df, from_date=PREDICT_FROM_DATE,
                value_col=query)
            params = model.get_params()
            model_params.append([query, str(params['order'])])
        if count == 0:
            continue
        stayhome_order_date = place.get('ClosedFrom'
            ) if place else SOCIAL_DISTANCE_ORDER_DATE
        df = df[(df['date'] >= data_start_date) & (df['date'] <= data_end_date)
            ]
        df_before = df[df['date'] <= reopen_date]
        df_after = df[df['date'] >= reopen_date_minus_1]
        df_prediction = df[df['is_predicted'] == 1]
        if config.TRENDS_APPLY_NORMALIZATION:
            max_value = df[query].max()
            baseline = df_before[query].median()
            df['value'] = df[query].apply(lambda x: (x - baseline) / max_value)
            df_before['value'] = df_before[query].apply(lambda x: (x -
                baseline) / max_value)
            df_after['value'] = df_after[query].apply(lambda x: (x -
                baseline) / max_value)
            baseline = 0
            value_range = [-1, 1]
        else:
            max_value = df[query].max()
            baseline = df_before[query].median()
            df['value'] = df[query]
            df_before['value'] = df_before[query]
            df_after['value'] = df_after[query]
        query_text = query.split('+')[0].strip(
            ) + ' + ...' if '+' in query else query
        actual_mean, actual_meanCI95min, actual_meanCI95max = (
            mean_confidence_interval(df_prediction[query]))
        predict_mean = df_prediction['prediction'].mean()
        diff = round(100 * (actual_mean - predict_mean) / predict_mean, 1)
        diffCI95min = round(100 * (actual_meanCI95min - predict_mean) /
            predict_mean, 1)
        diffCI95max = round(100 * (actual_meanCI95max - predict_mean) /
            predict_mean, 1)
        x_date = list(df['date'])[int(df['date'].count() / 2)]
        diff_annot = go.layout.Annotation(text=
            f'<b>{query_text}</b><br><sub><b style="color:{config.COLOR_UPTREND if diff >= 0 else config.COLOR_DOWNTREND}">{diff}%</b>; 95%CI, [{diffCI95min}%, {diffCI95max}%]</sub>'
            , showarrow=False, xanchor='center', yanchor='top', x=x_date, y
            =0.0, xshift=0, yshift=-5, xref=
            f"x{'' if idx == 0 else idx + 1}", yref=
            f"y{'' if idx == 0 else idx + 1}")
        annotations.append(diff_annot)
        max_y = max(df[query].max(), abs(df[query].min()))
        min_y = -max_y
        shape_lockdown = go.layout.Shape(**{'type': 'rect', 'y0': 100, 'y1':
            -100, 'x0': COVID_START_DATE, 'x1': REOPEN_DATE, 'xref': 'x1',
            'yref': 'y1', 'layer': 'below', 'fillcolor': '#eeeeee', 'line':
            dict(width=0), 'line_width': 0})
        fig.add_shape(shape_lockdown, row=row, col=col)
        shape = go.layout.Shape(**{'type': 'line', 'y0': baseline, 'y1':
            baseline, 'x0': str(df['date'].values[0]), 'x1': str(df['date']
            .values[-1]), 'xref': 'x1', 'yref': 'y1', 'layer': 'below',
            'line': {'color': 'rgb(200, 200, 200)', 'width': 1.5}})
        fig.add_shape(shape, row=row, col=col)
        if stayhome_order_date:
            shape_stayhome_order = go.layout.Shape(**{'type': 'line', 'y0':
                -0.25, 'y1': 0.25, 'x0': stayhome_order_date, 'x1':
                stayhome_order_date, 'xref': 'x1', 'yref': 'y1', 'line': {
                'color': 'blue', 'width': 1.5, 'dash': 'dot'}})
            fig.add_shape(shape_stayhome_order, row=row, col=col)
        subplot_before = go.Scatter(x=df_before['date'], y=df_before[
            'value'], mode='lines', name='Before Lockdown', line=dict(width
            =1, color=config.LINE_COLOR_BEFORE), line_shape='linear',
            showlegend=False)
        subplot_after = go.Scatter(x=df_after['date'], y=df_after['value'],
            mode='lines', name='Actual Queries', line=dict(width=1.5, color
            =config.LINE_COLOR_AFTER), line_shape='linear', showlegend=
            showlegend)
        subplot_prediction = go.Scatter(x=df_prediction['date'], y=
            df_prediction['prediction'], mode='lines', name=
            'Expected Queries', line=dict(width=2, color=config.
            LINE_COLOR_BEFORE, dash='dot'), line_shape='linear', showlegend
            =showlegend)
        subplot_lockdown_legend = go.Bar(x=[reopen_date], y=[0], name=
            'Early Lockdown Phase', showlegend=showlegend, marker_color=
            '#eeeeee')
        fig.add_trace(subplot_before, row=row, col=col)
        fig.add_trace(subplot_after, row=row, col=col)
        fig.add_trace(subplot_prediction, row=row, col=col)
        if idx == 0:
            fig.add_trace(subplot_lockdown_legend, row=row, col=col)
    fig_title = ''
    fig.update_layout(title={'text': fig_title, 'x': 0.5, 'xanchor':
        'center'}, title_font=dict(size=12), height=50 + n_rows * 175,
        width=250 * n_cols, coloraxis=dict(colorscale='Bluered_r'),
        showlegend=True, plot_bgcolor='rgb(255,255,255)', titlefont={'size':
        30}, margin={'t': 50}, annotations=annotations, legend=dict(
        orientation='v', yanchor='bottom', y=0, xanchor='right', x=1,
        bgcolor='white', bordercolor='#333', borderwidth=1))
    fig.update_xaxes(showgrid=False, showticklabels=False, showline=False)
    fig.update_yaxes(showgrid=False, showticklabels=False, showline=True,
        range=value_range)
    mkdir_if_not_exist(config.TRENDS_OUTPUT_DIR)
    df_params = pd.DataFrame(model_params, columns=['Query', 'Order'])
    df_params.to_csv('%s/ARIMA_orders_%s.csv' % (config.TRENDS_OUTPUT_DIR,
        group), index=False)
    url = py.iplot(fig, filename=group, file_id=group)
    print('URL:', url.src)
    if config.TRENDS_EXPORT_FIGURES:
        mkdir_if_not_exist(config.TRENDS_FIGURES_DIR)
        fig.write_image('%s/%s_%s_%s.jpg' % (config.TRENDS_FIGURES_DIR,
            country, state or 'All', group))
    else:
        fig.show()

import config
import pandas as pd
import chart_studio
from plotly.subplots import make_subplots
import plotly.graph_objects as go
import chart_studio.plotly as py
from datetime import datetime

from scripts.trends import get_data_filename, get_group_queries
from scripts.trends.predict import arima_predict, mean_confidence_interval
from utils.io import mkdir_if_not_exist

# COVID_START_DATE = "2020-01-11"
COVID_START_DATE = "2020-03-15" # "2020-01-12"
REOPEN_DATE = "2020-06-09"
REOPEN_DATE_MINUS_1 = "2020-06-07"

DATA_START_DATE = "2020-01-01"
DATA_END_DATE = "2020-08-31"
SOCIAL_DISTANCE_ORDER_DATE = "2020-03-16"
PREDICT_FROM_DATE = REOPEN_DATE


def plot_trends(group, country="US", state=None, place=None, predictive_method="ARIMA"):
    """ Plot trends of queries
    @param group: Query group
    @param country: Country code
    @param state: State code
    @param place: Place data in dictionary
    @param predictive_method: ARIMA or GAM
    """
    print(f"* Plotting Google Trends of `{group}` for {country} - {state or 'All'}")
    group_queries = get_group_queries(group, only_root=True)

    n_queries = len(group_queries)
    n_cols = 3
    n_rows = int(n_queries / n_cols) + (1 if n_queries % n_cols else 0)

    # Annotations
    annotations = []

    # Initialize figure with subplots
    subplot_titles = ["%s..." % t[:22] if len(t) >= 22 else t for t in group_queries]
    fig = make_subplots(
        rows=n_rows, cols=n_cols, subplot_titles=subplot_titles,
        shared_yaxes=True,
        print_grid=True
    )

    # Marked Dates
    covid_start_date = COVID_START_DATE
    reopen_date = REOPEN_DATE
    reopen_date_minus_1 = REOPEN_DATE_MINUS_1
    data_start_date = DATA_START_DATE
    data_end_date = DATA_END_DATE

    # Figure variable
    baseline = 0
    value_range = [0, 100]

    # Model params
    model_params = []

    for idx, query in enumerate(group_queries):
        row = int(idx / n_cols) + 1
        col = idx % n_cols + 1
        showlegend = idx == 0

        query_file_path = get_data_filename(group, query, country=country, state=state, full=True)
        df = pd.read_csv(query_file_path, parse_dates=True)
        count = df["date"].count()

        # ARIMA Model
        if query in df.columns:
            print("Query: ", query)
            # get_arima_params(df[query])
            df, model = arima_predict(df, from_date=PREDICT_FROM_DATE, value_col=query)
            params = model.get_params()
            model_params.append([query, str(params["order"])])
            # return False
        
        # No data
        if count == 0:
            continue

        # Process
        stayhome_order_date = place.get("ClosedFrom") if place else SOCIAL_DISTANCE_ORDER_DATE

        df = df[(df["date"] >= data_start_date) & (df["date"] <= data_end_date)]
        df_before = df[(df["date"] <= reopen_date)]
        df_after = df[(df["date"] >= reopen_date_minus_1)]
        df_prediction = df[df["is_predicted"] == 1]

        # Normalize
        if config.TRENDS_APPLY_NORMALIZATION:
            max_value = df[query].max()
            baseline = df_before[query].median()
            df["value"] = df[query].apply(lambda x: (x - baseline) / max_value)
            df_before["value"] = df_before[query].apply(lambda x: (x - baseline) / max_value)
            df_after["value"] = df_after[query].apply(lambda x: (x - baseline) / max_value)
            baseline = 0
            value_range = [-1, 1]
        else:
            max_value = df[query].max()
            baseline = df_before[query].median()
            df["value"] = df[query]
            df_before["value"] = df_before[query]
            df_after["value"] = df_after[query]

        # Compute difference
        query_text = query.split("+")[0].strip() + " + ..." if "+" in query else query
        actual_mean, actual_meanCI95min, actual_meanCI95max = mean_confidence_interval(df_prediction[query])
        predict_mean = df_prediction["prediction"].mean()
        diff = round(100 * (actual_mean - predict_mean) / predict_mean, 1)
        diffCI95min = round(100 * (actual_meanCI95min - predict_mean) / predict_mean, 1)
        diffCI95max = round(100 * (actual_meanCI95max - predict_mean) / predict_mean, 1)
        x_date = list(df['date'])[int(df["date"].count()/2)]
        diff_annot = go.layout.Annotation(
            text=f'<b>{query_text}</b><br><sub><b style="color:{config.COLOR_UPTREND if diff >= 0 else config.COLOR_DOWNTREND}">{diff}%</b>; 95%CI, [{diffCI95min}%, {diffCI95max}%]</sub>',
            showarrow=False, xanchor="center", yanchor="top", 
            x=x_date,
            y=0.0,
            xshift=0,
            yshift=-5,
            xref=f"x{'' if idx == 0 else idx + 1}",
            yref=f"y{'' if idx == 0 else idx + 1}"
        )
        annotations.append(diff_annot)

        # Lockdown period
        max_y = max(df[query].max(), abs(df[query].min()))
        min_y = -max_y
        shape_lockdown = go.layout.Shape(**{"type": "rect","y0":100,"y1": -100,"x0":COVID_START_DATE, 
                    "x1":REOPEN_DATE,"xref":"x1","yref":"y1","layer":"below",
                    "fillcolor":"#eeeeee", "line":dict(width=0), "line_width": 0})
        fig.add_shape(shape_lockdown, row=row, col=col)

        # Horizontal line 
        shape = go.layout.Shape(**{"type": "line","y0":baseline,"y1": baseline,"x0":str(df["date"].values[0]), 
                    "x1":str(df["date"].values[-1]),"xref":"x1","yref":"y1","layer":"below",
                    "line": {"color": "rgb(200, 200, 200)","width": 1.5}})
        fig.add_shape(shape, row=row, col=col)

        # Stay home order
        if stayhome_order_date:
            shape_stayhome_order = go.layout.Shape(**{"type": "line","y0":-0.25,"y1": 0.25,"x0":stayhome_order_date, 
                    "x1":stayhome_order_date,"xref":"x1","yref":"y1",
                    "line": {"color": "blue","width": 1.5, "dash": "dot"}})
            fig.add_shape(shape_stayhome_order, row=row, col=col)

        # Plot
        subplot_before = go.Scatter(x=df_before["date"], y=df_before["value"], 
                            mode="lines", name="Before Lockdown",
                            line=dict(width=1, color=config.LINE_COLOR_BEFORE), 
                            line_shape="linear", showlegend=False) # linear or spline 
        subplot_after = go.Scatter(x=df_after["date"], y=df_after["value"], 
                            mode="lines", name="Actual Queries",
                            line=dict(width=1.5, color=config.LINE_COLOR_AFTER), 
                            line_shape="linear", showlegend=showlegend) # linear or spline 
        subplot_prediction = go.Scatter(x=df_prediction["date"], y=df_prediction["prediction"], 
                            mode="lines", name="Expected Queries",
                            line=dict(width=2, color=config.LINE_COLOR_BEFORE, dash="dot"), 
                            line_shape="linear", showlegend=showlegend) # linear or spline 
        subplot_lockdown_legend = go.Bar(x=[reopen_date,], y=[0,], 
                            name="Early Lockdown Phase", 
                            showlegend=showlegend,
                            marker_color="#eeeeee")
        fig.add_trace(subplot_before, row=row, col=col)
        fig.add_trace(subplot_after, row=row, col=col)
        fig.add_trace(subplot_prediction, row=row, col=col)
        if idx == 0:
            fig.add_trace(subplot_lockdown_legend, row=row, col=col)

        # break

    # Caption
    # caption = go.layout.Annotation(
    #     showarrow=False,
    #     text="",
    #     xanchor="center",
    #     x=0.5,
    #     yanchor="top",
    #     y=0.0,
    #     yshift=0,
    # )

    # Layout
    # location = f"{country}.{state}" if state else country
    # fig_title = f"""Term: {group}. Location: {location}<br>
    # <span style="font-size: 14px;line-height:1">Period: {data_start_date} - {data_end_date}
    # <br>Lockdown Period: {covid_start_date} - {PREDICT_FROM_DATE}</span>"""
    fig_title = ""
    fig.update_layout(title={"text": fig_title, "x":0.5, "xanchor": "center"}, 
                    title_font=dict(size=12),
                    height=50 + n_rows * 175, width=250 * n_cols, coloraxis=dict(colorscale="Bluered_r"), 
                    showlegend=True, plot_bgcolor="rgb(255,255,255)", titlefont={"size": 30},
                    margin={"t": 50},
                    annotations=annotations,
                    legend=dict(
                        orientation="v",
                        yanchor="bottom",
                        y=0,
                        xanchor="right",
                        x=1,
                        bgcolor="white",
                        bordercolor="#333",
                        borderwidth=1
                    )
                )
    fig.update_xaxes(showgrid=False, showticklabels=False, showline=False)
    fig.update_yaxes(showgrid=False, showticklabels=False, showline=True, range=value_range)

    # Store model parameters
    mkdir_if_not_exist(config.TRENDS_OUTPUT_DIR)
    df_params = pd.DataFrame(model_params, columns=["Query", "Order"])
    df_params.to_csv("%s/ARIMA_orders_%s.csv" % (config.TRENDS_OUTPUT_DIR, group), index=False)

    # Create online URL
    url = py.iplot(fig, filename=group, file_id=group)
    print("URL:", url.src)

    if config.TRENDS_EXPORT_FIGURES:
        # Save
        mkdir_if_not_exist(config.TRENDS_FIGURES_DIR)
        fig.write_image("%s/%s_%s_%s.jpg" % (config.TRENDS_FIGURES_DIR, country, state or "All", group))
        # fig.show()
    else:
        # Show
        fig.show()