from sklearn.linear_model import Ridge
import sys
from social_features import get_social_features
from social_features import get_reply_only_social_features
import time
from collections import *
import re
import pickle
import numpy as np
from sklearn.preprocessing import Imputer

def main():

    input_file = sys.argv[1]
    outfile = sys.argv[2]

    MIN_FREQ = 5
    if len(sys.argv) > 3:
        MIN_FREQ = int(sys.argv[3])

    reply_only = False
    if len(sys.argv) > 4:
        reply_only = True

    max_lines = -1
    if len(sys.argv) > 5:
        max_lines = int(sys.argv[5])

    y = []
    X_dicts = []

    time_sum = 0

    with open(input_file) as f:
        for line_no, line in enumerate(f):
            if line_no == 0:
                continue

            if max_lines > 0 and line_no >= max_lines:
                break

            cols = line[:-1].split('\t')
            # All four social scores

            post = normalize(cols[4])
            reply = normalize(cols[5])

            if len(post) == 0 or len(reply) == 0:
                continue

            start = time.time()
            if not reply_only:
                feats = get_social_features(post, reply)
            else:
                feats = get_reply_only_social_features(post, reply)
            end = time.time()
            time_sum += end-start

            if len(feats) == 0:
                continue

            y.append(cols[0:4])
            X_dicts.append(feats)
            if len(X_dicts) % 50 == 0:
                print(('Saw %d instances, average time: %f' \
                      % (len(X_dicts), time_sum / len(X_dicts))))
                # break

    # Count features to remove rare ones
    feature_counts = Counter()
    for d in X_dicts:
        for k in d.keys():
            feature_counts[k] += 1

    # Convert to feature matrix and dump into a classifier
    feature_to_index = {}
    for d in X_dicts:
        for k in d.keys():
            if k not in feature_to_index and feature_counts[k] >= MIN_FREQ:
                feature_to_index[k] = len(feature_to_index)

    num_feats = len(feature_to_index)
    print(('Using %d features (out of %d)' % (num_feats, len(feature_counts))))

    with open(outfile + '.features.tsv', 'w') as outf:
        for k,v in feature_to_index.items():
            outf.write('%s\t%d\n' % (k, v))

    # TODO: Use median feature value for unknowns somehow
    print("creating X")
    X = []
    for j, d in enumerate(X_dicts):
        social = '\t'.join(y[j])

        vec = num_feats * [ float('nan') ]
        for k, v in d.items():
            # Skip rare features
            if k not in feature_to_index:
                continue
            i = feature_to_index[k]
            vec[i] = v
        X.append(np.array(vec))

    X = np.array(X)
    print(len(y), len(X))

    with open(outfile + '.zero.tsv', 'w') as outf:
        for j, row in enumerate(np.nan_to_num(X)):
            social = '\t'.join(y[j])
            outf.write("%s\t%s\n" % (social, " ".join([str(x) for x in row])))

def normalize(text):
    text = text.strip()
    if len(text) > 1:
        text = text[0].upper() + text[1:]

    text = re.sub(r'\bi\b', 'I', text)
    text = re.sub(r'\b([yY])oure\b', r"\1ou're", text)
    text = re.sub(r'\b([yY])ouve\b', r"\1ou've", text)
    text = re.sub(r'\b([wWdD])ont\b', r"\1on't", text)
    text = re.sub(r'\b([tTwW])hats\b', r"\1hat's", text)
    text = re.sub(r'\b[iI]m\b', "I'm", text)
    text = re.sub(r'\bu\b', 'you', text)

    return text

if __name__ == '__main__':
    main()


#!/usr/bin/env python
"""
This module is a Keras/Tensorflow based implementation of Pyramid Scene Parsing Networks.

Original paper & code published by Hengshuang Zhao et al. (2017)
"""
from __future__ import print_function
from __future__ import division
from os.path import splitext, join, isfile
from os import environ
from math import ceil
import argparse
import numpy as np
from scipy import misc, ndimage

from keras.applications.resnet50 import ResNet50
from keras.layers import Input, Dense, Flatten
from keras.optimizers import SGD
from keras import backend as K
from keras.models import Model, load_model
import tensorflow as tf

from utils import image_utils
from utils.data import open_file

class Discriminator(object):
    """Discriminator for classes"""

    def __init__(self, lr=1e-4, checkpoint=None):
        print("checkpoint %s" % checkpoint)
        """Instanciate a Resnet discriminator"""

        if checkpoint is None:
            print("Building Resnet discriminator")
            self.model = self.build_model(lr)
        else:
            print("Loading from checkpoint %s" % checkpoint)
            self.model = load_model(checkpoint)

        self.input_shape = (473,473)

    def build_model(self, lr):
        inp = Input((473,473,4))
        resnet = ResNet50(input_tensor=inp, weights=None, include_top=False)
        x = Flatten()(resnet.outputs[0])
        output = Dense(1, activation='sigmoid')(x)
        model = Model(inputs=inp, outputs=output)

        sgd = SGD(lr=lr, momentum=0.9, nesterov=True)
        model.compile(optimizer=sgd,
                        loss="binary_crossentropy",
                        metrics=['accuracy'])
        return model

    def predict(self, img, prediction, category):
        """
        Predict segmentation for an image.

        Arguments:
            img: must be rowsxcolsx3
            prediction: must be rowsxcolsxN-1
            category: must be 1 ... N
        """
        img_resized = misc.imresize(img, self.input_shape)
        img_preprocessed = image_utils.preprocess_image(img_resized)

        input_data = prepare_disc_data(img_preprocessed, prediction, category)
        input_data = input_data[np.newaxis, :, :, :]  # Append sample dimension for keras

        prediction = self.model.predict(input_data)[0]
        return prediction

def prepare_disc_data(img, prediction, category):
    s = prediction[category-1]
    s = image_utils.scale(s, img.shape[:2])
    s = s > 0.5
    s = s * 255
    data = np.concatenate((img, s[:,:,np.newaxis]), axis=2)
    return data

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('-i', '--input_path', type=str, default='example_images/ade20k.jpg',
                        help='Path the input image')
    parser.add_argument('-p', '--prediction', type=str, default='example_results/ade20k.h5',
                        help='Path to output')
    parser.add_argument('-c', '--category', type=int, default='1',
                        help='Model/Weights to use')
    parser.add_argument('--id', default="0")
    args = parser.parse_args()

    environ["CUDA_VISIBLE_DEVICES"] = args.id

    sess = tf.Session()
    K.set_session(sess)

    with sess.as_default():
        img = misc.imread(args.input_path)
        prediction = open_file(args.prediction, ftype="ap")
        print(args)

        disc = Discriminator()

        prob = disc.predict(img, prediction, args.category)
        print(prob)


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

import logging
import math
import numpy as np
#from PIL import ImageGrab #if windows or os X
import pyscreenshot as ImageGrab  #if linux
from PIL import Image
import cv2 #opencv
import io
import time
#from matplotlib import pyplot as plt
#plt.rcParams['figure.figsize'] = (30, 30)
#import seaborn as sns
import pandas as pd
import numpy as np
from random import randint
import os
from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.common.keys import Keys
#%matplotlib inline
from keras.models import model_from_json
from keras.models import Sequential
from keras.layers.core import Dense, Dropout, Activation, Flatten
from keras.layers.convolutional import Conv2D, MaxPooling2D
from keras.optimizers import SGD , Adam
from keras.callbacks import TensorBoard
from collections import deque
import random
import pickle
import json

