<jupyter_start><jupyter_text>**Colab to design an LC experiment with Trimmed Match**<jupyter_code>#@title **Getting Started**

#@markdown * Connect to the hosted runtime and run each cell after updating the necessary inputs
#@markdown * Download the file "example_data_for_design.csv" from the folder "example_datasets" in github.
#@markdown * Upload the csv file to your Google Drive and open it with Google Sheets
#@markdown * In the cell below, copy and paste the url of the sheet.<jupyter_output><empty_output><jupyter_text>Data input<jupyter_code>#@title Load the libraries needed for the design     

BAZEL_VERSION = '3.0.0'
!wget https://github.com/bazelbuild/bazel/releases/download/{BAZEL_VERSION}/bazel-{BAZEL_VERSION}-installer-linux-x86_64.sh
!chmod +x bazel-{BAZEL_VERSION}-installer-linux-x86_64.sh
!./bazel-{BAZEL_VERSION}-installer-linux-x86_64.sh
!sudo apt-get install python3-dev python3-setuptools git
!git clone https://github.com/google/trimmed_match
!python3 -m pip install ./trimmed_match
!pip install colorama
!pip install gspread-dataframe

"""Loading the necessary python modules."""
import datetime
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import numpy as np
import pandas as pd
import re
from scipy import stats
import warnings

from IPython.display import display
from IPython.core.interactiveshell import InteractiveShell

import gspread
from colorama import Fore, Style
from gspread_dataframe import set_with_dataframe
from google.colab import auth
from google.colab import data_table
from google.colab import widgets
from oauth2client.client import GoogleCredentials
from google.colab import drive
from trimmed_match.design.common_classes import GeoXType, TimeWindow
from trimmed_match.design.common_classes import GeoAssignment
from trimmed_match.design.trimmed_match_design import TrimmedMatchGeoXDesign
from trimmed_match.design.util import find_days_to_exclude, overlap_percent
from trimmed_match.design.util import check_time_periods, check_input_data
from trimmed_match.design.util import human_readable_number
from trimmed_match.design.util import expand_time_windows
from trimmed_match.design.util import CalculateMinDetectableIroas
from trimmed_match.design.util import format_design_table, create_output_table



warnings.filterwarnings('ignore')
InteractiveShell.ast_node_interactivity = "all"
#@markdown ---
#@markdown ### Enter the trix url for the sheet file containing the Client Sales Data:                               
#@markdown The spreadsheet should contain the mandatory columns:
#@markdown * date: date in the format YYYY-MM-DD
#@markdown * geo: the number which identifies the geo
#@markdown * response: variable on which you want to measure incrementality
#@markdown (e.g. sales, transactions)
#@markdown * cost: variable used as spend proxy (e.g. ad spend)
#@markdown * (optional) other columns can be present in the spreadsheet.

#@markdown ---

## load the trix in input

#@markdown Spreadsheet URL
client_sales_table = "add your url here, which should look like https://docs.google.com/spreadsheets/d/???/edit#gid=???" #@param {type:"string"}
auth.authenticate_user()
gc = gspread.authorize(GoogleCredentials.get_application_default())
def read_trix(url: str):
  wks = gc.open_by_url(url).sheet1
  data = wks.get_all_values()
  headers = data.pop(0)
  return pd.DataFrame(data, columns=headers)
geo_level_time_series = read_trix(client_sales_table)

#@markdown ### [OPTIONAL] Enter the trix url for the sheet file containing the geo pairs:
#@markdown The spreadsheet should contain the mandatory columns:
#@markdown * geo: the number which identifies the geo
#@markdown * pair: the number which identifies the pair for the corresponding
#@markdown geo. Pair numbers should be 1,...,N for a pairing with N pairs.
#@markdown * assignment: a string, where "Filtered" indicates pairs excluded. 
#@markdown Use any other value for the assignment column for pairs that
#@markdown should be included.

#@markdown NOTE: Avoid using the data from the chosen evaluation_period below
#@markdown for generating the pairing (if passed here) as this could lead to
#@markdown overfitting.


pairs_table = "add your url here, which should look like https://docs.google.com/spreadsheets/d/???/edit#gid=???" #@param {type:"string"}

