import requests
from tests.helpers import NetworkTest
import time
CONTROLLER = '127.0.0.1'
KYTOS_API = 'http://%s:8181/api' % CONTROLLER

class TestE2ESDNTraceInvertedCircuit:
    net = None
    circuit = None
    @staticmethod
    def setup_class(cls):
        """run at the beginning for all tests"""
        cls.net = NetworkTest(CONTROLLER, topo_name='linear10')
        cls.net.start()
        cls.net.restart_kytos_clean()
        cls.net.wait_switches_connect()

    @classmethod
    def teardown_class(cls):
        print('somente na terminação')
        cls.net.stop()

    @staticmethod
    def create_evc(vlan_id, interface_a="00:00:00:00:00:00:00:01:1", interface_z="00:00:00:00:00:00:00:0a:1"):
        payload = {
            "name": "Vlan_%s" % vlan_id,
            "enabled": True,
            "dynamic_backup_path": True,
            "uni_a": {
                "interface_id": interface_a,
                "tag": {"tag_type": 1, "value": vlan_id}
            },
            "uni_z": {
                "interface_id": interface_z,
                "tag": {"tag_type": 1, "value": vlan_id}
            }
        }
        api_url = KYTOS_API + '/kytos/mef_eline/v2/evc/'
        response = requests.post(api_url, json=payload)
        assert response.status_code == 201, response.text
        data = response.json()
        return data['circuit_id']

        response = requests.get(api_url + "switches")
        switches = response.json()
        for switch in switches["switches"]:
            response = requests.post(api_url + f"switches/{switch}/enable")
            response = requests.post(api_url + f"interfaces/switch{switch}/enable")

        response = requests.get(api_url + "links")
        links = response.json()
        for link in links["links"]:
            response = requests.post(api_url + f"links/{link}/enable")

        time.sleep(10)

    @staticmethod
    def get_evc(circuit_id):
        api_url = KYTOS_API + '/kytos/mef_eline/v2/evc/'
        response = requests.get(api_url + circuit_id)
        assert response.status_code == 200, response.text
        data = response.json()
        return data

    def test_100_run_sdntrace_cp(self):
        """Run SDNTrace-CP (Control Plane)."""
        # Trace from UNI_A
        payload = {
            "trace": {
                "switch": {"dpid": "00:00:00:00:00:00:03:01", "in_port": 1},
                "eth": {"dl_type": 33024, "dl_vlan": 400}
            }
        }

        api_url = KYTOS_API + '/amlight/sdntrace_cp/v1/trace'
        response = requests.put(api_url, json=payload)
        assert response.status_code == 200, response.text
        data = response.json()
        assert "result" in data, data
        assert len(data["result"]) == 10, data

        expected = [
            (
                l['endpoint_b']['switch'],
                l['endpoint_b']['port_number'],
                l['metadata']['s_vlan']['value']
            )
            for l in self.circuit['current_path']
        ]
        expected.insert(0, ('00:00:00:00:00:00:03:1', 1, 400))

        actual = [
            (step['dpid'], step['port'], step['vlan'])
            for step in data["result"]
        ]

        assert expected == actual, f"Expected {expected}. Actual: {actual}"

        # Trace from UNI_Z
        payload = {
            "trace": {
                "switch": {"dpid": "00:00:00:00:00:00:01:1", "in_port": 1},
                "eth": {"dl_type": 33024, "dl_vlan": 400}
            }
        }
        api_url = KYTOS_API + '/amlight/sdntrace_cp/v1/trace'
        response = requests.put(api_url, json=payload)
        assert response.status_code == 200, response.text
        data = response.json()
        assert "result" in data
        assert len(data["result"]) == 10, data

        expected = [
            (
                l['endpoint_a']['switch'],
                l['endpoint_a']['port_number'],
                l['metadata']['s_vlan']['value']
            )
            for l in reversed(self.circuit['current_path'])
        ]
        expected.insert(0, ('00:00:00:00:00:00:01:1', 1, 400))

        actual = [
            (step['dpid'], step['port'], step['vlan'])
            for step in data["result"]
        ]

        assert expected == actual, f"Expected {expected}. Actual: {actual}"


import torch
import torch.nn as nn
import torch.nn.functional as F
import math
import time

def timeit(x, func, iter=10):
    torch.cuda.synchronize()
    start = time.time()
    for _ in range(iter):
        y = func(x)
    torch.cuda.synchronize()
    runtime = (time.time()-start)/iter
    return runtime

class HOGLayer(nn.Module):
    def __init__(self, nbins=10, pool=8, max_angle=math.pi, stride=1, padding=1, dilation=1):
        super(HOGLayer, self).__init__()
        self.nbins = nbins
        self.stride = stride
        self.padding = padding
        self.dilation = dilation
        self.pool = pool
        self.max_angle = max_angle
        mat = torch.FloatTensor([[1, 0, -1],
                                [2, 0, -2],
                                [1, 0, -1]])
        mat = torch.cat((mat[None], mat.t()[None]), dim=0)
        self.register_buffer("weight", mat[:,None,:,:])
        self.pooler = nn.AvgPool2d(pool, stride=pool, padding=0, ceil_mode=False, count_include_pad=True)

    def forward(self, x):
        with torch.no_grad():
            gxy = F.conv2d(x, self.weight, None, self.stride,
                            self.padding, self.dilation, 1)
            #2. Mag/ Phase
            mag = gxy.norm(dim=1)
            norm = mag[:,None,:,:]
            phase = torch.atan2(gxy[:,0,:,:], gxy[:,1,:,:])

            #3. Binning Mag with linear interpolation
            phase_int = phase / self.max_angle * self.nbins
            phase_int = phase_int[:,None,:,:]

            n, c, h, w = gxy.shape
            out = torch.zeros((n, self.nbins, h, w), dtype=torch.float, device=gxy.device)
            out.scatter_(1, phase_int.floor().long()%self.nbins, norm)
            out.scatter_add_(1, phase_int.ceil().long()%self.nbins, 1 - norm)

            return self.pooler(out)

class HOGLayerMoreComplicated(nn.Module):
    def __init__(self, nbins=10, pool=8, max_angle=math.pi, stride=1, padding=1, dilation=1,
                 mean_in=False, max_out=False):
        super(HOGLayerMoreComplicated, self).__init__()
        self.nbins = nbins
        self.stride = stride
        self.padding = padding
        self.dilation = dilation
        self.pool = pool
        self.max_angle = max_angle
        self.max_out = max_out
        self.mean_in = mean_in

        mat = torch.FloatTensor([[1, 0, -1],
                                 [2, 0, -2],
                                 [1, 0, -1]])
        mat = torch.cat((mat[None], mat.t()[None]), dim=0)
        self.register_buffer("weight", mat[:,None,:,:])
        self.pooler = nn.AvgPool2d(pool, stride=pool, padding=0, ceil_mode=False, count_include_pad=True)

    def forward(self, x):
        if self.mean_in:
            return self.forward_v1(x)
        else:
            return self.forward_v2(x)

    def forward_v1(self, x):
        if x.size(1) > 1:
            x = x.mean(dim=1)[:,None,:,:]

        gxy = F.conv2d(x, self.weight, None, self.stride,
                       self.padding, self.dilation, 1)
        # 2. Mag/ Phase
        mag = gxy.norm(dim=1)
        norm = mag[:, None, :, :]
        phase = torch.atan2(gxy[:, 0, :, :], gxy[:, 1, :, :])

        # 3. Binning Mag with linear interpolation
        phase_int = phase / self.max_angle * self.nbins
        phase_int = phase_int[:, None, :, :]

        n, c, h, w = gxy.shape
        out = torch.zeros((n, self.nbins, h, w), dtype=torch.float, device=gxy.device)
        out.scatter_(1, phase_int.floor().long() % self.nbins, norm)
        out.scatter_add_(1, phase_int.ceil().long() % self.nbins, 1 - norm)

        return self.pooler(out)

    def forward_v2(self, x):
        batch_size, in_channels, height, width = x.shape
        weight = self.weight.repeat(3, 1, 1, 1)
        gxy = F.conv2d(x, weight, None, self.stride,
                        self.padding, self.dilation, groups=in_channels)

        gxy = gxy.view(batch_size, in_channels, 2, height, width)

        if self.max_out:
            gxy = gxy.max(dim=1)[0][:,None,:,:,:]

        #2. Mag/ Phase
        mags = gxy.norm(dim=2)
        norms = mags[:,:,None,:,:]
        phases = torch.atan2(gxy[:,:,0,:,:], gxy[:,:,1,:,:])

        #3. Binning Mag with linear interpolation
        phases_int = phases / self.max_angle * self.nbins
        phases_int = phases_int[:,:,None,:,:]

        out = torch.zeros((batch_size, in_channels, self.nbins, height, width),
                          dtype=torch.float, device=gxy.device)
        out.scatter_(2, phases_int.floor().long()%self.nbins, norms)
        out.scatter_add_(2, phases_int.ceil().long()%self.nbins, 1 - norms)

        out = out.view(batch_size, in_channels * self.nbins, height, width)
        out = torch.cat((self.pooler(out), self.pooler(x)), dim=1)

        return out

if __name__ == '__main__':
    import cv2
    cuda = True

    path = '/home/etienneperot/Pictures/pencils.jpg'
    im = cv2.imread(path, cv2.IMREAD_GRAYSCALE)

    x = torch.from_numpy(im)[None, None]
    if cuda:
        x = x.cuda().float()
    else:
        x = x.float()

    hog = HOGLayer(nbins=12, pool=2)
    if cuda:
        hog = hog.cuda()

    y = hog(x)

    y2 = y.cpu().numpy()
    bin = 0
    while 1:
        im = y2[0, bin]
        im = (im-im.min())/(im.max()-im.min())
        print('bin: ', bin, ' ang: ', float(bin)/hog.nbins * 180.0)
        cv2.imshow('bin', im)
        cv2.waitKey()
        bin = (bin + 1) % hog.nbins


import os
import os.path as osp
import argparse
import numpy as np
import pandas as pd
from tqdm import tqdm
from sklearn import metrics

import torch

from dataset import create_data_loader, preprocess_df
from model import ClassificationModel
from utils import set_seeds, load_config, str2bool

def inference(model, test_loader, device, cls_last_sigmoid):
    model.eval()
    preds = []

    with torch.no_grad():
        for img, tab in tqdm(iter(test_loader)):
            img = img.float().to(device)
            tab = tab.float().to(device)

            model_pred = model(img, tab)
            if not cls_last_sigmoid:
                model_pred = torch.sigmoid(model_pred)
            model_pred = model_pred.squeeze(1).to('cpu')

            preds += model_pred.tolist()

    return preds

def get_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument('--seed', type=int, default=42)
    parser.add_argument('--df', type=str, default='test')
    parser.add_argument('--mode', type=str, default='test') # test hard valid
    parser.add_argument('--hflip', action='store_true')
    parser.add_argument('--vflip', action='store_true')
    parser.add_argument('--ckpt_list', nargs='+') # exp0/best.pt
    parser.add_argument('--img_size', type=int, default=512)
    parser.add_argument('--resize_by_split', action='store_true')
    parser.add_argument('--transform_type', type=str, default='resize')
    parser.add_argument('--pred_thres', type=float, default=0.5)
    parser.add_argument('--remark', type=str, default='')
    args = parser.parse_args()

    if args.mode == 'valid':
        assert args.df in ['train', 'train_5fold', 'train_heuristic_5fold']
    else:
        assert args.df in ['test', 'test_heuristic']

    args.base_dir = './'
    args.data_dir = osp.join(args.base_dir, 'data')
    args.submit_dir = osp.join(args.base_dir, 'submit')
    os.makedirs(args.submit_dir, exist_ok=True)

    return args