game_url = "./snake.html"
chrome_driver_path = "../chromedriver.exe"
loss_file_path = "./data/loss_df.csv"
actions_file_path = "./data/actions_df.csv"
scores_file_path = "./data/scores_df.csv"
time_file_path = "./data/time_df.csv"
log_file_path = "./data/game_logs.log"

for handler in logging.root.handlers[:]:
    logging.root.removeHandler(handler)

logging.basicConfig(filename=log_file_path, level=logging.INFO)

if os.path.isfile(loss_file_path):
    loss_df = pd.read_csv(loss_file_path)
else:
    loss_df = pd.DataFrame(columns =['loss'])
    f = open(loss_file_path, "w+")
    loss_df.to_csv(loss_file_path, index=False)

if os.path.isfile(scores_file_path):
    scores_df = pd.read_csv(scores_file_path)
else:
    scores_df = pd.DataFrame(columns = ['scores'])
    open(scores_file_path, "x")
    scores_df.to_csv(scores_file_path, index=False)

if os.path.isfile(actions_file_path):
    actions_df = pd.read_csv(actions_file_path)
else:
    actions_df = pd.DataFrame(columns = ['left', 'right', 'up', 'down'])
    open(actions_file_path, "x")
    actions_df.to_csv(actions_file_path, index=False)

if os.path.isfile(time_file_path):
    time_df = pd.read_csv(time_file_path)
else:
    time_df = pd.DataFrame(columns = ['time'])
    open(time_file_path, "x")
    time_df.to_csv(time_file_path, index=False)

#sudo apt-get install chromium-chromedriver

img_rows, img_cols = 64,64
SPEED = 360
WIDTH = 8
HEIGHT = 8
ACTIONS = 4
GAMMA = 0.99
OBSERVATION = 65.
EXPLORE = 70
FINAL_EPSILON = 0.1
INITIAL_EPSILON = 1
REPLAY_MEMORY = 50000
BATCH = 64
FRAME_PER_ACTION = 1
LEARNING_RATE = 1e-4
img_channels = 4

class Game:
    def __init__(self,speed, width, height):
        self.speed = speed
        self.width = width
        self.height = height
        chrome_options = Options()
        chrome_options.add_argument("disable-infobars")
        self._driver = webdriver.Chrome("/usr/lib/chromium-browser/chromedriver")
        # self._driver = webdriver.Chrome(executable_path = "chromedriver.exe",chrome_options=chrome_options)
        self._driver.set_window_position(x=-10,y=0)
        self._driver.set_window_size(img_rows*5+50,img_cols*5+50)
        self._driver.get("file://"+os.path.abspath(game_url))
        self._driver.execute_script("Init.instance_.speed="+str(self.speed))
        print(self.width)
        self._driver.execute_script("Init.instance_.width="+str(self.width))
        self._driver.execute_script("Init.instance_.height="+str(self.height))
    def get_crashed(self):
        return self._driver.execute_script("return Init.instance_.crashed")
    def get_playing(self):
        return self._driver.execute_script("return Init.instance_.playing")
    def restart(self):
        self._driver.find_element_by_tag_name("body").send_keys(Keys.ENTER)

        time.sleep(0.25)# no actions are possible
                        # for 0.25 sec after game starts,
                        # skip learning at this time and make the model wait
    def press_enter(self):
        self._driver.find_element_by_tag_name("body").send_keys(Keys.ENTER)
    def press_up(self):
        self._driver.find_element_by_tag_name("body").send_keys(Keys.ARROW_UP)
    def press_down(self):
        self._driver.find_element_by_tag_name("body").send_keys(Keys.ARROW_DOWN)
    def press_left(self):
        self._driver.find_element_by_tag_name("body").send_keys(Keys.ARROW_LEFT)
    def press_right(self):
        self._driver.find_element_by_tag_name("body").send_keys(Keys.ARROW_RIGHT)
    def get_time(self):
        time = self._driver.execute_script("return Init.instance_.time")
        return int(time)
    def get_score(self):
        score = self._driver.execute_script("return Init.instance_.score")
        return int(score)
    def get_just_eaten(self):
        just_eaten = self._driver.execute_script("return Init.instance_.just_eaten")
        return just_eaten
    def get_distance(self):
        distance = self._driver.execute_script("return Init.instance_.distance")
        return distance
    def get_food_loc(self):
        foodx = self._driver.execute_script("return Init.instance_.foodx")
        foody = self._driver.execute_script("return Init.instance_.foody")
        return foodx, foody
    def get_head_loc(self):
        headx = self._driver.execute_script("return Init.instance_.headx")
        heady = self._driver.execute_script("return Init.instance_.heady")
        return headx, heady
    def end(self):
        self._driver.close()

class Agent:
    def __init__(self,game):
        self._game = game
        self.start()
        time.sleep(.5)
    def is_running(self):
        return self._game.get_playing()
    def is_crashed(self):
        return self._game.get_crashed()
    def is_just_eaten(self):
        return self._game.get_just_eaten()
    def up(self):
        self._game.press_up()
    def start(self):
        self._game.press_enter()
    def down(self):
        self._game.press_down()
    def left(self):
        self._game.press_left()
    def right(self):
        self._game.press_right()

class State:
    def __init__(self,agent,game):
        self._agent = agent
        self._game = game
        self._display = show_img()
        self._display.__next__()
        self.prev_dist = 0
    def get_state(self,actions):
        actions_df.loc[len(actions_df)] = list(actions)
        score = self._game.get_score()
        time = self._game.get_time()
        distance = self._game.get_distance()
        reward = 0
        is_over = False
        if actions[0] == 1:
            self._agent.left()
        elif actions[1] == 1:
            self._agent.right()
        elif actions[2] == 1:
            self._agent.up()
        elif actions[3] == 1:
            self._agent.down()
        image = grab_screen()
        self._display.send(image)
        # is_eaten = self._game.get_just_eaten()
        reward = 0.1*(self.prev_dist - distance)
        self.prev_dist = distance
        if distance == 0:
            reward = 1
        if self._agent.is_crashed():
            scores_df.loc[len(scores_df)] = score
            time_df.loc[len(time_df)] = time
            reward = -1
            self._game.restart()
            is_over = True
        reward = round(reward,4)
        # print('reward: ' +str(reward))
        return image, reward, is_over

def save_obj(obj, name ):
    with open('data/'+ name + '.pkl', 'wb') as f:
        pickle.dump(obj, f, pickle.HIGHEST_PROTOCOL)

def load_obj(name ):
    with open('data/' + name + '.pkl', 'rb') as f:
        return pickle.load(f)

def grab_screen(_driver = None):
    x = 80
    y = 240
    screen =  np.array(ImageGrab.grab(bbox=(x+10,y,x+WIDTH*13,y+HEIGHT*13))) #bbox = region of interset on the entire screen
    image = process_img(screen)
    return image