geo_level_time_series = check_input_data(geo_level_time_series)

if pairs_table == "":
  pairing = None
  filtered_geos = []
else:
  pairs = read_trix(pairs_table)
  if not set(["pair", "geo"]).issubset(pairs.columns):
    raise ValueError("The sheet in input must have the columns " +
                     f"['pair', 'geo'], got {pairs.columns}")
  for colname in ["pair", "geo"]:
    pairs[colname] = pd.to_numeric(pairs[colname])

  filtered_geos = pairs.loc[pairs["assignment"] == "Filtered", "geo"].to_list()

  pairs.sort_values(by="pair", inplace=True)
  pairing = pd.DataFrame({"geo1": pairs["geo"].values[::2],
                          "geo2": pairs["geo"].values[1::2],
                          "pair": pairs["pair"].values[::2]})

## set parameters in other cells of the colab based on the loaded data
number_of_weeks_test = 4 # length of the experiment in weeks
number_of_days_test = number_of_weeks_test * 7<jupyter_output><empty_output><jupyter_text>Select the parameters for the design<jupyter_code>#@title Select the parameters for the design of the experiment  

use_cross_validation = True

#@markdown Specification of the GeoXType
geox_type = 'HOLD_BACK' #@param['HOLD_BACK', 'HEAVY_UP', 'GO_DARK'] {type:'string'}
geox_type = GeoXType[geox_type]

#@markdown Minimum detectable iROAS
minimum_detectable_iROAS =  3 #@param{type: "number"}

#@markdown Average value per unit response: 1 if the response is sales/revenue, else the average value (e.g. 80 USD) per transactions/footfall/contracts/etc.
average_order_value =  1#@param{type: "number"}

#@markdown Design framework based on hypothesis testing H0: iROAS = 0 vs. H1: iROAS >= minimal_detectable_iROAS
significance_level = 0.10 #@param {type:"number"}
power_level = 0.80 #@param {type:"number"}
calc_min_detectable_iroas = CalculateMinDetectableIroas(significance_level,
                                                        power_level)

#@markdown Configuration of the test duration and pre-test data for the design
experiment_duration_in_weeks = 4 #@param {type:"integer"}
design_start_date = "\"2020-01-01\"" #@param {type:"date"}
design_end_date = "\"2020-12-29\"" #@param {type:"date"}
eval_start_date = "\"2020-12-02\"" #@param {type:"date"}
coverage_test_start_date = "\"2020-11-04\"" #@param {type:"date"}

#@markdown List the maximum budget for the experiment e.g. 300000
experiment_budget = "300000" #@param{type: "string"}
experiment_budget = float(experiment_budget)

#@markdown List any alternative budget which you would like to test separated
#@markdown by a comma, e.g. 125000, 150000
alternative_budget = "125000" #@param{type: "string"}
additional_budget = [float(re.sub(r"\W+", "", x)) for x in
                     alternative_budget.split(',') if alternative_budget != ""]

## Additional constraints which will be flagged in red if not met in
## the design

# upper bound on the minimal detectable relative lift
minimum_detectable_lift_in_response_metric = 0.1 * 100
# lower bound on the baseline revenue covered by the treatment group
minimum_revenue_covered_by_treatment = 0.05 * 100

#@markdown List the geo_id of the geos you want to exclude separated by
#@markdown a comma e.g. 100,200. Leave empty to select all geos.
geos_exclude = "2,13,14" #@param {type: "string"}
geos_exclude = [] if geos_exclude == "" else [re.sub(r"\W+", "", x) for x in
                                              geos_exclude.split(',')]

#@markdown List the days and time periods that you want to exclude separated by
#@markdown a comma e.g. 2019/10/10, 2010/10/11, 2018/10/20-2018/11/20. The format for time periods
#@markdown is "YYYY/MM/DD - YYYY/MM/DD", where the two dates specify the
#@markdown start and end date for the period. The format for
#@markdown day is "YYYY/MM/DD". Leave empty to use all days/weeks.
day_week_exclude = "2020/03/02-2020/03/22, 2020/08/10, 2020/08/14" #@param {type: "string"}
day_week_exclude = [] if day_week_exclude == "" else [
    re.sub(r"\s+", "", x) for x in day_week_exclude.split(",")
]