def main(args, train_args, hflip=False, vflip=False):
    args.device = torch.device("cuda:0")

    df = pd.read_csv(osp.join(args.data_dir, f'{args.df}.csv'))
    if args.mode == 'valid':
        df = preprocess_df(df, train_args.df_ver, drop_row=True)
        df = df[df["kfold"] == train_args.fold].reset_index(drop=True)
    else:
        df = preprocess_df(df, train_args.df_ver, drop_row=False)

    test_loader = create_data_loader(
        df, 'infer', args.img_size, data_dir=args.data_dir,
        hflip=hflip, vflip=vflip, norm_type=train_args.norm_type,
        transform_type=args.transform_type, resize_by_split=args.resize_by_split,
    )

    # exp235
    # args.tab_model = 'drop40'
    model = ClassificationModel(train_args).to(args.device)
    model.load_state_dict(torch.load(osp.join(args.work_dir, args.ckpt)))

    preds = inference(model, test_loader, args.device, train_args.cls_last_sigmoid)
    if args.mode in ['valid', 'hard']:
        preds = np.where(np.array(preds) > args.pred_thres, 1, 0)
    if args.mode == 'valid':
        labels = df[df['kfold'] == train_args.fold]['N_category'].tolist()
        valid_f1 = metrics.f1_score(y_true=labels, y_pred=preds, average='macro', zero_division=1)
        print(f'hflip={bool(hflip)} vflip={bool(vflip)} valid_f1:{valid_f1:.4f}')

    return preds

if __name__ == '__main__':
    args = get_parser()

    submit_file_name = args.mode
    if args.hflip:
        submit_file_name += '_hflip'
    if args.vflip:
        submit_file_name += '_vflip'
    if args.resize_by_split:
        submit_file_name += '_resizeBySplit'
    if args.transform_type != 'resize':
        submit_file_name += f'_{args.transform_type}'
    if args.img_size != 512:
        submit_file_name += f'_img{args.img_size}'
    if args.remark != '':
        submit_file_name += f'_{args.remark}'

    # preds
    ensemble = []
    ensemble_cnt = 0
    for idx in range(len(args.ckpt_list)):
        args.exp, args.ckpt = args.ckpt_list[idx].split('/')
        args.work_dir = osp.join('work_dirs', args.exp) # work_dirs/exp0

        submit_file_name += f'_{args.exp}'
        if args.ckpt.split('.')[0] != 'best':
            submit_file_name += args.ckpt.split('.')[0]

        set_seeds(args.seed)
        train_args = load_config(osp.join(args.work_dir, 'config.yaml'))

        for hflip in range(args.hflip + 1):
            for vflip in range(args.vflip + 1):
                preds = main(args, train_args, hflip, vflip)
                ensemble.append(preds)
                ensemble_cnt += 1

    if args.mode in ['test', 'hard']:
        # ensemble
        if args.mode == 'hard':
            ensemble = np.array(ensemble)
            preds = [np.argmax(np.bincount(ensemble[:,i])) for i in range(ensemble.shape[1])]
        elif args.mode == 'test':
            if ensemble_cnt > 1:
                ensemble = np.sum(ensemble, axis=0)
            elif ensemble_cnt == 1:
                ensemble = ensemble[0]
            preds = np.where(ensemble > args.pred_thres * ensemble_cnt, 1, 0)

        # output
        submit = pd.read_csv(osp.join(args.data_dir, 'sample_submission.csv'))
        submit['N_category'] = preds
        submit_path = osp.join(args.submit_dir, f'{submit_file_name}.csv')
        submit.to_csv(submit_path, index=False)


import numpy as np
import pickle

from abc import ABC
from keras.layers import Input, Dense, TimeDistributed
from keras.layers import LeakyReLU, Dropout, Bidirectional
from keras.layers import GRU, LSTM
from keras.models import Model
from keras.models import load_model

from modules.constants import Constants
from modules.word_embedding import WordEmbedding
from modules.summarizer.metrics import precision, recall, f1
from modules.summarizer.preprocessor import Preprocessor
from modules.utils.text_util import split_to_sentences

class BaseModel(ABC):
    def __init__(self):
        self._model = None
        print('[SUMMARIZER] Initializing Model')

    def train(self, x_train: np.ndarray, y_train: np.ndarray, x_val: np.ndarray, y_val: np.ndarray, log_path: str):
        history = self._model.fit(
            x_train,
            y_train,
            epochs=100,
            validation_data=(x_val, y_val)
        )
        logfile = open(Constants.MODEL_PATH + log_path, 'wb+')
        pickle.dump(history.history, logfile)

    def predict(self, x: np.ndarray) -> np.ndarray:
        return self._model.predict(x)

    def summarize(self, sentences: [str], word_embedding: WordEmbedding, n_sentence: int = 0) -> [str]:
        last_sentence_idx = len(sentences) - 1

        sentences_vector = Preprocessor.preprocess_text(sentences, word_embedding)
        model_predictions = self.predict(sentences_vector)

        predictions = model_predictions[0]

        index_confidence_list = []
        for idx, prediction in enumerate(predictions):
            index_confidence_list.append((idx, prediction[1]))

        selected_indices = []

        if n_sentence > 0:
            n = 0
            index_confidence_list.sort(key=lambda x: x[1], reverse=True)
            while n < n_sentence:
                if index_confidence_list[n][0] <= last_sentence_idx:
                    selected_indices.append(index_confidence_list[n][0])
                    n += 1
                else:
                    n += 1
                    n_sentence += 1
            selected_indices.sort()
        else:  # Use threshold
            for idx, confidence in index_confidence_list:
                if confidence > Constants.MODEL_THRESHOLD and idx < last_sentence_idx:
                    selected_indices.append(idx)

        summary = []
        for selected_index in selected_indices:
            summary.append(sentences[selected_index])

        return summary

    def evaluate(self, x_test: np.ndarray, y_test: np.ndarray) -> (float, float, float):
        loss, accuracy_score, precision_score, recall_score, f1_score = self._model.evaluate(x_test, y_test)
        return loss, accuracy_score, precision_score, recall_score, f1_score

    def save(self, model_name: str):
        self._model.save(Constants.MODEL_PATH + '{}.h5'.format(model_name))
        print('[SUMMARIZER] Model saved to {}'.format(Constants.MODEL_PATH + '{}.h5'.format(model_name)))

    def load(self, model_name: str):
        self._model = load_model(Constants.MODEL_PATH + '{}.h5'.format(model_name), custom_objects={'precision': precision, 'recall': recall, 'f1': f1})
        print('[SUMMARIZER] Model {} loaded'.format(model_name))

class FFNNModel(BaseModel):
    def __init__(self):
        super().__init__()
        inputs = Input(shape=(Constants.WORD_EMBEDDING_DIMENSION,))

        hidden_1 = Dense(256)(inputs)
        actvfn_1 = LeakyReLU(alpha=0.1)(hidden_1)
        drpout_1 = Dropout(0.5)(actvfn_1)

        hidden_2 = Dense(128)(drpout_1)
        actvfn_2 = LeakyReLU(alpha=0.1)(hidden_2)
        drpout_2 = Dropout(0.5)(actvfn_2)

        predictions = Dense(2, activation='softmax')(drpout_2)

        self._model = Model(inputs=inputs, outputs=predictions)
        self._model.summary()

        self._model.compile(
            optimizer='adam',
            loss='categorical_crossentropy',
            metrics=['accuracy', precision, recall, f1]
        )

class FFNNSeq(BaseModel):
    def __init__(self):
        super().__init__()
        inputs = Input(shape=(Constants.MAXIMUM_SENTENCE_LENGTH, Constants.WORD_EMBEDDING_DIMENSION,))

        input_dropout = Dropout(0.5)(inputs)

        hidden_1 = Dense(256)(input_dropout)
        actvfn_1 = LeakyReLU(alpha=0.1)(hidden_1)
        drpout_1 = Dropout(0.5)(actvfn_1)

        hidden_2 = Dense(128)(drpout_1)
        actvfn_2 = LeakyReLU(alpha=0.1)(hidden_2)
        drpout_2 = Dropout(0.5)(actvfn_2)

        predictions = TimeDistributed(Dense(
            2,
            activation='softmax'
        ))(actvfn_2)

        self._model = Model(inputs=inputs, outputs=predictions)
        self._model.summary()

        self._model.compile(
            optimizer='adam',
            loss='categorical_crossentropy',
            metrics=['accuracy', precision, recall, f1]
        )

class GRUModel(BaseModel):
    def __init__(self):
        super().__init__()
        inputs = Input(shape=(Constants.MAXIMUM_SENTENCE_LENGTH, Constants.WORD_EMBEDDING_DIMENSION,))

        input_dropout = Dropout(0.5)(inputs)

        gru_1 = GRU(
            64,
            dropout=0.5,
            recurrent_dropout=0.5,
            return_sequences=True
        )(input_dropout)

        gru_2 = GRU(
            64,
            dropout=0.5,
            recurrent_dropout=0.5,
            return_sequences=True
        )(gru_1)

        predictions = TimeDistributed(Dense(
            2,
            activation='softmax'
        ))(gru_2)

        self._model = Model(inputs=inputs, outputs=predictions)
        self._model.summary()

        self._model.compile(
            optimizer='adam',
            loss='categorical_crossentropy',
            metrics=['accuracy', precision, recall, f1]
        )

class LSTMModel(BaseModel):
    def __init__(self):
        super().__init__()
        inputs = Input(shape=(Constants.MAXIMUM_SENTENCE_LENGTH, Constants.WORD_EMBEDDING_DIMENSION,))

        input_dropout = Dropout(0.5)(inputs)

        lstm_1 = LSTM(
            64,
            dropout=0.5,
            recurrent_dropout=0.5,
            return_sequences=True
        )(input_dropout)

        lstm_2 = LSTM(
            64,
            dropout=0.5,
            recurrent_dropout=0.5,
            return_sequences=True
        )(lstm_1)

        predictions = TimeDistributed(Dense(
            2,
            activation='softmax'
        ))(lstm_2)

        self._model = Model(inputs=inputs, outputs=predictions)
        self._model.summary()

        self._model.compile(
            optimizer='adam',
            loss='categorical_crossentropy',
            metrics=['accuracy', precision, recall, f1]
        )

class BiGRUModel(BaseModel):
    def __init__(self):
        super().__init__()
        inputs = Input(shape=(Constants.MAXIMUM_SENTENCE_LENGTH, Constants.WORD_EMBEDDING_DIMENSION,))

        input_dropout = Dropout(0.5)(inputs)

        bigru_1 = Bidirectional(GRU(
            64,
            dropout=0.5,
            recurrent_dropout=0.5,
            return_sequences=True
        ))(input_dropout)

        bigru_2 = Bidirectional(GRU(
            64,
            dropout=0.5,
            recurrent_dropout=0.5,
            return_sequences=True
        ))(bigru_1)

        predictions = TimeDistributed(Dense(
            2,
            activation='softmax'
        ))(bigru_2)

        self._model = Model(inputs=inputs, outputs=predictions)
        self._model.summary()

        self._model.compile(
            optimizer='adam',
            loss='categorical_crossentropy',
            metrics=['accuracy', precision, recall, f1]
        )