def process_img(image):
    image = cv2.resize(image, (img_rows,img_cols))
    image = cv2.Canny(image, threshold1 = 100, threshold2 = 200)
    return  image

def show_img(graphs = False):
    while True:
        screen = (yield)
        window_title = "logs" if graphs else "game_play"
        cv2.namedWindow(window_title, cv2.WINDOW_NORMAL)
        imS = cv2.resize(screen, (200, 100))
        cv2.imshow(window_title, screen)
        if (cv2.waitKey(1) & 0xFF == ord('q')):
            cv2.destroyAllWindows()
            break

def init_cache():
    save_obj(INITIAL_EPSILON,"epsilon")
    t = 0
    save_obj(t,"time")
    D = deque()
    save_obj(D,"D")

def buildmodel():
    print("Now we build the model")
    model = Sequential()
    model.add(Conv2D(32, (7, 7), strides=(4, 4), padding='same',input_shape=(img_cols,img_rows,img_channels)))  #20*40*4
    model.add(Activation('relu'))
    model.add(Conv2D(64, (5, 5), strides=(2, 2), padding='same'))
    model.add(Activation('relu'))
    model.add(Conv2D(64, (3, 3), strides=(2, 3), padding='same'))
    model.add(Activation('relu'))
    model.add(Flatten())
    model.add(Dense(512))
    model.add(Activation('relu'))
    model.add(Dense(ACTIONS))
    adam = Adam(lr=LEARNING_RATE)
    model.compile(loss='mse',optimizer=adam)
    print("We finish building the model")
    return model

# buildmodel().summary()

def trainNetwork(model,game_state,observe=False):
    last_time = time.time()
    D = deque()
    do_nothing = np.zeros(ACTIONS)
#     do_nothing[0] = 1 #0 =>do nothing, 1=>left, 2=>right, 3=>up, 4=>down

    x_t, r_0, terminal = game_state.get_state(do_nothing)
#     print('x_t: {}'.format(x_t.shape))
    s_t = np.stack((x_t, x_t, x_t, x_t), axis=2)
#     print('s_t: {}'.format(s_t.shape))
    s_t = s_t.reshape(1, s_t.shape[0], s_t.shape[1], s_t.shape[2])
#     print('s_t reshaped: {}'.format(s_t.shape))
    initial_state = s_t
    model.save_weights("model_final.h5")

    if observe :
        OBSERVE = 999999999
        epsilon = FINAL_EPSILON
        print ("Now we load weight")
        model.load_weights("model_final.h5")
        adam = Adam(lr=LEARNING_RATE)
        model.compile(loss='mse',optimizer=adam)
        print ("Weight load successfully")
    else:                       #We go to training mode
        OBSERVE = OBSERVATION
        epsilon = INITIAL_EPSILON
        model.load_weights("model_final.h5")
        adam = Adam(lr=LEARNING_RATE)
        model.compile(loss='mse',optimizer=adam)
    t=0
    t = load_obj("time") # resume from the previous time step stored in file system
    print(t)
    while (True): #endless running

        loss = 0
        Q_sa = 0
        action_index = 0
        r_t = 0
        a_t = np.zeros([ACTIONS])

        #choose an action epsilon greedy
        if t % FRAME_PER_ACTION == 0:
            if  random.random() <= epsilon:
                #print("----------Random Action----------")
                action_index = random.randrange(ACTIONS)
                a_t[action_index] = 1
            else: # predict the output
                q = model.predict(s_t)
                max_Q = np.argmax(q)
                action_index = max_Q
                a_t[action_index] = 1

        #Reduced the epsilon
        if epsilon > FINAL_EPSILON and t > OBSERVE:
            epsilon -= (INITIAL_EPSILON - FINAL_EPSILON) / EXPLORE

        #run the selected action and observed next state and reward
        x_t1, r_t, terminal = game_state.get_state(a_t)
        #print('reward: {}'.format(r_t))
        #print('loop took {} seconds'.format(time.time()-last_time))
        last_time = time.time()
        x_t1 = x_t1.reshape(1, x_t1.shape[0], x_t1.shape[1], 1)
#         print('x_t1: {}'.format(x_t1.shape))
        s_t1 = np.append(x_t1, s_t[:, :, :, :3], axis=3)
#         print('s_t1: {}'.format(s_t1))

        # store the transition in D
        D.append((s_t, action_index, r_t, s_t1, terminal))
        if len(D) > REPLAY_MEMORY:
            D.popleft()

        if t > OBSERVE:

            #sample a minibatch to train on
            print(len(D), BATCH)
            minibatch = random.sample(D, BATCH)
            inputs = np.zeros((BATCH, s_t.shape[1], s_t.shape[2], s_t.shape[3]))
            targets = np.zeros((inputs.shape[0], ACTIONS))

            #Now we do the experience replay
            for i in range(0, len(minibatch)):
                state_t = minibatch[i][0]    # 2D stack of images
                action_t = minibatch[i][1]   #This is action index
                reward_t = minibatch[i][2]   #reward at state_t due to action_t
                state_t1 = minibatch[i][3]   #next state
                terminal = minibatch[i][4]   #wheather the agent died or survided due the action

                inputs[i:i + 1] = state_t

                targets[i] = model.predict(state_t)  # predicted q values
                Q_sa = model.predict(state_t1)      # predict q values for next step

                if terminal:
                    targets[i, action_t] = reward_t
                else:
                    targets[i, action_t] = reward_t + GAMMA * np.max(Q_sa)

            loss += model.train_on_batch(inputs, targets)
            loss_df.loc[len(loss_df)] = loss
        else:
            time.sleep(0.10)
        s_t = initial_state if terminal else s_t1
        t = t + 1

        if t % 1000 == 0:
            print("TimeStep: ", t, "Save model", end='')
            model.save_weights("model_final.h5", overwrite=True)
            save_obj(D,"D")
            save_obj(t,"time")
            save_obj(epsilon,"epsilon")
            loss_df.to_csv(loss_file_path,index=False)
            #print(scores_df.tail())
            scores_df.to_csv(scores_file_path,index=False)
            time_df.to_csv(time_file_path,index=False)
            actions_df.to_csv(actions_file_path,index=False)
            with open("model.json", "w") as outfile:
                json.dump(model.to_json(), outfile)

        state = ""
        if t <= OBSERVE:
            state = "observe"
        elif t > OBSERVE and t <= OBSERVE + EXPLORE:
            state = "explore"
        else:
            state = "train"

        # logging.info('TIMESTEP ' + str(t) +
        #               '| STATE ' + str(state) +
        #               '| EPSILON ' + str(epsilon) +
        #               '| ACTION ' + str(action_index) +
        #               '| Q_MAX ' + str(np.max(Q_sa)) +
        #               '| REWARD ' + str(r_t) +
        #               '| Loss ' + str(loss)
        #             )
        # print('.', end='')
        print("TIMESTEP", t, "/ STATE", state, "/ REWARD", r_t,"/ EPSILON", epsilon, "/ ACTION", action_index, "/ Q_MAX " , np.max(Q_sa), "/ Loss ", loss)

    print("Episode finished!")
    print("************************")