## convert input dates to datetimes (needed due to interactive parameters)
design_start_date = pd.to_datetime(design_start_date.replace("\"",""))
design_end_date = pd.to_datetime(design_end_date.replace("\"",""))
eval_start_date = pd.to_datetime(eval_start_date.replace("\"",""))
coverage_test_start_date = pd.to_datetime(coverage_test_start_date.replace("\"",
                                                                           ""))

number_of_days_test = experiment_duration_in_weeks * 7
eval_end_date = eval_start_date + datetime.timedelta(days=number_of_days_test-1)
coverage_test_end_date = coverage_test_start_date + datetime.timedelta(
    days=number_of_days_test - 1)

design_start_date = min(design_start_date, coverage_test_start_date,
                        eval_start_date)
design_end_date = max(design_end_date, coverage_test_end_date, eval_end_date)

## Find all the days we should exclude from the analysis from the input
periods_to_exclude = find_days_to_exclude(day_week_exclude)
days_exclude = expand_time_windows(periods_to_exclude)

## remove the excluded days from the rest of the analysis
geo_time_series = geo_level_time_series.copy()
geo_time_series = geo_time_series[~geo_time_series["date"].isin(days_exclude)]


## check that the user doesn't attempt by mistake to remove
## days/weeks in the evaluation or AA test periods.
days_in_eval = [
    x for x in geo_level_time_series["date"].drop_duplicates()
    if x in pd.Interval(eval_start_date, eval_end_date, closed="both")
]

days_in_coverage_test = [
    x for x in geo_level_time_series["date"].drop_duplicates()
    if x in pd.Interval(coverage_test_start_date, coverage_test_end_date,
                        closed="both")]

percentage_overlap_eval = overlap_percent(days_exclude, days_in_eval)
if percentage_overlap_eval > 0:
  raise ValueError((f'{Fore.RED}WARNING: {percentage_overlap_eval:.2} % of  ' +
                    f'the evaluation time period overlaps with days/weeks ' +
                    f'excluded in input. Please change eval_start_date.' +
                    f'\n{Style.RESET_ALL}'))

percentage_overlap_coverage_test = overlap_percent(days_exclude,
                                                   days_in_coverage_test)
if percentage_overlap_coverage_test > 0:
  raise ValueError(
      f'{Fore.RED}WARNING: {percentage_overlap_coverage_test:.2} % of '
      f'the aa test  time period overlaps with days/weeks ' +
      f'excluded in input. Please change coverage_test_start_date.' +
      f'\n{Style.RESET_ALL}')


## check that the evaluation and AA test periods do not
## overlap (if the user has changed them)
percentage_overlap_eval_coverage_test = overlap_percent(days_in_eval,
                                                        days_in_coverage_test)
if percentage_overlap_eval_coverage_test > 0:
  raise ValueError(f'{Fore.RED}WARNING: part of the evaluation time period ' +
                   f'overlaps with the coverage test period. Please change ' +
                   f'eval_start_date.\n{Style.RESET_ALL}')

try:
  pass_checks = check_time_periods(geox_data=geo_level_time_series,
                       start_date_eval=eval_start_date,
                       start_date_aa_test=coverage_test_start_date,
                       experiment_duration_weeks=experiment_duration_in_weeks,
                       frequency="infer")
except Exception as e:
  print(f'{Fore.RED} ERROR: ' + str(e) + f'\n{Style.RESET_ALL}')
  error_raised = e
  pass_checks = False

## check that the number of geos is even
geos_exclude = [int(x) for x in geos_exclude]
all_geos = set(geo_level_time_series["geo"].to_list())
non_existing_geos = set(geos_exclude) - set(all_geos)
if non_existing_geos:
  geos_exclude = [x for x in geos_exclude if x not in non_existing_geos]
  print(f'{Fore.RED}WARNING: Attempting to exclude the geos ' +
        f'{non_existing_geos} which do not exist in ' +
        f'the input trix.\n{Style.RESET_ALL}')
