import datetime as dt

import pandas as pd
import numpy as np

from wmfdata import spark, mariadb

## Load the RPython library so we can use R for graphs

%load_ext rpy2.ipython

%%R
library(ggplot2)
library(hrbrthemes)
library(tidyr)
library(lubridate)
library(zoo)
library(dplyr)
import::from(polloi, compress)

%%R

## Options
options(mc.cores = 4, digits = 3, scipen = 500)

## Defining a custom theme for all plots
commons_theme = function() {
    theme_ipsum_rc(
        base_size = 14, axis_title_size = 12, subtitle_size = 16,
        axis_title_just = 'cm'
    )
}

## We're operating with different format for timestamps, so we'll have to be able
## to parse them both with and without milliseconds.

def parse_dt(ts):
    try:
        return(dt.datetime.strptime(ts, '%Y-%m-%dT%H:%M:%S.%fZ'))
    except ValueError: ## no microseconds
        return(dt.datetime.strptime(ts, '%Y-%m-%dT%H:%M:%SZ'))

today = dt.datetime.now(dt.timezone.utc).date()
last_week = today - dt.timedelta(days = 7)

# Query to count fulltext and autocomplete searches on Commons

search_count_query = '''
WITH ac AS (
    SELECT TO_DATE(coalesce(meta.dt, client_dt, dt)) AS log_date,
           COUNT(DISTINCT event.searchsessionid, event.pageviewid) AS n_autocomp
    FROM event.searchsatisfaction
    WHERE year = 2020
    AND month = 8
    AND wiki = "commonswiki"
    AND useragent.is_bot = false
    AND event.subTest IS NULL
    AND event.action = "searchResultPage"
    AND event.isforced IS NULL -- only include non-test users
    AND event.source = "autocomplete"
    GROUP BY TO_DATE(coalesce(meta.dt, client_dt, dt))
), ft AS (
    SELECT TO_DATE(coalesce(meta.dt, client_dt, dt)) AS log_date,
           SUM(IF(event.hitsReturned > 0 , 1, 0)) AS n_fulltext_successful,
           SUM(IF(event.hitsReturned IS NULL , 1, 0)) AS n_fulltext_zeroresults
    FROM event.searchsatisfaction
    WHERE year = 2020
    AND month = 8
    AND wiki = "commonswiki"
    AND useragent.is_bot = false
    AND event.subTest IS NULL
    AND event.action = "searchResultPage"
    AND event.isforced IS NULL -- only include non-test users
    AND event.source = "fulltext"
    GROUP BY TO_DATE(coalesce(meta.dt, client_dt, dt))
)
SELECT ac.log_date, n_autocomp, n_fulltext_successful, n_fulltext_zeroresults
FROM ac
LEFT JOIN ft
ON ac.log_date = ft.log_date
'''

commons_searches_daily = spark.run(search_count_query)

autocomp_session_query = '''
SELECT event.searchsessionid,
    MIN(TO_DATE(coalesce(meta.dt, client_dt, dt))) AS session_start_date,
    SUM(1) AS num_searches
FROM event.searchsatisfaction
WHERE year = 2020
AND month = 8
AND wiki = "commonswiki"
AND useragent.is_bot = false
AND event.subTest IS NULL
AND event.action = "searchResultPage"
AND event.isforced IS NULL -- only include non-test users
AND event.source = "autocomplete"
GROUP BY event.searchsessionid
'''

autocomp_session_metrics = spark.run(autocomp_session_query)

fulltext_session_query = '''
SELECT event.searchsessionid,
    MIN(TO_DATE(coalesce(meta.dt, client_dt, dt))) AS session_start_date,
    SUM(1) AS num_searches
FROM event.searchsatisfaction
WHERE year = 2020
AND month = 8
AND wiki = "commonswiki"
AND useragent.is_bot = false
AND event.subTest IS NULL
AND event.action = "searchResultPage"
AND event.isforced IS NULL -- only include non-test users
AND event.source = "fulltext"
GROUP BY event.searchsessionid
'''

fulltext_session_metrics = spark.run(fulltext_session_query)

## Drop the session ID columns, we don't really need those
autocomp_session_metrics.drop(columns = 'searchsessionid', inplace = True)
fulltext_session_metrics.drop(columns = 'searchsessionid', inplace = True)

session_info_query = '''
WITH cs AS (
    SELECT event.searchsessionid,
           MIN(coalesce(meta.dt, client_dt, dt)) AS session_start_ts
    FROM event.searchsatisfaction
    WHERE year = 2020
    AND month = 8
    AND wiki = "commonswiki"
    AND useragent.is_bot = false
    AND event.subTest IS NULL
    AND event.action = "searchResultPage"
    AND event.isforced IS NULL -- only include non-test users
    GROUP BY event.searchsessionid
    HAVING SUM(1) < 50
),
se AS (
    SELECT event.searchsessionid,
           MAX(coalesce(meta.dt, client_dt, dt)) AS session_end_ts
    FROM event.searchsatisfaction
    WHERE year = 2020
    AND month = 8
    AND wiki = "commonswiki"
    AND useragent.is_bot = false
    AND event.subTest IS NULL
    AND event.isforced IS NULL -- only include non-test users
    GROUP BY event.searchsessionid
),
ct AS (
    SELECT event.searchsessionid, event.position
    FROM event.searchsatisfaction
    WHERE year = 2020
    AND month = 8
    AND wiki = "commonswiki"
    AND useragent.is_bot = false
    AND event.subTest IS NULL
    AND event.action = "visitPage"
    AND event.isforced IS NULL -- only include non-test users
),
dw AS (
    SELECT event.searchsessionid, max(event.checkin) AS last_checkin
    FROM event.searchsatisfaction
    WHERE year = 2020
    AND month = 8
    AND wiki = "commonswiki"
    AND useragent.is_bot = false
    AND event.subTest IS NULL
    AND event.action = "checkin"
    AND event.isforced IS NULL -- only include non-test users
    GROUP BY event.searchsessionid
)
SELECT cs.searchsessionid, session_start_ts, se.session_end_ts,
       IF(ct.searchsessionid IS NOT NULL, 1, 0) AS clicked_through,
       coalesce(ct.position, -1) AS position,
       IF(dw.last_checkin IS NOT NULL, dw.last_checkin, -1) AS last_checkin
FROM cs
JOIN se
ON cs.searchsessionid = se.searchsessionid
LEFT JOIN ct
ON cs.searchsessionid = ct.searchsessionid
LEFT JOIN dw
ON cs.searchsessionid = dw.searchsessionid
'''

session_info = spark.run(session_info_query)

session_info.drop(columns = 'searchsessionid', inplace = True)

r_session_info = session_info[['session_start_ts', 'session_end_ts']]

session_info['session_start'] = session_lengths['session_start_ts'].apply(parse_dt)
session_info['session_end'] = session_lengths['session_end_ts'].apply(parse_dt)
session_info['session_length'] = session_lengths['session_end'] - session_lengths['session_start']
session_info['session_start_date'] = session_lengths['session_start'].apply(lambda x: x.date())