class BiLSTMModel(BaseModel):
    def __init__(self):
        super().__init__()
        inputs = Input(shape=(Constants.MAXIMUM_SENTENCE_LENGTH, Constants.WORD_EMBEDDING_DIMENSION,))

        input_dropout = Dropout(0.5)(inputs)

        bilstm_1 = Bidirectional(LSTM(
            64,
            dropout=0.5,
            recurrent_dropout=0.5,
            return_sequences=True
        ))(input_dropout)

        bilstm_2 = Bidirectional(LSTM(
            64,
            dropout=0.5,
            recurrent_dropout=0.5,
            return_sequences=True
        ))(bilstm_1)

        predictions = TimeDistributed(Dense(
            2,
            activation='softmax'
        ))(bilstm_2)

        self._model = Model(inputs=inputs, outputs=predictions)
        self._model.summary()

        self._model.compile(
            optimizer='adam',
            loss='categorical_crossentropy',
            metrics=['accuracy', precision, recall, f1]
        )


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import json
import logging
import numpy as np
import os
import uuid

from pycocotools.cocoeval import COCOeval
from pycocotools import mask as COCOmask

from core.config import cfg
from utils.io import save_object
import utils.boxes as box_utils

logger = logging.getLogger(__name__)

def evaluate_masks(
    json_dataset,
    all_boxes,
    all_segms,
    output_dir,
    use_salt=True,
    cleanup=False
):
    res_file = os.path.join(
        output_dir, 'segmentations_' + json_dataset.name + '_results'
    )
    if use_salt:
        res_file += '_{}'.format(str(uuid.uuid4()))
    res_file += '.json'
    _write_coco_segms_results_file(
        json_dataset, all_boxes, all_segms, res_file)
    coco_eval = _do_segmentation_eval(json_dataset, res_file, output_dir)
    # Optionally cleanup results json file
    if cleanup:
        os.remove(res_file)
    return coco_eval

def _write_coco_segms_results_file(
    json_dataset, all_boxes, all_segms, res_file
):
    # [{"image_id": 42,
    #   "category_id": 18,
    #   "segmentation": [...],
    #   "score": 0.236}, ...]
    results = []
    for cls_ind, cls in enumerate(json_dataset.classes):
        if cls == '__background__':
            continue
        if cls_ind >= len(all_boxes):
            break
        cat_id = json_dataset.category_to_id_map[cls]
        results.extend(_coco_segms_results_one_category(
            json_dataset, all_boxes[cls_ind], all_segms[cls_ind], cat_id))
    logger.info(
        'Writing segmentation results json to: {}'.format(
            os.path.abspath(res_file)))
    with open(res_file, 'w') as fid:
        json.dump(results, fid)

def _coco_segms_results_one_category(json_dataset, boxes, segms, cat_id):
    results = []
    image_ids = json_dataset.image_index
    #image_ids = json_dataset.COCO.getImgIds()
    #image_ids.sort()
    assert len(boxes) == len(image_ids)
    assert len(segms) == len(image_ids)
    for i, image_id in enumerate(image_ids):
        dets = boxes[i]
        rles = segms[i]

        if isinstance(dets, list) and len(dets) == 0:
            continue

        dets = dets.astype(np.float)
        scores = dets[:, -1]

        results.extend(
            [{'image_id': image_id,
              'category_id': cat_id,
              'segmentation': rles[k],
              'score': scores[k]}
              for k in range(dets.shape[0])])

    return results

def _do_segmentation_eval(json_dataset, res_file, output_dir):
    tmp = [i-1 for i in json_dataset.list]
    coco_dt = json_dataset.COCO.loadRes(str(res_file))
    #coco_eval = COCOeval(json_dataset.COCO, coco_dt, 'segm')
    coco_eval = customCOCOeval(json_dataset.COCO, coco_dt, 'segm')
    coco_eval.params.imgIds = json_dataset.image_index
    coco_eval.evaluate()
    coco_eval.accumulate()
    #_log_detection_eval_metrics(json_dataset, coco_eval)
    coco_eval.summarize(class_index=tmp)
    eval_file = os.path.join(output_dir, 'segmentation_results.pkl')
    save_object(coco_eval, eval_file)
    logger.info('Wrote json eval results to: {}'.format(eval_file))
    return coco_eval

def evaluate_boxes(
    json_dataset, all_boxes, output_dir, use_salt=True, cleanup=False
):
    res_file = os.path.join(
        output_dir, 'bbox_' + json_dataset.name + '_results'
    )
    if use_salt:
        res_file += '_{}'.format(str(uuid.uuid4()))
    res_file += '.json'
    _write_coco_bbox_results_file(json_dataset, all_boxes, res_file)
    coco_eval = _do_detection_eval(json_dataset, res_file, output_dir)
    # Optionally cleanup results json file
    if cleanup:
        os.remove(res_file)
    return coco_eval

def _write_coco_bbox_results_file(json_dataset, all_boxes, res_file):
    # [{"image_id": 42,
    #   "category_id": 18,
    #   "bbox": [258.15,41.29,348.26,243.78],
    #   "score": 0.236}, ...]
    results = []
    for cls_ind, cls in enumerate(json_dataset.classes):
        if cls == '__background__':
            continue
        if cls_ind >= len(all_boxes):
            break
        cat_id = json_dataset.category_to_id_map[cls]
        results.extend(_coco_bbox_results_one_category(
            json_dataset, all_boxes[cls_ind], cat_id))
    logger.info(
        'Writing bbox results json to: {}'.format(os.path.abspath(res_file)))
    with open(res_file, 'w') as fid:
        json.dump(results, fid)

def _coco_bbox_results_one_category(json_dataset, boxes, cat_id):
    results = []
    image_ids = json_dataset.image_index
    #image_ids = json_dataset.COCO.getImgIds()
    #image_ids.sort()
    assert len(boxes) == len(image_ids)
    for i, image_id in enumerate(image_ids):
        dets = boxes[i]
        if isinstance(dets, list) and len(dets) == 0:
            continue
        dets = dets.astype(np.float)
        scores = dets[:, -1]
        xywh_dets = box_utils.xyxy_to_xywh(dets[:, 0:4])
        xs = xywh_dets[:, 0]
        ys = xywh_dets[:, 1]
        ws = xywh_dets[:, 2]
        hs = xywh_dets[:, 3]
        results.extend(
            [{'image_id': image_id,
              'category_id': cat_id,
              'bbox': [xs[k], ys[k], ws[k], hs[k]],
              'score': scores[k]} for k in range(dets.shape[0])])
    return results

def _do_detection_eval(json_dataset, res_file, output_dir):
    tmp = [i-1 for i in json_dataset.list]
    coco_dt = json_dataset.COCO.loadRes(str(res_file))
    #coco_eval = COCOeval(json_dataset.COCO, coco_dt, 'bbox')
    coco_eval = customCOCOeval(json_dataset.COCO, coco_dt, 'bbox')
    coco_eval.evaluate()
    coco_eval.accumulate()
    #_log_detection_eval_metrics(json_dataset, coco_eval)
    coco_eval.summarize(class_index=tmp)
    eval_file = os.path.join(output_dir, 'detection_results.pkl')
    save_object(coco_eval, eval_file)
    logger.info('Wrote json eval results to: {}'.format(eval_file))
    return coco_eval

def _log_detection_eval_metrics(json_dataset, coco_eval):
    def _get_thr_ind(coco_eval, thr):
        ind = np.where((coco_eval.params.iouThrs > thr - 1e-5) &
                       (coco_eval.params.iouThrs < thr + 1e-5))[0][0]
        iou_thr = coco_eval.params.iouThrs[ind]
        assert np.isclose(iou_thr, thr)
        return ind

    IoU_lo_thresh = 0.5
    IoU_hi_thresh = 0.95
    ind_lo = _get_thr_ind(coco_eval, IoU_lo_thresh)
    ind_hi = _get_thr_ind(coco_eval, IoU_hi_thresh)
    # precision has dims (iou, recall, cls, area range, max dets)
    # area range index 0: all area ranges
    # max dets index 2: 100 per image
    precision = coco_eval.eval['precision'][ind_lo:(ind_hi + 1), :, :, 0, 2]
    ap_default = np.mean(precision[precision > -1])
    logger.info(
        '~~~~ Mean and per-category AP @ IoU=[{:.2f},{:.2f}] ~~~~'.format(
            IoU_lo_thresh, IoU_hi_thresh))
    logger.info('{:.1f}'.format(100 * ap_default))
    for cls_ind, cls in enumerate(json_dataset.classes):
        if cls == '__background__':
            continue
        # minus 1 because of __background__
        precision = coco_eval.eval['precision'][
            ind_lo:(ind_hi + 1), :, cls_ind - 1, 0, 2]
        ap = np.mean(precision[precision > -1])
        logger.info('{:.1f}'.format(100 * ap))
    logger.info('~~~~ Summary metrics ~~~~')
    coco_eval.summarize()

