# coding: utf-8

# # Anomalies in meal prices
#
# In the Chamber of Deputies' CEAP, there is a list of 1,000's of meal expenses made by congresspeople. The law says that the congressperson cannot pay for any other, even being her advisor or SO. We want to work on this analysis to find possibly illegal and immoral expenses. They may have happened when the politician spent more than needed (e.g. the whole menu costs X but the bill was 2X) or too much in an specific period of time. In the end, we also want to alert about too expensive reibursements, even with an explanation behind of it.
#
# Note: remember to correct prices with an inflation index (e.g. IPCA).

# In[1]:

get_ipython().magic('matplotlib inline')
import matplotlib.pyplot as plt
import seaborn as sns
sns.set(color_codes=True)

plt.rcParams['figure.figsize'] = (20, 10)

# In[2]:

from serenata_toolbox.datasets import fetch

fetch('2016-11-19-reimbursements.xz', '../data')
fetch('2016-09-03-companies.xz', '../data')
fetch('2016-11-29-yelp-companies.xz', '../data')
fetch('2016-12-02-foursquare-companies.xz', '../data')

# In[3]:

import numpy as np
import pandas as pd

dataset = pd.read_csv('../data/2016-11-19-reimbursements.xz',
                      dtype={'applicant_id': np.str,
                             'cnpj_cpf': np.str,
                             'congressperson_id': np.str,
                             'subquota_number': np.str},
                      low_memory=False)
dataset = dataset[dataset['congressperson_id'].notnull()]
dataset['issue_date'] = pd.to_datetime(dataset['issue_date'], errors='coerce')
dataset['issue_date_day'] = dataset['issue_date'].apply(lambda date: date.day)
dataset['issue_date_month'] = dataset['issue_date'].apply(lambda date: date.month)
dataset['issue_date_year'] = dataset['issue_date'].apply(lambda date: date.year)
dataset['issue_date_weekday'] = dataset['issue_date'].apply(lambda date: date.weekday())
dataset['issue_date_week'] = dataset['issue_date'].apply(lambda date: date.week)

# The `companies.xz` dataset has businesses placed outside Brazilian area. We intentionally disconsider them.

# In[4]:

is_in_brazil = '(-73.992222 < longitude < -34.7916667) & (-33.742222 < latitude < 5.2722222)'
companies = pd.read_csv('../data/2016-09-03-companies.xz',
                        dtype={'cnpj': np.str},
                        low_memory=False)
companies = companies.query(is_in_brazil)
companies['cnpj'] = companies['cnpj'].str.replace(r'\D', '')
dataset = pd.merge(dataset, companies,
                   how='left',
                   left_on='cnpj_cpf',
                   right_on='cnpj',
                   suffixes=('', '_company'))

# In[5]:

dataset =     dataset.query('subquota_description == "Congressperson meal"')
companies =     companies[companies['cnpj'].isin(dataset.loc[dataset['cnpj'].notnull(),
                                                 'cnpj'])]

# In[6]:

dataset['total_net_value'].describe()

# In[7]:

dataset['total_net_value'].median()

# In[8]:

sns.distplot(dataset['total_net_value'],
             bins=30,
             kde=False)

# In[9]:

bottom_99 = dataset['total_net_value'].quantile(0.99)
bottom_99

# In[10]:

dataset[dataset['total_net_value'] < bottom_99].shape

# In[11]:

sns.distplot(dataset.loc[dataset['total_net_value'] < bottom_99, 'total_net_value'],
             bins=30,
             kde=False)

# In[12]:

bottom_99_dataset = dataset.query('total_net_value > {}'.format(bottom_99))
ranking = bottom_99_dataset.groupby('state_company')['total_net_value']     .median().sort_values(ascending=False)

sns.boxplot(x='state_company',
            y='total_net_value',
            data=bottom_99_dataset,
            order=ranking.index)

# In[13]:

bottom_99_dataset.query('state_company == "CE"').shape

# In[14]:

dataset.query('state_company == "CE"').shape

# In[15]:

bottom_99_dataset['state_company'].isnull().sum()

# In[16]:

bottom_99_dataset.query('state_company == "CE"')     .sort_values('total_net_value', ascending=False)

# ## Using Yelp to improve prices information

# In[17]:

yelp = pd.read_csv('../data/2016-11-29-yelp-companies.xz',
                   low_memory=False)
yelp.head()

# We have data for just 8.6% of the companies which received from the "Congressperson meal" subquota.

# In[18]:

yelp['price'].notnull().sum()

# In[19]:

companies.shape

# In[20]:

yelp['price'].isnull().sum()

# In[21]:

yelp['price.int'] = yelp['price'].str.len()
states_with_records =     yelp[yelp['price'].notnull()].groupby('location.state')['location.state'].count() > 10
states_with_records = states_with_records[states_with_records].index

# In[22]:

yelp_just_significant_states =     yelp[yelp['price'].notnull() &
         yelp['location.state'].isin(states_with_records)]
yelp_just_significant_states['location.state'].value_counts()

# Yelp won't be that useful for now, since we don't have a lot of data. Will leave it for another analysis.

# ## Predict prices

# The idea here is to try to predict the "right range for prices" in a specific place. If we can have a good accuracy, everything far from the prediction could be considered an outlier.

# In[23]:

bottom_99_dataset.iloc[0, :57]

# **DummyRegressor with mean strategy as a baseline**

# In[24]:

from sklearn.dummy import DummyRegressor
from sklearn.model_selection import train_test_split

X = bottom_99_dataset[['year']]
y = bottom_99_dataset['total_net_value']
X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=0)

model = DummyRegressor(strategy='mean')
model.fit(X_train, y_train)
model.score(X_test, y_test)

# In[25]:

from sklearn.preprocessing import LabelEncoder

le_state = LabelEncoder()
le_city = LabelEncoder()
factor_columns = ['state_company', 'city']
model_dataset = bottom_99_dataset.dropna(subset=factor_columns)
model_dataset['state_company'] = le_state.fit_transform(model_dataset['state_company'])
model_dataset['city'] = le_city.fit_transform(model_dataset['city'])

model_columns = ['cnpj',
                 'issue_date_day',
                 'issue_date_month',
                 'issue_date_year']
X = model_dataset[model_columns + factor_columns]
y = model_dataset['total_net_value']
X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=0)

# **LinearRegression**
#
# Not very good accuracy.

# In[26]:

from sklearn.linear_model import LinearRegression

model = LinearRegression(n_jobs=-1)
model.fit(X_train, y_train)
model.score(X_test, y_test)

# What if we could improve it using the type of business as a feature? e.g. restaurant, hotel, grill...

# In[27]:

import unicodedata

def normalize_string(string):
    if isinstance(string, str):
        nfkd_form = unicodedata.normalize('NFKD', string.lower())
        return nfkd_form.encode('ASCII', 'ignore').decode('utf-8')

# In[28]:

import nltk
from sklearn.feature_extraction.text import CountVectorizer

stopwords = nltk.corpus.stopwords.words('portuguese')
count_vect = CountVectorizer(stop_words=stopwords)
trade_names = dataset.loc[dataset['supplier'].notnull(),
                          'supplier'].unique()
trade_names = np.vectorize(normalize_string)(trade_names)
dataset_counts = count_vect.fit_transform(trade_names)

# In[29]:

frequent_words = sorted(list(zip(count_vect.get_feature_names(),
    np.asarray(dataset_counts.sum(axis=0)).ravel())), key=lambda x: -x[1])

# In[30]:

frequent_words[:20]

# In[31]:

frequent_words = dict(frequent_words)

excluded_keywords = ['ltda', 'cia', 'epp']
[frequent_words.pop(keyword) for keyword in excluded_keywords]

# In[32]:

def business_type(name):
    fun = np.vectorize(lambda x: normalize_string(x))
    keywords = set(fun(name.split(' '))) - set(stopwords)
    key_freqs = list(map(lambda x: (x, frequent_words.get(x)), list(keywords)))
    key_freqs = [key_freq for key_freq in key_freqs if key_freq[1] is not None]
    if key_freqs:
        key_freq = max(key_freqs, key=lambda x: x[1])
        return key_freq[0]

dataset['supplier_keyword'] = dataset['supplier'].apply(business_type)
bottom_99_dataset['supplier_keyword'] =     bottom_99_dataset['supplier'].apply(business_type)

# In[33]:

le_state = LabelEncoder()
le_city = LabelEncoder()
le_supplier_keyword = LabelEncoder()
factor_columns = ['state_company', 'supplier_keyword']
model_dataset = bottom_99_dataset.dropna(subset=factor_columns)
model_dataset['state_company'] = le_state.fit_transform(model_dataset['state_company'])
model_dataset['city'] = le_city.fit_transform(model_dataset['city'])
model_dataset['supplier_keyword'] = le_city.fit_transform(model_dataset['supplier_keyword'])

model_columns = ['cnpj',
                 'issue_date_day',
                 'issue_date_month',
                 'issue_date_year']
X = model_dataset[model_columns + factor_columns]
y = model_dataset['total_net_value']
X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=0)

# In[34]:

model = LinearRegression(n_jobs=-1)
model.fit(X_train, y_train)
model.score(X_test, y_test)

# Still not good enough. In fact, there's a serious mistake when considering this linear regression method for outlier detection: not just we're assuming that prices follow a linear regression (IMHO still good assumption, though), but outliers should be removed before running the regression. In other words, to find outliers, we must first remove them, having the regression score as just a accuracy estimator. May still be an interesting approach but we want to engage with something simpler first, to get the easier and most anomalous results.

# ## Common CNPJs
#
# Expenses in the same restaurant are expected to follow a normal distribution. Can we find outliers in companies with enough expenses to analyze?

# In[35]:

from scipy.stats import normaltest

def normaltest_pvalue(values):
    if len(values) >= 20:
        return normaltest(values).pvalue
    else:
        return 1

net_values_by_cnpj = dataset.groupby('cnpj_cpf')['total_net_value']     .agg([len, np.mean, np.std, normaltest_pvalue])     .sort_values('len', ascending=False)     .reset_index()
net_values_by_cnpj['threshold'] = net_values_by_cnpj['mean'] +     3 * net_values_by_cnpj['std']
applicants_per_cnpj = dataset.groupby('cnpj_cpf')['applicant_id']     .aggregate(lambda x: len(set(x))).reset_index()     .rename(columns={'applicant_id': 'congresspeople'})
net_values_by_cnpj = pd.merge(net_values_by_cnpj, applicants_per_cnpj)
net_values_by_cnpj.head()

# In[36]:

len(net_values_by_cnpj.query('normaltest_pvalue < .05')) / len(net_values_by_cnpj)

# In[37]:

data_with_threshold = pd.merge(dataset, net_values_by_cnpj, on='cnpj_cpf')     .sort_values('total_net_value', ascending=False)

# In[38]:

data_with_threshold['main_activity'] =     data_with_threshold['main_activity'].apply(normalize_string)

# Let's discard hotel reibursements. There's no method yet to discover for how long the congressperson stayed in the hotel, so we can expect a high standard deviation in their expenses. Even when detecting outliers, it's too hard to investigate manually given the same reason.
#
# Here, we just consider CNPJs where we have more than 20 reimbursements, from at least 3 congresspeople.