num_geos = len(all_geos - set(geos_exclude))
if num_geos % 2 != 0:
  geo_level_data = geo_level_time_series.groupby(
      "geo", as_index=False)["response"].sum()
  largest_geo = geo_level_data.loc[geo_level_data["response"].idxmax()]
  print(f'\nSince the number of geos is odd, we have removed the following' +
  f' geo (the one with largest response):')
  largest_geo
  geos_exclude.append(largest_geo["geo"])
#@title Summary of the possible designs 

if min([percentage_overlap_eval, percentage_overlap_coverage_test,
        percentage_overlap_eval_coverage_test]) > 0:
  raise ValueError(f'{Fore.RED}Either the evaluation time period or the AA ' +
                   f'test period overlaps with days/weeks excluded in input, ' +
                   f'or these two periods overlap. Please change them ' +
                   f'accordingly.\n{Style.RESET_ALL}')

if not pass_checks:
  raise ValueError(f'{Fore.RED} There is an error with the evaluation or' +
  f' aa test period in the previous colab cell.\nPlease correct that error' +
  f' and then rerun this cell.\n' +
  f'Previous error:\n' + str(error_raised) + '{Style.RESET_ALL}')

## set the number of excluded pairs to be tested
num_geos = len(set(geo_level_time_series["geo"].to_list()) - set(geos_exclude))
max_num_geopairs_trim = int(np.floor(num_geos/2 - 10))
# if any geo pairs was specified as filtered, we do not attempt further
# filtering
if pairs_table:
  candidate_num_pairs_excluded = [int(len(filtered_geos) / 2)]
else:
  candidate_num_pairs_excluded = range(0, max_num_geopairs_trim+1, 1)

## set the number of simulation used to compute the RMSE
number_of_simulations = 200

### Evaluate the RMSE of the possible designs ###

## remove the AA test period to make sure it's not used in the evaluation
## of the RMSE or in the training.
data_without_coverage_test_period = geo_time_series[
    (geo_time_series["date"] < coverage_test_start_date) |
    (geo_time_series["date"] > coverage_test_end_date)]

data_for_design = data_without_coverage_test_period.copy()
data_for_design = data_for_design[~data_for_design["geo"].isin(geos_exclude)]

time_window_for_design = TimeWindow(design_start_date, design_end_date)
time_window_for_eval = TimeWindow(eval_start_date, eval_end_date)
## initialize the TrimmedMatchGeoxDesign
pretest = TrimmedMatchGeoXDesign(
    geox_type=geox_type,
    pretest_data=data_for_design,
    response="response",
    spend_proxy="cost",
    matching_metrics={"response": 1.0, "cost": 0.01},
    time_window_for_design=time_window_for_design,
    time_window_for_eval=time_window_for_eval,
    pairs=pairing)

## run a first design with budget equal to the max. budget
preliminary_results, prel_results_detailed = pretest.report_candidate_designs(
    budget_list=[experiment_budget],
    iroas_list=[0],
    use_cross_validation=use_cross_validation,
    num_pairs_filtered_list=candidate_num_pairs_excluded,
    num_simulations=number_of_simulations)

## calculate the minimum detectable iROAS for a design with max. budget
chosen_design = preliminary_results.loc[
    preliminary_results["rmse_cost_adjusted"].idxmin()].squeeze()
lowest_detectable_iroas = calc_min_detectable_iroas.at(chosen_design["rmse"])

## two cases are possible:
##   1) if the minimum detectable iROAS with max. budget is greater than
##      the minimum_detectable_iROAS in input, then calculate the budget needed
##      to reach the minimum_detectable_iROAS in input and run a design with
##      such budget. This is the code in the if clause below;
##   2) if the minimum detectable iROAS with max. budget is smaller than
##      the minimum_detectable_iROAS in input, then run designs with
##      budgets equal to the max. budget plus/minus 20%.
##      This is the code in the else clause below;
minimum_iroas_aov = minimum_detectable_iROAS / average_order_value
if lowest_detectable_iroas > minimum_iroas_aov:
  budget_to_reach_min_det_iroas = (experiment_budget * lowest_detectable_iroas
                                   / minimum_iroas_aov)
  additional_results, results_detailed = pretest.report_candidate_designs(
    budget_list=[budget_to_reach_min_det_iroas] + additional_budget,
    iroas_list=[0],
    use_cross_validation=use_cross_validation,
    num_pairs_filtered_list=candidate_num_pairs_excluded,
    num_simulations=number_of_simulations)

  results = pd.concat([preliminary_results, additional_results], sort=False)

