# -*- coding: utf-8 -*-
"""
Capturing output representations
================================

This example demonstrates how the `capture_repr` configuration option
([Controlling what output is captured](https://sphinx-gallery.github.io/stable/configuration.html#capture-repr))
 works. The default `capture_repr` setting is
`('_repr_html_', '__repr__')` and was used to build this
Mkdocs-Gallery documentation. The output that is captured with this setting
is demonstrated in this example. Differences in outputs that would be captured
with other `capture_repr` settings are also explained.
"""
#%%
# Nothing is captured for the code block below because no data is directed to
# standard output and the last statement is an assignment, not an expression.

# example 1
a = 2
b = 10

#%%
# If you did wish to capture the value of `b`, you would need to use:

# example 2
a = 2
b = 10
b   # this is an expression

#%%
# Mkdocs-Gallery first attempts to capture the `_repr_html_` of `b` as this
# is the first 'representation' method in the `capture_repr` tuple. As this
# method does not exist for `b`, Mkdocs-Gallery moves on and tries to capture
# the `__repr__` method, which is second in the tuple. This does exist for
# `b` so it is captured and the output is seen above.
#
# A pandas dataframe is used in the code block below to provide an example of
# an expression with a `_repr_html_` method.

# example 3
import pandas as pd

df = pd.DataFrame(data = {'col1': [1, 2], 'col2': [3, 4]})
df

#%%
# The pandas dataframe `df` has both a `__repr__` and `_repr_html_`
# method. As `_repr_html_` appears first in the `capture_repr` tuple, the
# `_repr_html_` is captured in preference to `__repr__`.
#
# Statsmodels tables should also be styled appropriately:

# example 4
import numpy as np
import statsmodels.iolib.table
statsmodels.iolib.table.SimpleTable(np.zeros((3, 3)))

#%%
# For the example below, there is data directed to standard output and the last
# statement is an expression.

# example 5
print('Hello world')
a + b

#%%
# `print()` outputs to standard output, which is always captured. The
# string `'Hello world'` is thus captured. A 'representation' of the last
# expression is also captured. Again, since this expression `a + b` does not
# have a `_repr_html_` method, the `__repr__` method is captured.
#
#
# Matplotlib output
# -----------------
#
# Matplotlib function calls generally return a Matplotlib object as well as
# outputting the figure. For code blocks where the last statement is a
# Matplotlib expression, a 'representation' of the object will be captured, as
# well as the plot. This is because Matplotlib objects have a `__repr__`
# method and our `capture_repr` tuple contains `__repr__`. Note that
# Matplotlib objects also have a `__str__` method.
#
# In the example below, `matplotlib.pyplot.plot()` returns a list of
# `Line2D` objects representing the plotted data and the `__repr__` of the
# list is captured as well as the figure:

import matplotlib.pyplot as plt

plt.plot([1,2,3])

#%%
# To avoid capturing the text representation, you can assign the last Matplotlib
# expression to a temporary variable:

_ = plt.plot([1,2,3])

#%%
# Alternatively, you can add `plt.show()`, which does not return anything,
# to the end of the code block:

plt.plot([1,2,3])
plt.show()

#%%
# The `capture_repr` configuration
# --------------------------------
#
# The `capture_repr` configuration is `('_repr_html_', '__repr__')` by
# default. This directs Mkdocs-Gallery to capture 'representations' of the last
# statement of a code block, if it is an expression. Mkdocs-Gallery does
# this according to the order 'representations' appear in the tuple.
#
# With the default `capture_repr` setting, `_repr_html_` is attempted to be
# captured first. If this method does not exist, the `__repr__` method would be
# captured. If the `__repr__` also does not exist (unlikely for non-user
# defined objects), nothing would be captured. For example, if the the
# configuration was set to `'capture_repr': ('_repr_html_')` nothing would be
# captured for example 2 as `b` does not have a `_repr_html_`.
# You can change the 'representations' in the `capture_repr` tuple to finely
# tune what is captured in your example `.py` files.
#
# To only capture data directed to standard output you can set `capture_repr`
# to be an empty tuple: `capture_repr: ()`. With this setting, only data
# directed to standard output is captured. For the examples above, output would
# only be captured for example 4. Although the last statement is an expression
# for examples 2, 3 and 4 no 'representation' of the last expression would be
# output. You would need to add `print()` to the last expression to capture
# a 'representation' of it.
#
# The empty tuple setting imitates the behaviour of Sphinx-Gallery prior to
# v0.5.0, when this configuration was introduced.


#%%
import pandas as pd
ser_obj = pd.Series(range(10, 20))
print(ser_obj)
# 获取数据
print(ser_obj.values)
print('* '* 50)
# 获取索引
print(ser_obj.index)
print('* '* 50)
#%%
print(ser_obj * 2)
print('* '* 50)
print(ser_obj > 15)
print('* '* 50)
#%%
year_data = {2001: 17.8, 2005: 20.1, 2003: 16.5}
ser_obj2 = pd.Series(year_data)
print(ser_obj2)
print('* '* 50)
print(ser_obj2.index)
print('* '* 50)
print(ser_obj2[2001])
print('* '* 50)

#%%
print(ser_obj2.name)
print('* '* 50)
ser_obj2.name = 'temp'
print(ser_obj2.index.name)
print('* '* 50)
ser_obj2.index.name = 'year1'
print(ser_obj2.head())
print('* '* 50)

#%% md
# DataFrame
#%%
d2 =[{"name" : "xiaohong" ,"age" :32,"tel" :10010},{ "name": "xiaogang" ,"tel": 10000} ,{"name":"xiaowang" ,"age":22}]
df6=pd.DataFrame(d2)

print(df6)
print('* '* 50)
#%%
import pandas as pd
import numpy as np
dict_data = {'A': 1,
             'B': pd.Timestamp('20190926'),
             'C': pd.Series(1, index=list(range(4)),dtype='float32'),
             'D': np.array([3] * 4,dtype='int32'),
             'E': ["Python","Java","C++","C"],
             'F': 'wangdao' }
df_obj2 = pd.DataFrame(dict_data)
print(df_obj2)
print('* '* 50)

print(df_obj2.index)
print('* '* 50)
# df_obj2.index[0] = 2
#%%
dates = pd.date_range('20130101', periods=6)
df = pd.DataFrame(np.random.randn(6, 4), index=dates, columns=list('ABCD'))

print(df)
print('* '* 50)
#%%
print(df_obj2['A'])
print('* '* 50)
#把df的某一列取出来是series
print(type(df_obj2['A']))
print('* '* 50)
#%%
#增加列数据
df_obj2['G'] = df_obj2['D'] + 4
print(df_obj2.head())
print('* '* 50)
#%%
# 删除列
del(df_obj2['G'] )
print(df_obj2.head())
print('* '* 50)
#%% md
# 4 Pandas的索引操作
#%%
print(df_obj2.index)
print('* '* 50)
#%%
# 索引对象不可变（上面代码增加）
# df_obj2.index[0] = 2
#%% md
# 3 常见的Index种类
# •Index，索引
# •Int64Index，整数索引
# •MultiIndex，层级索引
# •DatetimeIndex，时间戳类型
#%%
ser_obj = pd.Series(range(5), index = list("abcde"))
print(ser_obj)
print('* '* 50)
ser_obj.index
#%%
# 行索引
print(ser_obj['b'])
print('* '* 50)
print(ser_obj[2])
print('* '* 50)
#%%
# 切片索引
print(ser_obj[1:3])
print('* '* 50)
print(ser_obj['b':'d'])
print('* '* 50)
#%%
# 不连续索引
print(ser_obj[[0, 2, 4]])
print('* '* 50)
print(ser_obj[['a', 'e']])
print('* '* 50)
#%%
# 布尔索引
ser_bool = ser_obj > 2
print(ser_bool)
print('* '* 50)
print(ser_obj[ser_bool])
print('* '* 50)

print(ser_obj[ser_obj > 2])
print('* '* 50)
#%% md
# 4 DataFrame索引
#%%
import numpy as np
df_obj = pd.DataFrame(np.random.randn(5,4), columns = ['a', 'b', 'c', 'd'])
print(df_obj.head())
print('* '* 50)
#%%
# 列索引
print(df_obj['a']) # 返回Series类型
print('* '* 50)
print(df_obj[['a']]) # 返回DataFrame类型
print('* '* 50)
print(type(df_obj[['a']])) # 返回DataFrame类型
print('* '* 50)
#%% md
# 1. loc 标签索引
#%%
# 标签索引 loc
# Series
print(ser_obj)
print(ser_obj['b':'d'])
print(ser_obj.loc['b':'d'])

# DataFrame
df_obj = pd.DataFrame(np.random.randn(5,4), columns = ['a', 'b', 'c', 'd'],index=list('abcde'))
print(df_obj['a'])
print('-'*50)
print(df_obj.loc['a'])
print('-'*50)
# 第一个参数索引行，第二个参数是列,loc或者iloc效率高于直接用取下标的方式
print(df_obj.loc['a':'c', 'a':'c'])
#%% md
# iloc 位置索引
#%%
# Series
print(ser_obj[1:3])
print(ser_obj.iloc[1:3])

# DataFrame
print(df_obj.iloc[0:2, 0:2]) # 注意和df_obj.loc[0:2, 'a']的区别
#%% md


from collections import Counter
from pytorch_transformers import GPT2Tokenizer, GPT2LMHeadModel
from pytorch_transformers import AdamW, WarmupLinearSchedule
import torch.nn.functional as F
import numpy as np
import torch
from torch.utils.data import Dataset
import time
import csv
import json