# In[39]:

is_hotel_reimbursement = data_with_threshold['main_activity']     .str.contains('hoteis').astype(np.bool)
outliers = data_with_threshold[~is_hotel_reimbursement]     .query('(congresspeople > 3) & (len >= 20) & (total_net_value > threshold)')
print(len(outliers), outliers['total_net_value'].sum())

# ## Foursquare
#
# Before totally discarding it, let's see how significant is the Foursquare dataset.

# In[40]:

foursquare = pd.read_csv('../data/2016-12-02-foursquare-companies.xz',
                         low_memory=False)
foursquare.head()

# In[41]:

foursquare.iloc[0]

# In[42]:

print(foursquare['price.tier'].notnull().sum(),
      foursquare['price.tier'].notnull().sum() / len(companies),
      foursquare.query('confirmed_match == True')['price.tier'].notnull().sum() / len(companies))

# ### Clustering for find the best group for a new restaurant

# In[43]:

companies.shape

# In[44]:

# is_cnpj = (dataset['cnpj_cpf'].str.len() == 14) & \
#     dataset['cnpj_cpf'].notnull() & \
#     dataset['document_type'] != 2
# cnpjs = dataset.sort_values('issue_date') \
#     .loc[is_cnpj, ['cnpj_cpf', 'supplier']] \
#     .drop_duplicates('cnpj_cpf', keep='last')
# cnpjs.head()

# In[45]:

is_cnpj = dataset['cnpj_cpf'].str.len() == 14
cnpj_list = dataset.loc[is_cnpj].groupby('cnpj')['total_net_value']     .agg([np.mean, np.std]).reset_index()
cnpj_list.shape

# In[46]:

cnpj_list.head()

# In[47]:

cnpj_list = pd.merge(cnpj_list,
                     dataset[['cnpj_cpf', 'supplier']].drop_duplicates('cnpj_cpf'),
                     how='left',
                     left_on='cnpj', right_on='cnpj_cpf')
del cnpj_list['cnpj_cpf']
cnpj_list.head()

# In[48]:

counts = dataset.loc[is_cnpj].groupby('cnpj')['applicant_id']     .agg({'congresspeople': (lambda x: len(np.unique(x))),
          'len': (lambda x: len(x))
         }).reset_index()

cnpj_list = pd.merge(cnpj_list, counts)

# **Calculate threshold for companies using their own receipts**

# In[49]:

threshold_for_cnpjs = cnpj_list.groupby('cnpj')     .apply(lambda x: x['mean'].mean() + 3 * x['std'].mean()).reset_index()     .rename(columns={0: 'threshold'})
threshold_for_cnpjs

cnpj_list = pd.merge(cnpj_list, threshold_for_cnpjs)
cnpj_list.head()

# In[50]:

HOTEL_REGEX = r'hote[l(eis)(ls)]'
is_hotel_reimbursement = (cnpj_list['supplier'].str.lower().str.contains(HOTEL_REGEX))

# Mark companies as having or not significant data.

# In[51]:

rows = (~is_hotel_reimbursement) &     (cnpj_list['congresspeople'] > 3) &     (cnpj_list['len'] > 20)
cnpj_list['has_significant_data'] = False
cnpj_list.loc[rows, 'has_significant_data'] = True

# In[52]:

print(cnpj_list['has_significant_data'].sum(),
      cnpj_list['has_significant_data'].sum() / len(cnpj_list['has_significant_data']))

# In[53]:

sns.lmplot('mean', 'std',
           data=cnpj_list.query('has_significant_data'),
           scatter_kws={'marker': 'D', 's': 100},
           size=10)

# **Predict threshold classifying companies in clusters by their price ranges**

# In[54]:

X = cnpj_list.loc[cnpj_list['has_significant_data'],
                  ['mean', 'std']]

# In[55]:

from sklearn.cluster import KMeans

model = KMeans(n_clusters=3, random_state=0)
model.fit(X)

# In[56]:

cnpj_list.loc[cnpj_list['has_significant_data'], 'y'] = model.predict(X)

# In[57]:

cnpj_list.query('y.notnull()').head()

# In[58]:

rows = (~cnpj_list['has_significant_data']) &     cnpj_list['std'].notnull() &     (~is_hotel_reimbursement)
X = cnpj_list.loc[rows, ['mean', 'std']]
cnpj_list.loc[rows, 'y'] = model.predict(X)

# In[59]:

threshold_for_groups = cnpj_list.groupby('y')     .apply(lambda x: x['mean'].mean() + 4 * x['std'].mean()).reset_index()     .rename(columns={0: 'threshold'})
threshold_for_groups

# In[60]:

group_thresholds = pd.merge(cnpj_list.query('~has_significant_data'),
                            threshold_for_groups,
                            on='y',
                            suffixes=('', '_group'))

cnpj_list = pd.merge(cnpj_list,
                     group_thresholds[['cnpj', 'threshold_group']],
                     how='left')
cnpj_list.loc[~cnpj_list['has_significant_data'], 'threshold'] =     cnpj_list['threshold_group']

# In[61]:

cnpj_list.query('(~has_significant_data) & std.notnull()').head()

# In[62]:

cnpj_list.query('has_significant_data').head()

# In[63]:

cnpj_list.query('threshold.notnull()').sample(5, random_state=10)

# In[64]:

del cnpj_list['threshold_group']

# In[65]:

merged = pd.merge(dataset, cnpj_list,
                  how='left',
                  left_on='cnpj_cpf',
                  right_on='cnpj',
                  suffixes=('', '_company'))

# In[66]:

merged['supplier'] = merged['supplier'].astype(np.str)
is_hotel_reimbursement =     (merged['supplier'].str.lower().str.contains(HOTEL_REGEX))

merged[~is_hotel_reimbursement].query('total_net_value > threshold').shape

# In[67]:

keys = ['year',
        'congressperson_name',
        'document_id',
        'total_net_value',
        'threshold',
        'cnpj_cpf',
        'has_significant_data',
        'name']

merged['diff'] = merged['threshold'] - merged['total_net_value']
merged[~(is_hotel_reimbursement | merged['has_significant_data'])]     .query('(total_net_value > threshold)')     .sort_values('diff', ascending=False).head(10)[keys]

# In[68]:

merged[~is_hotel_reimbursement].shape

# In[69]:

merged[~is_hotel_reimbursement]     .query('(total_net_value > threshold)')['total_net_value'].shape

# In[70]:

merged[~is_hotel_reimbursement]     .query('(total_net_value > threshold)')['total_net_value'].sum()

# In[71]:

merged[~is_hotel_reimbursement]     .query('(total_net_value > threshold) & (has_significant_data == False)')['total_net_value'].shape

# ## Conclusions
#
# For companies with significant data (defined by us as a company which received money at least 20x, from at least 3 distinct congresspeople), we use mean + 3 * std to detect outliers. Does not return all the suspect cases, but all of them, after some sampling investigation, seem to be very suspect.
#
# Since there's "significant data" just for 4% of the companies, we need a way for extrapolating the results for not so known ones. For doing so, we classify companies in 3 clusters using K-Means, considering mean and standard deviation of their prices as features. Once classified, we consider their threshold mean + 4 * stds of their clusters (one extra std compared to places where we have enough reimbursements to know better).
#
# Reimbursements made for expenses in hotels are discarded from this classifier, since they usually contain much more than meals (and we don't know for how long the congressperson was hosted in the place, not yet trustable for legal reports).

# In[ ]:


r"""
Module that creates and modifies parse trees of well formed boolean formulas.

A parse tree of a boolean formula is a nested list, where each branch is either
a single variable, or a formula composed of either two variables and a binary
operator or one variable and a unary operator. The function parse() produces
a parse tree that is simplified for the purposes of more efficient truth value
evaluation. The function polish_parse() produces the full parse tree of a boolean
formula which is used in functions related to proof and inference.  That is,
parse() is meant to be used with functions in the logic module that perform
semantic operations on a boolean formula, and polish_parse() is to be used with
functions that perform syntactic operations on a boolean formula.

AUTHORS:

- Chris Gorecki (2007): initial version

- Paul Scurek (2013-08-01): added polish_parse, cleaned up python code,
  updated docstring formatting

EXAMPLES:

Find the parse tree and variables of a string representation of a boolean formula::

    sage: import sage.logic.logicparser as logicparser
    sage: s = 'a|b&c'
    sage: t = logicparser.parse(s)
    sage: t
    (['|', 'a', ['&', 'b', 'c']], ['a', 'b', 'c'])

Find the full syntax parse tree of a string representation of a boolean formula::

    sage: import sage.logic.logicparser as logicparser
    sage: s = '(a&b)->~~c'
    sage: logicparser.polish_parse(s)
    ['->', ['&', 'a', 'b'], ['~', ['~', 'c']]]

Find the tokens and distinct variables of a boolean formula::

    sage: import sage.logic.logicparser as logicparser
    sage: s = '~(a|~b)<->(c->c)'
    sage: logicparser.tokenize(s)
    (['(', '~', '(', 'a', '|', '~', 'b', ')', '<->', '(', 'c', '->', 'c', ')', ')'], ['a', 'b', 'c'])

Find the parse tree of a boolean formula from a list of the formula's tokens::

    sage: import sage.logic.logicparser as logicparser
    sage: t = ['(', 'a', '->', '~', 'c', ')']
    sage: logicparser.tree_parse(t)
    ['->', 'a', ['~', 'c', None]]
    sage: r = ['(', '~', '~', 'a', '|', 'b', ')']
    sage: logicparser.tree_parse(r)
    ['|', 'a', 'b']

Find the full syntax parse tree of a boolean formula from a list of tokens::

    sage: import sage.logic.logicparser as logicparser
    sage: t = ['(', 'a', '->', '~', 'c', ')']
    sage: logicparser.tree_parse(t, polish = True)
    ['->', 'a', ['~', 'c']]
    sage: r = ['(', '~', '~', 'a', '|', 'b', ')']
    sage: logicparser.tree_parse(r, polish = True)
    ['|', ['~', ['~', 'a']], 'b']

"""
#*****************************************************************************
#       Copyright (C) 2007 Chris Gorecki <chris.k.gorecki@gmail.com>
#       Copyright (C) 2013 Paul Scurek <scurek86@gmail.com>
#
#  Distributed under the terms of the GNU General Public License (GPL)
#  as published by the Free Software Foundation; either version 2 of
#  the License, or (at your option) any later version.
#                  http://www.gnu.org/licenses/
#*****************************************************************************

from types import *
import string

__symbols = '()&|~<->^'
__op_list = ['~', '&', '|', '^', '->', '<->']

def parse(s):
    r"""
    Return a parse tree from a boolean formula s.

    INPUT:

    - ``s`` -- a string containing a boolean formula.

    OUTPUT:

    A list containing the prase tree and a list containing the
    variables in a boolean formula in this order:

    1. the list containing the pase tree
    2. the list containing the variables

    EXAMPLES:

    This example illustrates how to produce the parse tree of a boolean formula s.

    ::

        sage: import sage.logic.logicparser as logicparser
        sage: s = 'a|b&c'
        sage: t = logicparser.parse(s)
        sage: t
        (['|', 'a', ['&', 'b', 'c']], ['a', 'b', 'c'])
    """
    toks, vars_order = tokenize(s)
    tree = tree_parse(toks)
    # special case of tree == single variable
    if isinstance(tree, StringType):
        return ['&', tree, tree], vars_order
    return tree, vars_order