else:
  optimal_budget = (experiment_budget * lowest_detectable_iroas /
                    minimum_iroas_aov)
  lower_budget = optimal_budget *  0.8
  upper_budget = optimal_budget * 1.2
  list_of_budgets = [lower_budget, optimal_budget, upper_budget
                    ] + additional_budget
  results, results_detailed = pretest.report_candidate_designs(
      budget_list=list_of_budgets,
      iroas_list=[0],
      use_cross_validation=use_cross_validation,
      num_pairs_filtered_list=candidate_num_pairs_excluded,
      num_simulations=number_of_simulations)


# these are numerical identifier used in the table in input to identify the two
# groups
group_treatment = 1
group_control = 2
group_filtered = GeoAssignment.EXCLUDED

filtered_pairs = preliminary_results.loc[
    preliminary_results["rmse_cost_adjusted"].idxmin(), "num_pairs_filtered"]
axes_paired = pretest.output_chosen_design(filtered_pairs, base_seed=0,
                                           confidence=1-significance_level,
                                           group_control=group_control,
                                           group_treatment=group_treatment)
plt.close()

## assign geos to treatment and control groups
geopairs = pretest.geo_level_eval_data
geopairs.sort_values(by=["pair"], inplace=True)

geo_treatment = geopairs[geopairs["assignment"]==group_treatment]
geo_control = geopairs[geopairs["assignment"]==group_control]
treatment_geo = geo_treatment["geo"].to_list()
control_geo = geo_control["geo"].to_list()

budgets_for_design = results["budget"].drop_duplicates().to_list()

### AA test to check the coverage probability of the confidence interval ###

## remove the evaluation period to make sure it's not used in the evaluation
## of the RMSE or in the training.
data_without_eval_period = geo_time_series[
    (geo_time_series["date"] < eval_start_date) |
    (geo_time_series["date"] > eval_end_date)]


data_for_coverage_test = data_without_eval_period.copy()
data_for_coverage_test = data_for_coverage_test[
    ~data_for_coverage_test["geo"].isin(geos_exclude)]

time_window_for_coverage_test = TimeWindow(coverage_test_start_date,
                                           coverage_test_end_date)
## initialize the TrimmedMatchGeoxDesign
coverage_test_class = TrimmedMatchGeoXDesign(
    geox_type=geox_type,
    pretest_data=data_for_coverage_test,
    response="response",
    spend_proxy="cost",
    matching_metrics={"response": 1.0, "cost": 0.01},
    time_window_for_design=time_window_for_design,
    time_window_for_eval=time_window_for_coverage_test,
    pairs=pairing)

## calculate the point estimate for each simulation
aa_results, aa_results_detailed = coverage_test_class.report_candidate_designs(
    budget_list=[budgets_for_design[0]],
    iroas_list=[0],
    use_cross_validation=use_cross_validation,
    num_pairs_filtered_list=candidate_num_pairs_excluded,
    num_simulations=number_of_simulations)


## The code below this line only takes care of formatting the output

total_response = geo_level_time_series.loc[
    geo_level_time_series["date"].between(eval_start_date, eval_end_date),
    "response"].sum()
total_spend = geo_level_time_series.loc[
    geo_level_time_series["date"].between(eval_start_date, eval_end_date),
    "cost"].sum()

designs = create_output_table(results=results,
                              total_response=total_response,
                              total_spend=total_spend,
                              geo_treatment=geo_treatment,
                              budgets_for_design=budgets_for_design,
                              average_order_value=average_order_value,
                              num_geos=num_geos,
                              confidence_level=1-significance_level,
                              power_level=power_level)


designs_table = format_design_table(
    designs=designs,
    minimum_detectable_iroas=minimum_detectable_iROAS,
    minimum_lift_in_response_metric=minimum_detectable_lift_in_response_metric,
    minimum_revenue_covered_by_treatment=minimum_revenue_covered_by_treatment)

designs_table


#@title Select the design to be used in the experiment  
#@markdown Select the design using the number as displayed in the table in
#@markdown the cell called "Summary of the possible designs".