class customCOCOeval(COCOeval):

    def __init__(self, cocoGt=None, cocoDt=None, iouType="segm"):
        super().__init__(cocoGt, cocoDt, iouType)
        self.params.maxDets = [1, 10, 100, 200, 500]
        self.params.areaRng = [[0 ** 2, 1e5 ** 2], [0 ** 2, 16 ** 2], [16 ** 2, 64 ** 2], [64 ** 2, 1e5 ** 2]]
        self.params.areaRngLbl = ['all', 'small', 'medium', 'large']

    def computeIoU(self, imgId, catId):
        p = self.params
        if p.useCats:
            gt = self._gts[imgId,catId]
            dt = self._dts[imgId,catId]
        else:
            gt = [_ for cId in p.catIds for _ in self._gts[imgId,cId]]
            dt = [_ for cId in p.catIds for _ in self._dts[imgId,cId]]
        if len(gt) == 0 and len(dt) ==0:
            return []
        dt = sorted(dt, key=lambda x: -x['score'])
        if len(dt) > p.maxDets[-1]:
            dt=dt[0:p.maxDets[-1]]

        if p.useSegm:
            g = [g['segmentation'] for g in gt]
            d = [d['segmentation'] for d in dt]
        else:
            g = [g['bbox'] for g in gt]
            d = [d['bbox'] for d in dt]

        # compute iou between each dt and gt region
        ious = COCOmask.iou(d,g,[])
        return ious

    def evaluateImg(self, imgId, catId, aRng, maxDet):
        '''
        perform evaluation for single category and image
        :return: dict (single image results)
        '''
        #
        p = self.params
        if p.useCats:
            gt = self._gts[imgId,catId]
            dt = self._dts[imgId,catId]
        else:
            gt = [_ for cId in p.catIds for _ in self._gts[imgId,cId]]
            dt = [_ for cId in p.catIds for _ in self._dts[imgId,cId]]
        if len(gt) == 0 and len(dt) ==0:
            return None

        for g in gt:
          #if 'ignore' not in g:
              #g['ignore'] = 0
          #if g['ignore'] or (g['area']<aRng[0] or g['area']>aRng[1]):
          #    g['_ignore'] = 1
          if (g['area']<aRng[0] or g['area']>aRng[1]):
              g['_ignore'] = 1
          else:
              g['_ignore'] = 0

        # sort dt highest score first, sort gt ignore last
        # gt = sorted(gt, key=lambda x: x['_ignore'])
        gtind = [ind for (ind, g) in sorted(enumerate(gt), key=lambda ind_g: ind_g[1]['_ignore']) ]

        gt = [gt[ind] for ind in gtind]
        dt = sorted(dt, key=lambda x: -x['score'])[0:maxDet]
        # load computed ious
        N_iou = len(self.ious[imgId, catId])
        ious = self.ious[imgId, catId][0:maxDet, np.array(gtind)] if N_iou >0 else self.ious[imgId, catId]

        T = len(p.iouThrs)
        G = len(gt)
        D = len(dt)
        gtm  = np.zeros((T,G))
        dtm  = np.zeros((T,D))
        gtIg = np.array([g['_ignore'] for g in gt])
        dtIg = np.zeros((T,D))
        if not len(ious)==0:
            for tind, t in enumerate(p.iouThrs):
                for dind, d in enumerate(dt):
                    # information about best match so far (m=-1 -> unmatched)
                    iou = min([t,1-1e-10])
                    m   = -1
                    for gind, g in enumerate(gt):
                        # if this gt already matched, continue
                        if gtm[tind,gind]>0:
                            continue
                        # if dt matched to reg gt, and on ignore gt, stop
                        if m>-1 and gtIg[m]==0 and gtIg[gind]==1:
                            break
                        # continue to next gt unless better match made
                        if ious[dind,gind] < iou:
                            continue
                        # match successful and best so far, store appropriately
                        iou=ious[dind,gind]
                        m=gind
                    # if match made store id of match for both dt and gt
                    if m ==-1:
                        continue
                    dtIg[tind,dind] = gtIg[m]
                    dtm[tind,dind]  = gt[m]['id']
                    gtm[tind,m]     = d['id']
        # set unmatched detections outside of area range to ignore
        a = np.array([d['area']<aRng[0] or d['area']>aRng[1] for d in dt]).reshape((1, len(dt)))
        dtIg = np.logical_or(dtIg, np.logical_and(dtm==0, np.repeat(a,T,0)))
        # store results for given image and category
        return {
                'image_id':     imgId,
                'category_id':  catId,
                'aRng':         aRng,
                'maxDet':       maxDet,
                'dtIds':        [d['id'] for d in dt],
                'gtIds':        [g['id'] for g in gt],
                'dtMatches':    dtm,
                'gtMatches':    gtm,
                'dtScores':     [d['score'] for d in dt],
                'gtIgnore':     gtIg,
                'dtIgnore':     dtIg,
            }

    def summarize(self, class_index=None, verbose=1):
        '''
        Compute and display summary metrics for evaluation results.
        Note this functin can *only* be applied on the default parameter setting
        '''
        def _summarize( ap=1, iouThr=None, areaRng='all', maxDets=500 ):
            p = self.params
            iStr = ' {:<18} {} @[ IoU={:<9} | area={:>6s} | maxDets={:>3d} ] = {:0.3f}'
            titleStr = 'Average Precision' if ap == 1 else 'Average Recall'
            typeStr = '(AP)' if ap==1 else '(AR)'
            iouStr = '{:0.2f}:{:0.2f}'.format(p.iouThrs[0], p.iouThrs[-1]) \
                if iouThr is None else '{:0.2f}'.format(iouThr)

            aind = [i for i, aRng in enumerate(p.areaRngLbl) if aRng == areaRng]
            mind = [i for i, mDet in enumerate(p.maxDets) if mDet == maxDets]
            if ap == 1:
                # dimension of precision: [TxRxKxAxM]
                s = self.eval['precision']
                # IoU
                if iouThr is not None:
                    t = np.where(iouThr == p.iouThrs)[0]
                    s = s[t]
                if not class_index is None:
                    s = s[:,:,class_index,aind,mind]
                else:
                    s = s[:,:,:,aind,mind]
            else:
                # dimension of recall: [TxKxAxM]
                s = self.eval['recall']
                if iouThr is not None:
                    t = np.where(iouThr == p.iouThrs)[0]
                    s = s[t]
                if not class_index is None:
                    s = s[:,class_index,aind,mind]
                else:
                    s = s[:,:,aind,mind]
            if len(s[s>-1])==0:
                mean_s = -1
            else:
                mean_s = np.mean(s[s>-1])
            if verbose > 0:
                print(iStr.format(titleStr, typeStr, iouStr, areaRng, maxDets, mean_s))
            return mean_s
        def _summarizeDets():
            stats = np.zeros((14,))
            stats[0] = _summarize(1)
            stats[1] = _summarize(1, iouThr=.5, maxDets=self.params.maxDets[4])
            stats[2] = _summarize(1, iouThr=.75, maxDets=self.params.maxDets[4])
            stats[3] = _summarize(1, areaRng='small', maxDets=self.params.maxDets[4])
            stats[4] = _summarize(1, areaRng='medium', maxDets=self.params.maxDets[4])
            stats[5] = _summarize(1, areaRng='large', maxDets=self.params.maxDets[4])
            stats[6] = _summarize(0, maxDets=self.params.maxDets[0])
            stats[7] = _summarize(0, maxDets=self.params.maxDets[1])
            stats[8] = _summarize(0, maxDets=self.params.maxDets[2])
            stats[9] = _summarize(0, maxDets=self.params.maxDets[3])
            stats[10] = _summarize(0, maxDets=self.params.maxDets[4])
            stats[11] = _summarize(0, areaRng='small', maxDets=self.params.maxDets[4])
            stats[12] = _summarize(0, areaRng='medium', maxDets=self.params.maxDets[4])
            stats[13] = _summarize(0, areaRng='large', maxDets=self.params.maxDets[4])
            return stats
        def _summarizeKps():
            stats = np.zeros((10,))
            stats[0] = _summarize(1, maxDets=20)
            stats[1] = _summarize(1, maxDets=20, iouThr=.5)
            stats[2] = _summarize(1, maxDets=20, iouThr=.75)
            stats[3] = _summarize(1, maxDets=20, areaRng='medium')
            stats[4] = _summarize(1, maxDets=20, areaRng='large')
            stats[5] = _summarize(0, maxDets=20)
            stats[6] = _summarize(0, maxDets=20, iouThr=.5)
            stats[7] = _summarize(0, maxDets=20, iouThr=.75)
            stats[8] = _summarize(0, maxDets=20, areaRng='medium')
            stats[9] = _summarize(0, maxDets=20, areaRng='large')
            return stats
        if not self.eval:
            raise Exception('Please run accumulate() first')
        iouType = self.params.iouType
        if iouType == 'segm' or iouType == 'bbox':
            summarize = _summarizeDets
        elif iouType == 'keypoints':
            summarize = _summarizeKps
        self.stats = summarize()

    def __str__(self, cass_index=None):
        self.summarize(class_index)

def evaluate_box_proposals(
    json_dataset, roidb, thresholds=None, area='all', limit=None
):
    """Evaluate detection proposal recall metrics. This function is a much
    faster alternative to the official COCO API recall evaluation code. However,
    it produces slightly different results.
    """
    # Record max overlap value for each gt box
    # Return vector of overlap values
    areas = {
        'all': 0,
        'small': 1,
        'medium': 2,
        'large': 3,
        '96-128': 4,
        '128-256': 5,
        '256-512': 6,
        '512-inf': 7}
    area_ranges = [
        [0**2, 1e5**2],    # all
        [0**2, 32**2],     # small
        [32**2, 96**2],    # medium
        [96**2, 1e5**2],   # large
        [96**2, 128**2],   # 96-128
        [128**2, 256**2],  # 128-256
        [256**2, 512**2],  # 256-512
        [512**2, 1e5**2]]  # 512-inf
    assert area in areas, 'Unknown area range: {}'.format(area)
    area_range = area_ranges[areas[area]]
    gt_overlaps = np.zeros(0)
    num_pos = 0
    for entry in roidb:
        gt_inds = np.where(entry['gt_classes'] > 0)[0]
        gt_boxes = entry['boxes'][gt_inds, :]
        gt_areas = entry['seg_areas'][gt_inds]
        valid_gt_inds = np.where(
            (gt_areas >= area_range[0]) & (gt_areas <= area_range[1]))[0]
        gt_boxes = gt_boxes[valid_gt_inds, :]
        num_pos += len(valid_gt_inds)
        non_gt_inds = np.where(entry['gt_classes'] == 0)[0]
        boxes = entry['boxes'][non_gt_inds, :]
        if boxes.shape[0] == 0:
            continue
        if limit is not None and boxes.shape[0] > limit:
            boxes = boxes[:limit, :]
        overlaps = box_utils.bbox_overlaps(
            boxes.astype(dtype=np.float32, copy=False),
            gt_boxes.astype(dtype=np.float32, copy=False))
        _gt_overlaps = np.zeros((gt_boxes.shape[0]))
        for j in range(min(boxes.shape[0], gt_boxes.shape[0])):
            # find which proposal box maximally covers each gt box
            argmax_overlaps = overlaps.argmax(axis=0)
            # and get the iou amount of coverage for each gt box
            max_overlaps = overlaps.max(axis=0)
            # find which gt box is 'best' covered (i.e. 'best' = most iou)
            gt_ind = max_overlaps.argmax()
            gt_ovr = max_overlaps.max()
            assert gt_ovr >= 0
            # find the proposal box that covers the best covered gt box
            box_ind = argmax_overlaps[gt_ind]
            # record the iou coverage of this gt box
            _gt_overlaps[j] = overlaps[box_ind, gt_ind]
            assert _gt_overlaps[j] == gt_ovr
            # mark the proposal box and the gt box as used
            overlaps[box_ind, :] = -1
            overlaps[:, gt_ind] = -1
        # append recorded iou coverage level
        gt_overlaps = np.hstack((gt_overlaps, _gt_overlaps))

    gt_overlaps = np.sort(gt_overlaps)
    if thresholds is None:
        step = 0.05
        thresholds = np.arange(0.5, 0.95 + 1e-5, step)
    recalls = np.zeros_like(thresholds)
    # compute recall for each iou threshold
    for i, t in enumerate(thresholds):
        recalls[i] = (gt_overlaps >= t).sum() / float(num_pos)
    # ar = 2 * np.trapz(recalls, thresholds)
    ar = recalls.mean()
    return {'ar': ar, 'recalls': recalls, 'thresholds': thresholds,
            'gt_overlaps': gt_overlaps, 'num_pos': num_pos}

def evaluate_keypoints(
    json_dataset,
    all_boxes,
    all_keypoints,
    output_dir,
    use_salt=True,
    cleanup=False
):
    res_file = os.path.join(
        output_dir, 'keypoints_' + json_dataset.name + '_results'
    )
    if use_salt:
        res_file += '_{}'.format(str(uuid.uuid4()))
    res_file += '.json'
    _write_coco_keypoint_results_file(
        json_dataset, all_boxes, all_keypoints, res_file)
    # Only do evaluation on non-test sets (annotations are undisclosed on test)
    if json_dataset.name.find('test') == -1:
        coco_eval = _do_keypoint_eval(json_dataset, res_file, output_dir)
    else:
        coco_eval = None
    # Optionally cleanup results json file
    if cleanup:
        os.remove(res_file)
    return coco_eval

def _write_coco_keypoint_results_file(
    json_dataset, all_boxes, all_keypoints, res_file
):
    results = []
    for cls_ind, cls in enumerate(json_dataset.classes):
        if cls == '__background__':
            continue
        if cls_ind >= len(all_keypoints):
            break
        logger.info(
            'Collecting {} results ({:d}/{:d})'.format(
                cls, cls_ind, len(all_keypoints) - 1))
        cat_id = json_dataset.category_to_id_map[cls]
        results.extend(_coco_kp_results_one_category(
            json_dataset, all_boxes[cls_ind], all_keypoints[cls_ind], cat_id))
    logger.info(
        'Writing keypoint results json to: {}'.format(
            os.path.abspath(res_file)))
    with open(res_file, 'w') as fid:
        json.dump(results, fid)