def polish_parse(s):
    r"""
    Return the full syntax parse tree from a boolean formula s.

    INPUT:

    - ``s`` -- a string containing a boolean expression

    OUTPUT:

    The full syntax parse tree as a nested list

    EXAMPLES:

    This example illustrates how to find the full syntax parse tree of a boolean formula.

    ::

        sage: import sage.logic.logicparser as logicparser
        sage: s = 'a|~~b'
        sage: t = logicparser.polish_parse(s)
        sage: t
        ['|', 'a', ['~', ['~', 'b']]]

    AUTHORS:

    - Paul Scurek (2013-08-03)
    """
    toks, vars_order = tokenize(s)
    tree = tree_parse(toks, polish = True)
    # special case where the formula s is a single variable
    if isinstance(tree, StringType):
        return vars_order
    return tree

def tokenize(s):
    r"""
    Return the tokens and the distinct variables appearing in a boolean formula s.

    INPUT:

    - ``s`` -- a string representation of a boolean formula

    OUTPUT:

    The tokens and variables as an ordered pair of lists in the following order:

    1. A list containing the tokens of s, in the order they appear in s
    2. A list containing the distinct variables in s, in the order they appearn in s

    EXAMPLES:

    This example illustrates how to tokenize a string representation of a boolean formula.

    ::

        sage: import sage.logic.logicparser as logicparser
        sage: s = 'a|b&c'
        sage: t = logicparser.tokenize(s)
        sage: t
        (['(', 'a', '|', 'b', '&', 'c', ')'], ['a', 'b', 'c'])
    """
    i = 0
    toks = ['(']
    vars_order = []

    while i < len(s):
        tok = ""
        skip = valid = 1
        if s[i] in '()~&|^':
            tok = s[i]
        elif s[i:i + 2] == '->':
            tok = '->'
            skip = 2
        elif s[i:i + 3] == '<->':
            tok = '<->'
            skip = 3
        # check to see if '-', '<' or '>' are used incorrectly
        elif s[i] in '<->':
            msg = "'%s' can only be used as part of the operators '<->' or '->'." % (s[i])
            raise SyntaxError, msg
        if len(tok) > 0:
            toks.append(tok)
            i += skip
            continue
        else:
            # token is a variable name
            if s[i] == ' ':
                i += 1
                continue

            while i < len(s) and s[i] not in __symbols and s[i] != ' ':
                tok += s[i]
                i += 1

            if len(tok) > 0:
                if tok[0] not in string.letters:
                    valid = 0
                for c in tok:
                    if c not in string.letters and c not in string.digits and c != '_':
                        valid = 0

            if valid == 1:
                toks.append(tok)
                if tok not in vars_order:
                    vars_order.append(tok)
            else:
                msg = 'invalid variable name ' + tok
                msg += ": identifiers must begin with a letter and contain only "
                msg += "alphanumerics and underscores"
                raise NameError, msg

    toks.append(')')
    return toks, vars_order

def tree_parse(toks, polish = False):
    r"""
    Return a parse tree from the tokens in toks.

    INPUT:

    - ``toks`` -- a list of tokens from a boolean formula

    - ``polish`` -- (default: False) a boolean.  When true, tree_parse will return
      the full syntax parse tree.

    OUTPUT:

    A parse tree in the form of a nested list that depends on ``polish`` as follows:

    polish == False -- Return a simplified parse tree.

    polish == True -- Return the full syntax parse tree.

    EXAMPLES:

    This example illustrates the use of tree_parse when polish == False.

    ::

        sage: import sage.logic.logicparser as logicparser
        sage: t = ['(', 'a', '|', 'b', '&', 'c', ')']
        sage: logicparser.tree_parse(t)
        ['|', 'a', ['&', 'b', 'c']]

    We now demonstrate the use of tree_parse when polish == True.

    ::

        sage: t = ['(', 'a', '->', '~', '~', 'b', ')']
        sage: logicparser.tree_parse(t)
        ['->', 'a', 'b']
        sage: t = ['(', 'a', '->', '~', '~', 'b', ')']
        sage: logicparser.tree_parse(t, polish = True)
        ['->', 'a', ['~', ['~', 'b']]]
    """
    stack = []
    for tok in toks:
        stack.append(tok)
        if tok == ')':
            lrtoks = []
            while tok != '(':
                tok = stack.pop()
                lrtoks.insert(0, tok)
            branch = parse_ltor(lrtoks[1:-1], polish = polish)
            stack.append(branch)
    return stack[0]

def parse_ltor(toks, n = 0, polish = False):
    r"""
    Return a parse tree from toks, where each token in toks is atomic.

    INPUT:

    - ``toks`` -- a list of tokens. Each token is atomic.

    - ``n`` -- (default: 0) an integer representing which order of
      operations are occurring

    - ``polish`` -- (default: False) a boolean.  When true, double negations
      are not cancelled and negated statements are turned into list of length two.

    OUTPUT:

    The parse tree as a nested list that depends on ``polish`` as follows:

    polish == False - Return a simplified parse tree.

    polish == True - Return the full syntax parse tree.

    EXAMPLES:

    This example illustrates the use of parse_ltor when polish == False.

    ::

        sage: import sage.logic.logicparser as logicparser
        sage: t = ['a', '|', 'b', '&', 'c']
        sage: logicparser.parse_ltor(t)
        ['|', 'a', ['&', 'b', 'c']]

    ::

        sage: import sage.logic.logicparser as logicparser
        sage: t = ['a', '->', '~', '~', 'b']
        sage: logicparser.parse_ltor(t)
        ['->', 'a', 'b']

    We now repeat the previous example, but with polish == True.

    ::

        sage: import sage.logic.logicparser as logicparser
        sage: t = ['a', '->', '~', '~', 'b']
        sage: logicparser.parse_ltor(t, polish = True)
        ['->', 'a', ['~', ['~', 'b']]]

    """
    i = 0
    for tok in toks:
        if tok == __op_list[n]:
            if tok == '~':
                if not polish:
                    # cancel double negations
                    if toks[i] == '~' and toks[i + 1] == '~':
                        del toks[i]
                        del toks[i]
                        return parse_ltor(toks, n)
                    args = [toks[i], toks[i + 1], None]
                    toks[i] = args
                    del toks[i + 1]
                    return parse_ltor(toks, n)
                # This executes when creating the full syntax parse tree
                else:
                    j = i
                    while toks[j] == '~':
                        j += 1
                    while j > i:
                        args = [toks[j - 1], toks[j]]
                        toks[j - 1] = args
                        del toks[j]
                        j -= 1
                    return parse_ltor(toks, n = n, polish = polish)
            else:
                args = [toks[i - 1], toks[i], toks[i + 1]]
                toks[i - 1] = [args[1], args[0], args[2]]
                del toks[i]
                del toks[i]
                return parse_ltor(toks, n)
        i += 1
    if n + 1 < len(__op_list):
        return parse_ltor(toks, n + 1)
    if len(toks) > 1:
        raise SyntaxError
    return toks[0]

def apply_func(tree, func):
    r"""
    Apply func to each node of tree.  Return a new parse tree.

    INPUT:

    - ``tree`` -- a parse tree of a boolean formula

    - ``func`` -- a function to be applied to each node of tree.  This may
      be a function that comes from elsewhere in the logic module.

    OUTPUT:

    The new parse tree in the form of a nested list

    EXAMPLES:

    This example uses :func:`apply_func` where ``func`` switches two entries of tree.

    ::

        sage: import sage.logic.logicparser as logicparser
        sage: t = ['|', ['&', 'a', 'b'], ['&', 'a', 'c']]
        sage: f = lambda t: [t[0], t[2], t[1]]
        sage: logicparser.apply_func(t, f)
        ['|', ['&', 'c', 'a'], ['&', 'b', 'a']]
    """
    if type(tree[1]) is ListType and type(tree[2]) is ListType:
        lval = apply_func(tree[1], func)
        rval = apply_func(tree[2], func)
    elif type(tree[1]) is ListType:
        lval = apply_func(tree[1], func)
        rval = tree[2]
    elif type(tree[2]) is ListType:
        lval = tree[1]
        rval = apply_func(tree[2], func)
    else:
        lval = tree[1]
        rval = tree[2]
    return func([tree[0], lval, rval])


# -*- coding: utf-8 -*-
"""
Assignment Chaper 01 (Think Stats 2, Allen B. Downey)
Ex 1-2
Self-study on statistics using pyhton
@author: Github: @rafaelmm82
"""

from __future__ import print_function, division

import sys
import numpy as np
import thinkstats2

from collections import defaultdict

def ReadFemResp(dct_file='2002FemResp.dct',
                dat_file='2002FemResp.dat.gz',
                nrows=None):
    """Reads the NSFG respondent data.

    dct_file: string file name
    dat_file: string file name

    returns: DataFrame
    """
    dct = thinkstats2.ReadStataDct(dct_file)
    df = dct.ReadFixedWidth(dat_file, compression='gzip', nrows=nrows)
    CleanFemResp(df)
    return df

def CleanFemResp(df):
    """Recodes variables from the respondent frame.

    df: DataFrame
    """
    pass

def ReadFemPreg(dct_file='2002FemPreg.dct',
                dat_file='2002FemPreg.dat.gz'):
    """Reads the NSFG pregnancy data.

    dct_file: string file name
    dat_file: string file name

    returns: DataFrame
    """
    dct = thinkstats2.ReadStataDct(dct_file)
    df = dct.ReadFixedWidth(dat_file, compression='gzip')
    CleanFemPreg(df)
    return df

def CleanFemPreg(df):
    """Recodes variables from the pregnancy frame.

    df: DataFrame
    """
    # mother's age is encoded in centiyears; convert to years
    df.agepreg /= 100.0

    # birthwgt_lb contains at least one bogus value (51 lbs)
    # replace with NaN
    df.loc[df.birthwgt_lb > 20, 'birthwgt_lb'] = np.nan

    # replace 'not ascertained', 'refused', 'don't know' with NaN
    na_vals = [97, 98, 99]
    df.birthwgt_lb.replace(na_vals, np.nan, inplace=True)
    df.birthwgt_oz.replace(na_vals, np.nan, inplace=True)
    df.hpagelb.replace(na_vals, np.nan, inplace=True)

    df.babysex.replace([7, 9], np.nan, inplace=True)
    df.nbrnaliv.replace([9], np.nan, inplace=True)

    # birthweight is stored in two columns, lbs and oz.
    # convert to a single column in lb
    # NOTE: creating a new column requires dictionary syntax,
    # not attribute assignment (like df.totalwgt_lb)
    df['totalwgt_lb'] = df.birthwgt_lb + df.birthwgt_oz / 16.0

    # due to a bug in ReadStataDct, the last variable gets clipped;
    # so for now set it to NaN
    df.cmintvw = np.nan