selected_design =   1#@param {type:"integer"}

if selected_design not in designs.index:
  raise ValueError(f'the selected design must be one of {designs.index.to_list()}, got {selected_design}')

selected_design = int(selected_design)
final_design = designs[designs.index == selected_design]
selected_budget = final_design["Budget"].values[0]
filtered = final_design["Treatment/control/excluded geos"].values[0]
filtered_pairs = int(filtered.replace(" ", "").split('/')[2]) // 2


## Uncomment the following line to override the automatic choice for
## the number of filtered pairs. For example, using filtered_pairs = 5 will
## filter out the pairs 1, 2, 3, 4, 5.

# filtered_pairs = 5

###


axes_paired = pretest.output_chosen_design(filtered_pairs, base_seed=0,
                                           confidence=1-significance_level,
                                           group_control=group_control,
                                           group_treatment=group_treatment)
plt.close()
#@title Scatterplot and time series comparison of different metrics for treatment vs. control groups  

for ax in axes_paired[1]:
  ylim = ax.get_ylim()
  for period in periods_to_exclude:
    if period.first_day < period.last_day:
      useless=ax.fill_between([period.first_day, period.last_day], ylim[0], ylim[1],
                              facecolor="gray", alpha=0.5)
    else:
      useless=ax.vlines(period.first_day, ylim[0], ylim[1],
                         color="gray", alpha=0.5)

  handles, labels = ax.get_legend_handles_labels()
  patch = mpatches.Patch(color='grey', label='removed by the user')
  handles.append(patch)
  useless=ax.fill_between([coverage_test_start_date, coverage_test_end_date],
                          ylim[0], ylim[1], facecolor="g", alpha=0.5)
  patch = mpatches.Patch(color='g', label='left out for AA test')
  handles.append(patch)
  useless=ax.legend(handles=handles, loc='best')

axes_paired[1,1].figure
#@title Plot each pair of geos for comparison  
g = pretest.plot_pair_by_pair_comparison(group_control=group_control,
                                         group_treatment=group_treatment)<jupyter_output><empty_output><jupyter_text>Summary of the design and save the pretest data, the geopairs, treatment and control stores in a trix.<jupyter_code>#@title Summary and Results  

geopairs = pretest._geo_level_eval_data
geopairs.sort_values(by=["pair"], inplace=True)
included_pairs = geopairs["pair"] > filtered_pairs
treatment_geo = geopairs.loc[(
    included_pairs &
    (geopairs["assignment"] == group_treatment))]["geo"].to_list()
control_geo = geopairs.loc[(
    included_pairs &
    (geopairs["assignment"] == group_control))]["geo"].to_list()

temporary = geo_level_time_series[geo_level_time_series["geo"].isin(
    treatment_geo)]
treatment_time_series = temporary[
    (temporary["date"] >= design_start_date)
    & (temporary["date"] <= design_end_date)].groupby(
        "date", as_index=False)[["response", "cost"]].sum()

temporary = geo_level_time_series[geo_level_time_series["geo"].isin(
    control_geo)]
control_time_series = temporary[
    (temporary["date"] >= design_start_date)
    & (temporary["date"] <= design_end_date)].groupby(
        "date", as_index=False)[["response", "cost"]].sum()

eval_window = ((treatment_time_series['date'] >= eval_start_date) &
                (treatment_time_series['date'] <= eval_end_date))
baseline = treatment_time_series[eval_window]["response"].sum()

result_to_out = results[results["budget"] ==
                        budgets_for_design[selected_design]]

print("Data in input:\n")
print("-  {} Geos \n".format(
    len(geo_level_time_series["geo"].drop_duplicates().index)))

print("Output:\n")
print("The output contains two lists of geos: one for treatment" +
      " and the other for control\n")

human_baseline = human_readable_number(baseline)
cost_baseline = budgets_for_design[selected_design] * 100 / baseline
print("-  {} Geo pairs for the experiment\n".format(len(treatment_geo)))
print("    Baseline store response: ${} for treatment\n".format(human_baseline))
print("    Cost/baseline = ${} / ${} ~ {:.3}%\n".format(selected_budget,
                                                        human_baseline,
                                                        cost_baseline))