def playGame(observe=False):
    game = Game(speed=SPEED, width=WIDTH, height=HEIGHT)
    snake = Agent(game)
    game_state = State(snake,game)
    model = buildmodel()
    try:
        trainNetwork(model,game_state)
    except StopIteration:
        game.end()

init_cache()
playGame(observe=False);

# supervised_frames = np.load("training_data_final_working.npy")
# frame = supervised_frames[0][0]
# action_index = supervised_frames[0][1]
# plt.imshow(frame)
# print('Action taken at this frame : Action index = {} i.e. jump'.format(str(action_index)))

# supervised_actions = []

# for frame in supervised_frames:
#     supervised_actions.append(frame[1])

# fig, axs = plt.subplots(ncols=1,nrows =2,figsize=(15,15))
# sns.distplot(supervised_actions,ax=axs[0])
# axs[1].set_title('AI gameplay distribution')
# axs[0].set_title('Human gameplay distribution')
# actions_df = pd.read_csv("./data/actions_df.csv")
# sns.distplot(actions_df,ax=axs[1])


#!/usr/bin/env python3
# Extract wav2vec (2.0) features from wav files and segments
# Natalia Tomashenko

import torch
import fairseq
import argparse
import os
import numpy as np
from fairseq.models.wav2vec.wav2vec2 import Wav2Vec2Model
from torch import nn
import torch.nn.functional as F

from kaldiio import WriteHelper, ReadHelper

class PretrainedWav2VecModel(nn.Module):
    def __init__(self, fname):
        super().__init__()

        model, cfg, task = fairseq.checkpoint_utils.load_model_ensemble_and_task([fname])
        self.normalize = ('normalize' in cfg['task']) and cfg['task']['normalize']
        model = model[0]
        model.eval()
        self.model = model

    def forward(self, x):
        with torch.no_grad():
            self.model.eval()
            if self.normalize:
                x = F.layer_norm(x, x.shape)
            return self.model.extract_features(x, padding_mask=None, mask=False)

class Prediction:

    def __init__(self, fname, gpu=0):
        self.gpu = gpu
        self.model = PretrainedWav2VecModel(fname).cuda(gpu)

    def __call__(self, x):
        x = torch.from_numpy(x).float().cuda(self.gpu)
        with torch.no_grad():
            f, padding = self.model(x.unsqueeze(0))

        return f.squeeze(0).cpu().numpy()

def write_features(model, input, output):
    os.makedirs(output, exist_ok=True)
    with ReadHelper(f'ark:extract-segments scp:{input}/wav.scp {input}/segments ark:-|') as reader:
        with WriteHelper(f'ark,scp:{output}/feats.ark,{output}/feats.scp') as writer:
            for key, (sf, wav) in reader:
                wav = wav.astype(dtype=np.float32)
                feat = model(wav)
                feat = np.repeat(feat, 2, axis=0)
                writer(key, feat)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Usage: extract_wav2vec.py <model> <input> <output>')
    parser.add_argument('model', type=str)
    parser.add_argument('input', type=str)
    parser.add_argument('output', type=str)
    args = parser.parse_args()

    print('Options:')
    print('  model: {}'.format(args.model))
    print('  input: {}'.format(args.input))
    print('  output: {}'.format(args.output))

    print('  Loading model...')
    model = Prediction(args.model)

    print('  Writing Features...')
    write_features(model, args.input, args.output)
    print('Done.')


from pathlib import Path
import os
from PIL import Image, ImageFont, ImageDraw
import numpy as np
import pandas as pd
from math import *

p = Path("resources/graphics/Pokemon/Icons")

df = pd.read_csv(Path("resources/PBS/compressed/pokemon.csv"), index_col=0)
width = 64
height = ceil(len(df) / 64)

canvas = Image.new("RGBA", (width, height), "#00000000")
draw = ImageDraw.Draw(canvas)