def ValidatePregnum(resp, preg):
    """Validate pregnum in the respondent file.

    resp: respondent DataFrame
    preg: pregnancy DataFrame
    """
    # make the map from caseid to list of pregnancy indices
    preg_map = MakePregMap(preg)

    # iterate through the respondent pregnum series
    for index, pregnum in resp.pregnum.iteritems():
        caseid = resp.caseid[index]
        indices = preg_map[caseid]

        # check that pregnum from the respondent file equals
        # the number of records in the pregnancy file
        if len(indices) != pregnum:
            print(caseid, len(indices), pregnum)
            return False

    return True

def MakePregMap(df):
    """Make a map from caseid to list of preg indices.

    df: DataFrame

    returns: dict that maps from caseid to list of indices into `preg`
    """
    d = defaultdict(list)
    for index, caseid in df.caseid.iteritems():
        d[caseid].append(index)
    return d

def main():
    """Tests the functions in this module.

    script: string script name
    """
    # read and validate the respondent file
    resp = ReadFemResp()

    assert(len(resp) == 7643)
    assert(resp.pregnum.value_counts()[1] == 1267)

    # read and validate the pregnancy file
    preg = ReadFemPreg()
    print(preg.shape)

    # validate that the pregnum column in `resp` matches the number
    # of entries in `preg`
    assert(ValidatePregnum(resp, preg))

    print('Everything is ok!')

if __name__ == '__main__':
    main()


#!/usr/bin/env python
# coding: utf-8

# # Advanced Binary Image Segmentation for the Geo- and Eco-sciences, using Deep Learning
#
# ## Case Study: Detecting Intertidal Reefs
#
# #### Daniel Buscombe, MARDA Science
#
# ![](https://mardascience.com/wp-content/uploads/2019/06/cropped-MardaScience_logo-5.png)

# Before you do anything, go to `File > Save copy in Drive` so you can keep and work on your own copy
#
# ## Part 5: Train a oyster reef masker using a custom learning rate scheduler
#
# This jupyter notebook running on Google Colab is part of the "Advanced Binary Image Segmentation for the Geo- and Eco-sciences, using Deep Learning" course. The main course website can be accessed [here](https://mardascience.gitlab.io/binary_image_segmentation_for_geosciences/#/)
#
# Then we will construct, train and evaluate the model. In the [oysterNet paper](https://zslpublications.onlinelibrary.wiley.com/doi/full/10.1002/rse2.134), the authors used a bigger, more sophisticated model for `instance segmentation`, that is, semantic segmentation that is aware of all the different `instances` of the class (i.e. each individual piece of reef). The model they use is called `Mask RCNN`, the implementation of which is [here](https://github.com/matterport/Mask_RCNN). That is a large and very complicated model that is hard to experiment with. The research behind the [oysterNet paper](https://zslpublications.onlinelibrary.wiley.com/doi/full/10.1002/rse2.134) is state-of-the-art.
#
# Here, we use a simpler model with fewer parameters (namely, a residual UNet) and acheive acceptable results for a semantic segmentation (predicting the masks of where the reefs are, rather than by individual instances). The UNet is the same we used in a [previous set of tutorials](https://mardascience.gitlab.io/deep_learning_landscape_classification/#/) and is relatively simple to adapt and play with to demonstrate a few principles.
#
# **In this tutorial, we'll adapt our training strategies in search of a more satisfactory result. This may or may not be directly transferable with the same fidelity to your own dataset, but this is a good trick to learn and understand that might help for your data modelling purposes. This time we use a cyclical learning rate scheduler**
#
# According to [this](https://arxiv.org/abs/1506.01186) and [this](https://www.pyimagesearch.com/2019/07/29/cyclical-learning-rates-with-keras-and-deep-learning/), there are a few reasons why we would want to trial this strategy
# * a low learning rate may not be sufficient to break out of the non-optimal areas of the loss landscape and descend into areas of the loss landscape with lower loss.
# * our model and optimizer may be very sensitive to our initial learning rate choice
#
#
# This is designed to demonstrate a problem-solving strategy, and also a principle that often applies to natural imagery:
#
# > It's not just the model you choose, it's how you train it that counts
#
#
#

# ### Import libraries

# In[ ]:

import os #for accessing operating system utilities
from glob import glob #for finding files that match a certain string pattern
import matplotlib.pyplot as plt #for plotting
import numpy as np #for numerical operations
import random, string #for creating random strings
import tensorflow as tf #tensorflow
import json # for reading lable annotations in json format
import requests #for downloading files
from PIL import Image, ImageFilter #for reading and filtering imagery
import skimage.draw #for making masks (raster label images) from label annotations
from skimage.transform import resize #for resizing imagery
from psutil import virtual_memory #for interrogating our filesystem and RAM specifications
from imageio import imwrite

# In[ ]:

from tensorflow.keras.callbacks import ReduceLROnPlateau, EarlyStopping, ModelCheckpoint
from random import shuffle
from tensorflow.keras.layers import Input, Conv2D, MaxPooling2D, UpSampling2D, BatchNormalization
from tensorflow.keras.layers import Concatenate, Conv2DTranspose, Flatten, Activation, Add
from tensorflow.keras.models import Model

# ### Prepare the data

# #### Download the  imagery
#

# For the remainder of the tutorials, we'll be using this version of the data, consisting of imagery and image labels, with an augmented training set consisting of double the original number of imagery, half of which have been augmented with random flips, zoms and rotations

# In[ ]:

# from https://stackoverflow.com/questions/38511444/python-download-files-from-google-drive-using-url

def download_file_from_google_drive(id, destination):
    URL = "https://docs.google.com/uc?export=download"

    session = requests.Session()

    response = session.get(URL, params = { 'id' : id }, stream = True)
    token = get_confirm_token(response)

    if token:
        params = { 'id' : id, 'confirm' : token }
        response = session.get(URL, params = params, stream = True)

    save_response_content(response, destination)

def get_confirm_token(response):
    for key, value in response.cookies.items():
        if key.startswith('download_warning'):
            return value

    return None

def save_response_content(response, destination):
    """
    response = filename for input
    destination = filename for output
    """
    CHUNK_SIZE = 32768

    with open(destination, "wb") as f:
        for chunk in response.iter_content(CHUNK_SIZE):
            if chunk: # filter out keep-alive new chunks
                f.write(chunk)

# In[ ]:

file_id = '19RgkzaD9w-rvAF9uBpqMONnwF0lJVPh3'
destination = 'train_images.tar.gz'
download_file_from_google_drive(file_id, destination)

# In[ ]:

get_ipython().system('tar -xf train_images.tar.gz > tmp.txt')

# In[ ]:

file_id = '1gjWGXO7mtBhJuSOqrG6B1inzLF8TfTSX'
destination = 'val_images.tar.gz'
download_file_from_google_drive(file_id, destination)

# In[ ]:

get_ipython().system('tar -xf val_images.tar.gz > tmp.txt')

# In[ ]:

file_id = '1x8HrgEStkBCdBmFNo1xa0x6LCVE2GE1g'
destination = 'train_labels.tar.gz'
download_file_from_google_drive(file_id, destination)

# In[ ]:

get_ipython().system('tar -xf train_labels.tar.gz > tmp.txt')

# In[ ]:

file_id = '1v0kPgwZVQNrGAk0awJi31W83CYvHJbeX'
destination = 'val_labels.tar.gz'
download_file_from_google_drive(file_id, destination)

# In[ ]:

get_ipython().system('tar -xf val_labels.tar.gz > tmp.txt')

# In[ ]:

root = './content/'

# Check to see how many images we now have to work with

# In[ ]:

train_files = glob(root+"1kx1k_dataset/train_images/data/*.png")
val_files = glob(root+"1kx1k_dataset/val_images/data/*.png")

print("# train files: %i" % (len(train_files)))
print("# validation files: %i" % (len(val_files)))

# ### Setting up model training

# #### Custom batch generator

# We are going to feed images to the network in batches. The batch size (number of images and associated labels) will be ...

# In[ ]:

batch_size = 6

# The next function is an image batch generator. A `generator` is a special type of python function that `yields` a set of data. In our case, it will yield a set of `batch_size` images and labels drawn randomly from the entire set of `files` provided
#
# It opens the file, reads it into a numpy array with correct dimensions checked, and then does the same for the label image. It scales the image by dividing by 255, do turn the 8-bit data scaled between 0 and 255 to data scaled between 0 and 1. The labels are flattened from 3D (RGB) to 2D (greyscale integers)

# In[ ]:

def image_batch_generator(files, sz, batch_size = 4):

  while True: # this is here because it will be called repeatedly by the training function

    #extract a random subset of files of length "batch_size"
    batch = np.random.choice(files, size = batch_size)

    #variables for collecting batches of inputs (x) and outputs (y)
    batch_x = []
    batch_y = []

    #cycle through each image in the batch
    for f in batch:

        #preprocess the raw images
        raw = Image.open(f)
        raw = raw.resize(sz)
        raw = raw.filter(ImageFilter.UnsharpMask(radius=20, percent=100))
        raw = np.array(raw)

        #check the number of channels because some of the images are RGBA or GRAY
        if len(raw.shape) == 2:
            raw = np.stack((raw,)*3, axis=-1)

        else:
            raw = raw[:,:,0:3]

        #get the image dimensions, find the min dimension, then square the image off
        nx, ny, nz = np.shape(raw)
        n = np.minimum(nx,ny)
        raw = raw[:n,:n,:]

        raw[np.isnan(raw)] = 1e-5
        raw[np.isinf(raw)] = 1e-5

        batch_x.append(raw)

        #get the masks.
        maskfile = f.replace('_images','_labels').replace('.png','.jpg')
        mask = Image.open(maskfile)
        # the mask is 3-dimensional so get the max in each channel to flatten to 2D
        try:
           mask = np.max(np.array(mask.resize(sz)),axis=2)
        except:
           mask = np.array(mask.resize(sz))

        # water pixels are always greater than 170
        mask = (mask>170).astype('int') ##170 = (2/3)*255

        mask = mask[:n,:n]

        mask[np.isnan(mask)] = 1e-5
        mask[np.isinf(mask)] = 1e-5
        batch_y.append(mask)

    #preprocess a batch of images and masks
    batch_x = np.array(batch_x) #/255. #divide image by 255 to normalize
    batch_y = np.array(batch_y)
    batch_y = np.expand_dims(batch_y,1) #add singleton dimension to batch_y

    yield (batch_x, batch_y) #yield both the image and the label together

# We will specify input imagery of size `(512, 512, 3)`

# In[ ]:

sz = (512, 512)

# #### Build the model

#

# In[ ]:

def batchnorm_act(x):
    x = BatchNormalization()(x)
    return Activation("relu")(x)

def conv_block(x, filters, kernel_size=(3, 3), padding="same", strides=1):
    conv = batchnorm_act(x)
    return Conv2D(filters, kernel_size, padding=padding, strides=strides)(conv)

def bottleneck_block(x, filters, kernel_size=(3, 3), padding="same", strides=1):
    conv = Conv2D(filters, kernel_size, padding=padding, strides=strides)(x)
    conv = conv_block(conv, filters, kernel_size=kernel_size, padding=padding, strides=strides)

    bottleneck = Conv2D(filters, kernel_size=(1, 1), padding=padding, strides=strides)(x)
    bottleneck = batchnorm_act(bottleneck)

    return Add()([conv, bottleneck])