summary_rmse = result_to_out.loc[result_to_out["num_pairs_filtered"]==
                                    filtered_pairs, "rmse"].values[0]
summary_minimum_detectable_iroas = calc_min_detectable_iroas.at(summary_rmse)
summary_minimum_detectable_lift = (cost_baseline *
                                   summary_minimum_detectable_iroas)
summary_minimum_detectable_iroas_aov = (
    summary_minimum_detectable_iroas * average_order_value)
print(f'Minimum detectable iROAS = ' +
      f'{summary_minimum_detectable_iroas_aov:.3}')
print(f'Minimum detectable lift in % = ' +
      f'{summary_minimum_detectable_lift:.2f}')

print(f"The design has Power {100 * power_level:.3}+% with Type-I error " +
      f"{100 *(significance_level):.3}% for testing H0: iROAS=0 vs " +
      f"H1: iROAS >= {summary_minimum_detectable_iroas_aov:.3}")
#@title Report stores for treatment and control separately and write to trix 

#@markdown ###Insert the name google sheets in which we will save the data.
#@markdown The trix contains 4 worksheets, named:
#@markdown * "pretest data", containing the geo level time series;
#@markdown * "geopairs", containing the pairs of geos and their assignment.
#@markdown * "treatment geos", contains the list of geos in the treatment;
#@markdown * "control geos", contains the geos in the control groups.
Client_Name = "Client_Name" #@param {type:"string"}
filename_design = Client_Name + "_design.csv" #@param {type:"string"}

geopairs_formatted = geopairs.copy()
geopairs_formatted.loc[geopairs_formatted["pair"] <= filtered_pairs,
                       "assignment"] = group_filtered
geopairs_formatted["assignment"] = geopairs_formatted["assignment"].map({
    group_filtered: "Filtered",
    group_control: "Control",
    group_treatment: "Treatment"
})

geo_level_time_series["period"] = [
   # -3 indicates days excluded
   -3 if x in days_exclude else (
   # 0 indicates days in the evaluation period
   0 if x>=eval_start_date and x<=eval_end_date else (
   # -2 indicates days in the coverage_test period
   -2 if x >= coverage_test_start_date and x <= coverage_test_end_date
   # -1 indicates days in the training period
   else -1))
   for x in geo_level_time_series["date"]
]

tmp = geo_level_time_series[geo_level_time_series["geo"].isin(treatment_geo +
                                                              control_geo)]
design_data = tmp.merge(
    geopairs[["geo", "pair", "assignment"]], on="geo", how="left")

tmp_parameters = {
    "geox_type": str(geox_type).replace("GeoXType.", ""),
    "minimum_detectable_iROAS": minimum_detectable_iROAS,
    "average_order_value": average_order_value,
    "significance_level": significance_level,
    "power_level": power_level,
    "experiment_duration_in_weeks": experiment_duration_in_weeks,
    "design_start_date": design_start_date.strftime("%Y-%m-%d"),
    "design_end_date": design_end_date.strftime("%Y-%m-%d"),
    "eval_start_date": eval_start_date.strftime("%Y-%m-%d"),
    "coverage_test_start_date": coverage_test_start_date.strftime("%Y-%m-%d"),
    "experiment_budget": experiment_budget,
    "alternative_budget": alternative_budget,
    "geo_exclude": ", ".join(str(x) for x in geos_exclude),
    "day_week_exclude": ", ".join(day_week_exclude),
    "selected_design": selected_design,
    "filtered_pairs": filtered_pairs
}

parameters = {"parameter": list(tmp_parameters.keys()),
              "value": list(tmp_parameters.values())}


sh = gc.create(filename_design)
wid = sh.add_worksheet("pretest data", rows=1, cols=1)
set_with_dataframe(wid, design_data)
wid = sh.add_worksheet("geopairs", rows=1, cols=1)
set_with_dataframe(wid, geopairs_formatted)
wid = sh.add_worksheet("treatment geos", rows=1, cols=1)
set_with_dataframe(wid, pd.DataFrame({"geo": treatment_geo}))
wid = sh.add_worksheet("control geos", rows=1, cols=1)
set_with_dataframe(wid, pd.DataFrame({"geo": control_geo}))
wid = sh.add_worksheet("parameters used in the design", rows=1, cols=1)
set_with_dataframe(wid, pd.DataFrame(parameters))
out = sh.del_worksheet(sh.sheet1)<jupyter_output><empty_output><jupyter_text>Appendix:<jupyter_code>#@markdown The following cell is optional and show the graph/table behind the automatic designs presented above
#@title Plot of the RMSE as a function of the # of trimmed pairs  