class GPT2:
    """
        High level interface for PyTorch's GPT-2 implementation.
        Parameters
        ----------
        fileName: string
            The path to the file (txt, json, csv)

        taskToken: string
            GPT-2 expects a task name for generation. For example, for movie titles, this can be "Movie: "
            It will be treated as a fill in the blank task.

        epochs: int
            The number of cycles over the training data.

        variant: string, default "small"
            The variant of GPT-2 to use. This can be "small", "medium", "large"

        batchSize: int, default 16
            The number of training instances per iteration.

        eos: string, default "<|endoftext|>"
            The end-of-sentence token.

        instanceMxLen: int, default None
            Max length of each document.

        txtSeparator: str, optional, default '\n'
            When the data file is a txt file, this is the character serparating documents.
            By default, the assumption is that each document is on a separate line.

        csvIndex: int, optional, default None
            When the data file is a csv file, this is the index of the column to parse from.

        jsonKey: str, optional, default None
            When the data file is a csv file, this is the json key to parse from.
            Typically, it is 'body' or 'text'.

        seedParams: dict, optional, default {'N_first': 1, 'minFreq': 5}
            Parameters that will be used while selecting the random seed used to initalize prediction.
                N_first: get the first n_first tokens of each training document.
                minFreq: The min frequency a seed must have before it is included in the pool.

                By default random seed is enabled. If seedParams = {} then static seed will be applied.

        optimParams: dict, default {"lr" : 3e-4}
            The optimizer paramters.

        schedParams: dict, default {"warmup_steps" : 400}
            The scheduler paramters.

    """

    def __init__(self, fileName,
                        taskToken,
                        epochs,
                        variant = "small",
                        batchSize = 32,
                        eos = "<|endoftext|>",
                        instanceMxLen = None,
                        txtSeparator = '\n',
                        csvIndex = None,
                        jsonKey = None,
                        seedParams = {'N_first': 1, 'minFreq': 5},
                        optimParams = {"lr" : 3e-4},
                        schedParams = {"warmup_steps" : 400}):

        self.device =  torch.device('cuda' if torch.cuda.is_available() else 'cpu')

        # To be set by the user
        self.csvIndex = csvIndex
        self.jsonKey = jsonKey
        self.txtSeparator = txtSeparator
        self.eos = eos
        self.taskToken = taskToken
        self.optimParams = optimParams
        self.schedParams = schedParams

        self.epochs = self.__checkIntParam(epochs, "epochs")
        self.batchSize = self.__checkIntParam(batchSize, "batchSize")
        self.fileName = self.__checkFileName(fileName)
        self.instanceMxLen = instanceMxLen
        self.seedParams = self.__checkSeedParams(seedParams)
        self.variant = self.__checkVariant(variant)

        self.examples = None
        self.__text = None
        self.model = None
        self.tokenizer = None
        self.optimizer = None
        self.scheduler = None

        self.__readFile()
        self.__get_tokenizer(self.variant)
        print()
        self.__get_model(self.variant)
        self.__get_optimizer()
        self.__get_scheduler()

    def __checkIntParam(self, param, name):
        if sum([type(param) != int, param <= 0]) >= 1:
            raise Exception("Make sure %s is a positive int."%name)
        return param

    def __checkVariant(self, variant):
        avl = ["small", "medium", "large"]
        if variant not in avl:
            raise Exception("Available variants are", " ".join(avl), "got", variant, "instead")

        #in PyTorch implementation, small = gpt2, medium = gpt2-medium etc.
        return "-"+variant if variant != "small" else ""

    def __checkFileName(self, fileName):
        extension = fileName.split(".")[-1]

        if extension not in ['txt', 'csv', 'json']:
            raise Exception("Only works with txt, csv, and json files.")

        if extension == 'csv' and self.csvIndex == None:
            raise Exception("Please provide a 'csvIndex' to indicate the column to parse from.")
        elif extension == 'json' and self.jsonKey == None:
            raise Exception("Please provide a 'jsonKey' to indicate json key to parse from.")

        self.extension = extension
        return fileName

    def __checkSeedParams(self,params):
        attr_list = ['N_first', 'minFreq']

        if not len(params):
            return {}

        elif not all(attr in params for attr in attr_list):
            raise Exception("seedParams should contain the attributes {attrs}".format(attrs = attr_list))

        return params

    def __readFile(self):
        file_obj = open(self.fileName)

        if self.extension == "txt":
            self.__text = file_obj.read().split(self.txtSeparator)
        elif self.extension == "csv":
            reader = csv.reader(file_obj)
            self.__text = [row[self.csvIndex] for row in reader]
        else:
            json_data = json.loads(file_obj.read())
            self.__text = [row[self.jsonKey] for row in json_data]

        if self.instanceMxLen == None:
            self.instanceMxLen = len(max(self.__text, key=len))
        self.examples = [self.taskToken+" "+inst+" "+self.eos for inst in self.__text
                            if len(inst)>0 and len(inst)<self.instanceMxLen]
        self.seeds =  self.__getStartWords()

    def __getStartWords(self):
        n_first = 1 if 'N_first' not in self.seedParams else self.seedParams['N_first']

        start_words = [' '.join(instance.split()[:n_first])
                           for instance in self.__text
                           if len(instance.split()) > n_first]

        freq_start = Counter(start_words)
        if not len(self.seedParams):
            word, _ = freq_start.most_common(1)[0]
            print("The word {w} is the Static Seed. ".format(w=word))
            return {word: 1.0}

        return {word:freq/len(self.examples) for word,freq  in freq_start.items() if freq >= self.seedParams['minFreq'] }

    def getSeed(self):
        """
            return a weighted seed.
            In case static seed is enabled, then the most frequent token will be the seed.
        """
        seeds = list(self.seeds.keys())
        probs = list(self.seeds.values())
        return np.random.choice(seeds, 1 , probs).tolist()

    def __get_batches(self):
        num_batches = len(self.examples) // self.batchSize
        for i in range(0, num_batches*self.batchSize, self.batchSize):
            yield self.examples[i:i+self.batchSize]

    def __encode_batch(self, batch):
        encoded = torch.Tensor().long().to(self.device)
        for inst in batch:
            docTens = torch.tensor(self.tokenizer.encode(inst)).unsqueeze(0).to(self.device)
            encoded = torch.cat([encoded, docTens[:,1:]], dim=1)
        return encoded

    def __get_optimizer(self):
        self.optimizer = AdamW(self.model.parameters(), **self.optimParams)

    def __get_model(self, variant):
        self.model = GPT2LMHeadModel.from_pretrained('gpt2'+variant)
        self.model = self.model.to(self.device)
        self.model.train()

    def __get_tokenizer(self, variant):
        self.tokenizer = GPT2Tokenizer.from_pretrained('gpt2'+variant)

    def __get_scheduler(self):
        self.scheduler = WarmupLinearSchedule(self.optimizer, t_total = -1, **self.schedParams)

    def choose_from_top(self, probs, n=5):
        ind = np.argpartition(probs, -n)[-n:]
        top_prob = probs[ind]
        top_prob = top_prob / np.sum(top_prob) # Normalize
        choice = np.random.choice(n, 1, p = top_prob)
        token_id = ind[choice][0]
        return token_id

    def select_nucleus(self, out, p = 0.5):
        probs = F.softmax(out, dim=-1)
        idxs = torch.argsort(probs, descending = True)
        res, prob,  cumsum = [], [], 0.
        for idx in idxs:
            res.append(idx)
            prob.append(probs[idx].item())
            cumsum+=probs[idx]
            if cumsum>p:
               break
        nucleus_prob = prob / np.sum(prob)
        choice = np.random.choice(res , 1,p = nucleus_prob)[0]
        return choice

    def run(self):
        sum_loss = 0.0
        batch_count = 0
        last_loss = None
        batches_len = (len(self.examples) // self.batchSize) * self.epochs
        for e in range(self.epochs):
            batches = self.__get_batches()
            batch_times = []
            for batch in batches:
                start_time = time.time()

                encoded = self.__encode_batch(batch)

                outputs = self.model(encoded, labels=encoded)

                loss, logits = outputs[:2]
                loss.backward()
                sum_loss += loss.detach().data

                batch_count += 1
                self.optimizer.step()
                self.scheduler.step()
                self.optimizer.zero_grad()
                self.model.zero_grad()

                batch_time = time.time() - start_time
                batch_times.append(batch_time)

                mean_time = sum(batch_times)/len(batch_times)
                remaining_batches = batches_len - batch_count
                remaining_seconds = remaining_batches * mean_time
                remaining_time = time.strftime("%H:%M:%S",
                    time.gmtime(remaining_seconds))
                progress = "{:.2%}".format(batch_count/batches_len)
                print('Epoch: {}/{}'.format(e+1, self.epochs),
                      'Progress:', progress,
                      'Loss: {}'.format(last_loss),
                      'ETA:', remaining_time)
            last_loss = sum_loss
            sum_loss = 0.0

    def isRedundant(self, cur_ids, next_token):
        cur_ids_list = list(cur_ids.squeeze().to('cpu').numpy())
        next_id = next_token.item()
        if cur_ids_list[-1] == next_id: #prev word == next word
            return True
        else:
            return False

    def generate_document(self, n,
                       isNucleus=True,
                       instanceMxLen=None,
                       k=None, p=None,
                       uniq=True,
                       noRepetition = False):
        self.model.eval()
        res = set()
        max_len = instanceMxLen if instanceMxLen!=None else self.instanceMxLen
        with torch.no_grad():
            while len(res) < n:
                cur_ids = torch.tensor(self.tokenizer.encode(self.taskToken+" "+self.getSeed()[0])).unsqueeze(0).to(self.device)

                for i in range(max_len):
                    outputs = self.model(cur_ids, labels=cur_ids)
                    _, logits = outputs[:2]

                    if isNucleus:
                        if p!=None:
                            next_token_id = self.select_nucleus(logits[0,-1], p=p)
                        next_token_id = self.select_nucleus(logits[0,-1])

                    else: #topk
                        softmax_logits = torch.softmax(logits[0,-1], dim=0)
                        if k!=None:
                            next_token_id = self.choose_from_top(softmax_logits.to('cpu').numpy(), n=k)
                        next_token_id = self.choose_from_top(softmax_logits.to('cpu').numpy())

                    if noRepetition:
                        if self.isRedundant(cur_ids, next_token_id):
                            continue

                    cur_ids = torch.cat([cur_ids, torch.ones((1,1)).long().to(self.device) * next_token_id], dim = 1)
                    if next_token_id in self.tokenizer.encode(self.eos):
                        break

                doc = self.tokenizer.decode(list(cur_ids.squeeze().to('cpu').numpy())).strip()
                if uniq:
                    if doc not in self.examples:
                        res.add(doc)
                else:
                    res.add(doc)

        return list(res)


import skimage.io as io
import skimage.transform as skt
import numpy as np
from PIL import Image
from src.models.class_patcher import patcher
from src.utils.imgproc import *

class patcher(patcher):
    def __init__(self, body='./body/body_mishe.png', **options):
        super().__init__('ミーシェ', body=body, pantie_position=[910, 1929], **options)
        self.mask = io.imread('./mask/mask_mishe.png')
        self.sign_position = [933, 1482]
        try:
            self.add_sign = self.options['add_sign']
        except:
            self.add_sign = self.ask(question='Add immoral sign?', default=False)
        if self.add_sign:
            try:
                self.sign = Image.open(self.options['fsign'])
            except:
                self.sign = Image.open('./material/anna_sign.png')
            self.sign = self.sign.resize((369,746))

    def convert(self, image):
        pantie = np.array(image)
        pantie = np.bitwise_and(pantie, self.mask)
        [r, c, d] = pantie.shape

        # move from hip to front
        patch = np.copy(pantie[-140:-5, 546:, :])
        pantie[-115:, 546:, :] = 0
        patch = skt.resize(patch[::-1, ::-1, :], (patch.shape[0], 63), anti_aliasing=True, mode='reflect')
        [pr, pc, d] = patch.shape
        pantie[127 - 5:127 - 5 + pr, :pc, :] = np.uint8(patch * 255)

        # Affine transform matrix
        src_cols = np.linspace(0, c, 10)
        src_rows = np.linspace(0, r, 10)
        src_rows, src_cols = np.meshgrid(src_rows, src_cols)
        src = np.dstack([src_cols.flat, src_rows.flat])[0]
        shifter_row = np.zeros(src.shape[0])
        shifter_col = np.zeros(src.shape[0])
        shifter_row[30:-30] = (np.sin(np.linspace(0, 1 * np.pi, src.shape[0]) - np.pi / 32) * 100)[30:-30]
        shifter_row[:30] = (np.sin(np.linspace(0, 1 * np.pi, src.shape[0]) + np.pi / 2) * 60)[:30]
        shifter_row[-30:] = (np.sin(np.linspace(0, 1 * np.pi, src.shape[0]) - np.pi / 2) * 80)[-30:]
        shifter_col[13:-30] = -(np.sin(np.linspace(0, 1 * np.pi, src.shape[0]) + np.pi / 8) * 22)[13:-30]

        shifter_row = np.convolve(shifter_row, np.ones(20) / 20, mode='valid')
        shifter_col = np.convolve(shifter_col, np.ones(10) / 10, mode='valid')
        shifter_row = skt.resize(shifter_row, (100, 1), anti_aliasing=True, mode='reflect')[:, 0]
        shifter_col = skt.resize(shifter_col, (100, 1), anti_aliasing=True, mode='reflect')[:, 0]

        dst_rows = src[:, 1] + shifter_row - 110
        dst_cols = src[:, 0] + shifter_col
        dst = np.vstack([dst_cols, dst_rows]).T
        affin = skt.PiecewiseAffineTransform()
        affin.estimate(src, dst)
        pantie = np.uint8(skt.warp(pantie, affin) * 255)[:310, :, :]

        # Finalize
        pantie_ = skt.resize(pantie, (np.int(pantie.shape[0] * 2.05), np.int(pantie.shape[1] * 2.05)), anti_aliasing=True, mode='reflect')
        pantie = np.uint8(pantie_ * 255)
        return Image.fromarray(pantie)

    def patch(self, image, transparent=False):
        image = self.convert(image)
        if transparent:
            patched = Image.new("RGBA", self.body_size)
        else:
            patched = self.body.copy()

        if self.add_sign:
            self.paste(patched, self.sign, self.sign_position)
        patched = self.paste(patched, image, self.pantie_position)
        return patched


import numpy as np
import pytest

from quaternions import Quaternion
from quaternions.rotation import euler_rotation, rotation_axis_angle
from quaternions.utils import is_pair, is_point, is_scalar, is_vector

########################################
#    constructor
########################################

@pytest.mark.parametrize(
    "obj, expected",
    [
        # scalar
        (-1.5, Quaternion([-1.5, 0, 0, 0])),
        (5, Quaternion([5, 0, 0, 0])),
        # vector
        (([0.4, 1, 2.5]), Quaternion([0, 0.4, 1.0, 2.5])),
        (np.array([1, -2, 3]), Quaternion([0, 1, -2, 3])),
        # point
        ([-3, 0, 2, 0], Quaternion([-3, 0, 2, 0])),
        (np.array([1.2, 1, 0, -0.6]), Quaternion([1.2, 1, 0, -0.6])),
        # pair
        ((9, [1, 2, 3]), Quaternion([9, 1, 2, 3])),
        ((0, np.array([1, 2, 3])), Quaternion([0, 1, 2, 3])),
        # quaternion
        (Quaternion(-2.5), Quaternion(-2.5)),
        (Quaternion([-1, 2, -3, 4]), Quaternion([-1, 2, -3, 4])),
    ],
)
def test_quaternion_init(obj, expected):
    result = Quaternion(obj)

    assert result == expected

@pytest.mark.parametrize("obj", ["0.5", [0, 1, 2, 3, 4], list("letters"),])
def test_quaternion_init_error(obj):
    with pytest.raises(TypeError):
        Quaternion(obj)

########################################
#    properties
########################################

@pytest.mark.parametrize(
    "quat, expected",
    [
        (Quaternion(2), [2, 0, 0, 0]),
        (Quaternion([1, 2, 3]), [0, 1, 2, 3]),
        (Quaternion(-np.array([1, 2, 3, 4])), [-1, -2, -3, -4]),
        (Quaternion((9, [4, 5, 6])), np.array([9, 4, 5, 6])),
    ],
)
def test_point(quat, expected):
    point = quat.point

    assert is_point(point)
    assert all(x == y for x, y in zip(point, expected))

@pytest.mark.parametrize(
    "quat, expected",
    [
        (Quaternion(0), 0.0),
        (Quaternion([1, 2, 3]), 0.0),
        (Quaternion([-2.5, 1, 2, 3]), -2.5),
        (Quaternion((1, [5, 6, 7])), 1.0),
    ],
)
def test_scalar_part(quat, expected):
    scalar = quat.scalar_part

    assert is_scalar(scalar)
    assert scalar == expected

@pytest.mark.parametrize(
    "quat, expected",
    [
        (Quaternion(5), [0, 0, 0]),
        (Quaternion([1, 2, 3]), np.array([1, 2, 3])),
        (Quaternion([0, -1, 2, 3]), [-1, 2, 3]),
        (Quaternion((0, [-7, -8, -9])), -np.array([7, 8, 9])),
    ],
)
def test_vector_part(quat, expected):
    vector = quat.vector_part

    assert is_vector(vector)
    assert all(x == y for x, y in zip(vector, expected))

@pytest.mark.parametrize(
    "quat, expected",
    [
        (Quaternion(0), (0, [0, 0, 0])),
        (Quaternion(np.array([4, 9, 2])), (0.0, [4.0, 9, 2])),
        (Quaternion([1, 2, 3, 4]), (1, np.array([2, 3, 4]))),
        (Quaternion((-1, [1, -2, 3])), (-1.0, [1, -2, 3])),
    ],
)
def test_pair(quat, expected):
    pair = quat.pair

    assert is_pair(pair)
    s, v = pair
    es, ev = expected
    assert s == es
    assert all(x == y for x, y in zip(v, ev))

@pytest.mark.parametrize(
    "quat, expected",
    [
        (Quaternion(0), np.zeros([2, 2])),
        (Quaternion(1), np.eye(2)),
        (Quaternion([1, 2, 3]), np.array([[1j, -2 - 3j], [2 - 3j, -1j]])),
        (Quaternion([5, -1, 2, 3]), np.array([[5 - 1j, -2 - 3j], [2 - 3j, 5 + 1j]])),
        (Quaternion((4, [3, 2, 1])), np.array([[4 + 3j, -2 - 1j], [2 - 1j, 4 - 3j]])),
    ],
)
def test_complex_matrix(quat, expected):
    matrix = quat.complex_matrix

    assert np.allclose(matrix, expected)

@pytest.mark.parametrize(
    "quat, expected",
    [
        (Quaternion(-3), True),
        (Quaternion([1, 2, 3]), False),
        (Quaternion(np.array([0, 0, 0])), True),
        (Quaternion([1, 2, 3, 4]), False),
        (Quaternion([1, 0, 0, 0]), True),
        (Quaternion((-1, [1, 2, 3])), False),
        (Quaternion((-1, [0.0, 0.0, 0.0])), True),
    ],
)
def test_is_scalar(quat, expected):
    s = quat.is_scalar

    assert s is expected

@pytest.mark.parametrize(
    "quat, expected",
    [
        (Quaternion(8), False),
        (Quaternion(0), True),
        (Quaternion([1, 2, 3]), True),
        (Quaternion([1, 2, 3, 4]), False),
        (Quaternion([0, 1, 2, 3]), True),
        (Quaternion((-1, [1, 2, 3])), False),
        (Quaternion((0.0, [1.0, 2.0, 3.0])), True),
    ],
)
def test_is_vector(quat, expected):
    v = quat.is_vector

    assert v is expected

@pytest.mark.parametrize(
    "quat, expected",
    [
        (Quaternion(1), False),
        (Quaternion(0), True),
        (Quaternion([1, 2, 3]), False),
        (Quaternion([0, 0, 0]), True),
        (Quaternion([1, 2, 3, 4]), False),
        (Quaternion([0, 0, 0, 0]), True),
        (Quaternion((-1, [1, 2, 3])), False),
        (Quaternion((0.0, [0, 0, 0])), True),
    ],
)
def test_is_zero(quat, expected):
    z = quat.is_zero

    assert z is expected

@pytest.mark.parametrize(
    "quat, expected",
    [
        (Quaternion(1), True),
        (Quaternion(-1), True),
        (Quaternion(0), False),
        (Quaternion([1, 2, 3]), False),
        (Quaternion([0, -1, 0]), True),
        (Quaternion([1, 2, 3, 4]), False),
        (Quaternion([0, 1, 0, 0]), True),
        (Quaternion((-1, [1, 2, 3])), False),
        (Quaternion((-1 / 2, [-1 / 2, 1 / 2, 1 / 2])), True),
    ],
)
def test_is_unit(quat, expected):
    result = quat.is_unit

    assert result is expected

@pytest.mark.parametrize(
    "quat, expected",
    [
        (Quaternion(0), Quaternion(0)),
        (Quaternion(-2), Quaternion(-2.0)),
        (Quaternion([0, 0, 0]), Quaternion(0)),
        (Quaternion([1, 2, 3]), Quaternion(-np.array([1, 2, 3]))),
        (Quaternion(np.array([0, 0, 0, 0])), Quaternion(0)),
        (Quaternion([1, 2, 3, 4]), Quaternion([1, -2, -3, -4])),
        (Quaternion((0.0, [0, 0, 0])), Quaternion(0)),
        (Quaternion((1, [-2, -3, -4])), Quaternion([1, 2, 3, 4])),
    ],
)
def test_conjugate(quat, expected):
    c = quat.conjugate

    assert c == expected

def test_inverse_error():
    quat = Quaternion(0.0)
    with pytest.raises(ArithmeticError):
        quat.inverse

@pytest.mark.parametrize(
    "quat, expected",
    [
        (Quaternion(2), Quaternion(0.5)),
        (Quaternion([1, 2, 3]), Quaternion(-np.array([1, 2, 3]) / 14.0)),
        (Quaternion([1, 2, 3, 4]), Quaternion(np.array([1, -2, -3, -4]) / 30)),
        (Quaternion((1, [-2, -3, -4])), Quaternion(np.array([1, 2, 3, 4]) / 30.0)),
    ],
)
def test_inverse(quat, expected):
    i = quat.inverse

    assert i == expected

@pytest.mark.parametrize(
    "quat, norm, squared_norm",
    [
        (Quaternion(0), 0.0, 0.0),
        (Quaternion(-2), 2.0, 4),
        (Quaternion([0, 0, 0]), 0, 0),
        (Quaternion([1, 2, 3]), np.sqrt(14), 14.0),
        (Quaternion(np.array([0, 0, 0, 0])), 0.0, 0),
        (Quaternion([1, 2, 3, 4]), np.sqrt(30), 30),
        (Quaternion((0.0, [0, 0, 0])), 0, 0),
        (Quaternion((1, [-2, -3, -4])), np.sqrt(30), 30),
    ],
)
def test_norms(quat, norm, squared_norm):
    n = quat.norm
    sn = quat.squared_norm

    assert n == pytest.approx(norm, abs=1e-8)
    assert sn == pytest.approx(squared_norm, abs=1e-8)

##############################
# OPERATOR OVERLOADING
##############################

@pytest.mark.parametrize(
    "p, q",
    [
        (Quaternion(0), Quaternion(0)),
        (Quaternion([1, 2, 3]), Quaternion(np.array([1, 2, 3]))),
        (Quaternion([0, 1, 2, 3]), Quaternion(np.array([0, 1, 2, 3]))),
        (Quaternion((0, [1, 2, 3])), Quaternion((0, np.array([1, 2, 3])))),
    ],
)
def test_eq(p, q):
    assert p == q

@pytest.mark.parametrize(
    "p, q",
    [
        (Quaternion(2), Quaternion(-3)),
        (Quaternion([-2, 2, -2]), Quaternion(np.array([1, 2, 3]))),
        (Quaternion([-6, 3, 2, 3]), Quaternion(np.array([0, 1, 2, 3]))),
        (Quaternion((0, [-1, -2, 3])), Quaternion((0, np.array([1, 2, 3])))),
    ],
)
def test_ne(p, q):
    assert p != q

@pytest.fixture(
    scope="function",
    params=[
        # p
        # expected +p
        # expected -p
        # expected ~p
        # expected abs(p)
        (Quaternion(0), Quaternion(0), Quaternion(0), Quaternion(0), 0,),
        (Quaternion(-1.2), Quaternion(-1.2), Quaternion(1.2), Quaternion(-1.2), 1.2,),
        (
            Quaternion([1, 1, 0]),
            Quaternion([1, 1, 0]),
            Quaternion([-1, -1, 0]),
            Quaternion([-1, -1, 0]),
            np.sqrt(2),
        ),
        (
            Quaternion([-1, 1, -1, 1]),
            Quaternion([-1, 1, -1, 1]),
            Quaternion([1, -1, 1, -1]),
            Quaternion([-1, -1, 1, -1]),
            2.0,
        ),
    ],
)
def fxt_quat_unary(request):
    return request.param

@pytest.fixture(
    scope="function",
    params=[
        # p,  q
        # expected p + q
        # expected p - q
        # expected p * q
        # expected p / q
        (
            Quaternion(3),
            -1.5,
            Quaternion(1.5),
            Quaternion(4.5),
            Quaternion(-4.5),
            Quaternion(-2),
        ),
        (
            Quaternion(2),
            np.array([1, 0, 0]),
            Quaternion([2, 1, 0, 0]),
            Quaternion([2, -1, 0, 0]),
            Quaternion([2, 0, 0]),
            Quaternion((0, [-2, 0, 0])),
        ),
        (
            Quaternion(np.array([1, 0, 0])),
            Quaternion(np.array([0, 1, 0])),
            Quaternion(np.array([1, 1, 0])),
            Quaternion(np.array([1, -1, 0])),
            Quaternion(np.array([0, 0, 1])),
            Quaternion(np.array([0, 0, -1])),
        ),
        (
            Quaternion(3),
            np.array([1, 1, 1]),
            Quaternion([3, 1, 1, 1]),
            Quaternion([3, -1, -1, -1]),
            Quaternion(3 * np.array([1, 1, 1])),
            Quaternion(-np.array([1, 1, 1])),
        ),
        (
            Quaternion(np.array([1, 0, 1, 0])),
            [2, 0, 0, 1],
            Quaternion(np.array([3, 0, 1, 1])),
            Quaternion(np.array([-1, 0, 1, -1])),
            Quaternion(np.array([2, 1, 2, 1])),
            Quaternion(np.array([2, -1, 2, -1]) / 5),
        ),
        (
            Quaternion([-1, 1, 1, 0]),
            [0, 0, 1, 1],
            Quaternion(np.array([-1, 1, 2, 1])),
            Quaternion(np.array([-1, 1, 0, -1])),
            Quaternion(np.array([-1, 1, -2, 0])),
            Quaternion(np.array([1, -1, 2, 0]) / 2),
        ),
    ],
)
def fxt_bin_ops(request):
    return request.param

fxt_quat = [
    Quaternion(1.0),
    Quaternion(np.array([1, 2, 3])),
    Quaternion([1, 2, 3, 4]),
    Quaternion((-1, [-1, 0, 1])),
]

fxt_quat_struct = [
    1,
    [1, 0, 1],
    np.array([-1, 0, -2]),
    [-1, 2, -1, 0],
    np.array([-1, 0, -1, -2]),
    (0, [1, 2, 3]),
    (0, np.array([-1, 0, 0])),
]

fxt_quat_struct_list = [
    [[0, 1, -1], [-1, 1, -1], [2.3, 0.2, 1.5],],  # list of vectors
    np.array([[1, 2, 3], [-1, 0, 2], [-2, 1, -1],]),  # array of vectors
    [[5, 2, 1, 0], [0, -1, 1, -1], [1.1, 2.2, 1.5, -4.5],],  # list of points
    np.array([[0, 1.1, 0, 3], [0.2, 1.0, 0, 2], [-1.5, 1, 3, 1],]),  # array of points
    [(1, np.array([1, 2, 3])), (0, [-1, -2, 0]), (-1, [2, 0, -2]),],  # list of pairs
    [  # list of quaternions
        Quaternion(-3.0),
        Quaternion([0, 1, 0]),
        Quaternion(np.array([2, -1, 0])),
        Quaternion([5, 2, -1, 2]),
        Quaternion(np.array([-1, 0, 2, 2])),
        Quaternion((0, [1, 2, 3])),
    ],
]

##############################
# unary ops
##############################

def test_quaternion_pos(fxt_quat_unary):
    p, expected, *_ = fxt_quat_unary
    result = +p
    assert result == expected

def test_quaternion_neg(fxt_quat_unary):
    p, _, expected, *_ = fxt_quat_unary
    result = -p
    assert result == expected

def test_quaternion_invert(fxt_quat_unary):
    p, *_, expected, _ = fxt_quat_unary
    result = ~p
    assert result == expected

def test_quaternion_abs(fxt_quat_unary):
    p, *_, expected = fxt_quat_unary
    result = abs(p)
    assert result == pytest.approx(expected)

##############################
# binary ops
##############################

def test_quaternion_add(fxt_bin_ops):
    p, q, expected, *_ = fxt_bin_ops
    result = p + q
    assert expected == result

def test_quaternion_sub(fxt_bin_ops):
    p, q, _, expected, *_ = fxt_bin_ops
    result = p - q
    assert expected == result

def test_quaternion_mul(fxt_bin_ops):
    p, q, *_, expected, _ = fxt_bin_ops
    result = p * q
    assert expected == result

def test_quaternion_truediv(fxt_bin_ops):
    p, q, *_, expected = fxt_bin_ops
    result = p / q
    assert expected == result

def test_quaternion_iadd(fxt_bin_ops):
    p, q, expected, *_ = fxt_bin_ops
    p += q
    assert p == expected

def test_quaternion_isub(fxt_bin_ops):
    p, q, _, expected, *_ = fxt_bin_ops
    p -= q
    assert p == expected

def test_quaternion_imul(fxt_bin_ops):
    p, q, *_, expected, _ = fxt_bin_ops
    p *= q
    assert p == expected

def test_quaternion_itruediv(fxt_bin_ops):
    p, q, *_, expected = fxt_bin_ops
    p /= q
    assert p == expected

@pytest.mark.parametrize("p", fxt_quat)
@pytest.mark.parametrize("struct", fxt_quat_struct_list)
def test_quaternion_add_struct(p, struct):
    result = p + struct
    assert all(r - Quaternion(q) == p for r, q in zip(result, struct))

@pytest.mark.parametrize("p", fxt_quat)
@pytest.mark.parametrize("struct", fxt_quat_struct_list)
def test_quaternion_sub_struct(p, struct):
    result = p - struct
    assert all(r + Quaternion(q) == p for r, q in zip(result, struct))

@pytest.mark.parametrize("p", fxt_quat)
@pytest.mark.parametrize("struct", fxt_quat_struct_list)
def test_quaternion_mul_struct(p, struct):
    result = p * struct
    assert all(r / Quaternion(q) == p for r, q in zip(result, struct))

@pytest.mark.parametrize("p", fxt_quat)
@pytest.mark.parametrize("struct", fxt_quat_struct_list)
def test_quaternion_truediv_struct(p, struct):
    result = p / struct
    assert all(r * Quaternion(q) == p for r, q in zip(result, struct))

@pytest.mark.parametrize("p", fxt_quat)
@pytest.mark.parametrize("struct", fxt_quat_struct_list)
def test_quaternion_iadd_struct(p, struct):
    before = p
    p += struct
    assert all(r - Quaternion(q) == before for r, q in zip(p, struct))

@pytest.mark.parametrize("p", fxt_quat)
@pytest.mark.parametrize("struct", fxt_quat_struct_list)
def test_quaternion_isub_struct(p, struct):
    before = p
    p -= struct
    assert all(r + Quaternion(q) == before for r, q in zip(p, struct))

@pytest.mark.parametrize("p", fxt_quat)
@pytest.mark.parametrize("struct", fxt_quat_struct_list)
def test_quaternion_imul_struct(p, struct):
    before = p
    p *= struct
    assert all(r / Quaternion(q) == before for r, q in zip(p, struct))

@pytest.mark.parametrize("p", fxt_quat)
@pytest.mark.parametrize("struct", fxt_quat_struct_list)
def test_quaternion_itruediv_struct(p, struct):
    before = p
    p /= struct
    assert all(r * Quaternion(q) == before for r, q in zip(p, struct))

@pytest.mark.parametrize("p", fxt_quat_struct)
@pytest.mark.parametrize("q", fxt_quat)
def test_quaternion_radd(p, q):
    result = p + q
    assert result == Quaternion(p) + q

@pytest.mark.parametrize("p", fxt_quat_struct)
@pytest.mark.parametrize("q", fxt_quat)
def test_quaternion_rsub(p, q):
    result = p - q
    assert result == Quaternion(p) - q

@pytest.mark.parametrize("p", fxt_quat_struct)
@pytest.mark.parametrize("q", fxt_quat)
def test_quaternion_rmul(p, q):
    result = p * q
    assert result == Quaternion(p) * q

@pytest.mark.parametrize("p", fxt_quat_struct)
@pytest.mark.parametrize("q", fxt_quat)
def test_quaternion_rtruediv(p, q):
    result = p / q
    assert result == Quaternion(p) / q

@pytest.mark.parametrize("struct", fxt_quat_struct_list)
@pytest.mark.parametrize("q", fxt_quat)
def test_quaternion_radd_struct(struct, q):
    result = struct + q
    assert all(r - q == Quaternion(p) for r, p in zip(result, struct))

@pytest.mark.parametrize("struct", fxt_quat_struct_list)
@pytest.mark.parametrize("q", fxt_quat)
def test_quaternion_rsub_struct(struct, q):
    result = struct - q
    assert all(r + q == Quaternion(p) for r, p in zip(result, struct))

@pytest.mark.parametrize("struct", fxt_quat_struct_list)
@pytest.mark.parametrize("q", fxt_quat)
def test_quaternion_rmul_struct(struct, q):
    result = struct * q
    assert all(r / q == Quaternion(p) for r, p in zip(result, struct))

@pytest.mark.parametrize("struct", fxt_quat_struct_list)
@pytest.mark.parametrize("q", fxt_quat)
def test_quaternion_rtruediv_struct(struct, q):
    result = struct / q
    assert all(r * q == Quaternion(p) for r, p in zip(result, struct))

########################################
#    instance methods
########################################

@pytest.mark.parametrize("arg", fxt_quat_struct)
@pytest.mark.parametrize("q", fxt_quat)
def test_quaternion_group_conjugate(q, arg):
    expected = q * arg * q.inverse
    result = q.group_conjugate(arg)
    assert result == expected

@pytest.mark.parametrize("arg", fxt_quat_struct_list)
@pytest.mark.parametrize("q", fxt_quat)
def test_quaternion_group_conjugate_struct(q, arg):
    expected = q * arg * q.inverse
    result = q.group_conjugate(arg)
    assert all(r == e for r, e in zip(result, expected))

# 'AXphi' means rotation of angle phi around axis AX
UNIT_QUATERNION = {
    "x90": np.sqrt(2) / 2 * Quaternion([1, 1, 0, 0]),
    "y90": np.sqrt(2) / 2 * Quaternion([1, 0, 1, 0]),
    "z90": np.sqrt(2) / 2 * Quaternion([1, 0, 0, 1]),
}
NON_UNIT_QUATERNION = [
    Quaternion(0),
    Quaternion(-9),
    Quaternion([1, 2, 3]),
    Quaternion([1, 2, 3, 4]),
    Quaternion((-1, [3, 2, 1])),
]

EULER_SEQUENCE = ["xyz", "xzy", "yxz", "yzx", "zxy", "zyx"]

@pytest.mark.parametrize("q", NON_UNIT_QUATERNION)
def test_quaternion_axis_angle_fails(q):
    with pytest.raises(ArithmeticError):
        q.axis_angle()

@pytest.mark.parametrize(
    "u, exp_axis, exp_angle",
    [
        (Quaternion(1), np.array([1, 0, 0]), 0,),
        (
            1 / 2 * Quaternion([1, 1, 1, 1]),
            np.array([1, 1, 1]) / np.sqrt(3),
            2 / 3 * np.pi,
        ),
        (UNIT_QUATERNION["x90"], np.array([1, 0, 0]), 1 / 2 * np.pi,),
        (UNIT_QUATERNION["y90"], np.array([0, 1, 0]), 1 / 2 * np.pi,),
        (UNIT_QUATERNION["z90"], np.array([0, 0, 1]), 1 / 2 * np.pi,),
        (
            Quaternion((np.sqrt(3) / 2, np.array([1, 1, 0]) / 2 / np.sqrt(2))),
            np.array([1, 1, 0]) / np.sqrt(2),
            1 / 3 * np.pi,
        ),
    ],
)
def test_quaternion_axis_angle(u, exp_axis, exp_angle):
    axis, angle = u.axis_angle()

    assert np.allclose(axis, exp_axis)
    assert np.allclose(angle, exp_angle)

@pytest.mark.parametrize("q", NON_UNIT_QUATERNION)
def test_quaternion_rot_matrix_fails(q):

    with pytest.raises(ArithmeticError):
        q.rot_matrix()

@pytest.mark.parametrize(
    "quat, exp",
    [
        (Quaternion(1), np.eye(3)),
        (UNIT_QUATERNION["x90"], np.array([[1, 0, 0], [0, 0, -1], [0, 1, 0]]),),
        (UNIT_QUATERNION["y90"], np.array([[0, 0, 1], [0, 1, 0], [-1, 0, 0]]),),
        (UNIT_QUATERNION["z90"], np.array([[0, -1, 0], [1, 0, 0], [0, 0, 1]]),),
        (1 / 2 * Quaternion([1, 1, 1, 1]), np.array([[0, 0, 1], [1, 0, 0], [0, 1, 0]])),
    ],
)
def test_quaternion_rot_matrix(quat, exp):
    result = quat.rot_matrix()

    print(result)

    assert np.allclose(result, exp)

@pytest.mark.parametrize("q", NON_UNIT_QUATERNION)
def test_quaternion_euler_angles_fails_unit(q):
    with pytest.raises(ArithmeticError):
        q.euler_angles("x", "y", "z")

@pytest.mark.parametrize("axes", ["abc", "XYZ", "xxy", "yzy",])
def test_quaternion_euler_angles_fails_axes(axes):
    q = Quaternion(1)

    with pytest.raises(ValueError):
        q.euler_angles(*list(axes))

@pytest.mark.parametrize("axes", EULER_SEQUENCE)
def test_quaternion_euler_angles_identity(axes):

    result = Quaternion(1).euler_angles(*list(axes))
    assert result == (0, 0, 0)

@pytest.mark.parametrize(
    "quat, axes, expected",
    [
        (UNIT_QUATERNION["x90"], "xyz", (np.pi / 2, 0, 0)),
        (UNIT_QUATERNION["y90"], "xyz", (0, np.pi / 2, 0)),
        (UNIT_QUATERNION["z90"], "xyz", (0, 0, np.pi / 2)),
        (UNIT_QUATERNION["x90"], "xzy", (np.pi / 2, 0, 0)),
        (UNIT_QUATERNION["y90"], "xzy", (0, 0, np.pi / 2)),
        (UNIT_QUATERNION["z90"], "xzy", (0, np.pi / 2, 0)),
        (UNIT_QUATERNION["y90"], "yxz", (np.pi / 2, 0, 0)),
        (UNIT_QUATERNION["x90"], "yxz", (0, np.pi / 2, 0)),
        (UNIT_QUATERNION["z90"], "yxz", (0, 0, np.pi / 2)),
        (UNIT_QUATERNION["x90"], "yzx", (0, 0, np.pi / 2)),
        (UNIT_QUATERNION["y90"], "yzx", (np.pi / 2, 0, 0)),
        (UNIT_QUATERNION["z90"], "yzx", (0, np.pi / 2, 0)),
        (UNIT_QUATERNION["x90"], "zxy", (0, np.pi / 2, 0)),
        (UNIT_QUATERNION["y90"], "zxy", (0, 0, np.pi / 2)),
        (UNIT_QUATERNION["z90"], "zxy", (np.pi / 2, 0, 0)),
        (UNIT_QUATERNION["x90"], "zyx", (0, 0, np.pi / 2)),
        (UNIT_QUATERNION["y90"], "zyx", (0, np.pi / 2, 0)),
        (UNIT_QUATERNION["z90"], "zyx", (np.pi / 2, 0, 0)),
    ],
)
def test_quaternion_euler_angles(quat, axes, expected):

    result = quat.euler_angles(*list(axes))
    assert result == expected

@pytest.mark.parametrize("quat", list(UNIT_QUATERNION.values()))
@pytest.mark.parametrize("axes", EULER_SEQUENCE)
def test_quaternion_euler_angles_euler_rotations(quat, axes):
    axis3, axis2, axis1 = list(axes)
    angle3, angle2, angle1 = quat.euler_angles(axis3, axis2, axis1)
    composition = (
        euler_rotation(axis3, angle3)
        @ euler_rotation(axis2, angle2)
        @ euler_rotation(axis1, angle1)
    )
    rot_mat = quat.rot_matrix()

    assert np.allclose(rot_mat, composition)

########################################
#    class methods
########################################

@pytest.mark.parametrize(
    "arg",
    [
        -2,
        [1, 2, 3],
        np.array([-1, 0, -1]),
        [0, 1, 2, 3],
        np.array([3, -1, 0, -1]),
        (0, [1, 2, 3]),
        (0, np.array([1, 2, 3])),
        Quaternion([1, -1, 0, 1]),
    ],
)
def test_to_quaternion(arg):
    result = Quaternion.to_quaternion(arg)
    assert result == Quaternion(arg)

@pytest.mark.parametrize(
    "arg",
    [
        [[1, 2, 3], [-1, -1, 0], [2, 4, 7]],
        [np.array([1, 2, 3]), np.array([2, 3, 4])],
        np.array([[0, 1, 0], [-1, -2, -3], [1, 1, 1]]),
        np.repeat(-1.5, 30).reshape(-1, 3),
        [[0, 1, 2, 3], [-1, -1, 0, -2]],
        [np.array([1, 2, 4, 3]), np.array([2, -3, 3, 4])],
        np.array([[3, 0, 1, 0], [2, 1, 1, 1]]),
        np.repeat(-2.5, 16).reshape(-1, 4),
        [(0, [1, 2, 3]), (0, np.array([1, 2, 3])), (-2, [-1, -2, 0])],
        [Quaternion(0), Quaternion([1, 2, 3]), Quaternion([1, -1, 0, 1]),],
    ],
)
def test_to_quaternion_struct(arg):
    result = Quaternion.to_quaternion(arg)
    assert all(r == Quaternion(a) for r, a in zip(result, arg))

@pytest.mark.parametrize(
    "axis, angle, error",
    [
        ([1, 0], 1.0, ValueError),
        ([1, 0, 1], "1.5", ValueError),
        ([1e-17, 0, -1e-20], 1.2, ArithmeticError),
    ],
)
def test_rotation_quat_from_axis_angle_invalid(axis, angle, error):
    with pytest.raises(error):
        Quaternion.rotation_quat_from_axis_angle(axis, angle)

@pytest.mark.parametrize(
    "axis, angle, expected",
    [
        ([1, 0, 0], 0.0, Quaternion(1)),  # angle=0 -> identity rotation
        ([22, 0, 0], np.pi / 2, UNIT_QUATERNION["x90"],),  # axis=X-axis, angle=90
        (
            [0, -3, 0],
            np.pi,  # axis=neg. Y-axis, angle=180
            Quaternion(np.array([0, 0, -1, 0])),
        ),
        (
            [0, 0, 5],
            np.pi / 3,  # axis=Z-axis, angle=60
            Quaternion(np.array([np.sqrt(3) / 2, 0, 0, 1 / 2])),
        ),
    ],
)
def test_rotation_quat_from_axis_angle(axis, angle, expected):
    result = Quaternion.rotation_quat_from_axis_angle(axis, angle)
    assert result == expected

@pytest.mark.parametrize(
    "axis, angle, vectors",
    [
        ([1, 0, 0], 1.0, list("vectors")),
        ([1, 1, 1], 1.0, np.ones(6)),
        ([1, 0, 1], 1.0, np.ones(40).reshape(-1, 4)),
    ],
)
def test_rotate_by_axis_angle_invalid(axis, angle, vectors):
    with pytest.raises(ValueError):
        Quaternion.rotate_by_axis_angle(axis, angle, vectors)

@pytest.mark.parametrize(
    "axis, angle, vector",
    [
        ([1, 0, 0], 1.0, [1, 1, 1]),
        ([1, 0, -2], 1.2, np.array([0, 1, -1])),
        (np.array([1, 2, -3]), -0.4, [-1, 1.5, 2.5]),
    ],
)
def test_rotate_by_axis_angle(axis, angle, vector):
    expected = rotation_axis_angle(np.array(axis), angle) @ np.array(vector)
    result = Quaternion.rotate_by_axis_angle(axis, angle, vector)
    assert np.allclose(result, expected)

@pytest.mark.parametrize("n", [10, 20, 50, 100])
@pytest.mark.parametrize("angle", [1.2, np.pi / 3, 2.0, 3.3])
@pytest.mark.parametrize(
    "axis", [np.array([1, 1, 1]), np.array([2, 0, -1]), np.array([-1, 1.5, -1.3]),]
)
def test_rotate_by_axis_angle_vectors(n, axis, angle):
    # create vectors pointing in axis direction
    vectors = np.tile(axis, [n, 1])
    # vectors are fix points, we expect no transform
    expected = vectors
    # rotate them
    rotated = Quaternion.rotate_by_axis_angle(axis, angle, vectors)
    # test
    assert np.allclose(rotated, expected)

    # create vectors perpendicular to axis direction
    vectors = np.tile(np.array([axis[2], 0, -axis[0]]), [n, 1])
    # calculate the expected transform
    rot_matrix = rotation_axis_angle(axis, angle)
    expected = [rot_matrix @ p for p in vectors]
    # rotate them
    rotated = Quaternion.rotate_by_axis_angle(axis, angle, vectors)
    # test
    assert np.allclose(rotated, expected)


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

# In[ ]:

# This Python 3 environment comes with many helpful analytics libraries installed
# It is defined by the kaggle/python docker image: https://github.com/kaggle/docker-python
# For example, here's several helpful packages to load in

import numpy as np # linear algebra
import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)
import matplotlib.pyplot as plt
import seaborn as sns
import xgboost as xgb

