import os
import argparse
import torch
import random
from torch import optim
from math import exp

from model import Copynet, Normal, MoG, Vamp, GP_Full
from training import train, evaluate
from utils import reader
from decode import BeamSearch

def parse_arguments():
    p = argparse.ArgumentParser(description='Hyperparams')
    # Training config
    p.add_argument('-t', '--task', type=str, default="train",
                   help="specify the task to do: (train)ing, (eval)uation, decode")
    p.add_argument('-epochs', type=int, default=200,
                   help='number of epochs for train')
    p.add_argument('-bs', '--batch_size', type=int, default=64,
                   help='number of epochs for train')
    p.add_argument('-optim', type=int, default=2,
                   help="optimization method. 0: SGD; 1: Adagrad, 2: Adam")
    p.add_argument('-lr', type=float, default=0.001,
                   help='initial learning rate')
    p.add_argument('-grad_clip', type=float, default=5.0,
                   help='in case of gradient explosion')
    p.add_argument('-val_step', type=int, default=1,
                   help='the number of epoches between two validations')
    p.add_argument('-print_step', type=int, default=100,
                   help='the number of batches between two tf.summary')
    p.add_argument('-kw', type=float, default=0.,
                   help="scaler for kl annealing function")
    p.add_argument('-x0', type=int, default=25000,
                   help="mean of kl annealing function")
    # Model config
    p.add_argument('-mt', '--model_type', type=str, default='gp_full',
                   help='type of model: copynet/normal/mog/vamp/gp_full')
    p.add_argument('-e', '--embed_size', type=int, default=300,
                   help='embedding size')
    p.add_argument('-hs', '--hidden_size', type=int, default=512,
                   help='output dimension of the recognition network')
    p.add_argument('-k', '--latent_size', type=int, default=256,
                   help="dimension of the latent variable z")
    p.add_argument('-wd', '--word_dropout', type=float, default=0,
                   help='decoder input sentence dropout rate')
    p.add_argument('--embed_dropout', type=float, default=0.5,
                   help='word embedding dropout rate')
    p.add_argument('-c', '--components', type=int, default=10,
                   help="number of components for mixture gaussian")
    p.add_argument('-kld_sampled', type=int, default=0,
                   help="how to compute kld. 1: sampled kld; 0: analytical kld")
    p.add_argument('-v', '--kernel_v', type=float, default=65.0,
                   help="Hyper-parameter for prior kernel,  control the signal variance")
    p.add_argument('-r', '--kernel_r', type=float, default=0.0001,
                   help="Hyper-parameter for prior kernel.")
    p.add_argument('-vocab', '--vocab_size', type=int, default=20000,
                   help="vocab size")
    p.add_argument('-max_len', type=int, default=30,
                   help="max sequence length")
    p.add_argument('-using_cuda', type=bool, default=True,
                   help="whether to use cuda")
    # Generation - decoding strategy
    p.add_argument('-samn', '--sample_num', type=int, default=10,
                   help="Number of samples from latent codes.")
    p.add_argument('-bms', '--beam_size', type=int, default=10,
                   help="beam size, if 0, use greedy search; otherwise, using beam search.")
    p.add_argument('-df', '--data_file', type=str, default='../../data/GYAFC/em',
                   help='path to dataset')
    p.add_argument('-mf', '--model_file', type=str, default=None,
                   help='path to a pretrained model')
    # Generation - sampling latent codes
    p.add_argument('-topn', type=int, default=10,
                   help="number of generated samples for the decoder")
    p.add_argument('-std', type=float, default=1.,
                   help="Standard deviation for sampling z.")
    p.add_argument('-decf', '--decode_from', type=str, default='sample',
                   help="sample: decode from sampled latent codes; mean: decode from mean")
    p.add_argument('--seed', type=int, default=123)
    return p.parse_args()