def res_block(x, filters, kernel_size=(3, 3), padding="same", strides=1):
    res = conv_block(x, filters, kernel_size=kernel_size, padding=padding, strides=strides)
    res = conv_block(res, filters, kernel_size=kernel_size, padding=padding, strides=1)

    bottleneck = Conv2D(filters, kernel_size=(1, 1), padding=padding, strides=strides)(x)
    bottleneck = batchnorm_act(bottleneck)

    return Add()([bottleneck, res])

def upsamp_concat_block(x, xskip):
    u = UpSampling2D((2, 2))(x)
    return Concatenate()([u, xskip])

def res_unet(sz, f):
    inputs = Input(sz)

    ## downsample
    e1 = bottleneck_block(inputs, f); f = int(f*2)
    e2 = res_block(e1, f, strides=2); f = int(f*2)
    e3 = res_block(e2, f, strides=2); f = int(f*2)
    e4 = res_block(e3, f, strides=2); f = int(f*2)
    _ = res_block(e4, f, strides=2)

    ## bottleneck
    b0 = conv_block(_, f, strides=1)
    _ = conv_block(b0, f, strides=1)

    ## upsample
    _ = upsamp_concat_block(_, e4)
    _ = res_block(_, f); f = int(f/2)

    _ = upsamp_concat_block(_, e3)
    _ = res_block(_, f); f = int(f/2)

    _ = upsamp_concat_block(_, e2)
    _ = res_block(_, f); f = int(f/2)

    _ = upsamp_concat_block(_, e1)
    _ = res_block(_, f)

    ## classify
    outputs = Conv2D(1, (1, 1), padding="same", activation="sigmoid")(_)

    #model creation
    model = Model(inputs=[inputs], outputs=[outputs])
    return model

# This is a class imbalanced problem, with many more `non-reef` pixels compared to `reef` pixels (the target class)
#
# So, we will use (Sorensen-)Dice loss instead of the more common IoU score (Jaccard Index) which is more suitable for class balanced problems

# In[ ]:

smooth = 1.

def dice_coef(y_true, y_pred):
    y_true_f = tf.reshape(tf.dtypes.cast(y_true, tf.float32), [-1])
    y_pred_f = tf.reshape(tf.dtypes.cast(y_pred, tf.float32), [-1])
    intersection = tf.reduce_sum(y_true_f * y_pred_f)
    return (2. * intersection + smooth) / (tf.reduce_sum(y_true_f) + tf.reduce_sum(y_pred_f) + smooth)

def dice_coef_loss(y_true, y_pred):
    return 1.0 - dice_coef(y_true, y_pred)

# Next, we make and compile our model
#
# As we saw in lesson 1, model compilation is a necessary step, involving specifiying the 'optimizer' (we will use `rmsprop` but `adam` is also a good one to use, in my experience). The loss function is the Dice loss, and the metric we want to keep track of in the dice coefficient

# In[ ]:

model = res_unet(sz+(3,), batch_size)
model.compile(optimizer = 'adam', loss = dice_coef_loss, metrics = [dice_coef])

# Let's take a look at how many parameters we have to optimize

# In[ ]:

model.summary()

# We have over 600,000 trainable parameters. This is a lot, for any model, but tiny compared to the number of parameters in some deep learning models, with up to hundreds of millions of trainable parameters

# ### Train the model with a custom "learning rate scheduler"

# In[ ]:

callbacks = tf.keras.callbacks
backend = tf.keras.backend

class PlotLearning(callbacks.Callback):

    def on_train_begin(self, logs={}):
        self.i = 0
        self.x = []
        self.losses = []
        self.val_losses = []
        self.acc = []
        self.val_acc = []
        #self.fig = plt.figure()
        self.logs = []
    def on_epoch_end(self, epoch, logs={}):
        self.logs.append(logs)
        self.x.append(self.i)
        self.losses.append(logs.get('loss'))
        self.val_losses.append(logs.get('val_loss'))
        self.acc.append(logs.get('dice_coef'))
        self.val_acc.append(logs.get('val_dice_coef'))
        self.i += 1
        print('i=',self.i,'loss=',logs.get('loss'),'val_loss=',logs.get('val_loss'),'dice_coef=',logs.get('dice_coef'),'val_dice_coef=',logs.get('val_dice_coef'))

        #choose a random test image and preprocess
        f = np.random.choice(val_files)

        raw = np.array(Image.open(f).resize(sz))

        #predict the mask
        pred = 255*model.predict(np.expand_dims(raw, 0)).squeeze()
        print(np.max(pred))

        #mask post-processing
        msk  = (pred>170).astype('int')

        msk = np.stack((msk,)*3, axis=-1)

        #show the mask and the segmented image
        combined = np.concatenate([raw, msk, raw* msk], axis = 1)
        plt.axis('off')
        plt.imshow(combined)
        plt.show()
        #plt.savefig(str(self.i)+'.png', dpi=100, bbox_inches='tight')

# In Part 3 we used callback functions that adaptively change the pace of training (using an adaptive learning rate), called ["reduce loss on plateau"](https://www.tensorflow.org/api_docs/python/tf/keras/callbacks/ReduceLROnPlateau)
#
# This time we'll control precisely the learning rate scheduler to invoke a function that describes the learning rate as a function of training epoch
#
# To do so requires a new class that will feed the model the correct learning rate based on the epoch it is current training on
#
#

# In[ ]:

class LearningRateScheduler(callbacks.Callback):
    def __init__(self,
                 schedule,
                 learning_rate=None,
                 steps_per_epoch=None,
                 verbose=0):
        super(LearningRateScheduler, self).__init__()
        self.learning_rate = learning_rate
        self.schedule = schedule
        self.verbose = verbose
        self.warmup_epochs = 0
        self.warmup_steps = 0
        self.global_batch = 0

    def on_train_batch_begin(self, batch, logs=None):
        self.global_batch += 1
        if self.global_batch < self.warmup_steps:
            if not hasattr(self.model.optimizer, 'lr'):
                raise ValueError('Optimizer must have a "lr" attribute.')
            lr = self.learning_rate * self.global_batch / self.warmup_steps
            backend.set_value(self.model.optimizer.lr, lr)
            if self.verbose > 0:
                print('\nBatch %05d: LearningRateScheduler warming up learning '
                      'rate to %s.' % (self.global_batch, lr))

    def on_epoch_begin(self, epoch, logs=None):
        if not hasattr(self.model.optimizer, 'lr'):
            raise ValueError('Optimizer must have a "lr" attribute.')
        lr = float(backend.get_value(self.model.optimizer.lr))

        if epoch >= self.warmup_epochs:
            try:  # new API
                lr = self.schedule(epoch - self.warmup_epochs, lr)
            except TypeError:  # old API
                lr = self.schedule(epoch - self.warmup_epochs)
            if not isinstance(lr, (float, np.float32, np.float64)):
                raise ValueError('The output of the "schedule" function '
                                 'should be float.')
            backend.set_value(self.model.optimizer.lr, lr)

            if self.verbose > 0:
                print('\nEpoch %05d: LearningRateScheduler reducing learning '
                      'rate to %s.' % (epoch + 1, lr))

    def on_epoch_end(self, epoch, logs=None):
        logs = logs or {}
        logs['lr'] = backend.get_value(self.model.optimizer.lr)

# In[ ]:

# This is our cosine function implementation for a 'cyclical' learning rate between two specified bounds, `min_lr` (minimum learning rate) and `max_lr` (maximum learning rate)

# In[ ]:

def cosine_ratedecay(max_epochs, max_lr, min_lr=1e-6):
    """
    cosine scheduler.
    :param max_epochs: max epochs
    :param max_lr: max lr
    :param min_lr: min lr
    :return: current lr
    """
    max_epochs = max_epochs ##- 5 if warmup else max_epochs

    def ratedecay(epoch):
        lrate = min_lr + (max_lr - min_lr) * (
                1 + np.cos(np.pi*2 * epoch / max_epochs)) / 2

        return lrate

    return ratedecay

# We now have two new hyperparameters we didn't have before

# In[ ]:

# maximum learning rate (lambda)
max_lr = 1e-4

max_epochs = 200

# The "two bites at the cherry" learning rate scheduler. The intuition behind this is that the model searches progressively finer portions of the loss parameter space, up until about epoch 70, then the function increases again, using larger learning rates to explore other 'valleys' in the loss landscape.
#
# I informally call it the 'two bites at the cherry' function because it facilitates the following situation: the optimal low in the loss landscape was not effectively found by the first loop of progressively lower learning rates. Giving the model a 'second bite' allows it to find a new low by adjusting learning rates to be coarse, so it can hop to different areas of the loss landscape more effectively

# In[ ]:

min_lr = 1e-6
def ratedecay(epoch):
    lrate = min_lr + (max_lr - min_lr) * (
            1 + np.cos(np.pi*3 * epoch / max_epochs)) / 2

    return lrate

plt.plot(ratedecay(np.arange(max_epochs)))

# Now we have defined all these hyperparameters (parameters that we chose, not automatically determined by the model), we can make a function that builds all the various callbacks together into a list, which can then be passed to the model training (`.fit()` function)

# In[ ]:

def build_callbacks(filepath, lr_ratedecay, lr, steps_per_epoch):

    # set checkpoint file
    model_checkpoint = ModelCheckpoint(filepath, monitor='val_loss',
                                   verbose=0, save_best_only=True, mode='min',
                                   save_weights_only = True)

    # learning rate scheduler setting
    learning_rate_scheduler = LearningRateScheduler(lr_ratedecay, lr, steps_per_epoch,
                                                verbose=1)

    callbacks = [model_checkpoint, learning_rate_scheduler, PlotLearning()]

    return callbacks

# Next, set up a name for the `.h5` file that will be used to store model weights.
#
# Finally, we train the model by calling the `.fit()` command and providing all the generators and hyperparameters defined in the callbacks
#
# The number of training and validation steps is simply the number of respective files divided by the batch size

# In[ ]:

filepath = 'reefs_weights_uresnet'+str(batch_size)+'_'+str(max_epochs)+'epochs_cosine_lrscheduler.h5'

train_generator = image_batch_generator(train_files, sz, batch_size = batch_size)
val_generator  = image_batch_generator(val_files, sz, batch_size = batch_size)
train_steps = len(train_files) //batch_size
val_steps = len(val_files) //batch_size
steps_per_epoch = len(train_files) // batch_size

print(train_steps)
print(val_steps)
print(steps_per_epoch)

lr_ratedecay = cosine_ratedecay(max_epochs,max_lr)

# If you'd rather not wait the 3-6 hours it might take to train the model you should leave the following cell uncommented

# In[ ]:

file_id = '1LyjyOxuzYzmWrwh6kxPVDZZWEzZGehNq'
destination = 'weights_part5.h5'
download_file_from_google_drive(file_id, destination)
model.load_weights(destination)
max_epochs = 5 #just use 5 epochs

# If you wish to train the model your own model, uncomment the cells below

# In[ ]:

# hist = model.fit(train_generator,
#                 epochs = max_epochs, steps_per_epoch = train_steps,
#                 validation_data = val_generator, validation_steps = val_steps,
#                 callbacks = build_callbacks(filepath,lr_ratedecay, max_lr, steps_per_epoch))

# #### Plot the training history