# Input data files are available in the "../input/" directory.
# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory

import os
for dirname, _, filenames in os.walk('/kaggle/input'):
    for filename in filenames:
        print(os.path.join(dirname, filename))

# Any results you write to the current directory are saved as output.

# In[ ]:

df = pd.read_csv('/kaggle/input/quality-prediction-in-a-mining-process/MiningProcess_Flotation_Plant_Database.csv', decimal=',')
df['date'] = pd.to_datetime(df['date'])
df.head(2)

# In[ ]:

print('Shape: ', df.shape)
print('Columns: ')
print(df.columns)
print('Datatypes:')
print(df.dtypes)

# Show the hours with less than 180 records (missing data within the hour).

# In[ ]:

counts = df.groupby('date').count()
counts[counts['% Iron Feed'] < 180]

# Luckily only the first hour is missing 6 data points and one other hour is missing one. When creating an accurate time series index, we will arbitrarily take out the first couple 20 second intervals to match the amount of records found for those hours.
#
# We will now create the 20 second frequency Datetime Index.

# In[ ]:

# get a series of unique hourly timestamps
hours = pd.Series(df['date'].unique())
hours.index = hours
len(hours)

# In[ ]:

# create a date time index from the first to the last hour included in the date column
date_range = pd.date_range(start=df.iloc[0,0], end='2017-09-09 23:59:40', freq='20S')
# remove first couple observations consistent with the counts exploration above
date_range = date_range[6:]
date_range[-5:]