#@markdown The first graph below shows the RMSE of the iROAS estimator
#@markdown with respect to the number of excluded geo pairs,
#@markdown with the baseline store sales (treatment group + control group)
#@markdown displayed next to each point of (# excluded pairs, RMSE).

#@markdown The second graph below shows the proportion of confidence
#@markdown intervals that cointains zero in an A/A test.

coverage_test_result = []
for num_pairs_filtered in candidate_num_pairs_excluded:
  temp_df = aa_results_detailed[(budgets_for_design[0], 0, num_pairs_filtered)]
  temp_df["contains_zero"] = (temp_df["conf_interval_low"] <= 0) & (
    temp_df["conf_interval_up"] >= 0)
  coverage = temp_df['contains_zero'].mean()
  coverage_test_result.append({"coverage": coverage,
                         "num_pairs_filtered": num_pairs_filtered,
                         "budget": budgets_for_design[0],
                         "ci_level": temp_df["ci_level"][0]})

coverage_test_result = pd.DataFrame(coverage_test_result)

# call the function that creates all the axis
axes_dict = pretest.plot_candidate_design(results=results)
## display the results of the designs with different budgets in different tabs
list_of_str = ["budget = " + human_readable_number(budget)+"$"
               for budget in budgets_for_design]
tb = widgets.TabBar(list_of_str)
for ind in range(len(budgets_for_design)):

  result = results[results["budget"]==budgets_for_design[ind]].reset_index(
      drop=True)
  with tb.output_to(ind):
    ax = axes_dict[(budgets_for_design[ind], 0)].axes[0]
    labels = ax.get_yticks().tolist()
    labels = [str(round(float(x)*average_order_value, 2)) for x in labels]
    useless = ax.set_yticklabels(labels)
    if len(coverage_test_result) > 1:
      display(axes_dict[(budgets_for_design[ind], 0)])

    print_result = result.copy()
    print_result.rename(columns={"num_pairs_filtered": "# of excluded pairs"},
                        inplace=True)
    print_result["rmse"] = print_result["rmse"] * average_order_value
    print_result["rmse_cost_adjusted"] = (print_result["rmse_cost_adjusted"]
                                          * average_order_value)
    data_table.DataTable(print_result[["# of excluded pairs", "rmse",
                 "rmse_cost_adjusted", "experiment_response"]],
                 include_index=False)

if len(coverage_test_result) > 1:
  fig1 = plt.figure(figsize=(15, 7.5))
  ax1 = fig1.add_subplot(1, 1, 1)
  useless=ax1.plot(
      coverage_test_result['num_pairs_filtered'],
      coverage_test_result['coverage'], 'blue',
      label='Coverage')
  useless=ax1.hlines(
      y=coverage_test_result['ci_level'][0],
      xmin=min(coverage_test_result["num_pairs_filtered"]),
      xmax=max(coverage_test_result["num_pairs_filtered"]),
      colors='red',
      linestyles='dashed',
      label='Confidence level (nominal)')
  useless=ax1.set_xlabel('#(Excluded geo pairs)')
  useless=ax1.set_ylabel('Coverage')
  useless=ax1.set_title('A/A test confidence interval coverage')
  useless=ax1.legend()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10193788971756938, 'fraction_numerical': 0.010361902433060276, 'mean_word_length': 4.210543080108297, 'pattern_counts': {'":': 33, '<': 31, '<?xml version=': 0, '>': 39, 'https://': 4, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 40, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '328660', 'license': 'Apache-2.0', 'n_tokens_mistral': 10763, 'n_tokens_neox': 10154, 'n_words': 2455, 'path': 'trimmed_match/notebook/design_colab_for_trimmed_match.ipynb', 'repo_name': 'Praneethp09/Geo_Test', 'size': '32716'}