# In the above, we gave an output variable to the `.fit()` command. This contains the training histories. That is, losses and metrics as a function of epoch. You can access the variables in the dictionary like so

# In[ ]:

# hist.history.keys()

# Let's make a plot of the histories of both train and validation losses and dice coefficients, and also the history of the learning rate

# In[ ]:

# plt.figure(figsize=(20,10))
# plt.subplot(131)
# plt.plot(hist.history['dice_coef'], 'b', label='train Dice coefficient')
# plt.plot(hist.history['val_dice_coef'], 'k', label='validation Dice coefficient')
# plt.xlabel('Epoch number'); plt.ylabel('Dice coefficent')
# plt.legend()

# plt.subplot(132)
# plt.plot(hist.history['loss'], 'b', label='train loss')
# plt.plot(hist.history['val_loss'], 'k', label='validation loss')
# plt.xlabel('Epoch number'); plt.ylabel('Loss')
# plt.legend()

# plt.subplot(133)
# plt.plot(hist.history['lr'], 'g')
# plt.xlabel('Epoch number'); plt.ylabel('Learning rate')

# ### Test the model

# Get the test set of files and untar like we did the other sets

# In[ ]:

file_id = '1lL6cbUNhwAQsDl4P4LIYiOMOwKxrhq-c'
destination = 'test_images.tar.gz'
download_file_from_google_drive(file_id, destination)

# In[ ]:

get_ipython().system('tar -xf test_images.tar.gz > tmp.txt')

# In[ ]:

file_id = '1GVZgjzuVatp-OjU5s5DJOgAteUvN1yDC'
destination = 'test_labels.tar.gz'
download_file_from_google_drive(file_id, destination)

# In[ ]:

get_ipython().system('tar -xf test_labels.tar.gz > tmp.txt')

# Get a test generator

# In[ ]:

test_files = glob(root+"1kx1k_dataset/test_images/*.png")

test_generator = image_batch_generator(test_files, sz, batch_size = batch_size)

print("# test files: %i" % (len(test_files)))

# Use the `.evaluate()` function of the model to get average Dice scores and losses for the test set
#
#

# In[ ]:

# some other training parameters
steps = len(test_files) // batch_size

# testing
scores = model.evaluate(test_generator, steps=steps)

print('loss={loss:0.4f}, Mean Dice={dice_coef:0.4f}'.format(loss=scores[0], dice_coef=scores[1]))

# In the [oysterNet paper](https://zslpublications.onlinelibrary.wiley.com/doi/full/10.1002/rse2.134), the authors reeport a best precision of 0.771 and recall of 0.772. The Dice score is mathematically equivalent to the F1 score, the harmonic mean of precision and recall, or 77%
#
# In Part 4, we acheived an average dice score of ~72% using a Dice loss function and adaptive learning rates that were contingent on validation loss
#
# This time, with a cyclical learning rate, we acheive ~80% - approximately the same as the original paper (actually, a little higher!)


r""".. role:: html(raw)
   :format: html

Unitary Designs
===============

.. meta::
    :property="og:description": Learn about designs and their uses in quantum computing.

    :property="og:image": https://pennylane.ai/qml/_images/fano.png

.. related::

    tutorial_haar_measure Understanding the Haar measure

*Author: PennyLane dev team. Posted: 7 Sept 2021. Last updated: 7 Sept 2021.*

.. note::

   This demo is intended to be a sequel to the
   :doc:`demo about the Haar measure </demos/tutorial_haar_measure>`.
   If you are not familiar with the Haar measure, we recommend going through
   that demo first before exploring this one.

Take a close look at the following mathematical object:

.. figure:: /demonstrations/unitary_designs/fano_no_labels.svg
   :align: center
   :width: 30%

|

There are many things we can say about it: it consists of seven points and seven
lines (the circle counts as a line); each line contains three points, and each
point is contained in three lines. Furthermore, any pair of points occur
together in exactly one line. This object, called the `Fano plane
<https://en.wikipedia.org/wiki/Fano_plane>`__, is an instance of a mathematical
structure called a `projective plane
<https://en.wikipedia.org/wiki/Projective_plane>`__, which is just one example
of a `combinatorial design
<https://en.wikipedia.org/wiki/Combinatorial_design>`__. Designs are sets of
objects, and groups of those objects, that satisfy certain balance properties
and symmetries. They have been studied for hundreds of years in a huge variety
of contexts [#Handbook]_, from `error correcting codes
<http://citeseerx.ist.psu.edu/viewdoc/summary?doi=10.1.1.50.5465>`__, to `card
games <https://homepages.warwick.ac.uk/staff/D.Maclagan/papers/set.pdf>`__, and
even `agriculture
<http://www-groups.mcs.st-and.ac.uk/~rab/histLShand.pdf>`__. So, what about
quantum computing?

Designs are actually quite prevalent in quantum computing. You've almost
certainly come across one before, though you may not have realized it. At the
end of the Haar measure demo, we asked a very important question: "do we always
*need* to sample from the full Haar measure?". The answer to this is "no", and
the reasoning lies in the study of *unitary designs*.

In this demo, you'll learn the definition of :math:`t`-designs, what it means to
generalize them to unitary :math:`t`-designs, and you'll see some canonical
examples of designs in quantum computing. You'll also learn about their
connection with the Haar measure, what it means to *twirl* a quantum channel,
and explore how to leverage 2-designs in PennyLane to compute the average
fidelity of a noisy channel. You will experience directly a situation where we can
use a :math:`t`-design as a shortcut over the full Haar measure to greatly improve
the efficiency of a task 🎉.

From spheres to unitary :math:`t`-designs
-----------------------------------------

Spherical designs
^^^^^^^^^^^^^^^^^

Before diving into unitary designs, let's look at the sphere for some
intuition.  Suppose we have a polynomial in :math:`d` variables, and we would
like to compute its average over the surface of a real, :math:`d`-dimensional
unit sphere, :math:`S(R^d)`. We can do so by integrating that function over the
sphere (using the proper measure), but that would be a lot of parameters to
keep track of.

One could alternatively approximate the average by sampling thousands of points
uniformly at random on the sphere, evaluating the function at those points, and
computing their average value. That will always work, and it will get us close,
but it will not be exact.

In fact, both of those approaches may be overkill in some special cases---if the
terms in the polynomial have the same degree of at most :math:`t`, you can
compute the average **exactly** over the sphere using only a small set of points
rather than integrating over the entire sphere.  That set of points is called a
spherical :math:`t`-design. More formally [#Handbook]_, [#Delsarte]_:

.. admonition:: Definition
    :class: defn

    Let :math:`p_t: \mathcal{S}(R^d)\rightarrow R` be a polynomial in :math:`d`
    variables, with all terms homogeneous in degree at most :math:`t`. A
    set :math:`X = \{x: x \in \mathcal{S}(R^d)\}` is a spherical :math:`t`-design if

    .. math::

        \frac{1}{|X|} \sum_{x \in X} p_t(x) = \int_{\mathcal{S}(R^d)} p_t (u) d\mu(u)

    holds for all possible :math:`p_t`, where :math:`d\mu` is the uniform,
    normalized spherical measure. A spherical :math:`t`-design is also a
    :math:`k`-design for all :math:`k < t`.

Now this is a pretty abstract picture, so let's consider the 3-dimensional
sphere. This definition tells us that if we want to take the average of a
polynomial over a sphere where all terms have the same degree of at most 2, we
can do so using a small, representative set of points called a 2-design,
rather than the whole sphere. Similarly, if all terms of the polynomial have the
same degree of at most 3, we could use a 3-design.

But what are these representative sets of points?  Since we are using these
points as a stand-in for averaging over the whole sphere, we'd want the points
in the set to be distributed in a way that provides sufficient "coverage". In
the 3-dimensional case, the vertices of some familiar solids form :math:`t`-designs
[#Handbook]_, [#sph4design]_:

.. figure:: /demonstrations/unitary_designs/shapes.svg
   :align: center
   :width: 80%

|

We see from these illustrations that spherical designs are sets of
*evenly-spaced points*. As :math:`t` increases, the configurations become
increasingly sphere-like. Looking at this in a different way, the more complex a
function becomes as its degree increases, the closer the :math:`t`-design must
be to a sphere; we need to evaluate the function at more points in order to gain
sufficient information when a function is varying more quickly due to a higher
degree. In 3 dimensions, we can compute the average of a polynomial with degree
2 by evaluating it only at the points of a tetrahedron, despite the fact that it
doesn't look spherical at all. More complex functions require more points
and thus more intricate configurations for the design.  Spherical designs exist
for all :math:`t` and dimension :math:`d` [#Handbook]_. They are not always
unique, and may have varying numbers of points.

To show that this really works, let's look at an explicit example. Consider the
following polynomial in 3 variables:

.. math::

   f(x, y, z) = x^4 - 4 x^3 y + y^2 z^2

We can compute the average value of :math:`f` by integrating over a unit sphere:
the result is :math:`4/15 \approx 0.26667`. However, this integral is
non-trivial to evaluate by hand; the most straightforward way is to convert to
polar coordinates, and even then, it involves integrating functions with 4th and
5th powers of trigonometric functions.

Instead, this is a case where we can leverage the fact that all terms in the
polynomial have degree 4, and compute the average exactly using only a subset of
points that form a 4-design. We choose a dodecahedron for convenience; while
this is actually a 5 design, it also forms a 4-design, and is a more familiar
shape than the 4-design depicted above.

First, we define the set of points that comprise a dodecahedron:

"""

import numpy as np

# The golden ratio
g = (1 + np.sqrt(5)) / 2

# A dodecahedron has 20 points
dodecahedron = np.array([
    # 8 of them form a cube within the sphere
    [1, 1, 1], [1, 1, -1], [1, -1, 1], [1, -1, -1],
    [-1, 1, 1], [-1, 1, -1], [-1, -1, 1], [-1, -1, -1],

    # 4 of them form a rectangle within the y-z plane
    [0, g, 1/g], [0, g, -1/g], [0, -g, 1/g], [0, -g, -1/g],

    # 4 of them form a rectangle within the x-z plane
    [1/g, 0, g], [1/g, 0, -g], [-1/g, 0, g], [-1/g, 0, -g],

    # 4 of them form a rectangle within the x-y plane
    [g, 1/g, 0],[g, -1/g, 0], [-g, 1/g, 0], [-g, -1/g, 0],
])

# Normalize the points so they all fit in the unit sphere
dodecahedron = np.array(
   [point / np.linalg.norm(point) for point in dodecahedron]
)

######################################################################
# Now we define our function and compute the average over the dodecahedron:

def f(x, y, z):
    return (x ** 4) - 4 * (x ** 3) * y +  y ** 2 * z ** 2

dodeca_average = np.mean([f(*point) for point in dodecahedron])
print(dodeca_average)

######################################################################
# This is exactly the value we expect. What happens if we try to do this using
# only a 3-design, the cube?

# The first 8 points of the dodecahedron are a cube
cube = dodecahedron[:8]

cube_average = np.mean([f(*point) for point in cube])
print(cube_average)