# In[ ]:

# create lists from both the hours series and the new datetime index
hours_list = hours.index.format()
print(hours_list[:5])
seconds_list = date_range.format()
print(seconds_list[:5])

# In[ ]:

# match the new datetime index to the hours series and only append the timestamps if the datea and hour match the hours list
new_index = []
for idx in seconds_list:
    if (idx[:13] + ':00:00') in hours_list:
        new_index.append(idx)

#remove the one missing interval within the hour which we found earlier using the counts
new_index.remove('2017-04-10 00:00:00')
new_index[-20:]

# In[ ]:

print(len(new_index))
print(len(df))

# In[ ]:

df['index'] = new_index
df['index'] = pd.to_datetime(df['index'])
df.index = df['index']
df = df.loc[:, df.columns[:-1]]
df.rename(columns={'date': 'datetime hours'}, inplace=True)
df.head()

# ### Checking which variables have hourly vs 20-sec frequency
#
# We can determine the frequency of the variables by grouping the dataframe by hours and counting the number of unique values. For hourly variables it should be 1, for the higher frequency variables it should be close to 180.

# In[ ]:

unique_avg = []
for col in df.columns:
    unique_avg.append(df.groupby('datetime hours').apply(lambda x: len(x[col].unique())).mean())
plt.plot(np.arange(len(unique_avg)), unique_avg)
plt.title('Average Count of Unique Values per Hour for every Variable')
plt.ylabel('Count')
plt.xticks(list(range(len(unique_avg))), list(df.columns), rotation='vertical')
plt.show()