def _coco_kp_results_one_category(json_dataset, boxes, kps, cat_id):
    results = []
    image_ids = json_dataset.COCO.getImgIds()
    image_ids.sort()
    assert len(kps) == len(image_ids)
    assert len(boxes) == len(image_ids)
    use_box_score = False
    if cfg.KRCNN.KEYPOINT_CONFIDENCE == 'logit':
        # This is ugly; see utils.keypoints.heatmap_to_keypoints for the magic
        # indexes
        score_index = 2
    elif cfg.KRCNN.KEYPOINT_CONFIDENCE == 'prob':
        score_index = 3
    elif cfg.KRCNN.KEYPOINT_CONFIDENCE == 'bbox':
        use_box_score = True
    else:
        raise ValueError(
            'KRCNN.KEYPOINT_CONFIDENCE must be "logit", "prob", or "bbox"')
    for i, image_id in enumerate(image_ids):
        if len(boxes[i]) == 0:
            continue
        kps_dets = kps[i]
        scores = boxes[i][:, -1].astype(np.float)
        if len(kps_dets) == 0:
            continue
        for j in range(len(kps_dets)):
            xy = []

            kps_score = 0
            for k in range(kps_dets[j].shape[1]):
                xy.append(float(kps_dets[j][0, k]))
                xy.append(float(kps_dets[j][1, k]))
                xy.append(1)
                if not use_box_score:
                    kps_score += kps_dets[j][score_index, k]

            if use_box_score:
                kps_score = scores[j]
            else:
                kps_score /= kps_dets[j].shape[1]

            results.extend([{'image_id': image_id,
                             'category_id': cat_id,
                             'keypoints': xy,
                             'score': kps_score}])
    return results

def _do_keypoint_eval(json_dataset, res_file, output_dir):
    ann_type = 'keypoints'
    imgIds = json_dataset.COCO.getImgIds()
    imgIds.sort()
    coco_dt = json_dataset.COCO.loadRes(res_file)
    coco_eval = COCOeval(json_dataset.COCO, coco_dt, ann_type)
    coco_eval.params.imgIds = imgIds
    coco_eval.evaluate()
    coco_eval.accumulate()
    eval_file = os.path.join(output_dir, 'keypoint_results.pkl')
    save_object(coco_eval, eval_file)
    logger.info('Wrote json eval results to: {}'.format(eval_file))
    coco_eval.summarize()
    return coco_eval


"""Helper functions for decoding analysis."""
import warnings
warnings.simplefilter(action='ignore', category=FutureWarning)
warnings.simplefilter(action='ignore', category=DeprecationWarning)

import os
from pathlib import Path
import numpy as np

from scipy import stats

from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn import svm
from sklearn.model_selection import StratifiedKFold, RepeatedStratifiedKFold
from sklearn.linear_model import Perceptron
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis
from sklearn.linear_model import LogisticRegression
from sklearn.feature_selection import SelectKBest, f_classif

from skimage.measure import block_reduce

import matplotlib
import matplotlib.pyplot as plt
import matplotlib as mpl
from matplotlib.colors import ListedColormap
from matplotlib import cm, colors, colorbar

import time

import mne
from mne.stats import permutation_cluster_1samp_test, bootstrap_confidence_interval
from mne.decoding import (GeneralizingEstimator, SlidingEstimator, get_coef, LinearModel)
from mne.baseline import rescale

from general_helper_functions.data_general_utilities import moving_average, baseline_scaling

font = {'weight': 'bold',
        'size': 14}

matplotlib.rc('font', **font)

# figure(figsize=(8, 6), dpi=80)
SUPPORTED_CLASSIFIERS = ["Perceptron", "linear_svm"]

def sum_of_square(x):
    """Return sum of squares."""
    return np.sum(x ** 2, axis=0)

def plot_roi_specificity(root, rois, decoding_scores, cmap="Purples", vmin=0.5, vmax=1., filename=None):
    views = [{'azimuth': 0, 'elevation': -90}, {'azimuth': 180, 'elevation': -120}]
    surf = "inflated"
    parc = "aparc.a2009s"

    subjects_dir = Path(root, "derivatives", "fs")

    # Get the brain object:
    Brain = mne.viz.get_brain_class()

    # ====================================================================
    # Getting the surface labels:
    labels = mne.read_labels_from_annot("fsaverage", parc=parc, hemi='both', surf_name=surf,
                                        annot_fname=None, subjects_dir=subjects_dir, sort=True, verbose=None)

    # Plot the brain:
    brain = Brain('fsaverage', 'lh', surf, subjects_dir=subjects_dir, cortex='low_contrast',
                  background='white', size=(800, 400), alpha=1)

    # Generating the color scale for this group:
    norm = colors.Normalize(vmin=vmin, vmax=vmax)
    # Adding the color bar:
    # reindex the colormap so it starts deeper into the colorscale (e.g. color values = 0.1-1 over the scale of
    # vmin-vmax (ACC) of purple)
    cmap = ListedColormap(cm.get_cmap(cmap)(np.linspace(0.1, 1, 128)))

    # Looping through each label in the data_df:
    for ii, label in enumerate(rois):
        # Get the data of that specific label:
        label_val = decoding_scores[ii]
        if label_val.ndim > 0:
            label_val = np.mean(label_val)
        # Get the label object for the specific label:
        label_obj = [label_object for label_object in labels if label in label_object.name]
        # assert len(label_obj) == 1, "There was more than one label matching the name: {}!".format(label)
        # Plot this label according to the color scheme:
        print('%s: %f' % (label, label_val))

        brain.add_label(label_obj[0], color=cmap(norm(label_val)), alpha=1, hemi="lh")

    # Finally, getting the different views to plot and save:
    fig = plt.figure(figsize=(8, 4))
    for view_ind, view in enumerate(views):
        if isinstance(view, str):
            brain.show_view(view)
        elif isinstance(view, dict):
            brain.show_view(**view)
        else:
            raise Exception("The view must be a list of strings or of dict!")
        im = brain.screenshot()
        ax = fig.add_axes([(view_ind * 0.5), 0.2, 0.5, 0.6])
        ax.axis('equal')
        ax.autoscale(enable=True, axis='x', tight=True)
        ax.axis('off')
        # crop out whitespace
        nonwhite_pix = (im != 255).any(-1)
        nonwhite_row = nonwhite_pix.any(1)
        nonwhite_col = nonwhite_pix.any(0)
        cropped_screenshot = im[nonwhite_row][:, nonwhite_col]
        plt.tight_layout()
        # Plot the surface
        ax.imshow(cropped_screenshot)
    ax = fig.add_axes([0.47, 0.1, 0.015, 0.13])
    # norm = mpl.colors.Normalize(vmin=0.5, vmax=vmax)
    cb = mpl.colorbar.ColorbarBase(ax, cmap=cmap, norm=norm, spacing='proportional',
                                   orientation="vertical")
    cb.set_ticks((vmin, vmax))
    cb.ax.tick_params(size=0, labelsize=8)
    cb.ax.set_title('ACC', fontsize=8)

    if filename is not None:
        plt.savefig(filename, dpi=150)

def plot_decoding_data(x, y, times, save_path=None, roi="", file_prefix="", sampling_freq=512):
    """

    :param x:
    :param y:
    :param times:
    :param save_path:
    :param roi:
    :param file_prefix:
    :param sampling_freq:
    :return:
    """
    # Averaging the data per trial type:
    cond_data = []
    for cond in set(y):
        cond_data.append(np.mean(x[np.where(y == cond)[0]], axis=0))
    x_heatmap = np.vstack(cond_data)

    # Plot a heatmap of that:
    fig, ax = plt.subplots()
    # Plotting the comparison:
    ax.imshow(x_heatmap, extent=list(times[[0, -1]]) + [0, x_heatmap.shape[0]], aspect="auto")
    ax.hlines(y=x_heatmap.shape[0] / 2, linestyles="-.", xmin=times[0], xmax=times[-1], color="k")
    ax.vlines(x=0, ymin=0, ymax=x_heatmap.shape[0], color="k", linestyles="-.")
    ax.set_ylabel("channels")
    ax.set_xlabel("Time")
    plt.title("Channels mean activations across trials")
    ax.set_yticks([x_heatmap.shape[0] / 4, 3 * x_heatmap.shape[0] / 4])
    ax.set_yticklabels(set(y))
    if save_path is not None:
        file_name = Path(save_path, file_prefix
                         + "heatmap_{0}.png".format(roi))
        fig.savefig(file_name, transparent=True)
    plt.close(fig)

    # Create the events dict:
    event_dict = {cond: int(i) for i, cond in enumerate(np.unique(y))}
    # Group the data per conditions:
    data = []
    events_desc = []
    for cond in np.unique(y):
        data.extend(x[(np.where(y == cond)[0])])
        events_desc.extend([event_dict[cond]] * len(np.where(y == cond)[0]))
    events = np.column_stack((np.arange(0, len(data) * (times[-1] - times[0]) * sampling_freq,
                                        (times[-1] - times[0]) * sampling_freq, dtype=int),
                              np.zeros(len(events_desc), dtype=int),
                              np.array(events_desc)))
    # Create the info:
    ch_names = ["ch_{}".format(i) for i in range(x.shape[1])]
    info = mne.create_info(ch_names, ch_types=["ecog"] * x.shape[1], sfreq=sampling_freq)
    new_epochs = mne.EpochsArray(data, info, tmin=times[0], events=events,
                                 event_id=event_dict)

    # -------------------------------
    # Compute the global field power (following tutorial here:
    # https://mne.tools/dev/auto_examples/time_frequency/time_frequency_global_field_power.html):
    evoked = new_epochs.average()
    gfp = np.sum(evoked.data ** 2, axis=0)
    times = evoked.times * 1e3
    gfp = rescale(gfp, times, baseline=(None, 0))
    fig, ax = plt.subplots(figsize=(10, 4))
    ax.plot(times, gfp, linewidth=2.5)
    ax.axhline(0, linestyle='--', color='grey', linewidth=2)
    ci_low, ci_up = bootstrap_confidence_interval(evoked.data, random_state=0,
                                                  stat_fun=sum_of_square)
    ci_low = rescale(ci_low, evoked.times, baseline=(None, 0))
    ci_up = rescale(ci_up, evoked.times, baseline=(None, 0))
    ax.fill_between(times, gfp + ci_up, gfp - ci_low, alpha=0.3)
    ax.grid(True)
    ax.set_ylabel('GFP')
    ax.set_title("Global field power across all electrodes in roi {}".format(roi))
    if save_path is not None:
        file_name = Path(save_path, file_prefix
                         + "gfp_{0}.png".format(roi))
        plt.savefig(file_name, transparent=True)
    # Z-scoring the data:
    baseline_scaling(new_epochs, correction_method="zscore")

    # Plot the image:
    new_epochs.plot_image(scalings={"ecog": 1}, show=False, units=dict(ecog='zscore'), combine="mean")
    if save_path is not None:
        file_name = Path(save_path, file_prefix
                         + "plot_image_{0}.png".format(roi))
        plt.savefig(file_name, transparent=True)
    plt.close()
    # Now computing the evoked responses:
    evk = {cond: None for cond in np.unique(y)}
    for cond in event_dict.keys():
        evk[cond] = new_epochs[event_dict[cond]].average()
    # Plotting all the electrodes on top of another:
    fig, axs = plt.subplots(len(event_dict.keys()))
    for ind, cond in enumerate(event_dict.keys()):
        evk[cond].plot(titles=cond, axes=axs[ind], scalings={"ecog": 1}, ylim={"ecog": [-10, 10]}, show=False)
        axs[ind].set_ylabel("zscore")
    plt.tight_layout()
    if save_path is not None:
        file_name = Path(save_path, file_prefix
                         + "plot_evoked_{0}.png".format(roi))
        fig.savefig(file_name, transparent=True)
        plt.close(fig)
    fig, ax = plt.subplots()
    # Plotting the comparison:
    mne.viz.plot_compare_evokeds(evk, combine='mean', axes=ax, show=False)
    ax.set_ylabel("zscore")
    if save_path is not None:
        file_name = Path(save_path, file_prefix
                         + "plot_compared_evoked_{0}.png".format(roi))
        fig.savefig(file_name, transparent=True)
        plt.close(fig)
    plt.close()

    return None

