S3_PATH = "s3://net-mozaws-prod-us-west-2-pipeline-analysis/taarv2/cleaned_data/"
clean_data = sqlContext.read.parquet(S3_PATH).filter("discopane_loaded = true")

clean_data.groupby("branch").agg(F.countDistinct("client_id")).show()

"""
Explanation: Data Prep
Load in cleaned experiment data, generated from this notebook. Filter to clients that loaded the discopane (this is for the control group, since we have already cross referenced the TAAR logs. Clients who did not load the discopane never saw the control, serving as noise in the experiment.
End of explanation
"""


min_date = clean_data.select(F.min('submission_date_s3').alias('min_d')).collect()[0].min_d
max_date = clean_data.select(F.max('submission_date_s3').alias('max_d')).collect()[0].max_d
print min_date, max_date

"""
Explanation: Grab the min and max submission dates for filtering main_summary.
End of explanation
"""


ms = (
    sqlContext.read.option("mergeSchema", True)
     .parquet("s3://telemetry-parquet/main_summary/v4")
     .filter("submission_date_s3 >= '{}'".format(min_date))
     .filter("submission_date_s3 <= '{}'".format(date_plus_x_days(max_date, 7*N_WEEKS)))
     .filter("normalized_channel = 'release'")
     .filter("app_name = 'Firefox'")
)

# a client's enrollment date is determined by their first appearance in the experiment
enrollment_dates = (
    clean_data.groupby("client_id", "branch")
     .agg(F.min('submission_date_s3')
          .alias("enrollment_date"))
)

# join main_summary to exp data
joined = enrollment_dates.join(ms.select("submission_date_s3", "client_id"), on="client_id", how='left')

# verify join contains same number of distinct clients as the experiment data.
# this also initializes our cache
jc = joined.select("client_id").distinct().count()
cc = clean_data.select("client_id").distinct().count()

jc - cc

"""
Explanation: Load in main_summary, filtered to the min date of the experiment, and 42 days beyond its compleition to allow for 6-week Retention Analysis. We then join main_summary with the experiment data.
End of explanation
"""


joined = (
    joined.withColumn("period", get_period("enrollment_date", "submission_date_s3"))
          .filter("enrollment_date <= '{}'".format(max_date))
).distinct().cache()

joined.count()

ret_df  = get_retention(joined)

ret_df.to_csv("taar_v2_retention.csv", index=False)

"""
Explanation: Calculate Retention Data
Perform 12-week retention analysis based on this example. [1-12]-Week Retention are additionally included since we can get them at a low cost. We expand out to 12-week retention to better validate the data since the TAAR branches exhibit suspiciously similar retention values.
End of explanation
"""


%%bash
aws s3 cp taar_v2_retention.csv s3://net-mozaws-prod-us-west-2-pipeline-analysis/taarv2/

"""
Explanation: Write to s3 since this job is quite expensive and should only be run once.
End of explanation
"""


%%bash 
aws s3 cp s3://net-mozaws-prod-us-west-2-pipeline-analysis/taarv2/taar_v2_retention.csv .

ret_df = pd.read_csv("taar_v2_retention.csv")
ret_df.fillna(0, inplace=True)

plt.rcParams['figure.figsize'] = (12, 6)
fig, ax = plt.subplots()
for group, data in ret_df.groupby("branch"):
    (data.sort_values("period")
         .plot(x='period', 
               y='retention', 
               ax=ax, 
               label=group))
plt.ylabel("Retention")
plt.xlabel("Week (period)")
plt.title("1-12 Week Retention by Branch")
plt.show()

ret_df[ret_df.period >= 6.0].sort_values(["period", "branch"])

"""
Explanation: Load processed Retention Data
This section loads the data generated above without having to the re-run the entire notebook.
End of explanation
"""


day_over_day_retention = []
for i in range(40):
    d = date_plus_x_days("20180312", i)
    joinedx = joined.filter("enrollment_date = '{}'".format(d))
    ret_dfx = get_retention(joinedx)
    week6 = ret_dfx[ret_dfx.period == 6.0]
    for b, data in week6.groupby("branch"):
        x = {
            'branch': b,
            'ret': data['retention'].values[0],
            'date': d
        }
        day_over_day_retention.append(x)

day_over_day_retention_df = pd.DataFrame(day_over_day_retention)
day_over_day_retention_df.date = (
    pd.to_datetime(day_over_day_retention_df.date, format='%Y%m%d'))

plt.rcParams['figure.figsize'] = (12, 6)
fig, ax = plt.subplots()
for group, data in day_over_day_retention_df.groupby("branch"):
    data['ma'] = data.ret.rolling(window=6).mean()
    (data.sort_values("date")
         .plot(x='date', 
               y='ma', 
               ax=ax, 
               label=group))
plt.ylabel("Retention")
plt.xlabel("Enrollment Date")
plt.title("6 Week Retention by Enrollment Date, Branch (6-Day Moving Average)")
plt.show()

"""
Explanation: Investigate nearly identical retention lines for TAAR Branches
Let's look at 6-week retention over time by each enrollment date
End of explanation
"""


(joined.groupby("enrollment_date")
       .agg(F.countDistinct("client_id").alias("number of participants"))
       .sort("enrollment_date")
       .toPandas()
       .plot(x='enrollment_date'))

"""
Explanation: We see increased variability with time, which is most certainly due to the study being front-loaded with participants. Looking at enrollment:
End of explanation
"""


w6r = ret_df[ret_df.period==6]
w6r

(slinear, nlinear,
 sensemble, nensemble,
 scontrol, ncontrol) = [int(w6r[w6r.branch == b][i].values[0])
                        for b in ('linear-taar', 'ensemble-taar', 'control')
                        for i in ('n_week_clients', "total_clients")]



def get_effect(g1s, g2s, g1n, g2n):
    """
    Extracts the observed difference and p-value 
    for two proportions
    
    g1s: the number of successes for group 1
    g2s: the number of successes for group 2
    g1n: total trials for group 1
    g2n: total trials for group 2
    
    returns the effect and p-value as a tuple
    """
    # use counts to form a proportion and format
    effect = str(round((g1s*1.0 / g1n) - (g2s*1.0 / g2n), 4) * 100) + '%'
    
    # perform test of proportions
    pval = sm.stats.proportions_ztest(np.array([g1s, g2s]),
                               np.array([g1n, g2n]),
                               value=0.05)[1]
    return effect, pval

le, lp = get_effect(slinear, scontrol, nlinear, ncontrol)  
print "linear effect: {}\npvalue: {}\n".format(le, lp)

ee, ep = get_effect(sensemble, scontrol, nensemble, ncontrol)  
print "ensemble effect: {}\npvalue: {}".format(ee, ep)


"""
Explanation: we see most of our clients are enrolled before 2018-03-22. This is why the lines are so smooth for the first ~9 datapoints in the previous retention chart
Breaking retention down into shorter segments shows that there are indeed differences between the taar branches, however they track each other rather closely and are consistently higher than the control. This is evidence that the similarities we see in the lines for the first plot are in fact the true retention values and not a data-handling issue.
Results
End of explanation
"""