######################################################################
# This clearly differs from the true value. We need a design with :math:`t=4`
# or better in order to compute this average, and when such a design is
# available, we may save significant computational time.
#
# Unitary designs
# ^^^^^^^^^^^^^^^
#
# We've learned now that spherical designs are sets of evenly-spaced points, and
# saw how they can be used as a shortcut to evaluate the average of a
# polynomial up to a given degree :math:`t`. However, there was nothing quantum
# about this; there weren't even any complex numbers involved. A *unitary
# design* extends this concept from evenly-distributed points to
# evenly-distributed unitaries.  More formally, instead of averaging polynomials
# over spheres, we consider polynomials that are functions of the entries of
# unitary matrices [#Dankert]_, [#Gross]_.
#
# .. admonition:: Definition
#     :class: defn
#
#     Let :math:`P_{t,t}(U)` be a polynomial with homogeneous degree at most :math:`t` in
#     :math:`d` variables in the entries of a unitary matrix :math:`U`, and degree
#     :math:`t` in the complex conjugates of those entries. A unitary
#     :math:`t`-design is a set of :math:`K` unitaries :math:`\{U_k\}` such that
#
#     .. math::
#
#         \frac{1}{K} \sum_{k=1}^{K} P_{t,t}(U_k) = \int_{\mathcal{U}(d)}
#         P_{t,t} (U) d\mu(U)
#
#     holds for all possible :math:`P_{t,t}`, and where :math:`d\mu` is the
#     uniform *Haar measure*.
#
# We stress again that this expression is **exact**. The unitaries in a unitary
# design are a representative set of points that are "evenly spaced" across the
# unitary group. With just a subset of the full group, we can evaluate complex
# expressions that would be otherwise intractable.
#
# A surprising result about unitary designs is that they exist for all possible
# combinations of :math:`t` and :math:`d` [#Roy]_. There are some known lower
# bounds for the number of unitaries required; for example, a 2-design in
# dimension :math:`d` has at least :math:`d^4 - 2d^2 + 2` elements [#Gross]_,
# [#Roy]_.  However, actually finding the sets (and in particular, finding ones
# with minimal size), is a challenging problem [#Bannai]_, though very recently
# some constructions have been put forward [#Nakata]_.
#
#
# .. admonition:: Fun fact
#
#     Applying the elements of a unitary design to a fixed pure state produces a
#     set of vectors that form a *complex projective design* [#DankertThesis]_.
#     These are much like spherical designs, but they live in a complex vector
#     space. If you've ever studied the characterization of quantum systems, you
#     may have come across some special sets of measurements called mutually
#     unbiased bases (MUBs), or symmetric, informationally complete positive
#     operator valued measurements (SIC-POVMs). Both of these sets of vectors
#     are complex projective 2-designs [#Klappenecker]_.
#
#     .. figure:: /demonstrations/unitary_designs/sic-povm.svg
#        :align: center
#        :width: 80%
#
#        The vectors of the simplest SIC-POVM in dimension 2, plotted on a Bloch sphere.
#
# Unitary :math:`t`-designs in action
# -----------------------------------
#
# Unitary designs come into play in applications that require randomization, or
# sampling of random unitaries---essentially, they can be used as a stand-in for
# the Haar measure. The way in which the unitaries are used in the application may
# place restrictions on the value of :math:`t` that is required; arguably the most
# common is the unitary 2-design.
#
# While in general unitary designs are hard to construct, there are well known
# results for unitary 1-, 2-, and 3-designs based on familiar objects in quantum
# computing. Before we see what those are, let's explore an important use case.
#
# Average fidelity
# ^^^^^^^^^^^^^^^^
# A key application of unitary 2-designs is benchmarking quantum
# operations. Suppose we have a noisy quantum channel :math:`\Lambda` that should
# perform something close to the unitary operation :math:`V`.  What can we say
# about the performance of this channel?
#
# One metric of interest is the *fidelity*. Consider the state :math:`|0\rangle`.
# In an ideal case, we apply :math:`V` and obtain :math:`V|0\rangle`.  But applying the
# channel :math:`\Lambda` gives us something a little different. Since it's noisy,
# we must consider the state as a density matrix. The action of :math:`\Lambda` on
# our starting state is :math:`\Lambda(|0\rangle \langle 0|)`.  If :math:`\Lambda`
# was perfect, then :math:`\Lambda(|0\rangle \langle 0|) = V|0\rangle \langle
# 0|V^\dagger`, and the fidelity is
#
# .. math::
#
#     F(\Lambda, V) = \langle 0 | V^\dagger \cdot \Lambda(|0\rangle \langle 0|) \cdot V|0\rangle = 1.
#
# In reality, :math:`\Lambda` is not going to implement :math:`V` perfectly, and
# :math:`F < 1`. More importantly though, all we've computed so far is the fidelity when
# the initial state is :math:`|0\rangle`. What if the initial state is something
# different? What is the fidelity *on average*?
#
# To compute an average fidelity, we must do so with respect to the full set
# of Haar-random states. We usually generate random states by applying a
# Haar-random unitary :math:`U` to :math:`|0\rangle`. Thus to compute the average
# over all such :math:`U` we must evaluate
#
# .. math::
#
#     \bar{F}(\Lambda, V) = \int_{\mathcal{U}} d\mu(U) \langle 0 | U^\dagger V^\dagger \Lambda(U |0\rangle \langle 0| U^\dagger) V U |0\rangle.
#
# This is known as *twirling* the channel :math:`\Lambda`. Computing the average
# fidelity in this way would be a nightmare---we'd have to compute the fidelity
# with respect to an infinite number of states!
#
# However, consider the expression in the integral above. We have an inner product
# involving two instances of :math:`U`, and two instances of
# :math:`U^\dagger`. This means that the expression is a polynomial of degree 2 in
# both the elements of :math:`U` and its complex conjugates---this matches exactly
# the definition of a unitary 2-design. This means that if we can find a set of
# :math:`K` unitaries that form a 2-design, we can compute the average fidelity
# using only a finite set of initial states:
#
# .. math::
#
#     \frac{1}{K} \sum_{j=1}^K \langle 0 | U_j^\dagger V^\dagger \Lambda(U_j |0\rangle \langle 0|
#     U_j^\dagger) V^\dagger U_j |0\rangle = \int_{\mathcal{U}} d\mu(U) \langle 0
#     | U^\dagger V^\dagger \Lambda(U |0\rangle \langle 0| U^\dagger) V U |0\rangle
#
# This is great, but a question remains: what is the representative set of unitaries?
#
# The Clifford group
# ^^^^^^^^^^^^^^^^^^
#
# A beautiful result in quantum computing is that some special groups you may
# already be familiar with are unitary designs:
#
# - the Pauli group forms a unitary 1-design, and
# - the Clifford group forms a unitary 3-design.
#
# By the definition of designs, this means the Clifford group is also a 1-
# and 2-design.
#
# The :math:`n`-qubit Pauli group, :math:`\mathcal{P}(n)`, is the set of all tensor
# products of Pauli operations :math:`X`, :math:`Y`, :math:`Z`, and :math:`I`. The
# :math:`n`-qubit Clifford group, :math:`\mathcal{C}(n)`, is the *normalizer* of the
# Pauli group. In simpler terms, the Clifford group is the set of operations that
# send Paulis to Paulis (up to a phase) under conjugation i.e.,
#
# .. math::
#
#    C P C^\dagger = \pm P^\prime, \quad \forall P, P^\prime \in \mathcal{P}(n), \quad C \in \mathcal{C}(n).
#
# The Clifford group has some profoundly interesting properties and countless uses
# across quantum computing, from circuit compilation to error correcting
# codes. For a single qubit, the group is built from just two operations. One is
# the Hadamard:
#
# .. math::
#
#    H X H^\dagger = Z, \quad H Y H^\dagger = -Y, \quad H Z H^\dagger = X.
#
# This clearly maps Paulis to Paulis (up to a phase). The other is the phase gate :math:`S`:
#
# .. math::
#
#    S X S^\dagger = Y, \quad S Y S^\dagger = -X, \quad S Z S^\dagger = Z.
#
# If both :math:`H` and :math:`S` map Paulis to Paulis, then products of them do
# as well. In group theory terms, the single-qubit Clifford group is
# generated by :math:`H` and :math:`S`.  For example, consider the action of
# :math:`HS`:
#
# .. math::
#
#    (HS) X (HS)^\dagger = -Y, \quad (HS) Y (HS)^\dagger = -Z, \quad (HS) Z (HS)^\dagger = X.
#
# Since :math:`Y = iXZ`, it is enough to specify Clifford operations by how they
# act on :math:`X` and :math:`Z`.  For a particular Clifford, there are 6 possible
# ways it can transform :math:`X`, namely :math:`\pm X, \pm Y`, or :math:`\pm Z`.  Once
# that is determined, there are four remaining options for the transformation of
# :math:`Z`, leading to 24 elements total.
#
# It takes some work, but you can take combinations of :math:`H` and :math:`S`
# and evaluate their action on :math:`X` and :math:`Z` (or look at their matrix
# representations) until you find all 24 unique elements. The results of
# this endeavour are expressed below as strings:

single_qubit_cliffords = [
 '',
 'H', 'S',
 'HS', 'SH', 'SS',
 'HSH', 'HSS', 'SHS', 'SSH', 'SSS',
 'HSHS', 'HSSH', 'HSSS', 'SHSS', 'SSHS',
 'HSHSS', 'HSSHS', 'SHSSH', 'SHSSS', 'SSHSS',
 'HSHSSH', 'HSHSSS', 'HSSHSS'
]

######################################################################
# To see for yourself how this set of unitaries is evenly distributed, try
# applying each of the Cliffords to the initial state :math:`|0\rangle`, and
# plot the resulting states on the Bloch sphere. You'll find they are
# symmetric and evenly spaced; in fact, they are all eigenstates of :math:`X`,
# :math:`Y`, and :math:`Z`. Furthermore, under the full group action, the result
# is balanced in the sense that each eigenstate is obtained the same number of
# times.
#
# The multi-qubit Clifford group can also be
# specified by only a small set of generators (in fact, only one more
# than is needed for the single-qubit case). Together, :math:`H`, :math:`S`, and
# CNOT (on every possible qubit or pair of qubits) generate the :math:`n`-qubit
# group. Be careful though---the size of the group increases exponentially. The
# 2-qubit group alone has 11520 elements! The size can be worked out in a manner
# analogous to that we used above in the single qubit case: by looking at the
# combinatorics of the possible ways the gates can map Paulis with only
# :math:`X` and :math:`Z` to other Paulis.
#
# An experiment
# ^^^^^^^^^^^^^
# The whole idea of unitary designs may sound too good to be true. Can we
# *really* compute the exact average fidelity using just 24 operations? In this
# section, we put them to the test: we'll compute the average fidelity of an
# operation first with experiments using a large but finite amount of Haar-random
# unitaries, and then again with only the Clifford group.

import pennylane as qml

# Scipy allows us to sample Haar-random unitaries directly
from scipy.stats import unitary_group

# set the random seed
np.random.seed(42)

# Use the mixed state simulator
dev = qml.device("default.mixed", wires=1)

######################################################################
# Let's set up a noisy quantum channel. To keep things simple, assume it
# consists of applying :class:`~.pennylane.SX`, the square-root of
# :math:`X` gate, followed by a few different types of noise. First, write a
# quantum function for our ideal experiment:

def ideal_experiment():
    qml.SX(wires=0)
    return qml.state()