# Only the Iron and Silica Feed and Concentrate variables seem hourly, the rest seems to contain higher frequency measurements. Yet, the unique averages are much higher than 1 for Silica Concentrate especially, which could indicate some inconsistencies.

# ### Interpolation Cleaning

# We further checked individual variables to see if there are any outliers etc. We noticed that there seemed to be some interpolated values which can be detrimental to any modelling attempts.

# In[ ]:

# some values for Silica Concentration seem interpolated so we're removing the values for all those hours
#some imports
from pandas.plotting import register_matplotlib_converters
register_matplotlib_converters()

#get list of % Silica Concentrate values, see the first 5 that have more than one hourly value
silica_unique = df.groupby('datetime hours').apply(lambda x: len(x['% Silica Concentrate'].unique()))
print(silica_unique[silica_unique > 1][:5])

# plot before the interpolations are taken out of the dataframe
plt.plot(df['% Silica Concentrate'][df['datetime hours'] == silica_unique[silica_unique > 1].index[0]])

# take out interpolated hours for % Silica Concentrate
interpolated_hours = silica_unique[silica_unique >1].index.format()
clean_df=df[~df['datetime hours'].isin(interpolated_hours)]

#finish the graph
plt.title('Interpolated hour of % Silica Concentrate')
plt.xlabel('Datetime Index [Day Hour:Minute]')
plt.ylabel('Silica Concentrate [%]')
plt.legend(loc='best')
plt.show()

# Next, we will graph some of the input variables and check for more interpolations. Both Iron Feed and Silica Feed seem to have a large amount of interpolated values.

# In[ ]:

plt.plot(clean_df.index, clean_df['% Iron Feed'])
plt.plot(clean_df.index, clean_df['% Silica Feed'])
plt.title('Iron and Silica Percentage of Input Feed')
plt.legend(loc='best')
plt.ylabel('Iron/Silica Feed [%]')
plt.xlabel('Date [Year-Month]')
plt.show()

# It seems that there are some missing values for the Iron and Silica Feed as well, so let us investigate the most frequently occuring values.

# In[ ]:

#if_unique = clean_df.groupby('datetime hours').apply(lambda x: len(x['% Iron Feed'].unique()))
#sf_unique = clean_df.groupby('datetime hours').apply(lambda x: len(x['% Silica Feed'].unique()))
#print(if_unique[if_unique > 1][:5])
print('Count of unique hours in cleaned df: ', len(clean_df.groupby('datetime hours').mean()))
print('Count of unique % Iron Feed values: ',len(clean_df['% Iron Feed'].unique()))
print('Count of unique % Silica Feed values: ',len(clean_df['% Silica Feed'].unique()))
print('Reference: Count of unique % Silica Concentrate values: ',len(clean_df['% Silica Concentrate'].unique()))

# In[ ]:

# function to get unique values of a df column and their counts
def get_unique_counts(column):
    df = pd.DataFrame()

    uv_list, count_list = list(column.unique()), []

    for uv in uv_list:
        count_list.append(len(column[column == uv]))

    df['unique_values'] = uv_list
    df['count'] = count_list
    return df

# Less than

# In[ ]:

if_unique = get_unique_counts(clean_df['% Iron Feed']).sort_values('count',ascending=False)
sf_unique = get_unique_counts(clean_df['% Silica Feed']).sort_values('count',ascending=False)
print(if_unique.head(10))
print(sf_unique.head(10))

# The four highest frequencies show the same count for both variables, let's look at the graphs to confirm that these were interpolated.

# In[ ]:

for i in range(6):
    clean_df['% Silica Feed'][clean_df['% Silica Feed'] == sf_unique.iloc[i,0]].plot()
    clean_df['% Iron Feed'][clean_df['% Iron Feed'] == if_unique.iloc[i,0]].plot()
    plt.show()

# In[ ]:

clean_df.groupby([clean_df.index.date, clean_df.index.hour]).mean()

# I will remove the intervals that feature seemingly unclean data, i.e. the four highest frequency observations.

# In[ ]:

dirty_idx = []
for i in range(4):
    dirty_idx.extend(clean_df['% Silica Feed'][clean_df['% Silica Feed'] == sf_unique.iloc[i,0]].index.format())
dirty_idx
print(len(dirty_idx), len(clean_df))
clean_df=clean_df[~clean_df.index.isin(dirty_idx)]
print(clean_df.shape)
clean_df['% Silica Feed'].plot()
clean_df['% Iron Feed'].plot()
plt.show()

# ### Correlation Plots

# In[ ]:

pair_cols = list(df.columns[1:8])
pair_cols.extend(df.columns[-2:])
print(pair_cols)
#smol_df = clean_df.loc[:,pair_cols]
sns.pairplot(clean_df.loc[:,pair_cols])
plt.show()

# No apparent meaningful patterns besides between Iron and Silica Concentrate and Iron and Silica Feed (which are to be expected).
#
# We also decided to check which minute of the hour showed the highest correlation with % Silica Concentrate for each variable. Our hypothesis was that they should peak around when the measurements where usually taken.

# In[ ]:

# minute correlations
corr_df = pd.DataFrame(index=clean_df.columns[1:])
for minute in range(60):
    min_df = clean_df[clean_df.index.minute == minute]
    corr_df[str(minute)] = min_df.groupby([min_df.index.date, min_df.index.hour, min_df.index.minute]).mean().corr().iloc[:,-1]
corr_df = corr_df.transpose()

corr_df.iloc[:,:-2].plot(legend=False)
plt.title("Correlations of Variables vs Silica Concentrate Grouped by Minute")
plt.ylabel("Correlation")
plt.xlabel("Minute of the Hour")
plt.show()

# In[ ]:

def rmse(actual, preds):
    return np.sqrt(np.sum((np.array(actual)-np.array(preds))**2) / len(actual))
def mape(actual, preds):
    return np.sum(np.abs((np.array(actual)-np.array(preds))/(np.array(actual)))) / len(actual)
def mae(actual, preds):
    return np.sum(np.abs((np.array(actual)-np.array(preds)))) / len(actual)

# ### XDBoost

# In[ ]:

import xgboost as xgb

# In[ ]:

start=0
end=100

X = clean_df.iloc[start*24*180:end*24*180,1:-2]
y = clean_df.iloc[start*24*180:end*24*180,-1]
X_test = clean_df.iloc[(end)*24*180:,1:-2]
y_test = clean_df.iloc[(end)*24*180:,-1]

xgbr= xgb.XGBRegressor(max_depth=8, n_estimators=50, min_sample_split = 500, subsample=0.5, silent=True, colsample_bytree=0.8, gamma=100)

xgbr.fit(X,y)
print(xgbr.feature_importances_)
print(xgbr.score(X,y))

# In[ ]:

preds=xgbr.predict(X_test)
pred_df = pd.DataFrame(preds, columns=['predictions'])
print('Train RMSE: ' + str(rmse(y, xgbr.predict(X))))
print('RMSE: ' + str(rmse(y_test, preds)))
print('MAE: ' + str(mae(y_test, preds)))
print('MAPE: ' + str(mape(y_test, preds)))
plt.plot(y_test, label = 'Actual')
pred_df.index = y_test.index
plt.plot(pred_df, label = 'Prediction')
plt.legend()
plt.ylim(0,6)
plt.title('XGBoost Regressor Model Forecast')
plt.xticks(rotation='vertical')
plt.ylabel('Silica Concentrate [%]')
plt.xlabel('Time [Days]')
plt.show()

# ### Ridge Regression

# In[ ]:

from sklearn.linear_model import Ridge

start=0
end=100
X = clean_df.iloc[start*24*180:end*24*180,1:-2]
y = clean_df.iloc[start*24*180:end*24*180,-1]
X_test = clean_df.iloc[(end)*24*180:,1:-2]
y_test = clean_df.iloc[(end)*24*180:,-1]

rr= Ridge(alpha= 1, fit_intercept=False, normalize=True)

rr.fit(X,y)
print(rr.get_params)
print(rr.score(X,y))
print(rr.score(X_test,y_test))

# In[ ]:

preds=rr.predict(X_test)
pred_df = pd.DataFrame(preds, columns=['predictions'])
y_pred= df.iloc[140*24*180:147*24*180,-1]
print('Train RMSE: ', rmse(y, rr.predict(X)))
print('RMSE: ', rmse(y_test, preds))
print('MAPE: ', mape(y_test, preds))
print('MAE: ', mae(y_test, preds))
plt.plot(y_test, label = 'Actual')
pred_df.index = y_test.index
plt.plot(pred_df, label = 'Predictions')
plt.title('Ridge Regression Model Forecast')
plt.xticks(rotation='vertical')
plt.ylabel('Silica Concentrate [%]')
plt.xlabel('Time [Days]')
plt.ylim(0,6)
plt.legend()
plt.show()