def main():
    # --------------------------------------------
    # print("*** Parsing arguments ***")
    args = parse_arguments()
    args.using_cuda = args.using_cuda and torch.cuda.is_available()

    # --------------------------------------------
    print("*** Preparing dataset ... ***")
    train_iter, val_iter, test_iter, txtfield = reader(suffix=".tsv",
                                                       rpath=args.data_file,
                                                       batch_size=args.batch_size)
    print("[TRAIN]:%d (dataset:%d)\t[VAL]:%d (dataset:%d)\t[TEST]:%d (dataset:%d)"
          % (len(train_iter), len(train_iter.dataset),
             len(val_iter), len(val_iter.dataset),
             len(test_iter), len(test_iter.dataset)))
    print("[vocab]:%d" % (len(txtfield.vocab)))

    # --------------------------------------------
    print("*** Initializing the model ... ***")
    if args.model_type == "copynet":
        model = Copynet(args, txtfield)
    elif args.model_type == "normal":
        model = Normal(args, txtfield)
    elif args.model_type == "mog":
        model = MoG(args, txtfield)
    elif args.model_type == "vamp":
        model = Vamp(args, txtfield)
    elif args.model_type == "gp_full":
        model = GP_Full(args, txtfield)
    else:
        raise ValueError("Unrecognized model type: %s" % args.model_type)

    # --------------------------------------------
    if (args.model_file is not None) and os.path.isfile(args.model_file):
        print("Loading the pretrained model from: {}".format(args.model_file))
        if args.using_cuda:
            model.load_state_dict(torch.load(args.model_file))
        else:
            model.load_state_dict(torch.load(args.model_file, map_location='cpu'))
    else:
        print("Couldn't find the pretrained model from: {}".format(args.model_file))
    if args.using_cuda: model.cuda()

    # --------------------------------------------
    optimizer = None
    params = model.parameters()
    if args.optim == 0:
        print("*** Choosing SGD as the optimization method ... ***")
        optimizer = optim.SGD(params, lr=args.lr)
    elif args.optim == 1:
        print("*** Choosing Adagrad as the optimization method ... ***")
        optimizer = optim.Adagrad(params, lr=args.lr)
    elif args.optim == 2:
        print("*** Choosing Adam as the optimization method ... ***")
        optimizer = optim.Adam(params, lr=args.lr)
    else:
        raise ValueError("Unspecified optimization method")

    # --------------------------------------------
    if args.task == "train":
        print("*** Training with validation ... ***")
        train(model, optimizer, train_iter, val_iter, txtfield, args)
    elif args.task == "eval":
        print("*** Testing ... ***")
        test_nll, test_kld = evaluate(model, test_iter, txtfield, args)
        print("[TEST] NLL: %5.4f KLD: %5.4f PPLx: %5.4f" % (test_nll, test_kld, exp(test_nll+test_kld)))
    elif args.task == "decode":
        print("*** Decode ... ***")
        # set random seed
        random.seed(args.seed)
        torch.manual_seed(args.seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(args.seed)
        beam_searcher = BeamSearch(model, txtfield, args)
        if args.decode_from == 'sample':
            beam_searcher.decode_sample(test_iter, txtfield, args, num_samples=30000)
        else:
            beam_searcher.decode_beam(test_iter, txtfield, args, num_samples=30000)
    else:
        raise ValueError("Unrecognized task: {}".format(args.task))

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt as e:
        print("[STOP]", e)


# Ambient から定期的にデータを取得する

from AmbientDataHandler import AmbientDataHandler
import datetime
from get_filepath import get_filepath
import pandas as pd
import ambient
import os
from recieve_username_auth import recieve_username_auth

def main():
    adh = AmbientDataHandler()
    read_config_df = adh.fetch_ambient_read_config()
    chID_readKey_df = read_config_df[["chID", "readKey"]]
    # chID_set =  set(chID_readKey_df["chID"])

    all_read_data_df = pd.DataFrame(columns=["created","chID"] + list(set(read_config_df["dataID"])))
    datetime_now = datetime.datetime.now()

    # 最後に読み込みを行った日時を読み込む。
    # そこを起点として新しいデータを読み込む。
    file = "../resource/latest_read_datetime.txt"
    filepath = get_filepath(file)
    try:
        with open(filepath, encoding="shift-jis", mode = "r") as f:
            latest_read_datetime = f.read()
    except FileNotFoundError:
        # 最終読み込み時刻を保存したファイルが無ければ、読み込み間隔を取得し、その分を引いた時刻を起点とする
        # それが見つからなければエラーを返す
        file = "../resource/cron_setting.csv"
        filepath = get_filepath(file)
        try:
            cron_setting_df = pd.read_csv(filepath, encoding="shift-jis", index_col=0)
        except FileNotFoundError:
            message = "最後に読み込みを行った時刻を保存したファイルが存在しません"
            file = "../stderr.log"
            filepath = get_filepath(file)
            with open(filepath, mode = "a") as f:
                f.write(message + "\n")
            return
        data_getting_run_span = int(cron_setting_df.loc["data_getting", "run_span"])
        latest_read_datetime = (datetime_now - datetime.timedelta(minutes=data_getting_run_span)).strftime("%Y-%m-%d %H:%M:%S")

    start_date_time = latest_read_datetime
    end_date_time = datetime_now.strftime("%Y-%m-%d %H:%M:%S")

    for chID in set(chID_readKey_df["chID"]):
        # chIDが同じなら一意にreadKeyが決まるはず。
        readKey = set(chID_readKey_df.loc[chID_readKey_df["chID"] == chID, "readKey"]).pop()

        # WriteKeyは使わないので ""でよい。
        am = ambient.Ambient(chID,"", readKey)
        read_data = am.read(start = start_date_time, end = end_date_time)
        read_data_df = pd.DataFrame(read_data)
        read_data_df["chID"] = chID
        all_read_data_df = pd.concat([all_read_data_df, read_data_df], axis = 0)

    # all_read_data_dfが空 = 更新するデータが無い.
    # よって、以後の処理を行わない。
    if all_read_data_df.empty:
        print("No data to read")
        return

    all_read_data_df = all_read_data_df.rename(columns={"created": "datetime"})
    # Ambient から帰ってくる日付はUTC
    all_read_data_df["datetime"] = pd.to_datetime(all_read_data_df["datetime"]).dt.tz_convert('Asia/Tokyo')
    # ここから後でtimezoneの情報はいらないと思うので消す。
    all_read_data_df["datetime"] = list(map(lambda x: x.replace(tzinfo = None).strftime("%Y-%m-%d %H:%M:%S"), all_read_data_df["datetime"]))

    melted_df = all_read_data_df.melt(id_vars=["datetime", "chID"], var_name="dataID")
    # chIDとdataIDをキーとしてdatatyoeとhouse_nameの情報を列に追加する。
    update_dict = {
        (row["chID"], row["dataID"]): (row["datatype"], row["house_name"])
        for _, row in read_config_df.iterrows()
    }
    melted_df[["datatype", "house_name"]] = melted_df.apply(lambda x: update_dict.get((x["chID"], x["dataID"]), (None, None)), axis=1, result_type="expand")
    melted_df = melted_df.drop(columns=["chID", "dataID"])
    melted_df = melted_df.reindex(columns=["datetime", "house_name", "datatype", "value"])

    file = "../resource/I2C_data_period.csv"
    filepath = get_filepath(file)
    melted_df.to_csv(filepath, mode="a", encoding="shift-jis", header = not os.path.exists(filepath), index=False)

    # 読み込んだデータの中で最新のものを抜きだし、定期通知用に記録する。
    latest_df = pd.DataFrame()
    house_datatype_tuple_list = melted_df[["house_name", "datatype"]].apply(tuple, axis=1).drop_duplicates().tolist()
    for house_datatype_tuple in house_datatype_tuple_list:
        house_datatype_df = melted_df.loc[(melted_df["house_name"] == house_datatype_tuple[0]) & (melted_df["datatype"] == house_datatype_tuple[1]), :]
        latest_datetime = max(house_datatype_df["datetime"])
        concat_df = house_datatype_df.loc[house_datatype_df["datetime"] == latest_datetime, :]
        latest_df = pd.concat([latest_df, concat_df], axis = 0)

    file = "../resource/I2C_data_latest_period.csv"
    filepath = get_filepath(file)
    latest_df.to_csv(filepath, mode="w", encoding="shift-jis", header = True, index=False)

    # データの保存が無事終了してから、最終読み込み日時を書き込む
    # 読み込み日時が更新されない＝どこかでエラー発生
    # 次に読み込む時に、保存できなかった分のデータも一緒に読み込む
    file = "../resource/latest_read_datetime.txt"
    filepath = get_filepath(file)
    with open(filepath, encoding="shift-jis", mode = "w") as f:
        f.write(end_date_time)

if recieve_username_auth():
    main()


#User function Template for python3

from typing import List
from queue import Queue
class Solution:
    #Function to return Breadth First Traversal of given graph.
    def bfsOfGraph(self, V: int, adj: List[List[int]]) -> List[int]:
        # code here
        bfsTraversal = []
        visited = [False] * V  # To keep track of visited nodes
        queue = Queue()  # Queue to store nodes to be visited

        # Start BFS from node 0
        queue.put(0)
        visited[0] = True

        while not queue.empty():
            currentNode = queue.get()  # Get the front node from the queue
            bfsTraversal.append(currentNode)  # Add the current node to the result

            # Visit all adjacent nodes of the current node
            for neighbor in adj[currentNode]:
                if not visited[neighbor]:
                    queue.put(neighbor)
                    visited[neighbor] = True

        return bfsTraversal

#{
 # Driver Code Starts

if __name__ == '__main__':
	T=int(input())
	for i in range(T):
		V, E = map(int, input().split())
		adj = [[] for i in range(V)]
		for _ in range(E):
			u, v = map(int, input().split())
			adj[u].append(v)
		ob = Solution()
		ans = ob.bfsOfGraph(V, adj)
		for i in range(len(ans)):
		    print(ans[i], end = " ")
		print()

# } Driver Code Ends


# -*- coding: utf-8 -*-
"""
the learning module

Description
-----------

The learning module set of functions provide a framework to optimise and classify
EO data for both per pixel or object properties

"""

# This must go first or it causes the error:
#ImportError: dlopen: cannot load any more object with static TLS

from xgboost.sklearn import XGBClassifier, XGBRegressor
import xgboost as xgb

from tqdm import tqdm

import matplotlib.pyplot as plt
from collections import OrderedDict
import glob
from sklearn import svm
from osgeo import gdal, ogr#,osr
import numpy as np
from sklearn.pipeline import Pipeline
from sklearn.model_selection import (StratifiedKFold, GroupKFold, KFold,
                                     train_test_split,GroupShuffleSplit,
                                     StratifiedGroupKFold,
                                     PredefinedSplit)
from sklearn.ensemble import (RandomForestClassifier, ExtraTreesClassifier,
                              GradientBoostingClassifier,RandomForestRegressor,
                              GradientBoostingRegressor, ExtraTreesRegressor,
                              VotingRegressor, VotingClassifier, StackingRegressor,
                              StackingClassifier,
                              HistGradientBoostingRegressor,
                              HistGradientBoostingClassifier)
from sklearn.linear_model import LogisticRegression, SGDClassifier
from sklearn.preprocessing import (LabelEncoder, MaxAbsScaler, MinMaxScaler,
                                   Normalizer, PowerTransformer,StandardScaler,
                                   QuantileTransformer)
from sklearn.svm import SVC, SVR, NuSVC, NuSVR, LinearSVC, LinearSVR
from sklearn.feature_selection import VarianceThreshold, RFECV
from sklearn.inspection import permutation_importance
from sklearn import metrics
from sklearn.model_selection import GridSearchCV, RandomizedSearchCV
#from catboost import (CatBoostClassifier, CatBoostRegressor,Pool )
import lightgbm as lgb
# from autosklearn.classification import AutoSklearnClassifier
# from autosklearn.regression import AutoSklearnRegressor
from skorch import NeuralNetClassifier
from optuna.integration import LightGBMPruningCallback
from optuna import create_study
import torch.nn as nn
import joblib
import joblib as jb
from geospatial_learn.raster import array2raster
from geospatial_learn.shape import _bbox_to_pixel_offsets#, zonal_stats
#from tpot import TPOTClassifier, TPOTRegressor
import warnings
from geospatial_learn.raster import _copy_dataset_config
import geospatial_learn.handyplots as hp
import geopandas as gpd
import pandas as pd
from psutil import virtual_memory
import os
gdal.UseExceptions()
ogr.UseExceptions()

# for current issue with numpy
np.warnings = warnings

# TODO -NOT WORKING of course stupid object structure making feeding params
# impossible
def create_model_optuna(X_train, outModel, clf='erf', group=None, random=False,
                 cv=5, params=None, pipe=None, cores=-1, strat=True,
                 test_size=0.3, regress=False, return_test=True,
                 scoring=None, class_names=None, save=True):
    """
    Train a model using the Optuna framework

    Parameters
    ---------------

    X_train: np array
              numpy array of training data where the 1st column is labels.
              If the groupkfold is used, the last column will be the group labels

    outModel: string
               the output model path which is a gz file, if using keras it is
               h5

    params: dict
            a dict of model params (see scikit learn). If using a pipe(line)
            remember to prefix each param as follows with parent object
            and two underscores.
            param_grid ={"selector__threshold": [0, 0.001, 0.01],
             "classifier__n_estimators": [1075]}

    clf: string
          an sklearn or xgb classifier/regressor
          logit, sgd, linsvc, svc, svm, nusvm, erf, rf, gb, xgb,

    group: np.array
            array of group labels for train/test split and grid search
            useful to avoid autocorrelation

    random: bool
             if True, a random param search

    cv: int
         no of folds

    cores: int or -1 (default)
            the no of parallel jobs

    strat: bool
            a stratified grid search

    test_size: float
            percentage to hold out to test

    regress: bool
              a regression model if True, a classifier if False

    return_test: bool
              return X_test and y_test along with results (last two entries
              in list)

    scoring: string
              a suitable sklearn scoring type (see notes)

    class_names: list of strings
                class names in order of their numercial equivalents

    Returns
    -------
    A list of:

    [grid.best_estimator_, grid.cv_results_, grid.best_score_,
            grid.best_params_, classification_report, X_test, y_test]

    """

    bands = X_train.shape[1]-1

    #X_train = X_train.transpose()
    X_train = X_train[X_train[:,0] != 0]

    # Remove non-finite values
    if group is not None: #replace this later not good
        inds = np.where(np.isfinite(X_train).all(axis=1))
        group = group[inds]

    X_train = X_train[np.isfinite(X_train).all(axis=1)]
    # y labels
    y_train = X_train[:,0]

    # remove labels from X_train
    X_train = X_train[:,1:bands+1]

    #no_classes = len(np.unique(y_train))

    # this is not a good way to do this
    if group is not None:

        # maintain group based splitting from initial train/test split
        # to main train set
        # TODO - sep func?
        splitter = GroupShuffleSplit(test_size=test_size, n_splits=1, random_state=0)
        split = splitter.split(X_train, y_train, group)
        train_inds, test_inds = next(split)

        X_test = X_train[test_inds]
        y_test = y_train[test_inds]
        X_train = X_train[train_inds]
        y_train = y_train[train_inds]
        group_trn = group[train_inds]

        group_kfold = GroupKFold(n_splits=cv)
        # Create a nested list of train and test indices for each fold
        k_kfold = group_kfold.split(X_train, y_train, group_trn)

        train_ind2, test_ind2 = [list(traintest) for traintest in zip(*k_kfold)]

        cv = [*zip(train_ind2, test_ind2)]

    else:
        X_train, X_test, y_train, y_test = train_test_split(
            X_train, y_train, test_size=test_size, random_state=0)

        cv = StratifiedKFold(n_splits=cv, shuffle=True, random_state=0)

    if regress is True:
        scr = metrics.mean_squared_error
    else:
        scr = metrics.log_loss

    # issue is how to pass the dict as an arg in this parent function
    def objective(trial, X, y, cv, group, score=scr):

        # how can one pass this as an arg
        param_grid = {
            #"device_type": trial.suggest_categorical("device_type", ['gpu']),
           #'metric': 'rmse',
            'random_state': 48,
            'n_estimators': 20000,
            'reg_alpha': trial.suggest_loguniform('reg_alpha', 1e-3, 10.0),
            'reg_lambda': trial.suggest_loguniform('reg_lambda', 1e-3, 10.0),
            'colsample_bytree': trial.suggest_categorical('colsample_bytree', [0.3,0.4,0.5,0.6,0.7,0.8,0.9, 1.0]),
            'subsample': trial.suggest_categorical('subsample', [0.4,0.5,0.6,0.7,0.8,1.0]),
            'learning_rate': trial.suggest_categorical('learning_rate', [0.006,0.008,0.01,0.014,0.017,0.02]),
            'max_depth': trial.suggest_categorical('max_depth', [10,20,100]),
            'num_leaves' : trial.suggest_int('num_leaves', 1, 1000),
            'min_child_samples': trial.suggest_int('min_child_samples', 1, 300),
            'cat_smooth' : trial.suggest_int('min_data_per_groups', 1, 100)
            }

        cv_scores = np.empty(10)

        for idx, (train_idx, test_idx) in enumerate(cv.split(X, y, group)):

            # these next two need altered to fit my data structures above
            X_train, X_test = X[train_idx], X[test_idx]
            y_train, y_test = y[train_idx], y[test_idx]
            # had objective=binary but that is for class
            model = lgb.LGBMRegressor( **param_grid)
            model.fit(
                X_train,
                y_train,
                eval_set=[(X_test, y_test)],
                #eval_metric="l2",
                early_stopping_rounds=100,
                callbacks=[
                    LightGBMPruningCallback(trial, "rmse")
                ],  # Add a pruning callback - suspect this may be a source of problems
            )
            preds = model.predict(X_test)
            cv_scores[idx] = score(y_test, preds)

        return np.mean(cv_scores)

    study = create_study(direction="minimize", study_name="Model training")
    func = lambda trial: objective(trial, X_train, y_train, group_kfold, group_trn,
                                   score=metrics.mean_squared_error)
    study.optimize(func, n_trials=50)

    print(f"\tBest value (rmse or r2): {study.best_value:.5f}")
    print(f"\tBest params:")

def _group_cv(X_train, y_train, group, test_size=0.2, cv=10, strat=False):

    """
    Return the splits and and vars for a group grid search
    """

    # maintain group based splitting from initial train/test split
    # to main train set
    # TODO - sep func?
    splitter = GroupShuffleSplit(test_size=test_size, n_splits=1, random_state=0)
    split = splitter.split(X_train, y_train, group)
    train_inds, test_inds = next(split)

    X_test = X_train[test_inds]
    y_test = y_train[test_inds]
    X_train = X_train[train_inds]
    y_train = y_train[train_inds]
    group_trn = group[train_inds]

    if strat == True:
        group_kfold = StratifiedGroupKFold(n_splits=cv).split(X_train,
                                                              y_train,
                                                              group_trn)
    else:
        group_kfold = GroupKFold(n_splits=cv).split(X_train,
                                                    y_train,
                                                    group_trn)

    # all this not required produces same as above - keep for ref though
    # # Create a nested list of train and test indices for each fold
    # k_kfold = group_kfold.split(X_train, y_train, groups=group_trn)

    # train_ind2, test_ind2 = [list(traintest) for traintest in zip(*k_kfold)]

    # cv = [*zip(train_ind2, test_ind2)]

    return X_train, y_train, X_test, y_test, group_kfold

def rec_feat_sel(X_train, featnames, preproc=('scaler', None),  clf='erf',
                 group=None,
                 cv=5, params=None, cores=-1, strat=True,
                 test_size=0.3, regress=False, return_test=True,
                 scoring=None, class_names=None, save=True, cat_feat=None):

    """
    Recursive feature selection

    Parameters
    ---------------

    X_train: np array
              numpy array of training data where the 1st column is labels.
              If the groupkfold is used, the last column will be the group labels

    params: dict
            a dict of model params (see scikit learn). If using a pipe(line)
            remember to prefix each param as follows with parent object
            and two underscores.
            param_grid ={"selector__threshold": [0, 0.001, 0.01],
             "classifier__n_estimators": [1075]}

    clf: string
          an sklearn or xgb classifier/regressor
          logit, sgd, linsvc, svc, svm, nusvm, erf, rf, gb, xgb,

    group: np.array
            array of group labels for train/test split and grid search
            useful to avoid autocorrelation

    cv: int
         no of folds

    cores: int or -1 (default)
            the no of parallel jobs

    strat: bool
            a stratified grid search

    test_size: float
            percentage to hold out to test

    regress: bool
              a regression model if True, a classifier if False

    scoring: string
              a suitable sklearn scoring type (see notes)

    class_names: list of strings
                class names in order of their numercial equivalents

    Returns
    -------

    bool index of features, list of chosen feature names
    """
    #TODO need to make all this a func
    # Not woring with hgb....
    clfdict = {'rf': RandomForestClassifier(random_state=0),
               'erf': ExtraTreesClassifier(random_state=0),
               'gb': GradientBoostingClassifier(random_state=0),
               'xgb': XGBClassifier(random_state=0),
               'logit': LogisticRegression(),
               # 'catb': CatBoostClassifier(logging_level='Silent', # supress trees in terminal
               #                            random_seed=42),
               # 'catbgpu': CatBoostClassifier(logging_level='Silent', # supress trees in terminal
               #                            random_seed=42,
                                          # task_type="GPU",
                                          # devices='0:1'),
               'lgbm': lgb.LGBMClassifier(random_state=0),
 #use_best_model=True-needs non empty eval set
                'hgb': HistGradientBoostingClassifier(random_state=0)}

    regdict = {'rf': RandomForestRegressor(random_state=0),
               'erf': ExtraTreesRegressor(random_state=0),
               'gb': GradientBoostingRegressor(random_state=0),
               'xgb': XGBRegressor(random_state=0),
               # 'catb': CatBoostRegressor(logging_level='Silent',
               #                           random_seed=42),
               # 'catbgpu': CatBoostClassifier(logging_level='Silent', # supress trees in terminal
               #                            random_seed=42,
               #                            task_type="GPU",
                                          # devices='0:1'),
               'lgbm': lgb.LGBMRegressor(random_state=0),

               'hgb': HistGradientBoostingRegressor(random_state=0)}

    if regress is True:
        model = regdict[clf]
        if scoring is None:
            scoring = 'r2'
    else:
        model = clfdict[clf]
        cv = StratifiedKFold(cv)
        if scoring is None:
            scoring = 'accuracy'
    bands = X_train.shape[1]-1

    #X_train = X_train.transpose()
    X_train = X_train[X_train[:,0] != 0]

    # Remove non-finite values
    if group is not None: #replace this later not good
        inds = np.where(np.isfinite(X_train).all(axis=1))
        group = group[inds]

    X_train = X_train[np.isfinite(X_train).all(axis=1)]
    # y labels
    y_train = X_train[:,0]

    # remove labels from X_train
    X_train = X_train[:,1:bands+1]

    #no_classes = len(np.unique(y_train))

    # this is not a good way to do this
    # Does this matter for feature selection??
    if group is not None:

        X_train, y_train, X_test, y_test, cv = _group_cv(X_train, y_train,
                                                         group, test_size,
                                                         cv)

    else:
        X_train, X_test, y_train, y_test = train_test_split(
            X_train, y_train, test_size=test_size, random_state=0)

    rfecv = RFECV(estimator=model,
                  step=1,
                  cv=cv,
                  scoring=scoring,
                  n_jobs=cores) # suspect this is no of folds

    pipeline  = Pipeline([preproc,
                          ('selector', rfecv)])

    # VERY slow - but I guess it is grid searching feats first
    #rfecv
    pipeline.fit(X_train, y_train)

    # First experiment is to add this as a fixed part of the process at the start
    # as it will slow it down otherwise

    # featind = pipeline[1].support_ # gains the feat indices(bool)
    # featnmarr = np.array(featnames)
    # featnames_sel = featnmarr[featind==True].tolist()

    # as X_train has changed we cant select from it within here

    return pipeline

def create_model(X_train, outModel, clf='erf', group=None, random=False,
                 cv=5, params=None, pipe='default', cores=-1, strat=True,
                 test_size=0.3, regress=False, return_test=True,
                 scoring=None, class_names=None, save=True, cat_feat=None):

    """
    Brute force or random model creating using scikit learn.

    Parameters
    ---------------

    X_train: np array
              numpy array of training data where the 1st column is labels.
              If the groupkfold is used, the last column will be the group labels

    outModel: string
               the output model path which is a gz file, if using keras it is
               h5

    params: dict
            a dict of model params (see scikit learn). If using a pipe(line)
            remember to prefix each param as follows with parent object
            and two underscores.
            param_grid ={"selector__threshold": [0, 0.001, 0.01],
             "classifier__n_estimators": [1075]}

    pipe: str,dict,None
            if 'default' will include a preprocessing pipeline consisting of
            StandardScaler(), MinMaxScaler(), Normalizer(), MaxAbsScaler(),
            otherwise specify in this form
            pipe = {'scaler': [StandardScaler(), MinMaxScaler(),
                  Normalizer()]}
            or None will not preprocess the data

    clf: string
          an sklearn or xgb classifier/regressor
          logit, sgd, linsvc, svc, svm, nusvm, erf, rf, gb, xgb,

    group: np.array
            array of group labels for train/test split and grid search
            useful to avoid autocorrelation

    random: bool
             if True, a random param search

    cv: int
         no of folds

    cores: int or -1 (default)
            the no of parallel jobs

    strat: bool
            a stratified grid search

    test_size: float
            percentage to hold out to test

    regress: bool
              a regression model if True, a classifier if False

    return_test: bool
              return X_test and y_test along with results (last two entries
              in list)

    scoring: string
              a suitable sklearn scoring type (see notes)

    class_names: list of strings
                class names in order of their numercial equivalents

    Returns
    -------
    A list of:

    [grid.best_estimator_, grid.cv_results_, grid.best_score_,
            grid.best_params_, classification_report, X_test, y_test]

    Notes:
    ------
        Scoring types - there are a lot - some of which won't work for
        multi-class, regression etc - see the sklearn docs!

        'accuracy', 'adjusted_rand_score', 'average_precision', 'f1',
        'f1_macro', 'f1_micro', 'f1_samples', 'f1_weighted', 'neg_log_loss',
        'neg_mean_absolute_error', 'neg_mean_squared_error',
        'neg_median_absolute_error', 'precision', 'precision_macro',
        'precision_micro', 'precision_samples', 'precision_weighted',
        'r2', 'recall', 'recall_macro', 'recall_micro', 'recall_samples',
        'recall_weighted', 'roc_auc'

    """

    clfdict = {'rf': RandomForestClassifier(random_state=0),
               'erf': ExtraTreesClassifier(random_state=0),
               'gb': GradientBoostingClassifier(random_state=0),
               'xgb': XGBClassifier(random_state=0),
               'logit': LogisticRegression(),
               # 'catb': CatBoostClassifier(logging_level='Silent', # supress trees in terminal
               #                            random_seed=42),
               # 'catbgpu': CatBoostClassifier(logging_level='Silent', # supress trees in terminal
               #                            random_seed=42,
               #                            task_type="GPU",
               #                            devices='0:1'),
               'lgbm': lgb.LGBMClassifier(random_state=0),

                'hgb': HistGradientBoostingClassifier(random_state=0),
                'svm': SVC(),
                'nusvc': NuSVC(),
                'linsvc': LinearSVC()}

    regdict = {'rf': RandomForestRegressor(random_state=0),
               'erf': ExtraTreesRegressor(random_state=0),
               'gb': GradientBoostingRegressor(random_state=0),
               'xgb': XGBRegressor(random_state=0),
               # 'catb': CatBoostRegressor(logging_level='Silent',
               #                           random_seed=42),
               # 'catbgpu': CatBoostRegressor(logging_level='Silent', # supress trees in terminal
               #                            random_seed=42,
               #                            task_type="GPU",
               #                            devices='0:1'),
               'lgbm': lgb.LGBMRegressor(random_state=0),
               'hgb': HistGradientBoostingRegressor(random_state=0),
                'svm': SVR(),
                'nusvc': NuSVR(),
                'linsvc': LinearSVR()}

    if regress is True:
        model = regdict[clf]
        if scoring is None:
            scoring = 'r2'
    else:
        model = clfdict[clf]
        if group is None:
            cv = StratifiedKFold(cv)
        if scoring is None:
            scoring = 'accuracy'

    bands = X_train.shape[1]-1

    #X_train = X_train.transpose()
    X_train = X_train[X_train[:,0] != 0]

    # Remove non-finite values
    if group is not None: #replace this later not good
        inds = np.where(np.isfinite(X_train).all(axis=1))
        group = group[inds]

    X_train = X_train[np.isfinite(X_train).all(axis=1)]
    # y labels
    y_train = X_train[:,0]

    # remove labels from X_train
    X_train = X_train[:,1:bands+1]

    #no_classes = len(np.unique(y_train))

    # this is not a good way to do this
    if regress == True:
        strat = False # failsafe

    if group is not None: # becoming a mess

        X_train, y_train, X_test, y_test, cv = _group_cv(X_train, y_train,
                                                             group, test_size,
                                                             cv, strat=strat)
    else:
        X_train, X_test, y_train, y_test = train_test_split(
            X_train, y_train, test_size=test_size, random_state=0)
        #cv = StratifiedKFold(cv)

    if pipe == 'default':
        # the dict must be in order of proc to work hence this
        # none is included to ensure
        # non prep data is considered
        #  should add selector?? seems to produce a load of nan errors (or warnings)
        sclr = {'scaler': [StandardScaler(), MinMaxScaler(),
                           Normalizer(), MaxAbsScaler(),
                           QuantileTransformer(output_distribution='uniform'),
                           #PowerTransformer(),
                           None]}
        sclr.update(params)

    else:
        sclr = pipe
        sclr.update(params)

    sk_pipe = Pipeline([("scaler", StandardScaler()),
                        #("selector", None), lord knows why this fails on var thresh
                        ("classifier", model)])

    if random is True:
        # recall the model is in the pipeline
        grid = RandomizedSearchCV(sk_pipe, param_distributions=sclr,
                                  n_jobs=cores, n_iter=20,  verbose=2)
    else:
        grid = GridSearchCV(sk_pipe,  param_grid=sclr,
                                    cv=cv, n_jobs=cores,
                                    scoring=scoring, verbose=1)

    grid.fit(X_train, y_train)

    joblib.dump(grid.best_estimator_, outModel)

    testresult = grid.best_estimator_.predict(X_test)

    if regress == True:
        regrslt = regression_results(y_test, testresult)
        results = [grid]

    else:
        crDf = hp.plot_classif_report(y_test, testresult, target_names=class_names,
                                      save=outModel[:-3]+'._classif_report.png')

        confmat = metrics.confusion_matrix(testresult, y_test, labels=class_names)
        disp = metrics.ConfusionMatrixDisplay(confusion_matrix=confmat,
                                      display_labels=class_names)
        disp.plot()

        # confmat = hp.plt_confmat(X_test, y_test, grid.best_estimator_,
        #                          class_names=class_names,
        #                          cmap=plt.cm.Blues,
        #                          fmt="%d",
        #                          save=outModel[:-3]+'_confmat.png')

        results = [grid, crDf, confmat]

    if return_test == True:
        results.extend([X_test, y_test])
    return results

def combine_models(X_train, modelist, mtype='regress', method='voting', group=None,
                   test_size=0.3, outmodel=None, class_names=None, params=None,
                   final_est='xgb', cv=5):#, cores=1):

    """
    Combine models using either the voting or stacking methods in scikit-learn

    Parameters
    ----------

    X_train: np array
              numpy array of training data where the 1st column is labels.

    outmodel: string
               the output model path which is a gz file, if using keras it is
               h5

    modelist: dict
            a list of tuples of model type (str) and model (class) e.g.
            [('gb', gmdl), ('rf', rfmdl), ('lr', reg3)]

    mtype: string
            either clf or regress

    method: string
            either voting or k = clusterer.labels_stacking

    test_size: float
            percentage to hold out to test

    group: np.array
            array of group labels for train/test split and grid search
            useful to avoid autocorrelation

    class_names: list of strings
                class names in order of their numercial equivalents

    params: dict
            a dict of model params (If using stacking method for final estimator)

    final_est: string
             the final estimator one of (rf, gb, erf, xgb, logit)

    """
    bands = X_train.shape[1]-1

    X_train = X_train[X_train[:,0] != 0]

    if group is not None: #replace this later not good
        inds = np.where(np.isfinite(X_train).all(axis=1))
        group = group[inds]

    X_train = X_train[np.isfinite(X_train).all(axis=1)]
    # y labels
    y_train = X_train[:,0]

    # remove labels from X_train
    X_train = X_train[:,1:bands+1]

        # Remove non-finite values

    # this is not a good way to do this
    if group is not None:
        # this in theory should not make any difference at this stage...
        # ...included for now
        # maintain group based splitting from initial train/test split
        # to main train set
        # TODO - sep func?
        splitter = GroupShuffleSplit(test_size=test_size, n_splits=1, random_state=0)
        split = splitter.split(X_train, y_train, group)
        train_inds, test_inds = next(split)

        X_test = X_train[test_inds]
        y_test = y_train[test_inds]
        X_train = X_train[train_inds]
        y_train = y_train[train_inds]
        group_trn = group[train_inds]

        group_kfold = GroupKFold(n_splits=cv)
        # Create a nested list of train and test indices for each fold
        k_kfold = group_kfold.split(X_train, y_train, group_trn)

        train_ind2, test_ind2 = [list(traintest) for traintest in zip(*k_kfold)]

        cv = [*zip(train_ind2, test_ind2)]

    else:
        X_train, X_test, y_train, y_test = train_test_split(
            X_train, y_train, test_size=test_size, random_state=0)

    if method == 'voting':
        comb = VotingRegressor(estimators=modelist)#, n_jobs=cores)
        # we only wish to predict really - but  necessary
        # for sklearn model construct
    else:
        clfdict = {'rf': RandomForestClassifier(random_state=0),
                   'erf': ExtraTreesClassifier(random_state=0),
                   'gb': GradientBoostingClassifier(random_state=0),
                   'xgb': XGBClassifier(random_state=0),
                   'logit': LogisticRegression(),
                   'lgbm': lgb.LGBMClassifier(random_state=0),
                    'hgb': HistGradientBoostingClassifier(random_state=0),
                    'svm': SVC(),
                    'nusvc': NuSVC(),
                    'linsvc': LinearSVC()}

        regdict = {'rf': RandomForestRegressor(random_state=0),
                   'erf': ExtraTreesRegressor(random_state=0),
                   'gb': GradientBoostingRegressor(random_state=0),
                   'xgb': XGBRegressor(random_state=0),
                   'lgbm': lgb.LGBMRegressor(random_state=0),
                   'hgb': HistGradientBoostingRegressor(random_state=0),
                    'svm': SVR(),
                    'nusvc': NuSVR(),
                    'linsvc': LinearSVR()}

        if mtype == 'regress':
            # won't accept the dict even with the ** to unpack it
            fe = regdict[final_est]()
            fe.set_params(**params)

            comb = StackingRegressor(estimators=modelist,
                                final_estimator=fe)
        else:
            fe = clfdict[final_est]()
            fe.set_params(**params)
            cv = StratifiedKFold(cv)
            comb = StackingClassifier(estimators=modelist,
                                final_estimator=fe)

    comb.fit(X_train, y_train)

    # Since there is no train/test this is misleading...
    train_pred = comb.predict(X_train)
    test_pred = comb.predict(X_test)
    if mtype == 'regress':
        print('On the train split (not actually trained on this)')
        regression_results(y_train, train_pred)

        print('On the test split')
        regression_results(y_test, test_pred)

    else:
        crDf = hp.plot_classif_report(y_test, test_pred, target_names=class_names,
                              save=outmodel[:-3]+'._classif_report.png')

        confmat = hp.plt_confmat(X_test, y_test, comb,
                                 class_names=class_names,
                                 cmap=plt.cm.Blues,
                                 fmt="%d",
                                 save=outmodel[:-3]+'_confmat.png')

    if outmodel is not None:
        joblib.dump(comb, outmodel)
    return comb, X_test, y_test

def regression_results(y_true, y_pred):

    # Regression metrics
    explained_variance = metrics.explained_variance_score(y_true, y_pred)
    mean_absolute_error = metrics.mean_absolute_error(y_true, y_pred)
    mse=metrics.mean_squared_error(y_true, y_pred)
    #mean_squared_log_error = metrics.mean_squared_log_error(y_true, y_pred)
    median_absolute_error = metrics.median_absolute_error(y_true, y_pred)
    r2=metrics.r2_score(y_true, y_pred)

    print('explained_variance: ', round(explained_variance,4))
    #print('mean_squared_log_error: ', round(mean_squared_log_error,4))
    print('r2: ', round(r2,4))
    print('MAE: ', round(mean_absolute_error,4))
    print('MSE: ', round(mse,4))
    print('MedianAE', round(median_absolute_error, 4))
    print('RMSE: ', round(np.sqrt(mse),4))
    #TODO add when sklearn updated
    display = metrics.PredictionErrorDisplay.from_predictions(
        y_true=y_true,
        y_pred=y_pred,
        kind="actual_vs_predicted",
        #ax=ax,
        scatter_kwargs={"alpha": 0.2, "color": "tab:blue"},
        line_kwargs={"color": "tab:red"},
    )

def RF_oob_opt(model, X_train, min_est, max_est, step, group=None,
               regress=False):

    """
    This function uses the oob score to find the best parameters.

    This cannot be parallelized due to the warm start bootstrapping, so is
    potentially slower than the other cross val in the create_model function

    This function is based on an example from the sklearn site

    This function plots a graph diplaying the oob rate

    Parameters
    ---------------------

    model: string (.gz)
            path to model to be saved

    X_train: np array
              numpy array of training data where the 1st column is labels

    min_est: int
              min no of trees

    max_est: int
              max no of trees

    step: int
           the step at which no of trees is increased

    regress: bool
              boolean where if True it is a regressor

    Returns: tuple of np arrays
    -----------------------

    error rate, best estimator

    """
    # This is a bit slow at present, needs profiled
    #t0 = time()

    bands = X_train.shape[1]-1

    #X_train = X_train.transpose()
    X_train = X_train[X_train[:,0] != 0]

    # Remove non-finite values
    if group is not None: #replace this later not good
        inds = np.where(np.isfinite(X_train).all(axis=1))
        group = group[inds]

    X_train = X_train[np.isfinite(X_train).all(axis=1)]
    # y labels
    y_train = X_train[:,0]

    # remove labels from X_train
    X_train = X_train[:,1:bands+1]

    RANDOM_STATE = 123

    if group is not None:

        # maintain group based splitting from initial train/test split
        # to main train set
        # TODO - sep func?
        splitter = GroupShuffleSplit(test_size=0.2, n_splits=1, random_state=0)
        split = splitter.split(X_train, y_train, group)
        train_inds, test_inds = next(split)

        X_test = X_train[test_inds]
        y_test = y_train[test_inds]
        X_train = X_train[train_inds]
        y_train = y_train[train_inds]

    else:
        X_train, X_test, y_train, y_test = train_test_split(
            X_train, y_train, test_size=0.2, random_state=0)

    print('iterating estimators')
    if regress is True:
        max_feat = X_train.shape[1]-1
        ensemble_clfs = [
        ("RandomForestRegressor, max_features='no_features'",
                RandomForestRegressor(warm_start=True, oob_score=True,
                                       max_features=max_feat,
                                       random_state=RANDOM_STATE))]
    else:
        ensemble_clfs = [
        ("RandomForestClassifier, max_features='sqrt'",
                RandomForestClassifier(warm_start=True, oob_score=True,
                                       max_features="sqrt",
                                       random_state=RANDOM_STATE)),
        ("RandomForestClassifier, max_features='log2'",
                RandomForestClassifier(warm_start=True, max_features='log2',
                                       oob_score=True,
                                       random_state=RANDOM_STATE)),
        ]

    # Map a classifier name to a list of (<n_estimators>, <error rate>) pairs.
    error_rate = OrderedDict((label, []) for label, _ in ensemble_clfs)

    # Range of `n_estimators` values to explore.
    min_estimators = min_est
    max_estimators = max_est

    for label, clf in ensemble_clfs:
        for i in tqdm(range(min_estimators, max_estimators + 1, step)):
            clf.set_params(n_estimators=i)
            clf.fit(X_train, y_train)

            # Record the OOB error for each `n_estimators=i` setting.
            oob_error = 1 - clf.oob_score_
            error_rate[label].append((i, oob_error))
    # Generate the "OOB error rate" vs. "n_estimators" plot.
    for label, clf_err in error_rate.items():
        xs, ys = zip(*clf_err)
        plt.plot(xs, ys, label=label)

    # suspect a slow down after here, likely the plot...
    plt.xlim(min_estimators, max_estimators)
    plt.xlabel("n_estimators")
    plt.ylabel("OOB error rate")
    plt.legend(loc="best")
    plt.show()

    # regression option
    if regress is True:
        max_features = max_feat
        er = np.array(error_rate["RandomForestRegressor, max_features='no_features'"][0:max_estimators])
        bestscr = er[:,1].min()
        data = er[np.where(er[:,1] == bestscr)]
        best_param = ((max_features, data[0]))
        best_model = RandomForestRegressor(warm_start=True, oob_score=True,
                                       max_features= max_features,
                                       n_estimators=best_param[1][0].astype(int),
                                       n_jobs=-1)
    else:
        sqrt = np.array(error_rate["RandomForestClassifier, max_features='sqrt'"][0:max_estimators])
        log2 = np.array(error_rate["RandomForestClassifier, max_features='log2'"][0:max_estimators])
    #NoFeat = np.array(error["RandomForestClassifier, max_features='None'"][0:max_estimators])
        minsqrt = sqrt[:,1].min()
        minLog = log2[:,1].min()

        if minsqrt>minLog:
            minVal = minLog
            max_features = 'log2'
            data = log2[np.where(log2[:,1] == minVal)]
        else:
            minVal = minsqrt
            max_features = 'sqrt'
            data = sqrt[np.where(sqrt[:,1] == minVal)]

        best_param = ((max_features, data[0]))

        best_model = RandomForestClassifier(warm_start=True, oob_score=True,
                                       max_features= max_features,
                                       n_estimators=best_param[1][0].astype(int),
                                       n_jobs=-1)
    best_model.fit(X_train, y_train)
    joblib.dump(best_model, model)

    if regress is True:
        pred = best_model.predict(X_test)
        print('Test results are:')
        regression_results(np.expand_dims(y_test,1), pred)

    return error_rate, best_param

def plot_feature_importances(modelPth, featureNames,
                             rank=True):

    """
    Plot the feature importances of an ensemble classifier

    Parameters
    --------------------------

    modelPth : string
               A sklearn model path or the estimator itself

    featureNames : list of strings
                   a list of feature names

    """
    if modelPth is not str:
        model = modelPth
    else:
        model = joblib.load(modelPth)

    # if model_type=='scikit':
    #     n_features = model.n_features_
    # if model_type=='xgb':
    #     n_features = model.n_features_in_
    # plt.barh(range(n_features), model.feature_importances_, align='center')
    # plt.yticks(np.arange(n_features), featureNames)
    # plt.xlabel("Feature importance")
    # plt.ylabel("Feature")
    # plt.ylim(-1, n_features)
    # plt.show()

    # fimp = pd.Series(model.feature_importances_, index=featureNames)
    importances = model.feature_importances_
    std = np.std([tree.feature_importances_ for tree in model.estimators_], axis=0)

    fimp = pd.Series(importances, index=featureNames)

    if rank == True:
        fimp = fimp.sort_values(ascending=False)

    fig, ax = plt.subplots()
    fimp.plot.bar(yerr=std, ax=ax)
    ax.set_title("Feature importances using MDI")
    ax.set_ylabel("Mean decrease in impurity")
    fig.tight_layout()

    return fimp

def plot_feat_importance_permutation(modelPth, featureNames,  X_test, y_test,
                                     rank=True):

    """
    Plot the feature importances of an ensemble classifier

    Parameters
    --------------------------

    modelPth : string
               A sklearn model path or the estimator itself

    featureNames : list of strings
                   a list of feature names

    Returns
    -------

    pandas df of importances

    """

    if modelPth is not str:
        model = modelPth
    else:
        model = joblib.load(modelPth)

    # if model_type=='scikit':
    #     n_features = model.n_features_
    # if model_type=='xgb':
    #     n_features = model.n_features_in_

    result = permutation_importance(
        model, X_test, y_test, n_repeats=10, random_state=0, n_jobs=-1
    )

    fimp = pd.Series(result.importances_mean, index=featureNames)
    if rank == True:
        fimp = fimp.sort_values(ascending=False)

    fig, ax = plt.subplots()
    fimp.plot.bar(yerr=result.importances_std, ax=ax)
    ax.set_title("Feature importances using permutation on full model")
    ax.set_ylabel("Mean accuracy decrease")
    fig.tight_layout()
    plt.show()

    return fimp

def classify_pixel(model, inputDir, bands, outMap, probMap):
    """
    A function to classify an image using a pre-saved model - assumes
    a folder of tiled rasters for memory management - classify_pixel_block is
    recommended instead of this function

    Parameters

    ---------------

    model: sklearn model
            a path to a scikit learn model that has been saved

    inputDir: string
               a folder with images to be classified

    bands: int
            the no of image bands eg 8

    outMap: string
             path to output image excluding the file format 'pathto/mymap'

    probMap: string
              path to output prob image excluding the file format 'pathto/mymap'

    FMT: string
          optional parameter - gdal readable fmt

        """

    imageList = glob.glob(inputDir+'*tile*')

    numBands = bands
    n = len(imageList)
    for im in range (0,n):

        inputImage = gdal.Open(imageList[im])
        shapeB = ((numBands, inputImage.RasterXSize* inputImage.RasterYSize))

        X = np.zeros(shape = shapeB)
        for band in tqdm(range(1,bands+1)):
            band1 = inputImage.GetRasterBand(band)
            im1 = band1.ReadAsArray()
            X[band-1] = im1.flatten()
        X = X.transpose()
        X = np.where(np.isfinite(X),X,0)

        #RF_clf = grid.best_estimator_

       #RF_clf.fit(X_train, y_train)

        # We only want the band stats - must remove the training numbers
        #X = X[:,1:len(bands+1)]

        model1 = joblib.load(model+'.gz')
        predictClass = model1.predict(X)
        #if inputImage.RasterXSize < inputImage.RasterYSize
        predictClass = np.reshape(predictClass, (im1.shape[0], im1.shape[1]))
#        else:
#            predictClass = np.reshape(predictClass, (inputImage.RasterXSize, inputImage.RasterYSize))

        print('Saving outputs')
        array2raster(predictClass, 1, inputImage, outMap+str(im)+'.tif', gdal.GDT_Byte, 'Gtiff')

        Probs=model1.predict_proba(X)
        print('classification '+str(im)+' done')
        #Class 1 = Deforestattion
        #if inputImage.RasterXSize < inputImage.RasterYSize:
        classArr = np.unique(predictClass)
        ProbsFinal = np.empty((im1.shape[0], im1.shape[1], inputImage.RasterCount))

        for band in tqdm(classArr):
            ProbsFinal[:,:,band-1] = np.reshape(Probs[:,band-1], (im1.shape[0], im1.shape[1]))

        print('Prob_'+str(im)+' done')
        array2raster(ProbsFinal, 4, inputImage, probMap+str(im)+'.tif', gdal.GDT_Float64, 'Gtiff')

def classify_pixel_bloc(model, inputImage,  outMap, bands=[1,2,3], blocksize=None,
                        FMT=None, ndvi = None, dtype = gdal.GDT_Int32):
    """
    A block processing classifier for large rasters, supports KEA, HFA, & Gtiff
    formats. KEA is recommended, Gtiff is the default

    Parameters
    ------------------

    model: sklearn model / keras model
            a path to a model that has been saved

    inputImage: string
                 path to image including the file fmt 'Myimage.tif'

    bands: band
            list of band indices to be used eg [1,2,3]

    outMap: string
             path to output image excluding the file format 'pathto/mymap'

    FMT: string
          optional parameter - gdal readable fmt

    blocksize: int (optional)
                size of raster chunck in pixels 256 tends to be quickest
                if you put None it will read size from gdal (this doesn't always pay off!)

    dtype: int (optional - gdal syntax gdal.GDT_Int32)
            a gdal dataype - default is int32

    Notes
    -----

    Block processing is sequential, but quite a few sklearn models are parallel
    so that has been prioritised rather than raster IO
    """
    if FMT == None:
        FMT = 'Gtiff'
        fmt = '.tif'
    if FMT == 'HFA':
        fmt = '.img'
    if FMT == 'KEA':
        fmt = '.kea'
    if FMT == 'Gtiff':
        fmt = '.tif'

    inDataset = gdal.Open(inputImage)

    outDataset = _copy_dataset_config(inDataset, outMap=outMap,
                                     dtype=dtype, bands=1)
    band = inDataset.GetRasterBand(1)
    cols = int(inDataset.RasterXSize)
    rows = int(inDataset.RasterYSize)
    outBand = outDataset.GetRasterBand(1)
    # So with most datasets blocksize is a row scanline
    if blocksize == None:
        blocksize = band.GetBlockSize()
        blocksizeX = blocksize[0]
        blocksizeY = blocksize[1]
    else:
        blocksizeX = blocksize
        blocksizeY = blocksize

    if os.path.splitext(model)[1] == ".h5":
        model1 = load_model(model)
    else:
        model1 = joblib.load(model)

    noBands = inDataset.RasterCount

    # minus the bands to zero start
    bands = list(np.array(bands)-1)

    if blocksizeY==1:
        rows = np.arange(cols, dtype=np.int)
        for row in tqdm(rows):
            i = int(row)
            j = 0
            #X = np.zeros(shape = (bands, blocksizeX))
            #for band in range(1,bands+1):

            X = inDataset.ReadAsArray(j, i, xsize=blocksizeX, ysize=blocksizeY)
            X.shape = ((noBands,blocksizeX))

            X = X[bands, :]

            if X.max() == 0:
                predictClass = np.zeros_like(rows, dtype = np.int32)
            else:
                X = np.where(np.isfinite(X),X,0) # this is a slower line
                X = X.transpose()
                #Xs = csr_matrix(X)
                predictClass = model1.predict(X)
                predictClass[X[:,0]==0]=0

            outBand.WriteArray(predictClass.reshape(1, blocksizeX),j,i)
                    #print(i,j)

    # else it is a block
    else:
        for i in tqdm(range(0, rows, blocksizeY)):
            if i + blocksizeY < rows:
                numRows = blocksizeY
            else:
                numRows = rows -i

            for j in range(0, cols, blocksizeX):
                if j + blocksizeX < cols:
                    numCols = blocksizeX
                else:
                    numCols = cols - j

                X = inDataset.ReadAsArray(j,i, xsize=numCols, ysize=numRows)
#                X = np.zeros(shape = (bands, numCols*numRows))
#                for band in range(1,bands+1):
#                    band1 = inDataset.GetRasterBand(band)
#                    data = band1.ReadAsArray(j, i, numCols, numRows)
                with warnings.catch_warnings():
                            warnings.simplefilter("ignore")
                            if X.max() == 0:
                                continue
                            else:

                                X.shape = ((noBands,numRows*numCols))
                                X = X[bands, :]
                                X = X.transpose()
                                X = np.where(np.isfinite(X),X,0)

                                predictClass = model1.predict(X)
                                # this deletes vals betweem 0 and 1 so it is scrubbed
                                # not sure why it was here
                                #predictClass[X[:,0]==0]=0
                                predictClass = np.reshape(predictClass, (numRows, numCols))
                                outBand.WriteArray(predictClass,j,i)
                #print(i,j)
    outDataset.FlushCache()
    outDataset = None

def prob_pixel_bloc(model, inputImage, bands, outMap, classes, blocksize=None,
                    FMT=None, one_class = None):
    """
    A block processing classifier for large rasters that produces a probability,
    output.

    Supports KEA, HFA, & Gtiff formats -KEA is recommended, Gtiff is the default

    Parameters
    ----------
    model: string
            a path to a scikit learn model that has been saved

    inputImage : string
                 path to image including the file fmt 'Myimage.tif'

    bands: int
            the no of image bands eg 8

    outMap: string
             path to output image excluding the file format 'pathto/mymap'

    classes: int
              no of classes

    blocksize: int (optional)
                size of raster chunck 256 tends to be quickest if you put None it
                will read size from gdal (this doesn't always pay off!)

    FMT: string
          optional parameter - gdal readable fmt eg 'Gtiff'

    one_class: int
                choose a single class to produce output prob raster

    -------------------------------------------------

    Block processing is sequential, but quite a few sklearn models are parallel
    so that has been prioritised rather than raster IO
    """
    if FMT == None:
        FMT = 'Gtiff'
        fmt = 'tif'
    if FMT == 'HFA':
        fmt = '.img'
    if FMT == 'KEA':
        fmt = '.kea'
    if FMT == 'Gtiff':
        fmt = '.tif'

    # TODO - a list of classes would be better eliminating the need for the one
    # class param
    if one_class != None:
        classes = one_class

    inDataset = gdal.Open(inputImage)

    outDataset = _copy_dataset_config(inputImage, outMap = outMap,
                                     bands = bands)
    band = inDataset.GetRasterBand(1)
    cols = inDataset.RasterXSize
    rows = inDataset.RasterYSize
    #outBand = outDataset.GetRasterBand(1)
    # So with most datasets blocksize is a row scanline
    if blocksize == None:
        blocksize = band.GetBlockSize()
        blocksizeX = blocksize[0]
        blocksizeY = blocksize[1]
    else:
        blocksizeX = blocksize
        blocksizeY = blocksize

    model1 = joblib.load(model)
    if blocksizeY==1:
        rows = np.arange(cols, dtype=np.int)
        for row in tqdm(rows):

            i = int(row)
            j = 0
            X = np.zeros(shape = (bands, blocksizeX))
            for band in range(1,bands+1):
                band1 = inDataset.GetRasterBand(band)
                data = band1.ReadAsArray(j, i, blocksizeX, 1)
                X[band-1] = data.flatten()
            X = X.transpose()
            X = np.where(np.isfinite(X),X,0) # this is a slower line
            Probs=model1.predict_proba(X)

            ProbsFinal = np.empty((data.shape[0], data.shape[1], classes))
            for band in range(1,classes+1):
                ProbsFinal[:,:,band-1] = np.reshape(Probs[:,band-1], (1, blocksizeX))
                outBand = outDataset.GetRasterBand(band)
                outBand.WriteArray(ProbsFinal[:,:,band-1],j,i)
                #print(i,j)

    # else it is a block
    else:
        model1 = joblib.load(model)
        for i in tqdm(range(0, rows, blocksizeY)):
            if i + blocksizeY < rows:
                numRows = blocksizeY
            else:
                numRows = rows -i

            for j in range(0, cols, blocksizeX):
                if j + blocksizeX < cols:
                    numCols = blocksizeX
                else:
                    numCols = cols - j
                X = np.zeros(shape = (bands, numCols*numRows))
                for band in range(1,bands+1):
                    band1 = inDataset.GetRasterBand(band)
                    data = band1.ReadAsArray(j, i, numCols, numRows)
                    X[band-1] = data.flatten()
                X = X.transpose()
                X = np.where(np.isfinite(X),X,0)  # this is a slower line
                Probs=model1.predict_proba(X)

                if one_class != None:
                    ProbsFinal = np.reshape(Probs[:,one_class-1],
                                            (data.shape[0], data.shape[1]))
                    outBand = outDataset.GetRasterBand(one_class)
                    outBand.WriteArray(ProbsFinal,j,i)
                else:
                    ProbsFinal = np.empty((data.shape[0], data.shape[1], classes))
                    for band in range(1,classes+1):
                        ProbsFinal[:,:,band-1] = np.reshape(Probs[:,band-1], (data.shape[0], data.shape[1]))
                        outBand = outDataset.GetRasterBand(band)
                        outBand.WriteArray(ProbsFinal[:,:,band-1],j,i)
                    #print(i,j)

    #outBand.FlushCache()
    outDataset.FlushCache()
    outDataset = None

def classify_object(model, inShape, attributes, field_name=None, write='gpd'):

    """
    Classify a polygon/point file attributes ('object based') using an sklearn
    model

    Parameters
    ------------------
    model: string
            path to input model

    inShape: string
              input shapefile path (must be .shp for now....)

    attributes: list of stings
                 list of attributes names

    field_name: string
                 name of classified label field (optional)

    write: string
                either gpd(geopandas) or ogr
    """

#    old method for ref - limited to .shp
#    dbf=simpledbf.Dbf5(inShape[:-4]+'.dbf')
#    csv = inShape[:-4]+'.csv'
#    dbf.to_csv(inShape[:-4]+'.csv')
#    df = dbf.to_dataframe()

    # it seems rather ugly/inefficient to read in every row via ogr
    model1 = joblib.load(model)

    df = gpd.read_file(inShape)

    X = df[attributes].to_numpy()

    if write == 'ogr':
        del df

    X[np.where(np.isnan(X))]=0
    X = X[np.isfinite(X).all(axis=1)]

    #Now the classification itself - see sklearn for details on params

    predictClass = model1.predict(X)

    # clear redundant variables from memory

    del X

    predictClass = predictClass.transpose()

    if write == 'ogr':
        shp = ogr.Open(inShape, 1)
        lyr = shp.GetLayer()
        fldDef = ogr.FieldDefn(field_name, ogr.OFTInteger)
        lyr.CreateField(fldDef)

        labels = np.arange(lyr.GetFeatureCount())

        for label in tqdm(labels):
            val=predictClass[label]
            feat = lyr.GetFeature(label)
            feat.SetField(field_name, int(val))
            lyr.SetFeature(feat)

        lyr.SyncToDisk()
        shp.FlushCache()
        shp = None
        lyr = None
    else:
        df[field_name] = predictClass
        df.to_file(inShape)

def get_polars(inShp, polars=["VV", "VH"]):

    """
    Get list of fields containing polarisations from a polygon/point file

    Parameters
    ----------

    inShp: string
          the input polygon

    polars: list of strings
            the attributes headed with polarisations eg 'VV'

    """

    shp = ogr.Open(inShp)
    lyr = shp.GetLayer()
    lyrdefn = lyr.GetLayerDefn()

    ootlist = []

    for f in range(lyrdefn.GetFieldCount()):
        defn = lyrdefn.GetFieldDefn(f)
        ootlist.append(defn.name)

    final = [o for o in ootlist if "VV" in o or "VH" in o]

    return final

def get_training_shp(inShape, label_field, feat_fields,  outFile = None):
    """
    Collect training from a shapefile attribute table. Used for object-based
    classification (typically).

    Parameters
    --------------------

    inShape: string
              the input polygon

    label_field: string
                  the field name for the class labels

    feat_fields: list
                  the field names of the feature data

    outFile: string (optional)
              path to training data to be saved (.gz)

    Returns
    ---------------------
    training data as a dataframe, first column is labels, rest are features
    list of reject features

    """
    # TODO Could be done in a less verbose way by gpd
    outData = list()

    feat_fields.insert(0,label_field)

    print('Loading & prepping data')
    shp = ogr.Open(inShape)
    lyr = shp.GetLayer()
    labels = np.arange(lyr.GetFeatureCount())

    rejects = []

    print('calculating stats')
    for label in tqdm(labels):
        #print(label)
        feat = lyr.GetFeature(label)
        if feat == None:
            print('no geometry for feature '+str(label))
            rejects.append(label)
            continue
        row = [feat.GetField(entry) for entry in feat_fields]
        outData.append(row)

    df=pd.DataFrame(outData, columns = feat_fields)
    outArray = df.to_numpy()

    if outFile != None:
        jb.dump(outArray, outFile, compress=2)

    return df, rejects

def get_training(inShape, inRas, bands, field, outFile = None):
    """
    Collect training as an np array for use with create model function

    Parameters
    --------------

    inShape: string
              the input shapefile - must be esri .shp at present

    inRas: string
            the input raster from which the training is extracted

    bands: int
            no of bands

    field: string
            the attribute field containing the training labels

    outFile: string (optional)
              path to the training file saved as joblib format (eg - 'training.gz')

    Returns
    ---------------------

    A tuple containing:
    -np array of training data
    -list of polygons with invalid geometry that were not collected

    """
    #t0 = time()
    outData = list()
    print('Loading & prepping data')
    raster = gdal.Open(inRas)
    shp = ogr.Open(inShape)
    lyr = shp.GetLayer()
    labels = np.arange(lyr.GetFeatureCount())
    rb = raster.GetRasterBand(1)
    rgt = raster.GetGeoTransform()
    mem_drv = ogr.GetDriverByName('Memory')
    driver = gdal.GetDriverByName('MEM')
    rejects = []

    print('calculating stats')
    for label in tqdm(labels):
        #print(label)
        feat = lyr.GetFeature(label)
        if feat == None:
            print('no geometry for feature '+str(label))
            continue
        iD = feat.GetField(field)
        geom = feat.GetGeometryRef()

        # Get raster georeference info

        src_offset = _bbox_to_pixel_offsets(rgt, geom)

        # calculate new geotransform of the feature subset
        new_gt = (
        (rgt[0] + (src_offset[0] * rgt[1])),
        rgt[1],
        0.0,
        (rgt[3] + (src_offset[1] * rgt[5])),
        0.0,
        rgt[5])

        # Create a temporary vector layer in memory
        mem_ds = mem_drv.CreateDataSource('out')
        mem_layer = mem_ds.CreateLayer('poly', None, ogr.wkbPolygon)
        mem_layer.CreateFeature(feat.Clone())

        # Rasterize
        rvds = driver.Create('', src_offset[2], src_offset[3], 1, gdal.GDT_Byte)
        rvds.SetGeoTransform(new_gt)
        gdal.RasterizeLayer(rvds, [1], mem_layer, burn_values=[1])
        rv_array = rvds.ReadAsArray()

        # Mask the source data array with our current feature
        # Use the logical_not to flip 0<->1 to get the correct mask effect
        # we also mask out nodata values explictly
        rb = raster.GetRasterBand(1)
        src_array = rb.ReadAsArray(src_offset[0], src_offset[1], src_offset[2],
                           src_offset[3])
        if np.shape(src_array) == ():
            rejects.append(label)
            continue
        # Read raster as arrays
        for band in range(1,bands+1):

            rb = raster.GetRasterBand(band)
            src_array = rb.ReadAsArray(src_offset[0], src_offset[1], src_offset[2],
                               src_offset[3])
            if src_array is None:
                src_array = rb.ReadAsArray(src_offset[0]-1, src_offset[1], src_offset[2],
                                           src_offset[3])

            masked = np.ma.MaskedArray(src_array,
                                       mask=np.logical_or(src_array == 0,
                                                          np.logical_not(rv_array)))

            datafinal = masked.flatten()

            if band == 1:
                X = np.zeros(shape = (datafinal.shape[0], bands+1))
            X[:,0] = iD

            X[:,band] = datafinal

        outData.append(X)
    outData = np.asarray(outData)
    outData = np.concatenate(outData).astype(None)

    if outFile != None:
        jb.dump(outData, outFile, compress=2)

    return outData, rejects

def rmse_vector_lyr(inShape, attributes):

    """
    Using sklearn get the rmse of 2 vector attributes
    (the actual and predicted of course in the order ['actual', 'pred'])

    Parameters
    -----------

    inShape: string
              the input vector of OGR type

    attributes: list
           a list of strings denoting the attributes

    """

    #open the layer etc
    shp = ogr.Open(inShape)
    lyr = shp.GetLayer()
    labels = np.arange(lyr.GetFeatureCount())

    # empty arrays for att
    pred = np.zeros((1, lyr.GetFeatureCount()))
    true = np.zeros((1, lyr.GetFeatureCount()))

    for label in labels:
        feat = lyr.GetFeature(label)
        true[:,label] = feat.GetField(attributes[0])
        pred[:,label] = feat.GetField(attributes[1])

    error = np.sqrt(metrics.mean_squared_error(true, pred))

    return error


# -*- coding: utf-8 -*-
# @Author: WuLC
# @Date:   2016-05-19 10:17:15
# @Last modified by:   WuLC
# @Last Modified time: 2016-05-19 10:18:13
# @Email: liangchaowu5@gmail.com

# binary search in recurse mode
# deal with case like: [1,1,1,1,1,2,3,1,1]
class Solution(object):
    def search(self, nums, target):
        """
        :type nums: List[int]
        :type target: int
        :rtype: bool
        """
        if len(nums) == 0:
            return False
        n = len(nums)
        if self.helper(nums, 0, n-1,target):
            return True
        else:
            return False

    def helper(self,nums,left,right,target):
        if left > right:
            return False
        if left == right:
            return nums[left] == target
        else:
            mid = (left+right)/2
            if nums[mid] == target:
                return True
            elif  nums[left]<=target<nums[mid]:
                return self.helper(nums,left,mid-1,target)
            elif nums[mid]<target<=nums[right]:
                return self.helper(nums,mid+1,right,target)
            else:
                return self.helper(nums,mid+1,right,target) or self.helper(nums,left,mid-1,target)

Algorithm/Python/081. Search in Rotated Sorted Array II.py


#
# editTCXFile.py
# Reads in a .tcx file from strava and replaces altitude change with the elevation value in each event record
#
import os
import xmltodict
import xml.etree.ElementTree as ET  # ET GUIDE: https://stackoverflow.com/questions/16419673/how-to-find-specific-elements-in-xml-using-elementtree
import pandas as pd
from datetime import datetime

#--- input params ---#
filename = 'Tread_Hour_1_of_2 - UPDATE.tcx'  # name of tcx file (do not include path)

#--------------------#

# Read in XML data from current working directory
input_tcx_file =  str(os.getcwd()) + '\\' + filename
tree = ET.parse(input_tcx_file)
root = tree.getroot()

# Establish output filename
outputfilename = input_tcx_file.replace('.tcx','_adj.tcx')
output_tcx_file = outputfilename

# Define the TCX namespace (one default, one for extension data)
namespace = {'ns': 'http://www.garmin.com/xmlschemas/TrainingCenterDatabase/v2'} # this is the default namespace
namespace_ext = {'ns': 'http://www.garmin.com/xmlschemas/ActivityExtension/v2'} # this is a namespace unique to the extension branches

# Find all trackpoints in the TCX file
trackpoints = root.findall(".//ns:Trackpoint", namespaces=namespace)

# Iterate through each trackpoint and replace altitude with elevation
for trackpoint in trackpoints:
    altitude_element = trackpoint.find(".//ns:AltitudeMeters", namespaces=namespace)
    if altitude_element is not None:
        elevation_element = trackpoint.find(".//ns:Elevation", namespaces=namespace_ext)
        elevation = elevation_element.text
        altitude_element.text = elevation

# NEED TO Clean up the string of encoding issues
# updated_xml_data = updated_xml_data.replace('ns0:','').replace('ns1:','').replace('ns2:','') # issue with this approach is the namespaces get removed in some branches

# Save the modified TCX data to a new file
tree.write(output_tcx_file)
print(f"TCX file with altitude replaced by elevation saved to {output_tcx_file}")


import pandas as pd
import json
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sn
import os

from project_annotation import Annotation

def load_json(file_path):
    if os.path.exists(file_path):
        with open(file_path) as f:
            data = json.load(f)
            return data
    else:
        print("Error: can't find path",file_path)
        return None

LABELS_SHL = {
        0: "Null",
        1: "Still",
        2: "Walking",
        3: "Run",
        4: "Bike",
        5: "Car",
        6: "Bus",
        7: "Train",
        8: "Subway",
        }

class Anno:
    # Total number of samples
    samples_total_dict = {'A': 4620201, 'B': 4557401, '0': 172701}

    def __init__(self, path, scenario):
        self.path = path
        self.scenario = scenario

        self.anno_dic = load_json(self.path)
        self.anno_obj = Annotation(type_=f"anno_{self.scenario}", path_=None, points=self.anno_dic)
        self.anno_mat = self.anno_obj.get_array( last_sample=self.samples_total_dict[self.scenario], convert_to_idx=True )

    @property
    def samples_total(self): return self.samples_total_dict[self.scenario]

def calc_F1Score(label,pred):
    # Total number of samples
    N = label.shape[0]

    F1 = 0
    for i in range(9):
        ni_l = np.sum((label == i))
        ni_p = np.sum((pred == i))

        if ni_l == 0:
            continue

        # weight
        wi = ni_l/N

        TP = np.sum((label == i) & (pred == i))
        FP = np.sum((label != i) & (pred == i))
        FN = np.sum((label == i) & (pred != i))

        if TP == 0:
            continue

        pi = (TP)/(TP+FP)
        ri = (TP)/(TP+FN)

        Fi = 2*(pi*ri)/(pi+ri)
        #print(f"Class {i} {Fi=} {ni_l=} {ni_p=} {TP=} {FP=} {FN=} {pi=} {ri=}")

        F1 += wi * Fi

    return F1

CARPM_indeces = {i:v for i,v in enumerate(["TP","TN","Overfill","Underfill","Merge","Fragmentation","Insertion","Deletion","Substitution"])}
#print(CARPM_indeces)

def mark_null_puddles(signal_A, signal_B, results, fill_value):
    s = 0
    l = -1
    for i in range(1,results.shape[0]):

        c = signal_A[i-1]

        # New actual class
        if signal_A[i] != c:

            # Search coherent activity for fragmentations
            if c != 0:
                sp = None
                for j in range(s,i):
                    if signal_B[j] == c:
                        if signal_B[j-1] == 0:
                            if sp is not None:
                                results[sp+1:j] = fill_value
                        sp = j
                    elif signal_B[j] != 0:
                        sp = None

            s = i
    return results

def calc_CARPM(label,pred):

    res = np.empty_like(label)
    res.fill(-1)

    # Insertion
    res[np.logical_and.reduce((label!=pred,label==0,pred!=0))] = 6

    # Deletion
    res[np.logical_and.reduce((label!=pred,label!=0,pred==0))] = 7

    # ~~~~~~~~~~~~~~~ #
    #  Special Cases  #
    # ~~~~~~~~~~~~~~~ #

    # Merge
    res = mark_null_puddles(signal_A=pred, signal_B=label, results=res, fill_value=4)

    # Fragmentation
    res = mark_null_puddles(signal_A=label, signal_B=pred, results=res, fill_value=5)

    # Overfill/Underfill
    s, sp = 0, 0
    for i in range(1,res.shape[0]):

        # Synced start can be skipped
        if label[i] != label[i-1] and pred[i] != pred[i-1]:
            s = i
            sp = i
            continue

        # Only consider cases where match is reached and no merge/fragmentation is already marked
        if label[i] == pred[i] and res[i-1] not in (4,5):

            # New actual activity class (Overfill)
            if label[i] != label[i-1] and label[i] != 0:
                # print(i,"Overfill")
                res[sp:i] = 2

            # End actual activity class (Underfill) +
            if label[i] != label[i-1] and label[i] == 0 and pred[sp-1] == label[i-1]:
                # print(i,"Underfill+")
                res[sp:i] = 3

            # New predicted activity class (Underfill)
            if pred[i] != pred[i-1] and pred[i] != 0:
                # print(i,"Underfill")
                res[s:i] = 3

            # End predicted activity class (Overfill) +
            if pred[i] != pred[i-1] and pred[i] == 0 and label[s-1] == pred[i-1]:
                # print(i,"Overfill+")
                res[s:i] = 2

        # ~~~~~~~~~~~~~~~ #
        #  Start Indeces  #
        # ~~~~~~~~~~~~~~~ #

        # New actual class
        if label[i] != label[i-1]:
            s = i

        # New predicted class
        if pred[i] != pred[i-1]:
            sp = i

    # True Positives
    res[np.logical_and(label==pred,label!=0)] = 0

    # True Negatives
    res[np.logical_and(label==pred,label==0)] = 1

    # Substitution
    res[np.logical_and.reduce((label!=pred,label!=0,pred!=0))] = 8

    return res

class GroundTruth:

    def __init__(self, path, scenario):
        self.path = path
        self.scenario = scenario
        self.anno = Anno(self.path, self.scenario)

    # Calculate weighted F1 score
    def get_F1(self, anno_pred, max_sample):

        label = self.anno.anno_mat[:max_sample]
        pred = anno_pred.anno_mat[:max_sample]

        return calc_F1Score(label, pred)

    # Calculate Continuous Activity Recognition Performance Metrics
    def get_CARPM(self, anno_pred=None, carpm=None, max_sample=0):
        assert anno_pred is not None or carpm is not None

        if carpm is None:
            label = self.anno.anno_mat[:max_sample]
            pred = anno_pred.anno_mat[:max_sample]

            carpm = calc_CARPM(label,pred)

        total = {v:np.sum(carpm==i) for i,v in CARPM_indeces.items()}
        ratio = {k:v/max_sample for k,v in total.items()}

        return total, ratio, carpm

import re
from sklearn.metrics import confusion_matrix

def get_sec(time_str):
    """Get seconds from time."""
    h, m, s = time_str.split(':')
    return int(h) * 3600 + int(m) * 60 + int(s)

def get_cm(label,pred,labels):
    return np.round(confusion_matrix(y_true=label,y_pred=pred,labels=labels,normalize='all'), 3)

path_pred = "static/uploads/pred"
class Task:
    GT = {'A': GroundTruth(path=os.path.join(path_pred,"anno_270417.json"), scenario='A'),
          'B': GroundTruth(path=os.path.join(path_pred,"anno_270417.json"), scenario='B'),
          '0': GroundTruth(path=os.path.join(path_pred,"anno_SHL_test.json"), scenario='0')}

    def __init__(self, path, scenario, load_carpm=True):
        self.path = path
        self.scenario = scenario

        if scenario == '0':
            self.settings = load_json(os.path.join(path,"settings.json"))
            #self.tlx = load_json(os.path.join(path,f"TLX_{scenario}.json"))

            self.log_path = os.path.join(path,"log_js.txt")
            self.carpm_path = os.path.join(path,"carpm.npy")

            self.anno = Anno( os.path.join(path,"anno.json"), scenario )

        else:
            self.settings = load_json(os.path.join(path,scenario,"settings.json"))
            self.tlx = load_json(os.path.join(path,f"TLX_{scenario}.json"))

            self.log_path = os.path.join(path,scenario,"log_js.txt")
            self.carpm_path = os.path.join(path,scenario,"carpm.npy")

            self.anno = Anno( os.path.join(path,scenario,"anno.json"), scenario )

        self.comp = self.get_completion_rate()
        self.F1score = self.GT[self.scenario].get_F1(anno_pred=self.anno, max_sample=self.max_sample)

        if load_carpm and os.path.exists(self.carpm_path):
            self.CARPM = self.GT[self.scenario].get_CARPM(carpm=self.load_carpm(), max_sample=self.max_sample)
        else:
            self.CARPM = self.GT[self.scenario].get_CARPM(anno_pred=self.anno, max_sample=self.max_sample)
            self.save_carpm()

    @property
    def carpm(self):
        return self.CARPM[2]

    @property
    def final_time(self):
        return int(self.comp['end_time']/self.comp['ratio_sighted'])

    @property
    def max_sample(self):
        return self.comp['max_sample']

    @property
    def samples_total(self): return self.anno.samples_total

    @property
    def AI(self):
        return self.settings["modus"]["AI"]

    def save_carpm(self):
        np.save(file=self.carpm_path, arr=self.carpm, allow_pickle=False, fix_imports=False)

    def load_carpm(self):
        return np.load(file=self.carpm_path, mmap_mode=None, allow_pickle=False, fix_imports=False)

    # Ratio how much of the material was covered
    # To DO: Add time to complete?
    def get_completion_rate(self):

        # Find maximum sample viewed
        max_sample = 0

        # Establish time it took to sight maximum sample and do last annotation
        start_time = '00:00:00'
        sight_time = '00:00:00'
        anno_time = '00:00:00'
        end_time = '00:00:00'

        with open(self.log_path, 'r') as f:
            for i,line in enumerate(f):
                end_time = line.split()[1]
                if i==0:
                    start_time = line.split()[1]

                # Check for sighted material
                result = re.search("\[Canvas\] Start: .+ Stop: .+ Selected:",line)
                if result is not None:
                    stop = int(result.group().split()[-2])
                    if stop > max_sample:
                        max_sample = stop
                        sight_time = line.split()[1]

                # Check for annotation
                result = re.search("\[Anno\] Add",line)
                if result is not None:
                    anno_time = line.split()[1]

        sight_time = get_sec(sight_time) - get_sec(start_time)
        anno_time = get_sec(anno_time) - get_sec(start_time)
        end_time = get_sec(end_time) - get_sec(start_time)

        max_sample = min(max_sample,self.samples_total)

        return {
                'max_sample': max_sample,
                'ratio_sighted': max_sample / self.samples_total,
                'start_time': start_time,
                'sight_time': sight_time,
                'anno_time': anno_time,
                'end_time': min(1200,end_time),
               }

    @property
    def confusion_matrix(self):

        label = self.GT[self.scenario].anno.anno_mat[:self.max_sample]
        pred = self.anno.anno_mat[:self.max_sample]

        return get_cm(label,pred,labels=list(range(9)))

    # Generate the confusion matrix for substituted classes
    @property
    def substitution_matrix(self):

        label = self.GT[self.scenario].anno.anno_mat[:self.max_sample]
        pred = self.anno.anno_mat[:self.max_sample]
        carpm = self.carpm

        label = label[carpm==8]
        pred = pred[carpm==8]

        return get_cm(label,pred,labels=list(range(1,9)))

class Participant:
    def __init__(self, path):
        assert os.path.exists(path)

        self.task = { "A": Task(path,"A"),
                      "B": Task(path,"B") }

        self.signup = load_json(os.path.join(path,"signup.json"))
        self.quest = load_json(os.path.join(path,"questionnaire.json"))

    @property
    def AI_first(self):
        return self.task["A"].AI

    def get_att(self, att):
        return self.signup[att]

class Study:
    def __init__(self, result_path):
        self.parts = {name: Participant(os.path.join(result_path,name)) for name in os.listdir(result_path)}

    def get_att(self, att):
        return { k: v.get_att(att) for k,v in self.parts.items()  }

def plot_cm(cm,title="test", substitution=False):
    labels = [LABELS_SHL[idx] for idx in list(range(1 if substitution else 0, 9))]

    fig = plt.figure(figsize = (7,7))

    df_cm = pd.DataFrame(cm, index=labels, columns=labels)
    ax = sn.heatmap(df_cm, cmap=sn.cubehelix_palette(as_cmap=True), annot=True, cbar=False, square=True, fmt='.3f') #vmin=0, vmax=1,
    ax.set_xlabel("Predicted Activity")
    ax.set_ylabel("Actual Activity")
    ax.set_title(title)
    plt.tight_layout()
    plt.savefig(os.path.join("img","CM_"+title.replace(" ","_")+".pdf"), bbox_inches = 'tight')
    return fig


import os
import parameters
import numpy as np
from string import ascii_lowercase

# Se parameters and initialize
data = parameters.default()
parameters.dump(data)
locals().update(data)
np.random.seed(seed)

dat_filename = "data/population-gain.npy"
prm_filename = "data/population-gain-parameters.json"
fig_filename = "figs/monkey-agent-comparison.pdf"

if not os.path.exists(dat_filename):
    raise Exception("Datafile not found. Please run 'gain-analysis.py'")

# Load and normalize score
score = np.load(dat_filename)
data = parameters.load(prm_filename)
locals().update(data)
score = (score - score.min())/(score.max() - score.min())

# -----------------------------------------------------------------------------
import matplotlib.pyplot as plt

fig = plt.figure(figsize = (12, 6))
ax =  plt.subplot(1,2,1, aspect=1)

from scipy.ndimage.filters import gaussian_filter
gscore = gaussian_filter(score, 1.0)
median = np.median(score)
C = ax.contour(gscore, levels=[median,],
                extent=[pmin, pmax, vmin, vmax],
                vmin=100, vmax=140, origin="lower", colors="0.25", zorder=50)
ax.set_title("Monkeys fitted behaviour", weight="bold")
ax.set_xlabel("α (probability weighting function)" )
ax.set_ylabel("β (utility function)")

#ax.set_xlabel("Probability distortion")
#ax.set_ylabel("Value distortion")

ax.axhline(0, lw=0.75, ls="--", color="black")
ax.axvline(1, lw=0.75, ls="--", color="black")
plt.text(0.01, .75, "risk averse", rotation=90, transform=ax.transAxes,
         ha="left", va="center", fontsize="small")
plt.text(0.01, .25, "risk seeking", rotation=90, transform=ax.transAxes,
         ha="left", va="center", fontsize="small")
plt.text(0.25, 0.01, "low ++ / high --", rotation=0, transform=ax.transAxes,
         ha="center", va="bottom", fontsize="small")
plt.text(0.75, 0.01, "low -- / high ++", rotation=0, transform=ax.transAxes,
         ha="center", va="bottom", fontsize="small")

# Display monkeys
monkeys = np.genfromtxt('monkeys.csv', delimiter=',', dtype=None, names=True)
X, Y = monkeys["gaindistortion"], monkeys["gainrisk_aversion"]
ax.scatter(X, Y, s=25, facecolor="blue",
           edgecolor="blue", marker="o", label="gain")
ax.scatter([X.mean()], [Y.mean()], s=75, facecolor="blue",
           edgecolor="white", marker="X", lw=1)

for i, label in enumerate(monkeys["name"]):
    ax.text(X[i]-0.02, Y[i],
             label.decode(), size="small",
            va = "center", ha="right", color="blue")

X, Y = monkeys["lossdistortion"], monkeys["lossrisk_aversion"]
ax.scatter(X, Y, s=25, facecolor="red",
           edgecolor="red", marker="o", label="loss")
ax.scatter([X.mean()], [Y.mean()], s=75, facecolor="red",
           edgecolor="white", marker="X", lw=1)
for i, label in enumerate(monkeys["name"]):
    label = label.decode()
    dy = 0.0 if label != "Yoh" else 0.025
    ax.text(X[i]-0.02, Y[i]+dy, label, size="small",
            va = "center", ha="right", color="red")
ax.legend(frameon=False)
ax.set_xlim(pmin, pmax), ax.set_ylim(vmin, vmax)

ax.text(-0.1, 1.1, ascii_lowercase[0],
        transform=ax.transAxes, size=20, weight='bold')

ax = plt.subplot(1,2,2, aspect=1)

p = "(selection_rate=%.2f,mutation_rate=%.2f,mixture_rate=%.2f)"
p = p % (selection_rate, mutation_rate, mixture_rate)
filename1 = "data/simulation-initial-population" + p + ".npy"
filename2 = "data/simulation-final-population" + p + ".npy"
agent_i = np.load(filename1)
agent_f = np.load(filename2)

from scipy.ndimage.filters import gaussian_filter
gscore = gaussian_filter(score, 1.0)
median = np.median(score)
C = ax.contour(gscore, levels=[median,],
                extent=[pmin, pmax, vmin, vmax],
                vmin=100, vmax=140, origin="lower", colors="0.25", zorder=50)
ax.set_title("Agents' behaviour", weight="bold")
ax.set_xlabel("α (probability weighting function)" )

ax.axhline(0, lw=0.75, ls="--", color="black")
ax.axvline(1, lw=0.75, ls="--", color="black")
plt.text(0.01, .75, "risk averse", rotation=90, transform=ax.transAxes,
         ha="left", va="center", fontsize="small")
plt.text(0.01, .25, "risk seeking", rotation=90, transform=ax.transAxes,
         ha="left", va="center", fontsize="small")
plt.text(0.25, 0.01, "low ++ / high --", rotation=0, transform=ax.transAxes,
         ha="center", va="bottom", fontsize="small")
plt.text(0.75, 0.01, "low -- / high ++", rotation=0, transform=ax.transAxes,
         ha="center", va="bottom", fontsize="small")

# Display agents
X, Y = agent_i[:,0], agent_i[:,1]
ax.scatter(X, Y, s=10, facecolor="black", edgecolor="None", marker="o",
           label="initial", alpha=.125)

X, Y = agent_f[:,0], agent_f[:,1]
ax.scatter(X, Y, s=30, facecolor="black", edgecolor="None", marker="o")
ax.scatter(X, Y, s=20, facecolor="white", edgecolor="None", marker="o")
ax.scatter(X, Y, s=20, facecolor=".5", edgecolor="None",
           marker="o", label="final", alpha=.25)

ax.scatter([X.mean()], [Y.mean()], s=75,
           facecolor="black", edgecolor="white", marker="X", lw=1)
X, Y = monkeys["gaindistortion"], monkeys["gainrisk_aversion"]
ax.scatter([X.mean()], [Y.mean()], s=75,
           facecolor="blue", edgecolor="white", marker="X", lw=1)
X, Y = monkeys["lossdistortion"], monkeys["lossrisk_aversion"]
ax.scatter([X.mean()], [Y.mean()], s=75,
           facecolor="red", edgecolor="white", marker="X", lw=1)
ax.legend(frameon=True, loc="upper right")
ax.set_xlim(pmin, pmax), ax.set_ylim(vmin, vmax)
ax.set_yticks([])

ax.text(-0.1, 1.1, ascii_lowercase[1],
        transform=ax.transAxes, size=20, weight='bold')

plt.tight_layout()
plt.savefig(fig_filename)
plt.show()