def find_channels_in_roi(roi, channels_labels):
    """
    This function checks which channels are found in a given ROI. The format of the ROI must be a list of labels that
    constitute that roi. The labels are of the format outputted by mne.get_montage_volume_labels, a dict with a list
    of labels for a given channel according to a specific parcellation.
    :param roi: (dict) key of  the dictionary corresponds to the ROI in question, and within it there is a string
    with all the labels from the free surfer parcellation of interest
    For example:
    "occipital": ["ctx-rh-lateraloccipital", "ctx-lh-lateraloccipital", "ctx-rh-inferiorparietal",
    "ctx-lh-inferiorparietal", "ctx-rh-pericalcarine", "ctx-lh-pericalcarine", "ctx-rh-cuneus", "ctx-lh-cuneus"]
    :param channels_labels: (dict) output of the mne function: mne.get_montage_volume_labels
    :return: (list of strings) list of the channels found in the ROI of interest
    """
    roi_channels = []
    # Looping through each channel found in the channels_labels dict:
    for ch in channels_labels.keys():
        # Looping through each label of this specific channel:
        for label in channels_labels[ch]:
            if label in roi:
                roi_channels.append(ch)
                break

    # Return the channels
    return roi_channels

def compute_pseudotrials(x, y, groups=None, n_trials=5, pad_val=np.nan, permute_trials=True):
    """
    This function enables computation of pseudotrials. This consists in averaging n_trials together. The averaging
    of course happens separately for groups and y labels separately (i.e. only trials of matching groups are averaged).
    Additionally, the attribution to trial to a group of pseudotrial is randomized, such that we don't have the
    following trials being averaged together.
    :param x: (np array) contains the data to use for the decoding. The data must be of dimension: trials *
    channels * time points. The averaging will be done along the trials dimension
    :param y: (list or np array) decoding targets. Must be of the same dimension as the first dim of x
    :param groups: (list or np array) groups to which each trial belongs!
    :param n_trials: (int) number of trials to average together. So if you set that to 5, every 5 trials of the same
    y condition (and groups condition) will be averaged together
    :param pad_val: (int, float, nan...) what to pad the array with when doing the average. Say you want to average
    every 5 trials but you have 26. By setting to np.nan you will make up for the array to be of size 30 by padding with
    nan, having the last trial being picked being averaged with nan.
    :param permute_trials: (boolean) whether to shuffle the trials order before doing the averaging. If set to true,
    it averages the trials in the order they come in the arrays. But by setting it to true, the order is randomized,
    to avoid any sort of dependencies
    :return:
    new_x: (np.array) x data averaged together. So it will be of the size of x but with first dimension / n_trials
    new_y:(np.array) decoding targets matching with the new_x dimensions
    new_groups: (np.array) groups matching with the new_x dimensions
    """
    # Print the original shape of the data:
    print("-" * 40)
    print("Welcome to pseudotrials computations")
    print("Original data shape:")
    print("x: {0}".format(x.shape))
    print("y: {0}".format(y.shape))
    # Getting the unique labels and the counts:
    values, counts = np.unique(y, return_counts=True)
    [print("{0}: {1}".format(values[i], counts[i])) for i in range(len(values))]
    new_x = []
    new_y = []
    new_groups = []
    # Checking the size of the inputs to ensure that they match:
    if len(y) != x.shape[0]:
        raise Exception("The dimension of the targets labels doesn't match the size of the data matrix!")
    if groups is not None and len(groups) != x.shape[0]:
        raise Exception("The dimension of the groups labels doesn't match the size of the data matrix!")
    if groups is not None and len(groups) != len(y):
        raise Exception("The dimension of the targets labels doesn't match the size of the groups labels!")
    # Depending on whether or not we have groups, the computation of pseudotrials need to be accounting for it:
    if groups is not None:
        # Looping through each unique group:
        for group in set(groups):
            # Extracting only those trials and labels from this group:
            group_x, group_y = x[np.where(groups == group)[0]], y[np.where(groups == group)[0]]
            # Now, looping through each unique label from this subset of the data to perform the averaging:
            for label in set(group_y):
                # Extract the data:
                data = group_x[np.where(group_y == label)]
                if permute_trials:
                    data = np.take(data, np.random.permutation(data.shape[0]), axis=0)
                avg_x = block_reduce(data, block_size=tuple([n_trials, *[1] * len(data.shape[1:])]),
                                     func=np.nanmean, cval=pad_val)
                # Now generating the labels and group:
                new_x.append(avg_x)
                new_y += [label] * avg_x.shape[0]
                new_groups += [group] * avg_x.shape[0]
    else:
        # In case there are no groups, one level less to dig to:
        for label in set(y):
            # Get the data:
            data = x[np.where(y == label)]
            # Permute if needed:
            if permute_trials:
                data = np.take(data, np.random.permutation(data.shape[0]), axis=0)
            # Apply block reduce function:
            avg_x = block_reduce(data, block_size=tuple([n_trials, *[1] * len(data.shape[1:])]),
                                 func=np.nanmean, cval=pad_val)
            # Now generating the labels and group:
            new_x.append(avg_x)
            new_y += [label] * avg_x.shape[0]
        new_groups = None
    # Turning everything back to numpy arrays:
    new_x = np.concatenate(new_x, axis=0)
    new_y = np.array(new_y)
    if new_groups is not None:
        new_groups = np.array(new_groups)
    # Printing info about the pseudotrials resulting from the procedure:
    print("Pseudotrials data shape:")
    print("x: {0}".format(new_x.shape))
    print("y: {0}".format(new_y.shape))
    values, counts = np.unique(new_y, return_counts=True)
    [print("{0}: {1}".format(values[i], counts[i])) for i in range(len(values))]
    return new_x, new_y, new_groups