# # Classification

# In[ ]:

sns.distplot(clean_df['% Silica Concentrate'])
plt.title('Distribution Plot for % Silica Concentrate')
plt.ylabel('Relative Frequency')
plt.xlabel('Silica Concentrate [%]')
plt.show()

# In[ ]:

#create hour column
cdf = clean_df.copy(deep=True)
cdf['hour'] = cdf.index.hour

# get labels
cdf['label'] = 0
cdf['label'][cdf['% Silica Concentrate'] > 3] = 1
print(cdf['label'][cdf['label'] == 1].count())
print(cdf['label'][cdf['label'] == 0].count())
print(cdf['label'][cdf['label'] == 0].count() / cdf['label'][cdf['label'] == 1].count())

# In[ ]:

import random

random.seed(69)
#start=0
#end=138

mdf = cdf.drop(columns = ['datetime hours', '% Iron Concentrate', '% Silica Concentrate'])

#create and sample train set for equal class distribution
#train = mdf.iloc[start*24*180:end*24*180]
train = mdf.iloc[:-14*24*180]
zero_idx = train[train['label'] == 0].index
sample_idx = random.sample(list(zero_idx), train[train['label'] == 1].shape[0])
sample_idx.extend(list(train[train['label'] == 1].index))
sample_idx = pd.DatetimeIndex(sample_idx).sort_values()
train = train.reindex(sample_idx)

X = train.iloc[:,:-1]
y = train.iloc[:,-1]

#X_eval = mdf.iloc[(end)*24*180:(end+7)*24*180,:-1]
#y_eval = mdf.iloc[(end)*24*180:(end+7)*24*180,-1]
#X_test = mdf.iloc[(end+7)*24*180:(end+14)*24*180,:-1]
#y_test = mdf.iloc[(end+7)*24*180:(end+14)*24*180,-1]
#X_eval = mdf.iloc[-14*24*180:-7*24*180,:-1]
#y_eval = mdf.iloc[-14*24*180:-7*24*180,-1]
X_test = mdf.iloc[-14*24*180:,:-1]
y_test = mdf.iloc[-14*24*180:,-1]

print(y[y == 0].count() / y[y==1].count())

# ### Initial XGBoost attempt

# In[ ]:

xgbc= xgb.XGBClassifier(max_depth=4, n_estimators=5, subsample=0.5, eval_metric='logloss', colsample_bytree=0.8,
                        min_child_weight=100, gamma=50)

xgbc.fit(X,y)
print(xgbc.feature_importances_)
print(xgbc.score(X,y))

# In[ ]:

from sklearn.metrics import confusion_matrix, accuracy_score, classification_report
preds=xgbc.predict(X_test)
results = confusion_matrix(y_test, preds)
print('Test Set Results')
print('Confusion Matrix :')
print(results)
print('Accuracy Score :',accuracy_score(y_test, preds) )
print('Report : ')
print(classification_report(y_test, preds))

# In[ ]:

eval_preds=xgbc.predict(X_eval)
eval_results = confusion_matrix(y_eval, eval_preds)
print('Evaluation Set Results')
print('Confusion Matrix :')
print(eval_results)
print('Accuracy Score :',accuracy_score(y_eval, eval_preds) )
print('Report : ')
print(classification_report(y_eval, eval_preds))

# In[ ]:

train_preds=xgbc.predict(X)
train_results = confusion_matrix(y, train_preds)
print('Training Set Results')
print('Confusion Matrix :')
print(train_results)
print('Accuracy Score :',accuracy_score(y, train_preds) )
print('Report : ')
print(classification_report(y, train_preds))

# In[ ]:

from sklearn.utils.multiclass import unique_labels
from sklearn.metrics import confusion_matrix, accuracy_score, classification_report

def plot_confusion_matrix(y_true, y_pred, classes,
                          normalize=False,
                          title=None,
                          cmap=plt.cm.Blues):
    """
    This function prints and plots the confusion matrix.
    Normalization can be applied by setting `normalize=True`.
    """
    if not title:
        if normalize:
            title = 'Normalized confusion matrix'
        else:
            title = 'Confusion matrix, without normalization'

    # Compute confusion matrix
    cm = confusion_matrix(y_true, y_pred)
    # Only use the labels that appear in the data
    classes = classes[unique_labels(y_true, y_pred)]
    if normalize:
        cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
        #print("Normalized confusion matrix")
    else:
        print()
        #print('Confusion matrix, without normalization')

    fig, ax = plt.subplots()
    im = ax.imshow(cm, interpolation='nearest', cmap=cmap)
    ax.figure.colorbar(im, ax=ax)
    # We want to show all ticks...
    ax.set(xticks=np.arange(cm.shape[1]),
           yticks=np.arange(cm.shape[0]),
           # ... and label them with the respective list entries
           xticklabels=classes, yticklabels=classes,
           title=title,
           ylabel='True label',
           xlabel='Predicted label')
    # Rotate the tick labels and set their alignment.
    plt.setp(ax.get_xticklabels(), rotation=45, ha="right",
             rotation_mode="anchor")

    # Loop over data dimensions and create text annotations.
    fmt = '.2f' if normalize else 'd'
    thresh = cm.max() / 2.
    for i in range(cm.shape[0]):
        for j in range(cm.shape[1]):
            ax.text(j, i, format(cm[i, j], fmt),
                    ha="center", va="center",
                    color="white" if cm[i, j] > thresh else "black")
    fig.tight_layout()
    return ax

# In[ ]:

plot_confusion_matrix(np.array(y_test), np.array(preds), classes=np.array(['pure (below cutoff)', 'impure (above cutoff)']), normalize=False)
plt.show()

# ## Grid Search

# In[ ]:

import random

random.seed(69)
#start=0
end=-14

mdf = cdf.drop(columns = ['datetime hours', '% Iron Concentrate', '% Silica Concentrate'])

#create and sample train set for equal class distribution
#train = mdf.iloc[start*24*180:end*24*180]
train = mdf.iloc[:-14*24*180]
zero_idx = train[train['label'] == 0].index
sample_idx = random.sample(list(zero_idx), train[train['label'] == 1].shape[0])
sample_idx.extend(list(train[train['label'] == 1].index))
sample_idx = pd.DatetimeIndex(sample_idx).sort_values()
train = train.reindex(sample_idx)

X = train.iloc[:,:-1]
y = train.iloc[:,-1]

X_eval = mdf.iloc[(end)*24*180:(end+7)*24*180,:-1]
y_eval = mdf.iloc[(end)*24*180:(end+7)*24*180,-1]
X_test = mdf.iloc[(end+7)*24*180:,:-1]
y_test = mdf.iloc[(end+7)*24*180:,-1]
#X_test = mdf.iloc[-14*24*180:,:-1]
#y_test = mdf.iloc[-14*24*180:,-1]

print(y[y == 0].count() / y[y==1].count())

# In[ ]:

#max_depth_list = [5,7,10,15]
#n_trees_list  = [50, 75, 100, 150, 200]

#feature_1 = []
#feature_2 = []
#train_acc = []
#test_acc = []
#test_precision = []
#test_recall = []
#test_trueones = []

#for max_depth in max_depth_list:
#    for n_tree in n_trees_list:
#        xgbc= xgb.XGBClassifier(max_depth=max_depth, n_estimators=n_tree, subsample=0.5, eval_metric='logloss', colsample_bytree=0.8,
#                        min_child_weight=100, gamma=50)
#
#        xgbc.fit(X,y)
#        pred = xgbc.predict(X_test)
#
#        feature_1.append(max_depth)
#        feature_2.append(n_tree)
#        train_acc.append(xgbc.score(X,y))
#        test_acc.append(xgbc.score(X_test, y_test))
#        cm = confusion_matrix(y_test, pred)
#        test_trueones.append(cm[1,1])
#        test_precision.append((cm[1,1]) / (cm[0,1] + cm[1,1]))
#        test_recall.append((cm[1,1]) / (cm[1,0] + cm[1,1]))
#        print(max_depth,n_tree)

# In[ ]:

#result_df = pd.DataFrame()
#result_df['feature_1'] = feature_1
#result_df['feature_2'] = feature_2
#result_df['train_acc'] = train_acc
#result_df['test_acc'] = test_acc
#result_df['test_precision'] = test_precision
#result_df['test_recall'] = test_recall
#result_df['test_trueones'] = test_trueones
#result_df

# In[ ]:

#max_depth_list2 = [2,3,4,5,6,8]
#n_trees_list2  = [3,5,6,7,8,9,10]#

#feature_12 = []
#feature_22 = []
#train_acc2 = []
#test_acc2 = []
#test_precision2 = []
#test_recall2 = []
#test_trueones2 = []

#for max_depth in max_depth_list2:
#    for n_tree in n_trees_list2:
#        xgbc= xgb.XGBClassifier(max_depth=max_depth, n_estimators=n_tree, subsample=0.5, eval_metric='logloss', colsample_bytree=0.8,
#                        min_child_weight=100, gamma=50)
#
#        xgbc.fit(X,y)
#        pred = xgbc.predict(X_eval)
#
#        feature_12.append(max_depth)
#        feature_22.append(n_tree)
#        train_acc2.append(xgbc.score(X,y))
 #       test_acc2.append(xgbc.score(X_eval, y_eval))
#        cm = confusion_matrix(y_eval, pred)
#        test_trueones2.append(cm[1,1])
#        test_precision2.append((cm[1,1]) / (cm[0,1] + cm[1,1]))
#        test_recall2.append((cm[1,1]) / (cm[1,0] + cm[1,1]))
#        #print(max_depth,n_tree)

# In[ ]:

#result_df2 = pd.DataFrame()
#result_df2['max_depth'] = feature_12
#result_df2['n_trees'] = feature_22
#result_df2['train_acc'] = train_acc2
#result_df2['test_acc'] = test_acc2
#result_df2['test_precision'] = test_precision2
#result_df2['test_recall'] = test_recall2
#result_df2['test_trueones'] = test_trueones2
#result_df2

# In[ ]:

from sklearn.metrics import confusion_matrix, accuracy_score, classification_report
max_depth_list = [2,3,4,5]
n_trees_list  = [2,3,4,5,6,7,8]
min_child_weight_list = [10,50,100,200,300]
gamma_list = [1,10,30,50,100]

feature_1 = []
feature_2 = []
feature_3 = []
feature_4 = []
train_acc = []
test_acc = []
test_precision = []
test_recall = []
test_trueones = []

for max_depth in max_depth_list:
    for n_tree in n_trees_list:
        for min_child_weight in min_child_weight_list:
            for gamma in gamma_list:
                xgbc= xgb.XGBClassifier(max_depth=max_depth, n_estimators=n_tree, subsample=0.5, eval_metric='logloss', #colsample_bytree=0.8,
                        min_child_weight=min_child_weight, gamma=gamma)

                xgbc.fit(X,y)
                pred = xgbc.predict(X_eval)

                feature_1.append(max_depth)
                feature_2.append(n_tree)
                feature_3.append(min_child_weight)
                feature_4.append(gamma)
                train_acc.append(xgbc.score(X,y))
                test_acc.append(xgbc.score(X_eval, y_eval))
                cm = confusion_matrix(y_eval, pred)
                test_trueones.append(cm[1,1])
                test_precision.append((cm[1,1]) / (cm[0,1] + cm[1,1]))
                test_recall.append((cm[1,1]) / (cm[1,0] + cm[1,1]))
                #print(max_depth,n_tree)

# In[ ]:

result_df3 = pd.DataFrame()
result_df3['max_depth'] = feature_1
result_df3['n_trees'] = feature_2
result_df3['min_child_weight'] = feature_3
result_df3['gamma'] = feature_4
result_df3['train_acc'] = train_acc
result_df3['test_acc'] = test_acc
result_df3['test_precision'] = test_precision
result_df3['test_recall'] = test_recall
result_df3['test_trueones'] = test_trueones
result_df3.to_csv('/kaggle/working/xgb_grid_search_results.csv')
result_df3

# In[ ]:

result_df3.sort_values(['test_trueones'], ascending=False).head(50)

# In[ ]:

result_df3.groupby([result_df3['max_depth'], result_df3['n_trees']])['test_trueones'].max()

# In[ ]:

result_df3[(result_df3['max_depth'] == 4) & (result_df3['n_trees'] == 8)]

# In[ ]:

xgbc= xgb.XGBClassifier(max_depth=5, n_estimators=2, subsample=0.5, eval_metric='logloss', min_child_weight = 300)

xgbc.fit(X,y)
print(xgbc.feature_importances_)
print(xgbc.score(X,y))

# In[ ]:

from sklearn.metrics import confusion_matrix, accuracy_score, classification_report
preds=xgbc.predict(X_eval)
results = confusion_matrix(y_eval, preds)
print('Test Set Results')
print('Confusion Matrix :')
print(results)
print('Accuracy Score :',accuracy_score(y_eval, preds) )
print('Report : ')
print(classification_report(y_eval, preds))

# In[ ]:

plot_confusion_matrix(np.array(y_eval), np.array(preds), classes=np.array(['pure (below cutoff)', 'impure (above cutoff)']), normalize=False)
plt.show()

# In[ ]:

from sklearn.metrics import roc_curve, roc_auc_score
fpr, tpr, thresholds = roc_curve(np.array(y_eval), xgbc.predict_proba(X_eval)[:,1])
roc_auc = roc_auc_score(np.array(y_eval), xgbc.predict_proba(X_eval)[:,1])

plt.figure()
lw = 2
plt.plot(fpr, tpr, color='darkorange', lw=lw, label='ROC curve (area = %0.2f)' % roc_auc)
plt.plot([0,1], [0,1], linestyle="--")
plt.xlim([0.0, 1.0])
plt.ylim([0.0, 1.05])
plt.xlabel('False Positive Rate')
plt.ylabel('True Positive Rate')
plt.title('ROC Curve')
plt.legend(loc="lower right")
plt.show()

# In[ ]:

xgbc= xgb.XGBClassifier(max_depth=3, n_estimators=7, subsample=0.5, eval_metric='logloss', min_child_weight=300)

xgbc.fit(X,y)
print(xgbc.feature_importances_)
print(xgbc.score(X,y))

# In[ ]:

from sklearn.metrics import confusion_matrix, accuracy_score, classification_report
preds=xgbc.predict(X_eval)
results = confusion_matrix(y_eval, preds)
print('Test Set Results')
print('Confusion Matrix :')
print(results)
print('Accuracy Score :',accuracy_score(y_eval, preds) )
print('Report : ')
print(classification_report(y_eval, preds))
plot_confusion_matrix(np.array(y_eval), np.array(preds), classes=np.array(['pure (below cutoff)', 'impure (above cutoff)']), normalize=False)
plt.show()

# In[ ]:

from sklearn.metrics import roc_curve, roc_auc_score
fpr, tpr, thresholds = roc_curve(np.array(y_eval), xgbc.predict_proba(X_eval)[:,1])
roc_auc = roc_auc_score(np.array(y_eval), xgbc.predict_proba(X_eval)[:,1])

plt.figure()
lw = 2
plt.plot(fpr, tpr, color='darkorange', lw=lw, label='ROC curve (area = %0.2f)' % roc_auc)
plt.plot([0,1], [0,1], linestyle="--")
plt.xlim([0.0, 1.0])
plt.ylim([0.0, 1.05])
plt.xlabel('False Positive Rate')
plt.ylabel('True Positive Rate')
plt.title('ROC Curve')
plt.legend(loc="lower right")
plt.show()

# ### Logistic Reg

# In[ ]:

from sklearn.linear_model import LogisticRegression

lr = LogisticRegression(fit_intercept=False, C=0.1)

lr.fit(X,y)
print(lr.decision_function(X))
print(lr.score(X,y))

# In[ ]:

from sklearn.metrics import confusion_matrix, accuracy_score, classification_report
preds=lr.predict(X_eval)
results = confusion_matrix(y_eval, preds)
print('Test Set Results')
print('Confusion Matrix :')
print(results)
print('Accuracy Score :',accuracy_score(y_eval, preds) )
print('Report : ')
print(classification_report(y_eval, preds))

# #### Grid

# In[ ]:

intercept_list = [True, False]
C_list  = [10,1,0.1,0.01,0.001,0.0001]

feature_1 = []
feature_2 = []
train_acc = []
test_acc = []
test_precision = []
test_recall = []
test_trueones = []

for C in C_list:
    for intercept in intercept_list:
        lr = LogisticRegression(fit_intercept=intercept, C=C, solver='liblinear')

        lr.fit(X,y)
        pred = lr.predict(X_eval)

        feature_1.append(C)
        feature_2.append(intercept)
        train_acc.append(lr.score(X,y))
        test_acc.append(lr.score(X_eval, y_eval))
        cm = confusion_matrix(y_eval, pred)
        test_trueones.append(cm[1,1])
        test_precision.append((cm[1,1]) / (cm[0,1] + cm[1,1]))
        test_recall.append((cm[1,1]) / (cm[1,0] + cm[1,1]))
        #print(C,intercept)

# In[ ]:

result_df = pd.DataFrame()
result_df['C'] = feature_1
result_df['intercept'] = feature_2
#result_df['min_child_weight'] = feature_3
#result_df['gamma'] = feature_4
result_df['train_acc'] = train_acc
result_df['test_acc'] = test_acc
result_df['test_precision'] = test_precision
result_df['test_recall'] = test_recall
result_df['test_trueones'] = test_trueones
result_df

# In[ ]:

C_list  = [10, 7.5, 5, 4, 3, 2,1,0.75, 0.5, 0.3, 0.1,0.05, 0.01]

feature_1 = []
feature_2 = []
train_acc = []
test_acc = []
test_precision = []
test_recall = []
test_trueones = []

for C in C_list:
    lr = LogisticRegression(fit_intercept=True, C=C, solver='liblinear')

    lr.fit(X,y)
    pred = lr.predict(X_eval)

    feature_1.append(C)
    #feature_2.append(intercept)
    train_acc.append(lr.score(X,y))
    test_acc.append(lr.score(X_eval, y_eval))
    cm = confusion_matrix(y_eval, pred)
    test_trueones.append(cm[1,1])
    test_precision.append((cm[1,1]) / (cm[0,1] + cm[1,1]))
    test_recall.append((cm[1,1]) / (cm[1,0] + cm[1,1]))
    #print(C,intercept)

# In[ ]:

result_df2 = pd.DataFrame()
result_df2['C'] = feature_1
#result_df2['intercept'] = feature_2
result_df2['train_acc'] = train_acc
result_df2['test_acc'] = test_acc
result_df2['test_precision'] = test_precision
result_df2['test_recall'] = test_recall
result_df2['test_trueones'] = test_trueones
result_df2

# In[ ]:

import xgboost as xgb
from sklearn.linear_model import LogisticRegression

lr = LogisticRegression(fit_intercept=True, C=7.50, solver='liblinear')

lr.fit(X,y)

lr_preds= lr.predict(X_eval)

# In[ ]:

from sklearn.metrics import confusion_matrix, accuracy_score, classification_report
results = confusion_matrix(y_eval, lr_preds)
print('Logistic Regression Results')
print('Confusion Matrix :')
print(results)
print('Accuracy Score :',accuracy_score(y_eval, lr_preds) )
print('Report: ')
print(classification_report(y_eval, lr_preds))
plot_confusion_matrix(np.array(y_eval), np.array(lr_preds), classes=np.array(['pure (below cutoff)', 'impure (above cutoff)']), normalize=False)
plt.show()

# ### Putting it all together

# In[ ]:

import random

random.seed(69)

mdf = df.drop(columns = ['datetime', '% Iron Concentrate', '% Silica Concentrate'])

train = mdf.iloc[:-14*24*180,:]
zero_idx = train[train['label'] == 0].index
sample_idx = random.sample(list(zero_idx), train[train['label'] == 1].shape[0])
sample_idx.extend(list(train[train['label'] == 1].index))
sample_idx = pd.DatetimeIndex(sample_idx).sort_values()
train = train.reindex(sample_idx)

X = train.iloc[:,:-1]
y = train.iloc[:,-1]

X_eval = mdf.iloc[-14*24*180:-7*24*180,:-1]
y_eval = mdf.iloc[-14*24*180:-7*24*180,-1]
X_test = mdf.iloc[-7*24*180:,:-1]
y_test = mdf.iloc[-7*24*180:,-1]

print(y[y == 0].count() / y[y==1].count())

# In[ ]:

import xgboost as xgb
from sklearn.linear_model import LogisticRegression

lr = LogisticRegression(fit_intercept=True, C=3.0, solver='liblinear')
xgbc= xgb.XGBClassifier(max_depth=2, n_estimators=2, eval_metric='logloss', subsample=0.5)
xgbc2 = xgb.XGBClassifier(max_depth=4, n_estimators=8, eval_metric='logloss', subsample=0.5, min_child_weight=300)

lr.fit(X,y)
xgbc.fit(X,y)
xgbc2.fit(X,y)

lr_preds= lr.predict(X_eval)
xgb_pred = xgbc.predict(X_eval)
xgb2_pred = xgbc2.predict(X_eval)

# In[ ]:

from sklearn.metrics import confusion_matrix, accuracy_score, classification_report
results = confusion_matrix(y_eval, xgb_pred)
print('XGBoost Results')
print('Confusion Matrix :')
print(results)
print('Accuracy Score :',accuracy_score(y_eval, xgb_pred) )
print('Report: ')
print(classification_report(y_eval, xgb_pred))

# In[ ]:

from sklearn.metrics import confusion_matrix, accuracy_score, classification_report
results = confusion_matrix(y_eval, xgb2_pred)
print('XGBoost Results')
print('Confusion Matrix :')
print(results)
print('Accuracy Score :',accuracy_score(y_eval, xgb2_pred) )
print('Report: ')
print(classification_report(y_eval, xgb2_pred))

# In[ ]:

from sklearn.metrics import confusion_matrix, accuracy_score, classification_report
results = confusion_matrix(y_eval, lr_preds)
print('Logistic Regression Results')
print('Confusion Matrix :')
print(results)
print('Accuracy Score :',accuracy_score(y_eval, lr_preds) )
print('Report: ')
print(classification_report(y_eval, lr_preds))

# #### building an hourly classifier

# In[ ]:

eval_df = pd.DataFrame()
eval_df['LogReg'] = lr_preds
eval_df['XGBoost'] = xgb_pred
eval_df['XGBoost2'] = xgb2_pred
eval_df.index = y_eval.index
eval_df.head()

# In[ ]:

hour_counts = eval_df.groupby([eval_df.index.date, eval_df.index.hour]).sum() / 180
hour_counts['Actual'] = y_eval.groupby([y_eval.index.date, y_eval.index.hour]).mean()
hour_counts.plot()

# In[ ]:

sns.boxplot(x="Actual", y="LogReg", data=hour_counts)

# In[ ]:

from sklearn.metrics import roc_curve, roc_auc_score
fpr, tpr, thresholds = roc_curve(np.array(hour_counts["Actual"]), np.array(hour_counts["LogReg"]))
roc_auc = roc_auc_score(np.array(hour_counts["Actual"]), np.array(hour_counts["LogReg"]))

plt.figure()
lw = 2
plt.plot(fpr, tpr, color='darkorange', lw=lw, label='ROC curve (area = %0.2f)' % roc_auc)
plt.plot([0,1], [0,1], linestyle="--")
plt.xlim([0.0, 1.0])
plt.ylim([0.0, 1.05])
plt.xlabel('False Positive Rate')
plt.ylabel('True Positive Rate')
plt.title('ROC Curve')
plt.legend(loc="lower right")
plt.show()

# In[ ]:

sns.boxplot(x="Actual", y="XGBoost", data=hour_counts)

# In[ ]:

from sklearn.metrics import roc_curve, roc_auc_score
fpr, tpr, thresholds = roc_curve(np.array(hour_counts["Actual"]), np.array(hour_counts["XGBoost"]))
roc_auc = roc_auc_score(np.array(hour_counts["Actual"]), np.array(hour_counts["XGBoost"]))

plt.figure()
lw = 2
plt.plot(fpr, tpr, color='darkorange', lw=lw, label='ROC curve (area = %0.2f)' % roc_auc)
plt.plot([0,1], [0,1], linestyle="--")
plt.xlim([0.0, 1.0])
plt.ylim([0.0, 1.05])
plt.xlabel('False Positive Rate')
plt.ylabel('True Positive Rate')
plt.title('ROC Curve')
plt.legend(loc="lower right")
plt.show()

# In[ ]:

sns.boxplot(x="Actual", y="XGBoost2", data=hour_counts)

# In[ ]:

from sklearn.metrics import roc_curve, roc_auc_score
fpr, tpr, thresholds = roc_curve(np.array(hour_counts["Actual"]), np.array(hour_counts["XGBoost2"]))
roc_auc = roc_auc_score(np.array(hour_counts["Actual"]), np.array(hour_counts["XGBoost2"]))

plt.figure()
lw = 2
plt.plot(fpr, tpr, color='darkorange', lw=lw, label='ROC curve (area = %0.2f)' % roc_auc)
plt.plot([0,1], [0,1], linestyle="--")
plt.xlim([0.0, 1.0])
plt.ylim([0.0, 1.05])
plt.xlabel('False Positive Rate')
plt.ylabel('True Positive Rate')
plt.title('ROC Curve')
plt.legend(loc="lower right")
plt.show()

# In[ ]:

avg_eval_preds = (hour_counts['LogReg'] + hour_counts['XGBoost'] + hour_counts['XGBoost2']) / 3
sns.boxplot(x=hour_counts["Actual"], y=avg_eval_preds)

# In[ ]:

from sklearn.metrics import roc_curve, roc_auc_score
fpr, tpr, thresholds = roc_curve(np.array(hour_counts["Actual"]), np.array(avg_eval_preds))
roc_auc = roc_auc_score(np.array(hour_counts["Actual"]), np.array(avg_eval_preds))

plt.figure()
lw = 2
plt.plot(fpr, tpr, color='darkorange', lw=lw, label='ROC curve (area = %0.2f)' % roc_auc)
plt.plot([0,1], [0,1], linestyle="--")
plt.xlim([0.0, 1.0])
plt.ylim([0.0, 1.05])
plt.xlabel('False Positive Rate')
plt.ylabel('True Positive Rate')
plt.title('ROC Curve')
plt.legend(loc="lower right")
plt.show()

# In[ ]:

cutoff_list = [0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5, 0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9]
acc = []
truezeros = []
falsezeros = []
trueones = []
falseones = []
imp_precision = []
imp_recall = []