for i, row in df.iterrows():
    try:
        img = (
            Image.open(p / f"{row.internalname}.png")
            .convert("RGBA")
            .resize((64, 32), resample=Image.NEAREST)
            .crop((0, 0, 32, 32))
        )
        canvas.alpha_composite(img, ((i % 64) * 32, (i // 64) * 32))
    except Exception as e:
        continue

canvas.save(Path("resources/graphics/generated/battler_ldtk_list.png"))

# for pth in p.glob("*.png"):
#     img = (
#         Image.open(pth)
#         .convert("RGBA")
#         .resize((64, 32), resample=Image.NEAREST)
#         .crop((0, 0, 32, 32))
#     )


# Face comparison functions
# See: https://github.com/ageitgey/face_recognition/blob/master/face_recognition/api.py

import dlib
import face_recognition_models
import numpy as np
import warnings

predictor_5_point_model = face_recognition_models.pose_predictor_five_point_model_location()
pose_predictor_5_point = dlib.shape_predictor(predictor_5_point_model)

predictor_68_point_model = face_recognition_models.pose_predictor_model_location()
pose_predictor_68_point = dlib.shape_predictor(predictor_68_point_model)

face_recognition_model = face_recognition_models.face_recognition_model_location()
try:
    face_encoder = dlib.face_recognition_model_v1(face_recognition_model)
except Exception as e:
    print(e)
    warnings.warn("Some functionality may be limited!")
    face_encoder = None

def _css_to_rect(css):
    """
    Convert a tuple in (top, right, bottom, left) order to a dlib `rect` object
    :param css:  plain tuple representation of the rect in (top, right, bottom, left) order
    :return: a dlib `rect` object
    """
    return dlib.rectangle(css[3], css[0], css[1], css[2])

def _raw_face_landmarks(face_image, face_locations, model="large"):
    #face_locations = [_css_to_rect(face_location) for face_location in face_locations]

    pose_predictor = pose_predictor_68_point

    if model == "small":
        pose_predictor = pose_predictor_5_point

    return [pose_predictor(face_image, face_location) for face_location in face_locations]

def compare_faces(known_face_encodings, face_encoding_to_check, tolerance=0.6):
    """
    Compare a list of face encodings against a candidate encoding to see if they match.
    :param known_face_encodings: A list of known face encodings
    :param face_encoding_to_check: A single face encoding to compare against the list
    :param tolerance: How much distance between faces to consider it a match. Lower is more strict. 0.6 is typical best performance.
    :return: A list of True/False values indicating which known_face_encodings match the face encoding to check
    """
    return list(face_distance(known_face_encodings, face_encoding_to_check) <= tolerance)

def face_distance(face_encodings, face_to_compare):
    """
    Given a list of face encodings, compare them to a known face encoding and get a euclidean distance
    for each comparison face. The distance tells you how similar the faces are.
    :param faces: List of face encodings to compare
    :param face_to_compare: A face encoding to compare against
    :return: A numpy ndarray with the distance for each face in the same order as the 'faces' array
    """
    if len(face_encodings) == 0:
        return np.empty((0))

    return np.linalg.norm(face_encodings - face_to_compare, axis=1)

def face_encodings(face_image, known_face_locations, num_jitters=1):
    """
    Given an image, return the 128-dimension face encoding for each face in the image.
    :param face_image: The image that contains one or more faces
    :param known_face_locations: Optional - the bounding boxes of each face if you already know them.
    :param num_jitters: How many times to re-sample the face when calculating encoding. Higher is more accurate, but slower (i.e. 100 is 100x slower)
    :return: A list of 128-dimensional face encodings (one for each face in the image)
    """
    raw_landmarks = _raw_face_landmarks(face_image, known_face_locations, model="small")

    return [np.array(face_encoder.compute_face_descriptor(face_image, raw_landmark_set, num_jitters)) for raw_landmark_set in raw_landmarks]


import pytest
"""OLGA (Sethna et al., 2019) is a fast tool for estimating the probability of generations
(i.e. Pgen). In this example we are going to estimate Pgens for the dash.csv data.
"""
def test_pgen_mouse():
    import pandas as pd
    from tcrdist.repertoire import TCRrep
    from tcrdist.pgen import OlgaModel
    import numpy as np

    df = pd.read_csv("dash.csv")
    tr = TCRrep(cell_df = df,
                organism = 'mouse',
                chains = ['alpha','beta'],
                db_file = 'alphabeta_gammadelta_db.tsv',
                compute_distances = True)

    # Load OLGA model as a python object
    olga_beta  = OlgaModel(chain_folder = "mouse_T_beta", recomb_type="VDJ")
    olga_alpha = OlgaModel(chain_folder = "mouse_T_alpha", recomb_type="VJ")

    # An example computing a single Pgen
    olga_beta.compute_aa_cdr3_pgen(tr.clone_df['cdr3_b_aa'][0])
    olga_alpha.compute_aa_cdr3_pgen(tr.clone_df['cdr3_a_aa'][0])

    # An example computing multiple Pgens
    olga_beta.compute_aa_cdr3_pgens(tr.clone_df['cdr3_b_aa'][0:5])
    olga_alpha.compute_aa_cdr3_pgens(tr.clone_df['cdr3_a_aa'][0:5])

    # An example computing 1920 Pgens more quickly with multiple cpus
    import parmap
    tr.clone_df['pgen_cdr3_b_aa'] = \
        parmap.map(
            olga_beta.compute_aa_cdr3_pgen,
            tr.clone_df['cdr3_b_aa'],
            pm_pbar=True,
            pm_processes = 2)

    tr.clone_df['pgen_cdr3_a_aa'] = \
        parmap.map(
            olga_alpha.compute_aa_cdr3_pgen,
            tr.clone_df['cdr3_a_aa'],
            pm_pbar=True,
            pm_processes = 2)

    """
    We can do something else useful. We've tweaked the original
    generative code in OLGA, so that you can generate CDRs,
    given a specific TRV and TRJ.

    Note that unfortunately not all genes are recognized in default OLGA models,
    but many are. This gives you an idea of what you can do. Here are 10
    CDR3s generated at random given a particular V,J usage combination
    """
    np.random.seed(1)
    olga_beta.gen_cdr3s(V = 'TRBV14*01', J = 'TRBJ2-5*01', n =10)
    olga_alpha.gen_cdr3s(V ='TRAV4-3*02', J ='TRAJ31*01',  n =10)

    """
    Using this approach, we can synthesize an 100K background,
    with similar gene usage frequency to our actual repertoire.
    Note, however, that given data availability,
    this is currently likely the most reliable for human beta chain.

    After OLGA's publication, a default mouse alpha model (mouse_T_alpha)
    was added to the OLGA GitHub repository. We've included that here
    but it should be used with caution as it is missing
    a number of commonly seen V genes.
    """
    np.random.seed(1)
    tr.synthesize_vj_matched_background(chain = 'beta')
    """
              v_b_gene    j_b_gene            cdr3_b_aa        pV        pJ       pVJ   weights      source
    0        TRBV14*01  TRBJ2-3*01        CASSLASAETLYF  0.033721  0.092039  0.002989  0.065742  vj_matched
    1      TRBV13-2*01  TRBJ2-3*01    CASGDAPDRTGAETLYF  0.118785  0.092039  0.010331  0.271309  vj_matched
    2      TRBV13-3*01  TRBJ1-1*01  CASSDGFSRTGGVNTEVFF  0.074051  0.106146  0.006923  1.009124  vj_matched
    3      TRBV13-3*01  TRBJ2-1*01       CASSDVQGGAEQFF  0.074051  0.117684  0.008915  1.021244  vj_matched
    4      TRBV13-3*01  TRBJ2-7*01     CASSSGTGGYIYEQYF  0.074051  0.204898  0.015366  1.670224  vj_matched
    ...            ...         ...                  ...       ...       ...       ...       ...         ...
    99995    TRBV14*01  TRBJ2-3*01  CASSPTGGAPYASAETLYF  0.033721  0.092039  0.002989  0.065742  vj_matched
    99996    TRBV17*01  TRBJ2-5*01       CASSRDPTQDTQYF  0.028110  0.124712  0.004930  0.650360  vj_matched
    99997    TRBV14*01  TRBJ2-3*01       CASSSTGGAETLYF  0.033721  0.092039  0.002989  0.065742  vj_matched
    99998  TRBV13-1*01  TRBJ2-1*01      CASSDWGKDYAEQFF  0.106042  0.117684  0.013373  2.622194  vj_matched
    99999     TRBV4*01  TRBJ2-3*01      CASSYDRGSAETLYF  0.040749  0.092039  0.002989  0.068343  vj_matched
    """
    np.random.seed(1)
    tr.synthesize_vj_matched_background(chain = 'alpha')
    """
               v_a_gene   j_a_gene        cdr3_a_aa        pV        pJ       pVJ   weights      source
    0      TRAV12N-3*01  TRAJ34*02     CAIASNTNKVVF  0.000438  0.000088  0.000088  0.006059  vj_matched
    1       TRAV3D-3*02  TRAJ33*01  CAVSAGADSNYQLIW  0.000088  0.000088  0.000088  0.005122  vj_matched
    2        TRAV3-3*01  TRAJ27*01     CAVSTNTGKLTF  0.014029  0.042964  0.000877  0.277471  vj_matched
    3        TRAV3-3*01  TRAJ26*01    CAVSHNYAQGLTF  0.014029  0.040947  0.001052  0.009155  vj_matched
    4        TRAV3-3*01  TRAJ26*01   CAVSARNYAQGLTF  0.014029  0.040947  0.001052  0.009155  vj_matched
    ...             ...        ...              ...       ...       ...       ...       ...         ...
    99995   TRAV3D-3*02  TRAJ21*01    CAVSVSNYNVLYF  0.000088  0.039982  0.000088  0.003758  vj_matched
    99996    TRAV3-3*01  TRAJ43*01    CAVSENNNNAPRF  0.014029  0.022271  0.000526  0.071093  vj_matched
    99997   TRAV3D-3*02  TRAJ26*01    CAVSGNYAQGLTF  0.000088  0.040947  0.000088  0.000296  vj_matched
    99998    TRAV3-3*01  TRAJ26*01   CAVKGNNYAQGLTF  0.014029  0.040947  0.001052  0.009155  vj_matched
    99999   TRAV9N-2*01  TRAJ15*01      CTYQGGRALIF  0.000088  0.043840  0.000088  0.020438  vj_matched
    """

    """"
    tcrdist3's integration of Pgen estimates makes it very easy to look for
    PUBLIC clusters of TCRs (i.e. high number of neighbors) with unlikely V(D)J
    recombinations.
    """
    from tcrdist.public import _neighbors_fixed_radius
    from tcrdist.public import _K_neighbors_fixed_radius
    tr.clone_df['neighbors'] = _neighbors_fixed_radius(pwmat = tr.pw_beta, radius = 18)
    tr.clone_df['K_neighbors'] = _K_neighbors_fixed_radius(pwmat = tr.pw_beta , radius = 18)
    tr.clone_df['pgen_cdr3_b_aa_nlog10'] = tr.clone_df['pgen_cdr3_b_aa'].apply(lambda x : -1*np.log10(x))
    tr.clone_df['nsubject'] = tr.clone_df['neighbors'].apply(lambda x: len(tr.clone_df['subject'][x].unique()))
    # nsubject > 1 implies quasi-publicity
    tr.clone_df['qpublic'] = tr.clone_df['nsubject'].apply(lambda x: x >1)

    # Note one can find neighbors based on paired-chain distances.
    from tcrdist.public import _neighbors_fixed_radius
    from tcrdist.public import _K_neighbors_fixed_radius
    tr.clone_df['neighbors'] = _neighbors_fixed_radius(pwmat = tr.pw_beta + tr.pw_alpha, radius = 50)
    tr.clone_df['K_neighbors'] = _K_neighbors_fixed_radius(pwmat = tr.pw_beta + tr.pw_alpha , radius = 50)
    tr.clone_df['pgen_cdr3_b_aa_nlog10'] = tr.clone_df['pgen_cdr3_b_aa'].apply(lambda x : -1*np.log10(x))
    tr.clone_df['nsubject'] = tr.clone_df['neighbors'].apply(lambda x: len(tr.clone_df['subject'][x].unique()))
    # nsubject > 1 implies quasi-publicity)
    tr.clone_df['qpublic'] = tr.clone_df['nsubject'].apply(lambda x: x >1 )
    """
    #Code for plotting pgen vs. neighbors
    import plotnine
    from plotnine import ggplot, geom_point, aes, ylab, xlab
    (ggplot(tr.clone_df,
        aes(x= 'pgen_cdr3_b_aa_nlog10',
            y ='K_neighbors',
            color = 'nsubject')) +
        geom_point(size = .1) +
        plotnine.scales.xlim(0,20) +
        plotnine.facets.facet_wrap('epitope', scales = "free_y")+ ylab('Neighbors') + xlab('-Log10 Pgen'))
    """


"""Implementation of Forecaster for Downscale

Contain the classes compound_poisson.forecast.downscale.Forecaster

Note to future developers: only forecasting of the test set (future) has been
    implemented. Implementation of forecasting the training set should handle
    rng like time_series, a rng(s) for forecasting training set, another rng(s)
    for test set

compound_poisson.forecast.forecast_abstract.Forecaster
    <- Forecaster
        <- ForecasterGp

compound_poisson.forecast.time_series.Forecaster
    <- compound_poisson.forecast.downscale.TimeSeriesForecaster
because
compound_poisson.downscale.TimeSeriesDownscale
    <>1- compound_poisson.forecast.downscale.TimeSeriesForecaster

TODO: improve signature use of simulate_forecasts()
"""

from os import path

import numpy as np
from sklearn import gaussian_process

from compound_poisson.forecast import distribution_compare
from compound_poisson.forecast import forecast_abstract
from compound_poisson.forecast import roc
from compound_poisson.forecast import time_series

class Forecaster(forecast_abstract.Forecaster):
    """Contain Monte Carlo forecasts for Downscale

    Statistical note, method for forward simulation:
        Each location has a Forecaster object which simulates independently,
            in parallel and asynchronously. This means that each location will
            forward simulate using DIFFERENT MCMC samples. This was chosen
            because it is easier to parallise and it is faster. Effects average
            out but this means individual forecasts should not be used as each
            location would have different MCMC parameters.
        Please refer to commit 206aa73 to implement the forward forecast
            synchronously. That is: draw a MCMC sample, each location forward
                simulate using that MCMC sample. Draw another MCMC sample,
                simulate, draw, ...etc.

    Notes:
        self.data contains the test set, this includes the model fields AND the
            precipitation. This means the test set precipitation does not need
            to be passed when assessing the performance of the forecast

    Attributes:
        downscale: Downscale object, to forecast from
        data: the test set (dataset.Data object)
        forecast_array: memmap of forecasts
            dim 0: for each (unmasked) location
            dim 1: for each simulation
            dim 2: for each time point
    """

    def __init__(self, downscale, memmap_dir):
        self.downscale = downscale
        self.data = None
        super().__init__(memmap_dir)

    # override
    def make_memmap_path(self):
        super().make_memmap_path(type(self.downscale).__name__)

    # override
    # additional parameter data
    def start_forecast(self, n_simulation, data):
        """Start forecast simulations, to be called initially

        Args:
            n_simulation: number of simulations
            data: the test set (dataset.Data object)
        """
        self.data = data
        self.n_time = len(data)
        self.time_array = data.time_array
        super().start_forecast(n_simulation)

    # implemented
    def copy_to_memmap(self, memmap_to_copy):
        for i in range(len(self.forecast_array)):
            memmap_to_copy_i = memmap_to_copy[i]
            self.forecast_array[i, 0:len(memmap_to_copy_i)] = memmap_to_copy_i

    # implemented
    def simulate_forecasts(self, index_range, is_print=True):
        forecast_message = []
        for i_space, time_series_i in enumerate(
                self.downscale.generate_unmask_time_series()):

            # extract model fields for each unmasked time_series
            lat_i = time_series_i.id[0]
            long_i = time_series_i.id[1]
            x_i = self.data.get_model_field(lat_i, long_i)

            message = ForecastMessage(time_series_i,
                                      x_i,
                                      self.n_simulation,
                                      self.memmap_path,
                                      self.forecast_array.shape,
                                      i_space,
                                      is_print)
            forecast_message.append(message)

        time_series_array = self.downscale.pool.map(
            ForecastMessage.forecast, forecast_message)
        self.downscale.replace_unmask_time_series(time_series_array)

    # implemented
    def get_prob_rain(self, rain, index=None):
        """Get the probability if it will rain at least of a certian amount

        Args:
            rain: scalar, amount of rain to evaluate the probability
            index: time index (optional), otherwise, take ROC from all time
                points

        Return:
            matrix, dim 0 for each location, dim 1 for each time step
        """
        if index is None:
            index = slice(self.n_time)
        # for forecast_array...
        #     dim 0 is location
        #     dim 1 is for each simulation
        #     dim 2 is for each time point
        p_rain = np.mean(self.forecast_array[:, :, index] > rain, 1)
        return p_rain

    # override
    # to provide the shape of the memmap
    def load_memmap(self, mode):
        super().load_memmap(
            mode, (self.downscale.area_unmask, self.n_simulation, self.n_time))

    def load_locations_memmap(self, mode):
        """Call load_memmap() for each forecaster in self.downscale
        """
        for time_series_i in self.downscale.generate_unmask_time_series():
            time_series_i.forecaster.load_memmap(mode)

    def del_locations_memmap(self):
        """Call del_memmap() for each forecaster in self.downscale
        """
        for time_series_i in self.downscale.generate_unmask_time_series():
            time_series_i.forecaster.del_memmap()

    def generate_time_series_forecaster(self):
        """Generate the forecaster for every unmasked time series. Also load
            the memmap. Caution: ensure to call del_memmap() for each of the
            forecaster after use
        """
        for time_series_i in self.downscale.generate_unmask_time_series():
            forecaster = time_series_i.forecaster
            forecaster.load_memmap("r")
            yield time_series_i.forecaster

    def generate_forecaster_no_memmap(self):
        """Generate the forecaster for every unmasked time series, do not load
            memmap, used for parallel computation by delaying the calling of
            load_memap() at a later stage
        """
        for time_series_i in self.downscale.generate_unmask_time_series():
            yield time_series_i.forecaster

    # implemented
    def get_roc_curve_array(
            self, rain_warning_array, time_index=None, pool=None):
        """Get array of ROC curves

        Evaluate the ROC curve for different amounts of precipitation

        Args:
            rain_warning_array: array of amount of precipitation to be detected
            time_index: optional, a pointer (eg slice or array of indices) for
                time points to take ROC curve of
            pool: optional, used for parallel computing

        Return:
            array of roc.Roc objects which can be None if a value of
                precipitation in rain_warning_array was never observed
        """
        if time_index is None:
            time_index = slice(len(self.data))
        mask = self.data.mask
        observed_rain = self.data.rain[time_index, np.logical_not(mask)]
        # swap axes so that...
        #     dim 0: for each location
        #     dim 1: for each time point
        observed_rain = np.swapaxes(observed_rain, 0, 1)
        # when flatten, this is comparable with the return value from
        # self.get_prob_rain()
        observed_rain = observed_rain.flatten()

        roc_array = []
        # get roc curve for every rain_warning, else None if that amount of
        # rain was never observed in the test set
        for rain_warning in rain_warning_array:
            if np.any(rain_warning < observed_rain):
                p_rain = self.get_prob_rain(rain_warning, time_index).flatten()
                roc_curve = roc.Roc(rain_warning, p_rain, observed_rain, pool)
                roc_array.append(roc_curve)
            else:
                roc_array.append(None)
        return roc_array

    # implemented
    def compare_dist_with_observed(self, n_linspace=100):
        """Return an object from distribution_compare, used to compare the
            distribution of the precipitation of the forecast and the observed

        Args:
            observed_rain: numpy array of observed precipitation
            n_linspace: number of points to evaluate between 0 mm and max
                observed rain

        Return: distribution_compare.Downscale object
        """
        comparer = distribution_compare.Downscale()
        comparer.compare(self, n_linspace)
        return comparer

class ForecasterSynch(Forecaster):
    """Each forecast sample share the same MCMC sample

    Current implementation is RAM intensive. Could do with a bit of work...

    Attributes:
        rng: a single rng which draws random integer for selecting with mcmc
            sample to use
    """

    def __init__(self, downscale, memmap_dir):
        self.rng = downscale.spawn_rng(1)
        super().__init__(downscale, memmap_dir)
        for time_series_i in downscale.generate_unmask_time_series():
            # this class shall set from the mcmc manually, synchronously
            time_series_i.set_from_mcmc = False

    # override
    def simulate_forecasts(self, index_range, is_print=False):
        """Each location shares the same posterior sample
        """

        # for each forecast, all time series set mcmc sample
        for i_forecast in index_range:

            # set parameters and model fields from a mcmc sample
            # the setting of the parameters is done in the super class method
            mcmc_index = self.rng.randint(
                self.downscale.burn_in, self.downscale.n_sample)

            for time_series_i in self.downscale.generate_unmask_time_series():
                time_series_i.read_memmap()
                time_series_i.set_parameter_from_sample_i(mcmc_index)
                time_series_i.del_memmap()

            self.downscale.update_all_cp_parameters()

            # do one sample
            # hack: set the member variable n_simulation to do only one
            # forecast (could be improved here)
            self.n_simulation = i_forecast + 1
            super().simulate_forecasts([i_forecast], False)

            print("Predictive sample", i_forecast)

class ForecasterGp(Forecaster):
    """
    Do GP smoothing on the parameters for every forecast

    Take samples from the MCMC to construct a GP. For each forecast, the GP is
        sampled from and be used as parameters for the compound-Poisson model.

    Attributes:
        gp_input: array of spatial points (latitude, longitude) of the
            trained model
        gp_array: array of GaussianProcessRegressor() objects, one for each
            parameter
    """

    def __init__(self, downscale, memmap_dir, topo_key):
        """Constructor

        Args:
            downscale: parent Downscale object
            memmap_dir: location of the forecast memmap
            topo_key: array of topography keys to use as gp inputs, eg
                ["latitude", "longitude"]
        """
        super().__init__(downscale, memmap_dir)
        self.gp_input = None
        self.gp_array = []

        area_unmask = downscale.area_unmask
        # get the input variables of the GP from the topography information
        topo_dic = downscale.topography

        n_sample = 100  # number of posterior samples to use
        # each location has multiple samples of beta, fit gp onto a sample of
        # samples
        # gp_input:
        #     dim 0: for each location
        #     dim 1: for each topo key
        self.gp_input = np.zeros((area_unmask, len(topo_key)))
        # gp_output: array of numpy for each parameter (eg reg for temperature)
        # each element has dimensions, dim 0: for each location, dim 1: for
        # each gp sample
        gp_output = []
        for i_parameter in range(downscale.n_parameter):
            gp_output.append(np.zeros((area_unmask, n_sample)))

        # get topography information for each location
        for i_location, time_series_i in enumerate(
                downscale.generate_unmask_time_series()):
            coordinates = time_series_i.id
            for i_key, key in enumerate(topo_key):
                self.gp_input[i_location, i_key] = (
                    topo_dic[key][coordinates[0], coordinates[1]])
            # ensure time_series does not set from the posterior sample, these
            # are set manually from the gp
            time_series_i.set_from_mcmc = False

        # select random mcmc samples
        rng = downscale.rng
        mcmc_index_array = rng.choice(
            range(downscale.burn_in, downscale.n_sample), n_sample)

        # extract parameters to fit onto
        for i_mcmc, mcmc_index in enumerate(mcmc_index_array):
            # set mcmc sample
            for time_series_i in downscale.generate_unmask_time_series():
                time_series_i.read_memmap()
                time_series_i.set_parameter_from_sample_i(mcmc_index)
                time_series_i.del_memmap()
            # extract parameter and save it to gp_output
            parameter_vector = downscale.get_parameter_vector().copy()
            for i_parameter in range(downscale.n_parameter):
                slice_index = slice(
                    i_parameter*area_unmask, (i_parameter+1)*area_unmask)
                parameter_i = parameter_vector[slice_index]
                gp_output[i_parameter][:, i_mcmc] = parameter_i

        # fit gp for each parameter
        for i_parameter in range(downscale.n_parameter):
            gp = GaussianProcessRegressor()
            gp.fit(self.gp_input, gp_output[i_parameter])
            # do not save x_train, this is set in forecasting, prevents saving
            # duplicates of this variable onto disk
            gp.delete_x_train()
            self.gp_array.append(gp)

    # override
    def simulate_forecasts(self, index_range):
        # loop: get mcmc sample, forecast, ..etc
        # the forecast is done in parallel
        area_unmask = self.downscale.area_unmask
        n_parameter = self.downscale.n_parameter

        # set x_train for each gp, they all point to self.gp_input, prevents
        # deep copies
        for gp_i in self.gp_array:
            gp_i.set_x_train(self.gp_input)

        for i_forecast in index_range:

            # get the parameter and smooth it using GP
            parameter_vector = np.zeros(area_unmask * n_parameter)
            for i_parameter in range(n_parameter):
                slice_index = slice(
                    i_parameter*area_unmask, (i_parameter+1)*area_unmask)
                parameter_i = self.gp_array[i_parameter].sample_y_at_train(
                    self.downscale.rng)
                parameter_vector[slice_index] = parameter_i.flatten()

            # set the smoothed parameter
            self.downscale.set_parameter_vector(parameter_vector)
            self.downscale.update_all_cp_parameters()
            # do one sample
            #
            # hack: set the member variable n_simulation to do only one
            # forecast (could be improved here)
            self.n_simulation = i_forecast + 1
            super().simulate_forecasts([i_forecast], False)

            print("Predictive sample", i_forecast)

        # remove x_train so that duplicates are not saved
        for gp_i in self.gp_array:
            gp_i.delete_x_train()

class GaussianProcessRegressor(gaussian_process.GaussianProcessRegressor):
    """Modification of sklearn.gaussian_process.GaussianProcessRegressor

    Custom modification of sklearn.gaussian_process.GaussianProcessRegressor to
        accept multiple independent samples of the output for a given input.
        In other words, multiple observations of the response for a given
        explanatory variable. It is designed to reduce the size of the kernel
        matrix compared to stacking multiple observations.
    How to use: pass a design matrix X (n_sample x n_features) and a response
        design matrix (n_sample x n_observations) to the method fit(). Methods
        as such predict() and sample_y() now return a response vector, merging
        the n_targets (using notation in the orginial source in the method
        fit()) together.

    Modifications:
        - The method fit() uses the parameters differently
        - The method fit() modifies the member variables further
        - The member variable copy_X_train = False by default
        - Aims to elimate the member variable y_train_ because it appears it is
            not used anywhere in prediction
        - Added new methods

    Unmodified copyright notices from the originial source:
        Authors: Jan Hendrik Metzen <jhm@informatik.uni-bremen.de>
        Modified by: Pete Green <p.l.green@liverpool.ac.uk>
        License: BSD 3 clause
    See https://scikit-learn.org/stable/modules/generated/sklearn.gaussian_process.GaussianProcessRegressor.html#sklearn.gaussian_process.GaussianProcessRegressor
        for further information on the orginial source
    """

    def __init__(self):
        super().__init__()
        self.copy_X_train = False

    # override
    def fit(self, X, y):
        super().fit(X, y)
        self.alpha_ = np.mean(self.alpha_, axis=1)
        self.y_train_ = None

    def delete_x_train(self):
        """Set the member variable X_train_ to be none, this prevents saving
             duplicates of this member variable (there are multiple gps using
             the same trainig set) to disk. Designed as a precaution.
        """
        self.X_train_ = None

    def set_x_train(self, x_train):
        """Set the member variable X_train_ to something after calling
            delete_x_train(). X_train_ has dimensions, dim 0: each location,
            dim 1: each feature
        """
        self.X_train_ = x_train

    def sample_y_at_train(self, rng):
        """Sample the GP at the training set
        """
        if self.X_train_ is None:
            raise Exception("Must set X_train_")
        return super().sample_y(self.X_train_, random_state=rng)

class TimeSeriesForecaster(time_series.Forecaster):
    """Used by TimeSeriesDownscale class

    Extended to handle MCMC samples in memmaps. The member variable
        forecast_array is shared with all spatial points and is handled by the
        corresponding Downscale object.

    Attributes:
        i_space: pointer for space, or the 0th dimension for the forecast_array
        mememap_path: location of the forecast_array
        memmap_shape: shape of the forecast_array
    """

    def __init__(self, time_series, memmap_path, i_space):
        super().__init__(time_series, path.dirname(memmap_path))
        self.i_space = i_space
        self.memmap_path = memmap_path
        self.memmap_shape = None

    # override
    def start_forecast(self, n_simulation, model_field, memmap_shape):
        """Start forecast simulations, to be called initially

        Override as a memmap does not need to be created to store the
            forecasts, this has already been done by the corresponding
            Downscale object.

        Args:
            n_simulation: number of simulations
            model_field: model fields for test set
            memmap_shape: shape of the forecast_array
        """
        self.memmap_shape = memmap_shape
        super().start_forecast(n_simulation, model_field)

    # override
    def make_memmap_path(self):
        """Do nothing, memmap_path has already been provided
        """
        pass

    # override
    def simulate_forecasts(self, index_range):
        # do not print progress
        super().simulate_forecasts(index_range, False)

    # override
    def resume_forecast(self, n_simulation, memmap_shape):
        """Simulate more forecasts

        Override as the handling of forecast_array is done already by the
            corresponding Downscale object.

        Args:
            n_simulation: total amount of simulations, ie should be higher than
                previous
            memmap_shape: shape of the forecast_array
        """
        if n_simulation > self.n_simulation:
            self.memmap_shape = memmap_shape
            n_simulation_old = self.n_simulation
            self.n_simulation = n_simulation
            self.load_memmap("r+")
            # False in argument to not print progress
            self.simulate_forecasts(range(n_simulation_old, self.n_simulation))
            self.del_memmap()

    # override
    def load_memmap(self, mode):
        """Load the memmap file for forecast_array

        Override to use the forcast_array provided by Downscale. This is shared
            with all spatial points so extract the corresponding slice.

        Args:
            mode: not used, force "r+", prevent a "w+" because the memmap is
                already created
        """
        mode = "r+"
        super().load_memmap(mode, self.memmap_shape)
        self.forecast_array = self.forecast_array[self.i_space]

class ForecastMessage(object):
    """Message to forecast all spatial points in parallel
    """

    def __init__(self,
                 time_series,
                 model_field,
                 n_simulation,
                 memmap_path,
                 memmap_shape,
                 i_space,
                 is_print):
        self.time_series = time_series
        self.model_field = model_field
        self.n_simulation = n_simulation
        self.memmap_path = memmap_path
        self.memmap_shape = memmap_shape
        self.i_space = i_space
        self.is_print = is_print

    def forecast(self):
        self.time_series.forecast(
            self.model_field, self.n_simulation,
            self.memmap_path, self.memmap_shape, self.i_space)
        if self.is_print:
            print("Predictive location", self.i_space)
        return self.time_series