######################################################################
# Next, we apply some noise. We do so by making use of a relatively new feature
# in PennyLane called `quantum function transforms <https://pennylane.readthedocs.io/en/latest/code/qml_transforms.html>`__. Such transforms work by
# modifying the underlying, low-level quantum tapes which queue the quantum
# operations. Suppose the noisy channel is composed of the following:

def noisy_operations(damp_factor, depo_factor, flip_prob):
    qml.AmplitudeDamping(damp_factor, wires=0)
    qml.DepolarizingChannel(depo_factor, wires=0)
    qml.BitFlip(flip_prob, wires=0)

######################################################################
# Let's create a transform that applies this noise to any quantum function
# *after* the original operations, but before the measurements.  We use the
# convenient :func:`~.pennylane.transforms.qfunc_transform` decorator:

@qml.qfunc_transform
def apply_noise(tape, damp_factor, depo_factor, flip_prob):
    # Apply the original operations
    for op in tape.operations:
        qml.apply(op)

    # Apply the noisy sequence
    noisy_operations(damp_factor, depo_factor, flip_prob)

    # Apply the original measurements
    for m in tape.measurements:
        qml.apply(m)

######################################################################
# We can now apply this transform to create a noisy version of our ideal
# quantum function:

# The strengths of various types of noise
damp_factor = 0.02
depo_factor = 0.02
flip_prob = 0.01

noisy_experiment = apply_noise(damp_factor, depo_factor, flip_prob)(ideal_experiment)

######################################################################
# The last part of the experiment involves applying a random unitary matrix
# before all the operations, and its inverse right before the measurements.  We
# can write another transform here to streamline this process:

@qml.qfunc_transform
def conjugate_with_unitary(tape, matrix):
    qml.QubitUnitary(matrix, wires=0)

    for op in tape.operations:
        qml.apply(op)

    qml.QubitUnitary(matrix.conj().T, wires=0)

    for m in tape.measurements:
        qml.apply(m)

######################################################################
# Finally, in order to perform a comparison, we need a function to compute the
# `fidelity <https://en.wikipedia.org/wiki/Fidelity_of_quantum_states>`__
# compared to the ideal operation.

from scipy.linalg import sqrtm

def fidelity(rho, sigma):
    # Inputs rho and sigma are density matrices
    sqrt_sigma = sqrtm(sigma)
    fid = np.trace(sqrtm(sqrt_sigma @ rho @ sqrt_sigma))
    return fid.real

######################################################################
# Let's now compute the average fidelity, averaging over 50000 Haar-random unitaries:

n_samples = 50000

fidelities = []

for _ in range(n_samples):
    # Select a Haar-random unitary
    U = unitary_group.rvs(2)

    # Apply transform to construct the ideal and noisy quantum functions
    conjugated_ideal_experiment = conjugate_with_unitary(U)(ideal_experiment)
    conjugated_noisy_experiment = conjugate_with_unitary(U)(noisy_experiment)

    # Use the functions to create QNodes
    ideal_qnode = qml.QNode(conjugated_ideal_experiment, dev)
    noisy_qnode = qml.QNode(conjugated_noisy_experiment, dev)

    # Execute the QNodes
    ideal_state = ideal_qnode()
    noisy_state = noisy_qnode()

    # Compute the fidelity
    fidelities.append(fidelity(ideal_state, noisy_state))

fid_mean = np.mean(fidelities)
print(f"Mean fidelity = {fid_mean}")

######################################################################
# Now let's repeat the procedure using only Clifford group elements. First, we
# write a quantum function that performs a Clifford operation (or its inverse)
# based on its string representation.

def apply_single_clifford(clifford_string, inverse=False):
    for gate in clifford_string:
        if gate == 'H':
            qml.Hadamard(wires=0)
        else:
            sign = -1 if inverse else 1
            qml.PhaseShift(sign * np.pi/2, wires=0)

######################################################################
# Next, we write a transform that applies a Clifford in the context of the full
# experiment, i.e., apply the Clifford, then the operations, followed by the
# inverse of the Clifford.

@qml.qfunc_transform
def conjugate_with_clifford(tape, clifford_string):
    apply_single_clifford(clifford_string, inverse=False)

    for op in tape.operations:
        qml.apply(op)

    apply_single_clifford(clifford_string, inverse=True)

    for m in tape.measurements:
        qml.apply(m)

######################################################################
# You may have noticed this transform has exactly the same form as
# ``conjugate_with_unitary`` from above. Only the input type has changed, since
# the application of Cliffords here is specified by their string representation.
#
# It's now time to run the experiments:

fidelities = []

for C in single_qubit_cliffords:
    conjugated_ideal_experiment = conjugate_with_clifford(C)(ideal_experiment)
    conjugated_noisy_experiment = conjugate_with_clifford(C)(noisy_experiment)

    ideal_qnode = qml.QNode(conjugated_ideal_experiment, dev)
    noisy_qnode = qml.QNode(conjugated_noisy_experiment, dev)

    ideal_state = ideal_qnode()
    noisy_state = noisy_qnode()

    fidelities.append(fidelity(ideal_state, noisy_state))

######################################################################
# Let's see how our results compare to the earlier simulation:

clifford_fid_mean = np.mean(fidelities)

print(f"Haar-random mean fidelity = {fid_mean}")
print(f"Clifford mean fidelity    = {clifford_fid_mean}")

######################################################################
# Incredible 🤯 🤯 🤯 We were able to compute the average fidelity using only
# 24 experiments. Furthermore, the mean fidelity obtained from the Clifford
# experiments is **exact**; even with 50000 Haar-random experiments, we see
# deviations starting a few decimal places in. Consider the resources that would
# be saved if you were actually implementing this in a lab! It's not hard to see
# why the Clifford group plays such an important role in characterization
# procedures.
#
# Conclusion
# ----------
#
# In this demo, we've barely scratched the surface of designs and their
# applications in quantum computing. While benchmarking is a key application
# area, there are many others.  The Pauli group as a unitary 1-design has
# applications in the construction of private quantum channels [#PQC]_.
# *Approximate* unitary :math:`t`-designs (where the equality in the definition
# is replaced by approximately equal up to some finite precision) are also of
# interest, as there ways to construct them that are more efficient than those
# of exact designs [#Dankert]_. In particular, it has been shown that
# approximate complex projective 4-designs have applications to the state
# discrimination problem [#Ambainis]_.
#
# Furthermore, unitary designs are not the only designs that you'll encounter
# in quantum computing. The familiar Hadamard gate is just a 2-dimensional
# example of a broader family of *Hadamard designs*, on which there has been
# extensive research [#Seberry]_. Some sets of `mutually orthogonal Latin
# squares <https://en.wikipedia.org/wiki/Mutually_orthogonal_Latin_squares>`__
# have a direct correspondence with mutually unbiased bases, which are optimal
# quantum measurements [#Gaeta]_, as well as complex projective designs; and Latin
# squares themselves have direct correspondence with affine and projective
# planes, bringing us full circle back to the Fano plane from which we began.
#
#
# .. figure:: /demonstrations/unitary_designs/affine-latin.svg
#    :align: center
#    :width: 80%
#
#    An affine plane, Hadamard matrix, and a depiction of mutually orthogonal Latin squares.
#
# References
# ----------
#
# .. [#Handbook]
#
#     C. J. Colbourn and J. H. Dinitz (2006) *Handbook of Combinatorial Designs,
#     Second Edition*.  Chapman & Hall/CRC.
#
# .. [#Delsarte]
#
#    P. Delsarte, J.M. Goethals, J.J. Seidel (1977) *Spherical Codes and Designs*. Geometriae
#    Dedicata 6 363-388.
#
# .. [#sph4design]
#
#    R. H. Hardin and N. J. A. Sloane (1992) *New spherical 4-designs*. Discrete
#    Mathematics, 106-107 (255-264). `(PDF)
#    <https://www.sciencedirect.com/science/article/pii/0012365X9290552Q>`__.
#
# .. [#Ambainis]
#
#    A. Ambainis and J. Emerson (2007) *Quantum t-designs: t-wise independence
#    in the quantum world.* Twenty-Second Annual IEEE Conference on
#    Computational Complexity 129-140.
#    `(arXiv) <https://arxiv.org/abs/quant-ph/0701126>`__.
#
# .. [#Klappenecker]
#
#    A. Klappenecker and M. Roetteler (2005) *Mutually unbiased bases, spherical
#    designs, and frames.* Proceedings of SPIE Vol. 5914.
#
# .. [#Dankert]
#
#    C. Dankert, R. Cleve, J. Emerson, and E. Levine (2009) *Exact and
#    Approximate Unitary 2-Designs: Constructions and Applications.* Phys. Rev. A 80, 012304.
#    `(arXiv) <https://arxiv.org/abs/quant-ph/0606161>`__.
#
# .. [#DankertThesis]
#
#    C. Dankert (2005) *Efficient Simulation of Random Quantum States and
#    Operators.* MSc Thesis, University of Waterloo. `(arXiv)
#    <https://arxiv.org/abs/quant-ph/0512217>`__.
#
# .. [#Gross]
#
#    D. Gross, K. Audenaert, and J. Eisert (2007) *Evenly distributed unitaries:
#    on the structure of unitary designs*. J. Math. Phys. 48, 052104.
#    `(arXiv) <https://arxiv.org/abs/quant-ph/0611002>`__.
#
# .. [#Roy]
#
#    A. Roy and A. J. Scott (2009) *Unitary designs and codes*. Des. Codes Cryptogr. 53 13-31.
#    `(arXiv) <https://arxiv.org/abs/0809.3813>`__.
#
# .. [#Bannai]
#
#    E. Bannai, M. Nakahara, D. Zhao, and Y. Zhu (2019) *On the explicit constructions of
#    certain unitary t-designs.* J. Phys. A: Math. Theor. 52 495301.
#    `(arXiv) <https://arxiv.org/abs/1906.04583>`__.
#
# .. [#Nakata]
#
#    Y. Nakata et al. (2021) *Quantum circuits for exact unitary t-designs and
#    applications to higher-order randomized benchmarking.* `(arXiv)
#    <https://arxiv.org/abs/2102.12617>`__.
#
# .. [#PQC]
#
#    A. Ambainis, M. Mosca, A. Tapp, and R. de Wolf (2000) *Private Quantum Channels*. Proc.
#    41st FOCS, 547-553. `(PDF) <https://homepages.cwi.nl/~rdewolf/publ/qc/AMTW00.pdf>`__.
#
# .. [#Seberry]
#
#    J. Seberry and M. Yamada (1992) *Hadamard matrices, sequences, and block
#    designs.* Contemporary Design Theory -- A Collection of Surveys
#    (D. J. Stinson and J. Dinitz, Eds.), John Wiley and Sons, 431-560.
#    `(PDF) <http://mathscinet.ru/files/YamadaSeberry1992.pdf>`__.
#
# .. [#Gaeta]
#
#    M. Gaeta, O. Di Matteo, A. B. Klimov, and H. de Guise (2014) *Discrete phase-space
#    approach to mutually orthogonal Latin squares*. J. Phys. A: Math. Theor. 47 (43) 435303.
#    `(arXiv) <https://arxiv.org/abs/1408.6742>`__.