for cutoff in cutoff_list:
    ddf = pd.DataFrame(index=avg_eval_preds.index)
    ddf['class'] = 0
    ddf['class'][avg_eval_preds > cutoff] = 1

    cm = confusion_matrix(hour_counts["Actual"], ddf['class'])
    acc.append((cm[1,1]+cm[0,0]) / len(ddf))
    truezeros.append(cm[0,0])
    falsezeros.append(cm[1,0])
    trueones.append(cm[1,1])
    falseones.append(cm[0,1])
    imp_precision.append((cm[1,1]) / (cm[0,1] + cm[1,1]))
    imp_recall.append((cm[1,1]) / (cm[1,0] + cm[1,1]))
    #print(C,intercept)

# In[ ]:

cutoffs = pd.DataFrame()
cutoffs['Cutoff'] = cutoff_list
cutoffs['Accuracy'] = acc
cutoffs['True Zeros'] = truezeros
cutoffs['False Zeros'] = falsezeros
cutoffs['True Ones'] = trueones
cutoffs['False Ones'] = falseones
cutoffs['Impure Precision'] = imp_precision
cutoffs['Impure Recall'] = imp_recall
cutoffs

# In[ ]:

#for i in range(len(cutoff_list)):
#    ddf = pd.DataFrame(index=avg_eval_preds.index)
#    ddf['class'] = 0
#    ddf['class'][avg_eval_preds > cutoff_list[i]] = 1
#
#    print(cutoff_list[i])
#    plot_confusion_matrix(np.array(hour_counts["Actual"]), np.array(ddf['class']), classes=np.array(['pure (below cutoff)', 'impure (above cutoff)']), normalize=False)

# In[ ]:

ddf = pd.DataFrame(index=hour_counts['LogReg'].index)
ddf['class'] = 0
ddf['class'][hour_counts['LogReg'] > 0.45] = 1

plot_confusion_matrix(np.array(hour_counts["Actual"]), np.array(ddf['class']), classes=np.array(['pure (below cutoff)', 'impure (above cutoff)']), normalize=False)
plt.show()

# ### Apply method to Test Set

# In[ ]:

#import random #

#random.seed(69)

#mdf = df.drop(columns = ['datetime', '% Iron Concentrate', '% Silica Concentrate'])

#train = mdf.iloc[:-14*24*180]
#zero_idx = train[train['label'] == 0].index
#sample_idx = random.sample(list(zero_idx), train[train['label'] == 1].shape[0])
#sample_idx.extend(list(train[train['label'] == 1].index))
#sample_idx = pd.DatetimeIndex(sample_idx).sort_values()
#train = train.reindex(sample_idx)

#X = train.iloc[:,:-1]
#y = train.iloc[:,-1]

#X_eval = mdf.iloc[-42*24*180:-14*24*180,:-1]
#y_eval = mdf.iloc[-42*24*180:-14*24*180,-1]
#X_test = mdf.iloc[-14*24*180:,:-1]
#y_test = mdf.iloc[-14*24*180:,-1]

#print(y[y == 0].count() / y[y==1].count())

# In[ ]:

#lr = LogisticRegression(fit_intercept=True, C=3.0, solver='liblinear')
#xgbc= xgb.XGBClassifier(max_depth=2, n_estimators=2, eval_metric='logloss', subsample=0.5)
#xgbc2 = xgb.XGBClassifier(max_depth=4, n_estimators=8, eval_metric='logloss', subsample=0.5, min_child_weight=300)

#lr.fit(X,y)
#xgbc.fit(X,y)
#xgbc2.fit(X,y)

#lr_preds= lr.predict(X_eval)
#xgb_pred = xgbc.predict(X_eval)
#xgb2_pred = xgbc2.predict(X_eval)

# In[ ]:

#eval_df = pd.DataFrame()
#eval_df['LogReg'] = lr_preds
#eval_df['XGBoost'] = xgb_pred
#eval_df['XGBoost2'] = xgb2_pred
#eval_df.index = X_test.index
#hour_counts = eval_df.groupby([eval_df.index.date, eval_df.index.hour]).sum() / 180
#hour_counts['Actual'] = y_test.groupby([y_test.index.date, y_test.index.hour]).mean()
#hour_counts['Average'] = (hour_counts['LogReg'] + hour_counts['XGBoost'] + hour_counts['XGBoost2']) / 3
#hour_counts['Prediction'] = 0
#hour_counts['Prediction'][hour_counts['Average'] > 0.45] = 1
#plot_confusion_matrix(np.array(hour_counts["Actual"]), np.array(hour_counts["Prediction"]), classes=np.array(['pure (below cutoff)', 'impure (above cutoff)']), normalize=False)
#plt.show()

# In[ ]:

#from sklearn.metrics import confusion_matrix, accuracy_score, classification_report
#results = confusion_matrix(hour_counts["Actual"], hour_counts["Prediction"])
#print('Logistic Regression Results')
#print('Confusion Matrix :')
#print(results)
#print('Accuracy Score :',accuracy_score(hour_counts["Actual"], hour_counts["Prediction"]) )
#print('Report: ')
#print(classification_report(hour_counts["Actual"], hour_counts["Prediction"]))

# In[ ]:


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

# In[ ]:

import sys
sys.path.insert(0,"./../") #so we can import our modules properly

# In[ ]:

# iPhython
from IPython.core.display import display, HTML
display(HTML("<style>.container { width:90% !important; }</style>"))

from matplotlib import rcParams
import matplotlib.pyplot as plt
get_ipython().run_line_magic('matplotlib', 'notebook')
get_ipython().run_line_magic('matplotlib', 'notebook')

import numpy as np
import pandas as pd

# # House C - Energy

# In[ ]:

path_to_house = "./datasets/dfC_300s.hdf"
df_house = pd.read_hdf(path_to_house)
print('\n\nStart time: {}'.format(df_house.index[1]))
print('End time: {}'.format(df_house.index[-1]))
print(df_house.columns)

# ## Check that boiler power does not exceed total consumed power

# In[ ]:

meters = ['C_boiler_power']

# Investigate only points in time where all values are available
cols = ['C_total_cons_power']
cols.extend(meters)
df_house_noNAN = df_house.loc[:,cols].dropna(axis=0, how='any')

# Check if total consumed power is larger than total of appliances
delta = df_house_noNAN.loc[:,meters].sum(axis=1) - df_house_noNAN.loc[:,'C_total_cons_power']
tmt = delta > 0

if np.any(tmt):
    print('Boiler Power is larger than total consumed power for some data points.')

    # investigate problematic cases
    print("Found {} out of {} ({:.2}%) values to be problematic.".format(tmt.sum(), len(tmt), tmt.sum()/len(tmt)*100))
    print("")
    print("Statistics of problematic values:")
    print(delta[tmt].describe())
    print()
    print((delta[tmt]/df_house_noNAN.loc[tmt, 'C_total_cons_power']).describe())
else:
    print("Total of all appliances is smaller than total consumed energy for all measurement points.")

# In[ ]:

l = ['Total Appliances = Boiler', 'Total Consumed']
if np.any(tmt):
    print('Boiler Power is larger than total consumed power for some data points.')
    ## plot params ##
    ha = 6
    ncols = 8
    nrows = np.sum(tmt)//ncols+1
    fig, ax = plt.subplots(figsize=(17,3*nrows), ncols=ncols, nrows=nrows)
    idxs = np.nonzero(tmt)[0]
    for i, idx in enumerate(idxs):
        minidx = idx-ha
        maxidx = idx+ha
        df_house_noNAN.iloc[minidx:maxidx].loc[:,'C_boiler_power'].plot(ax=ax[i//ncols, i%ncols], label='Total Appliances')
        df_house_noNAN.iloc[minidx:maxidx].loc[:,'C_total_cons_power'].plot(ax=ax[i//ncols, i%ncols], label='Total Consumed')
    fig.legend(l, loc='upper center')
else:
    print("Total of all appliances is smaller than total consumed energy for all measurement points.")

# ## Visualize Boiler Variables

# In[ ]:

fig, ax = plt.subplots(ncols=1, nrows=3, figsize=(17,12), sharex=True)
df_house.loc[:,['C_boilertemp_top', 'C_boilertemp_bottom']].plot(ax=ax[0])
ax[0].legend(loc='upper left')
ax2 = ax[0].twinx()
df_house.loc[:,'C_boiler_power'].plot(ax=ax2)
# df_house.loc[:,['C_boiler_on_utility','C_boiler_on_thermostat', 'C_boiler_on_relay']].plot(ax=ax2)
ax2.legend(loc='upper right')
df_house.loc[:,['C_boiler_heater_1_on', 'C_boiler_heater_2_on', 'C_boiler_heater_3_on']].plot(ax=ax[1])
ax[1].legend()
df_house.loc[:,['C_boiler_on_utility', 'C_boiler_on_relay','C_boiler_on_thermostat']].plot(ax=ax[2])
ax[2].legend()

# In[ ]:

fig, ax = plt.subplots(ncols=1, nrows=1, figsize=(17,4), sharex=True)
df_house.loc[:,['C_boilertemp_top', 'C_boilertemp_bottom']].plot(ax=ax)
ax.legend(loc='upper left')
ax2 = ax.twinx()
# df_house.loc[:,'C_boiler_power'].plot(ax=ax2)
df_house.loc[:,['C_boiler_on_utility','C_boiler_on_thermostat', 'C_boiler_on_relay']].plot(ax=ax2)
ax2.legend(loc='upper right')

# ## Heat Pump Power and Total Consumed Power (without HP)

# In[ ]:

fig, ax = plt.subplots(figsize=(17,4))
df_house.loc[:,'C_total_cons_power'].plot(ax=ax)
df_house.loc[:,'C_hp_power'].plot(ax=ax)
ax.legend()

# ## Visualize `C_solarlog_radiation`

# In[ ]:

fix, ax = plt.subplots(figsize=(17,4))
df_house.loc[:,'C_solarlog_radiation'].plot(ax=ax)

# ## Check that `C_pv_prod_power` is the sum of `C_to_batt_power` + `C_direct_cons_power`+ `C_to_net_power`

# In[ ]:

# Investigate only points in time where all values are available
df_house_noNAN = df_house.loc[:,['C_pv_prod_power', 'C_to_batt_power', 'C_direct_cons_power', 'C_to_net_power']].dropna(axis=0, how='any')
delta = df_house_noNAN.loc[:,'C_pv_prod_power'] - df_house_noNAN.loc[:,['C_to_batt_power', 'C_direct_cons_power', 'C_to_net_power']].sum(axis=1)
print("Statistics of deviations:")
delta.describe()

# ## Check that `C_total_cons_power` is the sum of `C_direct_cons_power` + `C_from_batt_power` + `C_from_net_power`

# In[ ]:

# Investigate only points in time where all values are available
df_house_noNAN = df_house.loc[:,['C_total_cons_power', 'C_direct_cons_power', 'C_from_batt_power', 'C_from_net_power']].dropna(axis=0, how='any')
delta = df_house_noNAN.loc[:,'C_total_cons_power'] - df_house_noNAN.loc[:,['C_direct_cons_power', 'C_from_batt_power', 'C_from_net_power']].sum(axis=1)
print("Statistics of deviations:")
delta.describe()

# ## Visualize Batter Power Flows

# In[ ]:

fig, ax = plt.subplots(figsize=(17,4))
df_house.loc[:,['C_to_batt_power', 'C_from_batt_power']].plot(ax=ax)
ax.legend(loc='upper left')
ax2 = ax.twinx()
df_house.loc[:,'C_batt_state'].plot(ax=ax2, color='green')
ax2.legend(loc='lower left')

# # House C - Weather

# In[ ]:

path_to_house = "./datasets/dfC_3600s.hdf"
df_weather = pd.read_hdf(path_to_house)
print('\n\nStart time: {}'.format(df_weather.index[1]))
print('End time: {}'.format(df_weather.index[-1]))
print(df_weather.columns)

# In[ ]:

fig, ax = plt.subplots(figsize=(17,16), ncols=1, nrows=4, sharex=True)
df_weather.loc[:,'C_weather_pressure'].plot(ax=ax[0])
ax[0].legend()
df_weather.loc[:,'C_weather_rainfall'].plot(ax=ax[1], color='orange')
ax[1].legend()
df_weather.loc[:,['C_weather_temperature_out', 'C_weather_temperature_in']].plot(ax=ax[2])
ax[2].legend()
df_weather.loc[:,['C_weather_humidity_out', 'C_weather_humidity_in']].plot(ax=ax[3])
ax[3].legend()

# # House C - Raw Data

# In[ ]:

import os
base_path = "../rawData/C/"
files = os.listdir(path=base_path)

blacklisted = ["capPeriods.hdf"]
for file in files:
    if file in blacklisted: continue
    print(file)
    print(os.path.join(base_path, file))
    df_rawData = pd.read_hdf(os.path.join(base_path, file))
    print(df_rawData.iloc[1,0]-df_rawData.iloc[0,0])
    print(df_rawData.iloc[-1,0]-df_rawData.iloc[-2,0])
    print()

# In[ ]:


import random
from numbers import Number
from typing import Optional

import numpy as np
import torch

__all__ = ["set_seed"]

def set_seed(seed: Optional[int] = None) -> Optional[int]:
    """Set random seed for reproduction.

    Parameters
    ----------
    seed : Optional[int], optional
        random seed, by default None

    Returns
    -------
    Optional[int]
        the random seed.

    Example
    -------
    >>> from graphwar import set_seed
    >>> set_seed(42)
    42

    """
    assert seed is None or isinstance(seed, Number), seed
    np.random.seed(seed)
    random.seed(seed)
    if seed is not None:
        torch.manual_seed(seed)
        torch.cuda.manual_seed(seed)
        # torch.cuda.manual_seed_all(seed)
    return seed