def pad_with_nan(array, shape):
    """
    This function pads a given array to be of the shape specified
    :param array: (np array) array to pad
    :param shape: (list of int or np array shape object) shape that the passed array must have
    :return: (np array) padded array to the size expected
    """
    pad_per_dim = []
    for dim_ind, dim_size in enumerate(array.shape):
        # Compare each axis to the desired shape:
        dim_pad = shape[dim_ind] - dim_size
        pad_per_dim.append([dim_pad // 2, dim_pad // 2 + dim_pad % 2])
    return np.pad(array, pad_per_dim,
                  mode='mean')

def slice_per(source, step, axis=0, overlapping=True):
    """
    This function takes a numpy array and splits it in sub arrays by taking n data points (specified as step).
    So say your array is of like so:
    1, 2, 3, 4, 5, 6,
    3, 4, 2, 2, 1, 3
    And you set axis to 1 and step to 2, it will return:
    array 1:
    1, 3, 5
    3, 2, 1
    array 2:
    2, 4, 6
    4, 2, 3
    :param source: (np array) source data to slice
    :param step: (int) steps to slice the data per
    :param axis: (int) axis along which to slice
    :param overlapping: (boolean) whether or not to do overlapping
    :return: (list of arrays) list of each slice of the array
    """
    if overlapping:
        return [source.take(indices=range(i, i + step), axis=axis) for i in range(source.shape[axis] - step)]
    else:
        return [source.take(indices=range(i, i + step), axis=axis) for i in
                range(0, source.shape[axis] - step, step)]

def equate_array_sizes(list_of_arrays):
    """
    This function equates arrays of different sizes. This function compute the size of the different arrays and find the
    one of highest size. This info is then passed to the pad  with nan function to pad the array
    :param list_of_arrays: (list of arrays) list of arrays to equate in size.
    :return:
    """
    # Get the size of all arrays in the list:
    shapes_list = np.array([array.shape for array in list_of_arrays])
    # Finding the max across all the lists:
    max_dim = np.max(shapes_list, axis=0)
    # Equating all the arrays to the same size:
    return [pad_with_nan(array, max_dim) for array in list_of_arrays]

def bin_time_win(data, sr, axis=-1, overlapping=False, bins_duration_ms=10, average=True):
    """
    This function bins data according to a specific bins duration. If average is True, that data will be binned
    in the last dimension every n samples, essentially adding a dimension that will then be stacked back onto the
    channels dimension (this is a bit of a weird thing to do when you have several channels, but with only one,
    it means your feature are basically over time as opposed to be over electrodes). If the average option is set to
    true, the data will be averaged in each bin, essentially reducing the time dimension
    1, 2, 1, 2, 1, 2, 1, 2 and you want 2ms bins, it will return this:
    2, 2, 2, 2
    1, 1, 1, 1
    :param data: (numpy array) data to bin
    :param sr: (int) sampling rate of the data
    :param axis: (int) axis along which to bin
    :param overlapping: (boolean) whether or not to do the binning overlapping
    :param bins_duration_ms: (int) duration of each bin in milliseconds
    :param average: (boolean) whether or not to average in each bin. If true, the data will be averaged any n ms (as
    set by bins_duration_ms), and  the overlapping will also apply
    :return: (np array): binned data
    """
    # Compute how many samples the time window corresponds to:
    n_samples = int(np.floor(bins_duration_ms * sr / 1000))
    if n_samples < 2:
        print("=" * 80)
        print("WARNING: The time window used for binning is less than two samples. "
              "Therefore, binning is not possible! Increase window length or don't bin")
        return
    if average:
        data_new = moving_average(data, n_samples, axis=-1, overlapping=overlapping)
    else:
        # Splitting the data:
        list_of_arrays = slice_per(
            data, n_samples, axis=axis, overlapping=overlapping)
        # Reshaping to flatten the second and third dim:
        list_of_arrays_2d = [array.reshape(
            (array.shape[0], -1)) for array in list_of_arrays]
        # Then, stacking these arrays along the second dimension, to keep things 3D:
        data_new = np.stack(list_of_arrays_2d, axis=2)

    return data_new

def file_name_generator(save_path, file_prefix, description, file_extension, data_type="ieeg"):
    """
    This function generates full file names according to the cogitate naming conventions:
    :param save_path: (pathlib path object or path string) root path to where the data should be saved
    :param file_prefix: (string) prfix of the file name
    :param description: (string) what some after teh prefix in the file name
    :param file_extension: (string) what comes after the description, if anything
    :param data_type: (string) data type the data are from
    :return: full_file_name (string) file path + file name
    """
    full_file_name = os.path.join(
        save_path, file_prefix + description + "_" + data_type + file_extension)

    return full_file_name

def get_cross_validation_data(cross_validation_iterator, x, y):
    """
    This function extracts the data for the train and test sets of each cross validation folds.
    :param cross_validation_iterator: (scikit learn cross validation iterators) See doc here:
    https://scikit-learn.org/stable/modules/cross_validation.html
    :param x: (numpy array) array containing the data to be used for the decoding and that will be split by the
    cross validation iteror
    :param y: (numpy array or list) list of the decoding targets that will be split in the same way as the x
    :return: x_train, y_train, x_test, y_test: (lists of arrays) holding the x and y of the train and test sets
    respectively
    """
    print("-" * 40)
    print("Welcome to get_cross_validation_data")
    # Preparing empty lists to hold train and test data of each folds:
    x_train, x_test, y_train, y_test = [], [], [], []
    # Extracting the splitted test from the y test:
    ctr = 0
    for train_index, test_index in cross_validation_iterator.split(x, y):
        x_train.append(x[train_index])
        y_train.append(y[train_index])
        x_test.append(x[test_index])
        y_test.append(y[test_index])
        # Print the sizes of the test and train sets:
        print("Fold {0} train set trial number={1}, test set trial number={2}".format(ctr, len(train_index),
                                                                                len(test_index)))
        ctr = ctr + 1

    return x_train, y_train, x_test, y_test

def print_trials_counts(prefix, labels):
    """
    This function is a small utility to print trials counts
    :param labels: (list of strings) contains the targets to count
    :param prefix: (string) what to call the set we are printing the counts for
    :return:
    """
    print("{0} set trial number: {1}".format(prefix, len(labels)))
    targets_name, targets_counts = np.unique(labels, return_counts=True)
    print("{0} set trials counts:".format(prefix))
    [print("{0}: {1}".format(targets_name[i], targets_counts[i])) for i in range(len(targets_name))]
    return None

def temporal_generalization_decoding(clf, x, decoding_target, cross_validation_parameters, metric="accuracy",
                                     train_group=None, test_group=None,
                                     groups=None, n_pseudotrials=None, shuffle_labels=False, classifier_n_jobs=1,
                                     do_only_diag=False, average_scores=False, n_channel_subsample = None, verbose=False):
    """
    This function performs decoding in either a time resolved or temporal generalization fashion. This function can do
    so in a couple of different ways depending on the passed parameters.

    Option 1:
    Within group decoding:
        If train_group and test_group are both set to either None or both to the same string, the decoding will
        be performed within this group. The train and test_groups refers to a condition within which you want to train
        and test (so say task relevant and irrelevant). So if you set both to None, that mean that you want to
        train and test on all the trials. If you specify the same string in each, that means you want to train and test
        specifically within a given subset of your trials. For this option to work, you must set
        cross_validation_parameters["n_folds"] to not None and > 2. Otherwise you have no way to test your decoding
        accuracy.
    Cross task generalization decoding:
        If you pass two differrent groups under train_group and test_groups, then that means you want to train on a
        given group and test in the other. For this type of decoding, you must pass a list in groups of  the same
        length as decoding_target, identifying for each trial to which group it belongs to segregate them.
        When you use this function, you can skip cross_validation alltogether by setting
        cross_validation_parameters["n_folds"] to None, which mean training on all trials of the one condition and
        testing on all trials of the other. But if you do want to do cross fold validation, you can. And if you do
        want to do that, you have two other sub-options. In cross fold validation, the train set will be split in
        n_folds. You can then train on n-1 fold and test on all the trials of the other group to test your decoding.
        But you can also split the test set by the same amount as the training set if trial counts is something you
        are worried about
    :param clf: (scikit learn pipeline object) pipeline to be used for the decoding
    :param x: (np array) contains the data to be used for the decoding. The first dimension should represent the trials
    :param decoding_target: (list of strings) column of the epochs object meta data on which to perform the decoding
    :param cross_validation_parameters: (dict) parameters of the cross validation
        "n_folds": (int or None) how many folds to do the cross validation
        "split_generalization_set": (boolean) whether or not you want to split the test set as much as the train set
        in the case of cross task generalization decoding
    :param train_group: (string or None) groups on which to train the classifier. The string must match what is found
    in the groups list
    :param test_group: (string or None) groups on which to test the classifier. The string must match what is found
    in the groups list
    :param groups: (list) must be of same length as decoding_target, as it contains the group appartenance of each trial
    :param n_pseudotrials: (int or None) set to None if you do not want to use pseudotrials. If you use an int,
    then the trials will be averaged every n trials (so if = 2, the pseudotrials will be computed by gathering every
    two trials, and the
    :param shuffle_labels: (bool) whether or not to shuffle the labels before performing the decoding. This
    enables this function to be called many times to shuffle the labels and generate a null distribution.
    CAREFUL: this is not the same as do_test_label_shuffle. The latter consists of shuffling the labels of the test
    set after training to generate a null distribution. If you have set that to true, you should set the other to False
    because it wouldn't make any sense.
    :param classifier_n_jobs: (int) how many jobs to use to preprocessing the classifier in parallel
    :param do_only_diag: (boolean) whether to do only diagonal (time resolved) or temporal generalization decoding
    :param average_scores: (boolean) whether to average the decoding scores across folds
    :param verbose: (boolean) whether or not to print info to command line
    :return:
    """
    if verbose:
        print("-" * 40)
        print("Performing decoding")
    # Check the different inputs:
    if train_group is None and test_group is not None or train_group is not None and test_group is None:
        raise Exception("You have passed {0} as train condition, but {1} as test condition. This doesn't work! "
                        "\nThe train condition must either be both None or both strings!".
                        format(train_group, test_group))
    if isinstance(train_group, str) and groups is None:
        raise Exception("You have passed {0} as train and {1} as test groups, but you haven't passed a groups array to "
                        "\nidentify which trial belongs to which group".format(train_group, test_group))
    # First, specifiying the decoder:
    if do_only_diag is True:
        time_gen = SlidingEstimator(clf, n_jobs=classifier_n_jobs, scoring=metric,
                                    verbose="ERROR")
    else:
        time_gen = GeneralizingEstimator(clf, n_jobs=classifier_n_jobs, scoring=metric,
                                         verbose="ERROR")

    # randomly subsample the data (used for robustness checks)
    if n_channel_subsample is not None:
        x = x[:, np.random.choice(x.shape[1], n_channel_subsample, replace=False), :]
    # Shuffling the trials if needed (to generate null distributions for example):
    if shuffle_labels:
        decoding_target = decoding_target[np.random.permutation(len(decoding_target))]
    # Then, generating the pseudotrials if required:
    if n_pseudotrials is not None:
        x, decoding_target, groups = compute_pseudotrials(x, decoding_target, groups, n_trials=n_pseudotrials)

    # --------------------------------------------------------------------------------------------------------------
    # Option 1: train and test within the same condition
    # --------------------------------------------------------------------------------------------------------------
    if train_group == test_group:
        # Checking that the cross fold validation is set to true:
        if cross_validation_parameters["n_folds"] is None:
            raise Exception("If you are doing within task decoding, you must use cross fold validation to be able"
                            "\nto test your trained decoding on something!")
        # Get the classes:
        y = decoding_target
        # Setting the cross validation:
        if cross_validation_parameters["n_folds"] == "leave_one_out":
            n_folds = len(y)
        else:
            n_folds = cross_validation_parameters["n_folds"]
        # Prepare list to hold the results:
        scores = []
        # Creating cross val iterator:
        skf = StratifiedKFold(n_splits=n_folds)
        # Getting the indices of the test and train sets from cross folder validation:
        cv_iter = list(skf.split(x, y))
        # Performing the decoding:
        for ind, train_test_ind in enumerate(cv_iter):
            if verbose:
                print_trials_counts("train", y[train_test_ind[0]])
                print_trials_counts("test", y[train_test_ind[1]])
            # Initiate time:
            start = time.time()
            # Train on this split:
            time_gen.fit(X=x[train_test_ind[0]],
                         y=y[train_test_ind[0]])
            if verbose:
                print("Train time={:.2f}sec".format(time.time() - start))
            # Initiate time and CPU counters:
            start = time.time()
            # Test decoder within the same condition:
            scores.append(time_gen.score(X=x[train_test_ind[1]], y=y[train_test_ind[1]]))
            if verbose:
                print("Test time={:.2f}sec".format(time.time() - start))
        # Compute the coefficient regardless:
        try:
            time_gen.fit(x, y)
            coef = get_coef(time_gen, 'coef_', inverse_transform=True)
        except ValueError:
            coef = None
            print("WARNING: The coefficient could not be computed as the classifier was not linear")

        # Saving the scores in dict to keep things consistent between options
        scores = np.array(scores)
    # --------------------------------------------------------------------------------------------------------------
    # Option 2: Cross task generalization:
    # --------------------------------------------------------------------------------------------------------------
    else:
        # Prepare result of test:
        scores = []
        # Getting the data of the condition on which to perform the decoding:
        x_train_condition = x[np.where(groups == train_group)]
        # Get the targets:
        y_train_condition = decoding_target[np.where(groups == train_group)]

        # Get the test data:
        x_generalization = x[np.where(groups == test_group)]
        y_generalization = decoding_target[np.where(groups == test_group)]

        # In case of cross task generalization, there is the option of performing cross fold validation. If this is
        # expected, the data in the train group will be splitted in n_folds to test against the test set. The test
        # set can further be subdived into n folds or not:
        if cross_validation_parameters["n_folds"] is not None:
            # Create the kfold object:
            skf = StratifiedKFold(n_splits=cross_validation_parameters["n_folds"])
            # Get the labels and data of the different splits:
            x_train_condition_train, y_train_condition_train, x_train_condition_test, y_train_condition_test = \
                get_cross_validation_data(skf, x_train_condition, y_train_condition)

            # If the generalization set needs to be split:
            if cross_validation_parameters["split_generalization_set"]:
                _, _, x_generalization_sets, y_generalization_sets = \
                    get_cross_validation_data(
                        skf, x_generalization, y_generalization)

                # Performing the decoding:
                for ind, train_set in enumerate(x_train_condition_train):
                    if verbose:
                        print_trials_counts("train", y_train_condition_train[ind])
                        print_trials_counts("test", y_generalization_sets[ind])
                    # Train on this split:
                    start = time.time()
                    time_gen.fit(X=train_set,
                                 y=y_train_condition_train[ind])
                    if verbose:
                        print("Train time={:.2f}sec".format(time.time() - start))
                    # Test decoder generalization on other condition
                    start = time.time()
                    scores.append(time_gen.score(X=x_generalization_sets[ind],
                                                 y=y_generalization_sets[ind]))
                    if verbose:
                        print("Test time={:.2f}sec".format(time.time() - start))
            else:  # On each cross validation fold, testing the decoder on the entire generalization set:
                for ind, train_set in enumerate(x_train_condition_train):
                    print_trials_counts("train", y_train_condition_train[ind])
                    print_trials_counts("test", y_generalization)
                    # Train on this split:
                    start = time.time()
                    time_gen.fit(X=train_set, y=y_train_condition_train[ind])
                    if verbose:
                        print("Train time={:.2f}sec".format(time.time() - start))
                    start = time.time()
                    scores.append(time_gen.score(X=x_generalization,
                                                 y=y_generalization))
                    if verbose:
                        print("Test time={:.2f}sec".format(time.time() - start))
        else:  # If no cross fold validation is expected
            print_trials_counts("train", y_train_condition)
            print_trials_counts("test", y_generalization)
            # Train on this split:
            start = time.time()
            time_gen.fit(X=x_train_condition, y=y_train_condition)
            if verbose:
                print("Train time={:.2f}sec".format(time.time() - start))
            start = time.time()
            scores.append(time_gen.score(X=x_generalization,
                                         y=y_generalization))
            if verbose:
                print("Test time={:.2f}sec".format(time.time() - start))
        scores = np.array(scores)
        # Compute the coefficient regardless:
        try:
            time_gen.fit(x_train_condition, y_train_condition)
            coef = get_coef(time_gen, 'coef_', inverse_transform=True)
        except ValueError:
            coef = None
            print("WARNING: The coefficient could not be computed as the classifier was not linear")

    # Averaging the scores if necessary. If there were no cross fold validation or if only the first fold was tested
    # there will be only 1 matrix, so no need to average:
    if average_scores and cross_validation_parameters["n_folds"] is not None:
        scores = np.mean(scores, axis=0)

    return scores, coef

# compute variance-corrected ttests for model comparison
# see: Statistical comparison of models using grid search, https://scikit-learn.org/stable/auto_examples/model_selection/plot_grid_search_stats.html
def corrected_std(differences, n_train, n_test):
    """Corrects standard deviation using Nadeau and Bengio's approach.

    Parameters
    ----------
    differences : ndarray of shape (n_samples,)
        Vector containing the differences in the score metrics of two models.
    n_train : int
        Number of samples in the training set.
    n_test : int
        Number of samples in the testing set.

    Returns
    -------
    corrected_std : float
        Variance-corrected standard deviation of the set of differences.
    """
    # kr = k times r, r times repeated k-fold crossvalidation,
    # kr equals the number of times the model was evaluated
    kr = len(differences)
    corrected_var = np.var(differences, ddof=1) * (1 / kr + n_test / n_train)
    corrected_std = np.sqrt(corrected_var)
    return corrected_std

def compute_corrected_ttest(differences, df, n_train, n_test):
    """Computes right-tailed paired t-test with corrected variance.

    Parameters
    ----------
    differences : array-like of shape (n_samples,)
        Vector containing the differences in the score metrics of two models.
    df : int
        Degrees of freedom.
    n_train : int
        Number of samples in the training set.
    n_test : int
        Number of samples in the testing set.

    Returns
    -------
    t_stat : float
        Variance-corrected t-statistic.
    p_val : float
        Variance-corrected p-value.
    """
    mean = np.mean(differences)
    std = corrected_std(differences, n_train, n_test)
    t_stat = mean / std
    p_val = stats.t.sf(np.abs(t_stat), df)  # right-tailed t-test
    return t_stat, p_val


'''
A set of utility methods that are used in different parts of the framework.
'''

import logging

# Set of bots read from a bot tsv file
bots = None
filterBots = False
def setFilterBots(fb,botfile):
    global filterBots,bots

    if fb:
        try:
            bots = set(long(bot) for bot in open(botfile,'r'))
            filterBots = fb
            logging.info("%s Bots loaded from %s"%(len(bots),botfile))
        except:
            logging.error("Botlist (%s) could not be loaded, Bots will not be filtered"%botfile)

def isBot(u_id):
    '''
    Returns true if we filter for bots and  u_id is a known bot.

    :arg ints: Boolean, if True compares u_id as int (default is False)
    '''
    if filterBots:
        if u_id in bots:
            return True
    return False

def create_time_stamps_month(fromym='200101',toym='201012'):
    '''
    Helper data structures for time stamps
    List of all time unites, i.e. every month. yyyymm
    '''
    time_stamps = []
    # A dictionary that serves as a lookup for the index of atime stamp
    time_stamps_index = {}

    min_year = fromym[:-2]
    min_month = fromym[-2:]
    max_year = toym[:-2]
    max_month = toym[-2:]

    i = 0
    #add remaining months in the first year
    for m in range(int(min_month),13):
        ts = '%s%02d'%(min_year,m)
        time_stamps.append(ts)
        time_stamps_index[ts] = i
        i += 1
    #add intermediate years
    for y in range(int(min_year)+1,int(max_year)):
        for m in range(1,13):
            ts = '%s%02d'%(y,m)
            time_stamps.append(ts)
            time_stamps_index[ts] = i
            i += 1
    #last year
    for m in range(1,int(max_month)+1):
        ts = '%s%02d'%(max_year,m)
        time_stamps.append(ts)
        time_stamps_index[ts] = i
        i += 1

    return (time_stamps,time_stamps_index)

def create_time_stamps_day(fromymd='20010101',toymd='20101231'):
    '''
    Helper data structures for time stamps
    List of all time unites, i.e. every month. yyyymm
    '''

    import calendar

    time_stamps = []
    # A dictionary that serves as a lookup for the index of atime stamp
    time_stamps_index = {}

    min_year = int(fromymd[:4])
    min_month = int(fromymd[4:6])
    min_day = int(fromymd[6:8])
    max_year = int(toymd[:4])
    max_month = int(toymd[4:6])
    max_day = int(toymd[6:8])
    i = 0
    #add remaining months in the first year
    for m in range(min_month,13):
        # calender.monthrange returns a tuple (day of the week, number of days in the month)
        nds = calendar.monthrange(min_year,m)[1]
        for d in range(min_day,nds+1):
            ts = '%s%02d%02d'%(min_year,m,d)
            time_stamps.append(ts)
            time_stamps_index[ts] = i
            i += 1
    #add intermediate years
    for y in range(min_year+1,max_year):
        for m in range(1,13):
            # calender.monthrange returns a tuple (day of the week, number of days in the month)
            nds = calendar.monthrange(y,m)[1]
            for d in range(1,nds+1):
                ts = '%s%02d%02d'%(y,m,d)
                time_stamps.append(ts)
                time_stamps_index[ts] = i
                i += 1
    #to last month
    for m in range(1,max_month):
        nds = calendar.monthrange(max_year,m)[1]
        for d in range(1,nds+1):
            ts = '%s%02d%02d'%(max_year,m,d)
            time_stamps.append(ts)
            time_stamps_index[ts] = i
            i += 1
    #to last day
    for d in range(1,max_day+1):
        ts = '%s%02d%02d'%(max_year,max_month,d)
        time_stamps.append(ts)
        time_stamps_index[ts] = i
        i += 1

    return (time_stamps,time_stamps_index)

def computeMonthStartEndtime(ym):
    '''
    Returns the starting and end datetime object for the yyyymm passed. I.e. the first and last day of the month

    :arg ym: str, 'yyyymm' format
    :returns: tuple of datetime objects
    '''
    from datetime import datetime
    import calendar

    y = int(ym[:4])
    m = int(ym[4:])

    start = datetime(y, m, 1)
    # days in the given month
    # calender.monthrange returns a tuple (day of the week, number of days in the month)
    d = calendar.monthrange(y,m)[1]

    end = datetime(y, m, d)

    return (start,end)

def numberOfMonths(ymStart,ymEnd):
    '''Returns the number of months between the parameters.

    :arg ymStart: str, 'yyyymm' format
    :arg ymEnd: str, 'yyyymm' format
    :returns: int, number of month
    '''
    months = 0
    months += (int(ymEnd[:4])-(int(ymStart[:4])+1)) * 12
    if int(ymEnd[:4])==int(ymStart[:4]):
        #same year
        months += (int(ymEnd[4:])-int(ymStart[4:]))+1
    else:
        months += 12-int(ymStart[4:])+1
        months += int(ymEnd[4:])

    return months

def cmap_discretize(cmapName, N):
    """
    From http://www.scipy.org/Cookbook/Matplotlib/ColormapTransformations

    :arg cmap: colormap instance, eg. cm.jet.
    :arg N: Number of colors.

    :returns: a discrete colormap from the continuous colormap cmap.
    """

    try:
        from numpy import array, linspace, zeros, interp
        import matplotlib
        from matplotlib import pyplot as plt
    except:
        logging.error('matplotlib or numpy not installed')
        logging.error("cmap_discretize() returns invalid cmap; you can't plot without these packages anyways")
        return cmapName

    if cmapName is None:
        cmap = plt.cm.spectral
    else:
        cmap = plt.cm.get_cmap(cmapName)

    # try:
    #     from scipy import interpolate
    # except:
    #     logging.error('scipy not installed')
    #     logging.info("cmap_discretize() can't interpolate the colormap, returns cmap %s unchanged"%cmapName)
    #     return cmap

    cdict = cmap._segmentdata.copy()
    # N colors
    colors_i = linspace(0,1.,N)
    # N+1 indices
    indices = linspace(0,1.,N+1)
    for key in ('red','green','blue'):
        # Find the N colors
        D = array(cdict[key])

        # using scipy
        # I = interpolate.interp1d(D[:,0], D[:,1])
        # colors = I(colors_i)

        # using numpy
        colors = interp(colors_i,D[:,0], D[:,1])

        # Place these colors at the correct indices.
        A = zeros((N+1,3), float)
        A[:,0] = indices
        A[1:,1] = colors
        A[:-1,2] = colors
        # Create a tuple for the dictionary.
        L = []
        for l in A:
            L.append(tuple(l))
        cdict[key] = tuple(L)
    # Return colormap object.
    return matplotlib.colors.LinearSegmentedColormap('colormap',cdict,1024)

def movingAverage(array, WINDOW=5):
    try:
        import numpy as N
    except:
        logging.warning("Moving average can't be computed (Numpy not installed)")
        return

    weightings = N.repeat(1.0, WINDOW) / WINDOW
    return N.convolve(array, weightings)[WINDOW-1:-(WINDOW-1)]


from django.db import models

######################################################################################################################

class Status(models.Model):
    id = models.AutoField(
        primary_key=True,
    )
    title = models.CharField(
        verbose_name='Название стадии',
        max_length=64,
        default='',
    )
    order = models.SmallIntegerField(
        verbose_name='Порядок в списке стадий',
        default=0,
    )
    next_status = models.ForeignKey(
        'self',
        verbose_name='Следующая стадия',
        related_name='NextStatus',
        null=True,
        blank=True,
        default=None,
        on_delete=models.SET_NULL,
    )

    def __str__(self):
        return '{0}'.format(self.title)

    class Meta:
        ordering = 'order', 'title',
        verbose_name = 'Статус'
        verbose_name_plural = 'Статусы'
        managed = True

######################################################################################################################

class Contract(models.Model):
    id = models.AutoField(
        primary_key=True,
    )
    title = models.CharField(
        verbose_name='Наименование',
        max_length=124,
        default='',
    )
    amount = models.CharField(
        verbose_name='Плановая сумма',
        max_length=124,
        default='',
    )
    comment = models.TextField(
        verbose_name='Комментарий',
        blank=True,
        default='',
    )
    closed = models.BooleanField(
        verbose_name='Закрыт',
        default=False,
    )
    create_date = models.DateTimeField(
        verbose_name='Дата создания контракта',
        auto_now_add=True,
        null=True,
    )

    def __str__(self):
        return '{0}'.format(self.title)

    class Meta:
        ordering = 'title',
        verbose_name = 'Контракт'
        verbose_name_plural = 'Контракты'
        managed = True

######################################################################################################################

class Stage(models.Model):
    id = models.AutoField(
        primary_key=True,
    )
    contract = models.ForeignKey(
        Contract,
        verbose_name='Контракт',
        null=True,
        related_name='ContractStatus',
        on_delete=models.SET_NULL,
    )
    status = models.ForeignKey(
        Status,
        verbose_name='Статус контракта',
        null=True,
        related_name='StatusContract',
        on_delete=models.SET_NULL,
    )
    comment_stage = models.TextField(
        verbose_name='Комментарий',
        default='',
    )
    create_date = models.DateTimeField(
        verbose_name='Дата создания стадии',
        auto_now_add=True,
        null=True,
    )

    def __str__(self):
        return '{0} - {1}'.format(self.contract, self.status)

    class Meta:
        ordering = 'status', 'contract',
        verbose_name = 'Стадия'
        verbose_name_plural = 'Стадии'
        managed = True

######################################################################################################################

class Attache(models.Model):
    id = models.AutoField(
        primary_key=True,
    )
    name = models.CharField(
        verbose_name='Название файла',
        max_length=124,
        default='',
    )
    file = models.FileField(
        verbose_name='Файл',
        upload_to='contract/attache/%Y/%m/%d',
        null=True,
        blank=True,
    )
    stage = models.ForeignKey(
        Stage,
        verbose_name='Стадия',
        null=True,
        related_name='AttacheStage',
        on_delete=models.SET_NULL,
    )
    create_date = models.DateTimeField(
        verbose_name='Дата создания файла',
        auto_now_add=True,
        null=True,
    )

    def get_human_size(self):
        value = self.file.size
        if value < 512000:
            value = value / 1024.0
            ext = 'Кб'
        elif value < 4194304000:
            value = value / 1048576.0
            ext = 'Мб'
        else:
            value = value / 1073741824.0
            ext = 'Гб'
        return '{0} {1}'.format(str(round(value, 2)), ext)

    def __str__(self):
        return '{0}'.format(self.name)

    class Meta:
        ordering = 'create_date',
        verbose_name = 'Вложение'
        verbose_name_plural = 'Вложения'
        managed = True

######################################################################################################################