'''
@Author: ConghaoWong
@Date: 2019-12-20 09:39:34
LastEditors: Conghao Wong
LastEditTime: 2020-09-16 16:27:24
@Description: classes and methods of training model
'''
import os
import random

import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
from tensorflow import keras
from tqdm import tqdm

from GridRefine import SocialRefine_one
from helpmethods import calculate_ADE_FDE_numpy, dir_check, list2array
from sceneFeature import TrajectoryMapManager
from visual import TrajVisual

class Base_Model():
    """
    Base model for prediction.

    Following items should be given when using this model:
    ```
    self.create_model(self), # create prediction model
    self.loss(self, model_output, gt, obs='null'),  # loss function when training model
    self.loss_eval(self, model_output, gt, obs='null'), # loss function when test model
    self.forward_train(self, mode_inputs), # model result in training steps
    self.forward_test(self, test_tensor:list). # model result in test steps
    ```
    """
    def __init__(self, train_info, args):
        self.args = args
        self.train_info = train_info

    def run_commands(self):
        self.get_data()     # 获取与训练数据有关的信息

        if self.args.load == 'null':
            self.model, self.optimizer = self.create_model()
            self.model.summary()
            self.train()
        else:
            self.model, self.agents_test = self.load_from_checkpoint()
            self.model.summary()

            if self.args.test:
                self.test_batch(
                    self.agents_test,
                    test_on_neighbors=False,
                    batch_size=0.2,         # set 0.5 on toy exp
                    social_refine=self.args.sr_enable,
                    draw_results=self.args.draw_results,
                    save_agents=False,
                )

    def get_data(self):
        self.obs_frames = self.args.obs_frames
        self.pred_frames = self.args.pred_frames
        self.total_frames = self.obs_frames + self.pred_frames
        self.log_dir = dir_check(self.args.log_dir)

        if not self.args.load == 'null':
            return

        self.agents_train = self.train_info['train_data']
        self.agents_test = self.train_info['test_data']
        self.train_number = self.train_info['train_number']
        self.sample_time = self.train_info['sample_time']

    def load_from_checkpoint(self):
        base_path = self.args.load + '{}'
        if self.args.save_best:
            best_epoch = np.loadtxt(os.path.join(self.args.log_dir, 'best_ade_epoch.txt'))[1].astype(int)
            model = keras.models.load_model(base_path.format('_epoch{}.h5'.format(best_epoch)))
        else:
            model = keras.models.load_model(base_path.format('.h5'))

        agents_test = np.load(base_path.format('test.npy'), allow_pickle=True)
        return model, agents_test

    def create_model(self):
        raise 'MODEL is not defined!'
        return model, optimizer

    def loss(self, model_output, gt, obs='null'):
        """
        Train loss, using ADE by default
        """
        self.loss_namelist = ['ADE_t']
        loss_ADE = calculate_ADE(model_output[0], gt)
        loss_list = tf.stack([loss_ADE])
        return loss_ADE, loss_list

    def loss_eval(self, model_output, gt, obs='null'):
        """
        Eval metrics, using ADE and FDE by default.
        return: `np.array`
        """
        self.loss_eval_namelist = ['ADE', 'FDE']
        return calculate_ADE(model_output[0], gt).numpy(), calculate_FDE(model_output[0], gt).numpy()

    def prepare_model_inputs_all(self, input_agents):
        model_inputs = []
        gt = []
        agent_index = []
        for agent_index_current, agent in enumerate(tqdm(input_agents, desc='Prepare inputs...')):
            model_inputs.append(agent.get_train_traj())
            gt.append(agent.get_gt_traj())
            agent_index.append(agent_index_current)

        model_inputs = tf.cast(tf.stack(model_inputs), tf.float32)
        gt = tf.cast(tf.stack(gt), tf.float32)
        return [model_inputs, gt], agent_index

    def prepare_model_inputs_batch(self, train_tensor=0, batch_size=0, init=False):
        """
        Get batch data from all data
        """
        if init:
            self.batch_start = 0
            self.train_length = len(train_tensor[1])
            return self.train_length

        start = self.batch_start
        end = (self.batch_start + batch_size) % self.train_length
        # 每次最多取 1 epoch
        if end < start:
            if type(train_tensor[0]) == list:
                train_inputs = [
                    tf.concat([
                        train_input[start:],
                        train_input[:end],
                    ], axis=0) for train_input in train_tensor[0]
                ]
            else:
                train_inputs = tf.concat([
                    train_tensor[0][start:],
                    train_tensor[0][:end],
                ], axis=0)

            gt = tf.concat([
                train_tensor[1][start:],
                train_tensor[1][:end],
            ], axis=0)

        elif start + batch_size < self.train_length:
            if type(train_tensor[0]) == list:
                train_inputs = [train_input[start:end] for train_input in train_tensor[0]]
            else:
                train_inputs = train_tensor[0][start:end]
            gt = train_tensor[1][start:end]

        else:
            train_inputs = train_tensor[0]
            gt = train_tensor[1]

        self.batch_start = end
        return train_inputs, gt, len(gt)

    def forward_train(self, model_inputs):
        """
        Run a training implement
        """
        output = self.model(model_inputs)
        if not type(output) == list:
            output = [output]
        return output

    def forward_test(self, test_tensor:list):
        """
        Run test once.
        `test_tensor` is a `list`. `test_tensor[0]` is the inputs of model and `test_tensor[1]` are their grount truths.
        """
        model_inputs = test_tensor[0]
        gt = test_tensor[1]
        output = self.model(model_inputs)
        if not type(output) == list:
            output = [output]
        return output, gt, model_inputs

    def test_during_training(self, test_tensor, input_agents, test_index):
        """
        Run test during training.
        Results will NOT be written to inputs.
        """
        model_output, gt, obs = self.forward_test(test_tensor)
        loss_eval = self.loss_eval(model_output, gt, obs=obs)
        return model_output, loss_eval, gt, input_agents

    def train(self):
        """
        Train the built model `self.model`
        """
        batch_number = int(np.ceil(self.train_number / self.args.batch_size))
        summary_writer = tf.summary.create_file_writer(self.args.log_dir)

        print('\n-----------------dataset options-----------------')
        if self.args.train_percent[0] and self.args.train_type == 'all':
            print('Sampling data from training sets. ({}x)'.format(self.args.train_percent))
        if self.args.reverse:
            print('Using reverse data to train. (2x)')
        if self.args.add_noise:
            print('Using noise data to train. ({}x)'.format(self.args.add_noise))
        if self.args.rotate:
            print('Using rotate data to train. ({}x)'.format(self.args.rotate))
        print('train_number = {}, total {}x train samples.'.format(self.train_number, self.sample_time))

        print('-----------------training options-----------------')
        print('model_name = {}, \ndataset = {},\nbatch_number = {},\nbatch_size = {},\nlr={}'.format(
            self.args.model_name,
            self.args.test_set,
            batch_number,
            self.args.batch_size,
            self.args.lr,
        ))

        print('\nPrepare training data...')
        self.train_tensor, self.train_index = self.prepare_model_inputs_all(self.agents_train)
        self.test_tensor, self.test_index = self.prepare_model_inputs_all(self.agents_test)
        train_length = self.prepare_model_inputs_batch(self.train_tensor, init=True)

        if self.args.save_model:
            self.test_data_save_path = os.path.join(self.args.log_dir, '{}.npy'.format(self.args.model_name + '{}'))
            np.save(self.test_data_save_path.format('test'), self.agents_test)
            np.save(self.test_data_save_path.format('args'), self.args)

        test_results = []
        test_loss_dict = dict()
        test_loss_dict['-'] = 0

        batch_number = 1 + (train_length * self.args.epochs)// self.args.batch_size
        print(batch_number, train_length, self.args.epochs, self.args.batch_size)

        time_bar = tqdm(range(batch_number), desc='Training...')
        best_ade = 100.0
        best_epoch = 0
        for batch in time_bar:
            ADE = 0
            ADE_move_average = tf.cast(0.0, dtype=tf.float32)    # 计算移动平均
            loss_list = []

            obs_current, gt_current, train_sample_number = self.prepare_model_inputs_batch(self.train_tensor, self.args.batch_size)

            if train_sample_number < 20:
                continue

            with tf.GradientTape() as tape:
                model_output_current = self.forward_train(obs_current)
                loss_ADE, loss_list_current = self.loss(model_output_current, gt_current, obs=obs_current)
                ADE_move_average = 0.7 * loss_ADE + 0.3 * ADE_move_average

            ADE += loss_ADE
            grads = tape.gradient(ADE_move_average, self.model.trainable_variables)
            self.optimizer.apply_gradients(zip(grads, self.model.trainable_variables))

            loss_list.append(loss_list_current)
            loss_list = tf.reduce_mean(tf.stack(loss_list), axis=0).numpy()

            epoch = (batch * self.args.batch_size) // train_length

            if (epoch >= self.args.start_test_percent * self.args.epochs) and (epoch % self.args.test_step == 0):
                model_output, loss_eval, _, _ = self.test_during_training(self.test_tensor, self.agents_test, self.test_index)
                test_results.append(loss_eval)
                test_loss_dict = create_loss_dict(loss_eval, self.loss_eval_namelist)
                ade_current = loss_eval[0]
                if ade_current <= best_ade:
                    best_ade = ade_current
                    best_epoch = epoch

                    if self.args.save_best:
                        self.model.save(os.path.join(self.args.log_dir, '{}_epoch{}.h5'.format(self.args.model_name, epoch)))
                        np.savetxt(os.path.join(self.args.log_dir, 'best_ade_epoch.txt'), np.array([best_ade, best_epoch]))

            if epoch % 2 == 0:
                train_loss_dict = create_loss_dict(loss_list, self.loss_namelist)
                loss_dict = dict(train_loss_dict, **test_loss_dict) # 拼接字典
                time_bar.set_postfix(loss_dict)

            with summary_writer.as_default():
                for loss_name in loss_dict:
                    value = loss_dict[loss_name]
                    tf.summary.scalar(loss_name, value, step=epoch)

        print('Training done.')
        print('Tensorboard training log file is saved at "{}"'.format(self.args.log_dir))
        print('To open this log file, please use "tensorboard --logdir {} --port 54393"'.format(self.args.log_dir))

        latest_epochs = 10
        test_results = list2array(test_results)
        latest_results = np.mean(test_results[-latest_epochs-1:-1, :], axis=0)
        print('In latest {} test epochs, average test loss = {}'.format(
            latest_epochs,
            latest_results
        ))
        np.savetxt(os.path.join(self.args.log_dir, 'train_log.txt'), list2array(test_results))

        if self.args.save_model:
            self.model_save_path = os.path.join(self.args.log_dir, '{}.h5'.format(self.args.model_name))
            self.model.save(self.model_save_path)

            print('Trained model is saved at "{}".'.format(self.model_save_path.split('.h5')[0]))
            print('To re-test this model, please use "python main.py --load {}".'.format(self.model_save_path.split('.h5')[0]))

            model_name = self.model_save_path.split('.h5')[0].split('/')[-1]
            np.savetxt('./results/result-{}{}.txt'.format(model_name, self.args.test_set), latest_results)
            with open('./results/path-{}{}.txt'.format(model_name, self.args.test_set), 'w+') as f:
                f.write(self.model_save_path.split('.h5')[0])

    def test_batch(self, agents_test, test_on_neighbors=False, draw_results=True, batch_size=0.2, save_agents=False, social_refine=False):
        """
        Eval model on test sets.
        Results WILL be written to inputs.
        测试可以分段进行，并使用`batch_size`以百分比形式调节时间段长短;
        `test_on_neighbors`将会被自动打开当`social_refine == True`
        """
        print('-----------------Test options-----------------')
        print('model_name = {},\ndataset = {},\ntest_length= {} * length of test video.\n'.format(
            self.args.model_name,
            self.args.test_set,
            batch_size,
        ))

        start_frame = agents_test[0].obs_frame
        end_frame = agents_test[-1].obs_frame
        frame_length = end_frame - start_frame

        # sort by obs time
        agents_batch = dict()
        for agent in agents_test:
            batch_index = min(int((agent.obs_frame - start_frame)/(batch_size * frame_length)), int(1/batch_size)-1)
            if not batch_index in agents_batch:
                agents_batch[batch_index] = []
            else:
                agents_batch[batch_index].append(agent)

        if social_refine:
            test_on_neighbors = True

        agents_batch, test_index = self.prepare_test_agents_batch(agents_batch, test_on_neighbors)

        # run test
        all_loss = []
        all_loss_batch = []
        for batch_index in agents_batch:
            batch_loss = []
            [test_tensor, _], _ = self.prepare_model_inputs_all(agents_batch[batch_index], calculate_neighbor=test_on_neighbors)
            pred = self.forward_train(test_tensor)
            pred = pred[0].numpy()

            for agent_index, index in enumerate(test_index[batch_index]):
                current_pred = pred[index]
                agents_batch[batch_index][agent_index].write_pred(current_pred[0])
                if test_on_neighbors:
                    agents_batch[batch_index][agent_index].write_pred_neighbor(current_pred[1:])

                if social_refine:
                    agents_batch[batch_index][agent_index].write_pred_sr(SocialRefine_one(
                        agent=agents_batch[batch_index][agent_index],
                        args=self.args,
                        epochs=10,
                        save=False,
                    ))

                loss = agents_batch[batch_index][agent_index].calculate_loss(SR=social_refine)
                all_loss.append(loss)
                batch_loss.append(loss)

            all_loss_batch.append(np.mean(np.stack(batch_loss), axis=0))

        average_loss = np.mean(np.stack(all_loss), axis=0)
        print('test_loss={}\nTest done.'.format(create_loss_dict(average_loss, ['ADE', 'FDE'])))
        # print(all_loss_batch)

        if draw_results:
            result_agents = []
            for batch_index in agents_batch:
                result_agents += agents_batch[batch_index]

            # draw results only
            for index in range(len(result_agents)):
                result_agents[index].draw_results(self.log_dir, '{}.png'.format(index), draw_neighbors=False)

            # draw results on video frames
            # tv = TrajVisual(save_base_path=self.args.log_dir, verbose=True, draw_neighbors=False, social_refine=social_refine)
            # tv.visual(result_agents, dataset=self.args.test_set)

        if save_agents:
            result_agents = []
            for batch_index in agents_batch:
                result_agents += agents_batch[batch_index]
            np.save(os.path.join(self.log_dir, 'pred.npy'), result_agents)
            return result_agents

    def test(self, agents_test, test_on_neighbors=False, social_refine=True, draw_results=True, batch_size=0.2, save_agents=False):
        """
        Eval model on test sets.
        Results WILL be written to inputs.
        """
        all_loss = []
        loss_name_list = ['ADE', 'FDE']
        loss_function = calculate_ADE_FDE_numpy

        self.test_tensor, self.test_index = self.prepare_model_inputs_all(self.agents_test)
        pred = self.forward_train(self.test_tensor)

        for index in tqdm(range(len(agents_test)), desc='Testing...'):
            obs = agents_test[index].get_train_traj().reshape([1, agents_test[index].obs_length, 2])

            # if test_on_neighbors and agents_test[index].neighbor_number > 0:
            #     obs_neighbor = (np.stack(agents_test[index].get_neighbor_traj())).reshape([agents_test[index].neighbor_number, agents_test[index].obs_length, 2])
            #     obs = np.concatenate([obs, obs_neighbor], axis=0)

            agents_test[index].write_pred(pred[0].numpy()[index])
            # if test_on_neighbors:
            #     agents_test[index].write_pred_neighbor(pred[1:].numpy()[index])

            # if social_refine:
            #     agents_test[index].write_pred_sr(SocialRefine_one(agents_test[index], self.args_old))

            if draw_results:
                agents_test[index].draw_results(self.log_dir, '{}.png'.format(index), draw_neighbors=False # test_on_neighbors
                )

            all_loss.append(agents_test[index].calculate_loss())

        loss = np.mean(np.stack(all_loss), axis=0)

        print('test_loss={}'.format(create_loss_dict(loss, loss_name_list)))
        # for l in loss:
        #     print(loss, end='\t')
        print('\nTest done.')

        if save_agents:
            np.save(os.path.join(self.log_dir, 'pred.npy'), agents_test)
        return agents_test

    def prepare_test_agents_batch(self, agents_batch:dict, test_on_neighbors=False):
        """
        Prepare test agents and save test order. (When test on neighbors of current agent)
        returns: Test agents (in batch order) `agents_batch` and their order `test_index`.
        """
        # save test order
        test_index = dict()
        for batch_index in agents_batch:
            total_count = 0
            if not batch_index in test_index:
                test_index[batch_index] = []

            for agent_index, _ in enumerate(agents_batch[batch_index]):
                start_count = total_count
                total_count += 1
                if test_on_neighbors:
                    nei_len = agents_batch[batch_index][agent_index].neighbor_number
                    total_count += nei_len
                test_index[batch_index].append([i for i in range(start_count, total_count)])

        return agents_batch, test_index

class BGM(Base_Model):
    """
    `B`uilding a Dynamic `G`uidance `M`ap for Trajectory Prediction
    """
    def __init__(self, train_info, args):
        super().__init__(train_info, args)
        self.given_maps_when_test=False

    def create_model(self):
        positions = keras.layers.Input(shape=[self.obs_frames, 2])
        traj_maps = keras.layers.Input(shape=[self.args.gridmapsize, self.args.gridmapsize])
        start_point = tf.reshape(positions[:, -1, :], [-1, 1, 2])

        # sequence feature
        positions_n = positions - start_point
        positions_embadding_lstm = keras.layers.Dense(64)(positions_n)
        traj_feature = keras.layers.LSTM(64, return_sequences=True)(positions_embadding_lstm)
        feature_flatten = tf.reshape(traj_feature, [-1, self.obs_frames * 64])
        sequence_feature = keras.layers.Dense(self.obs_frames * 32, activation=tf.nn.tanh)(feature_flatten)

        # context feature
        traj_maps_r = tf.reshape(traj_maps, [-1, self.args.gridmapsize, self.args.gridmapsize, 1])
        average_pooling = keras.layers.AveragePooling2D([2, 2], padding='same')(traj_maps_r)
        cnn1 = keras.layers.Conv2D(32, [8, 8], activation=tf.nn.relu)(average_pooling)
        cnn2 = keras.layers.Conv2D(32, [5, 5], activation=tf.nn.relu)(cnn1)
        pooling2 = keras.layers.AveragePooling2D([2, 2])(cnn2)
        flatten = keras.layers.Flatten()(pooling2)
        context_feature = keras.layers.Dense(self.obs_frames * 32, activation=tf.nn.tanh)(flatten)

        # joint feature
        concat_feature = tf.concat([sequence_feature, context_feature], axis=-1)
        feature_fc = keras.layers.Dense(self.pred_frames * 64)(concat_feature)
        feature_reshape = tf.reshape(feature_fc, [-1, self.pred_frames, 64])
        output5 = keras.layers.Dense(2)(feature_reshape)
        output5 = output5 + start_point

        lstm = keras.Model(inputs=[positions, traj_maps], outputs=[output5])
        lstm.build(input_shape=[None, self.obs_frames, 2])
        lstm_optimizer = keras.optimizers.Adam(lr=self.args.lr)

        return lstm, lstm_optimizer

    def get_feature(self, inputs, layer_name):
        submodel = keras.Model(inputs=self.model.input, outputs=self.model.get_layer(layer_name).output)
        return submodel(inputs)

    def prepare_model_inputs_all(self, input_agents, calculate_neighbor=False):
        input_trajs = []
        input_maps = []
        gt = []
        agent_index = []
        for agent_index_current, agent in enumerate(tqdm(input_agents, desc='Prepare inputs...')):
            input_trajs.append(agent.get_train_traj())
            input_maps.append(agent.get_traj_map())
            gt.append(agent.get_gt_traj())
            agent_index.append(agent_index_current)

            if calculate_neighbor and agent.neighbor_number:
                for traj, traj_map in zip(agent.get_neighbor_traj(), agent.get_traj_map_for_neighbors()):
                    input_trajs.append(traj)
                    input_maps.append(agent.get_traj_map())
                    # No GT

        input_trajs = tf.cast(tf.stack(input_trajs), tf.float32)
        input_maps = tf.cast(tf.stack(input_maps), tf.float32)
        gt = tf.cast(tf.stack(gt), tf.float32)
        return [[input_trajs, input_maps], gt], agent_index

    def prepare_test_agents_batch(self, agents_batch:dict, test_on_neighbors=False):
        # create trajectory map for each batch
        if not type(self.given_maps_when_test) == np.ndarray:
            traj_maps = [TrajectoryMapManager(agents_batch[batch_index]) for batch_index in agents_batch]
        else:
            traj_maps = self.given_maps_when_test
            print('Using given maps')

        # write traj map and save batch order
        test_index = dict()
        for batch_index, traj_map in zip(agents_batch, traj_maps):
            total_count = 0
            if not batch_index in test_index:
                test_index[batch_index] = []

            for agent_index, _ in enumerate(agents_batch[batch_index]):
                agents_batch[batch_index][agent_index].write_traj_map(traj_map)
                start_count = total_count
                total_count += 1
                if test_on_neighbors:
                    agents_batch[batch_index][agent_index].write_traj_map_for_neighbors(traj_map)
                    nei_len = agents_batch[batch_index][agent_index].neighbor_number
                    total_count += nei_len
                test_index[batch_index].append([i for i in range(start_count, total_count)])

        return agents_batch, test_index

class Linear(Base_Model):
    def __init__(self, train_info, args):
        super().__init__(train_info, args)
        self.args.batch_size = 1
        self.args.epochs = 1
        self.args.draw_results = False
        self.args.train_percent = 0.0

    def run_commands(self):
        self.get_data()
        self.model, self.optimizer = self.create_model()
        self.test_batch(
            self.agents_test,
            batch_size=1.0,
            test_on_neighbors=False,
            social_refine=False,
            draw_results=self.args.draw_results,
            save_agents=False
        )

    def predict_linear(self, x, y, x_p, diff_weights=0):
        if diff_weights == 0:
            P = np.diag(np.ones(shape=[x.shape[0]]))
        else:
            P = np.diag(softmax([(i+1)**diff_weights for i in range(x.shape[0])]))

        A = tf.transpose(tf.stack([np.ones_like(x), x]))
        A_p = tf.transpose(tf.stack([np.ones_like(x_p), x_p]))
        Y = tf.transpose(y)

        P = tf.cast(P, tf.float32)
        A = tf.cast(A, tf.float32)
        A_p = tf.cast(A_p, tf.float32)
        Y = tf.reshape(tf.cast(Y, tf.float32), [-1, 1])

        B = tf.matmul(tf.matmul(tf.matmul(tf.linalg.inv(tf.matmul(tf.matmul(tf.transpose(A), P), A)), tf.transpose(A)), P), Y)
        Y_p = np.matmul(A_p, B)
        return Y_p, B

    def predict_linear_for_person(self, positions, diff_weights):
        t = np.array([t for t in range(self.obs_frames)])
        t_p = np.array([t + self.obs_frames for t in range(self.pred_frames)])
        x = tf.transpose(positions)[0]
        y = tf.transpose(positions)[1]

        x_p, _ = self.predict_linear(t, x, t_p, diff_weights=diff_weights)
        y_p, _ = self.predict_linear(t, y, t_p, diff_weights=diff_weights)

        return tf.transpose(tf.reshape(tf.stack([x_p, y_p]), [2, self.pred_frames]))

    def create_model(self):
        return self.predict_linear_for_person, 0

    def forward_train(self, train_tensor, index):
        input_trajs = train_tensor[0][index[0]:index[1]]
        gt = train_tensor[1][index[0]:index[1]]

        results = []
        for inputs_current in input_trajs:
            results.append(self.model(inputs_current, diff_weights=self.args.diff_weights))

        return [tf.stack(results)], gt, input_trajs

    def forward_test(self, test_tensor):
        input_trajs = test_tensor[0]
        gt = test_tensor[1]

        results = []
        for inputs_current in input_trajs:
            results.append(self.model(inputs_current, diff_weights=self.args.diff_weights))

        return output, gt, input_trajs

"""
helpmethods
"""

def create_loss_dict(loss, name_list):
    return dict(zip(name_list, loss))

def softmax(x):
    return np.exp(x)/np.sum(np.exp(x),axis=0)

def calculate_ADE(pred, GT):
    """input_shape = [batch, pred_frames, 2]"""
    pred = tf.cast(pred, tf.float32)
    GT = tf.cast(GT, tf.float32)
    return tf.reduce_mean(tf.linalg.norm(pred - GT, ord=2, axis=2))

def calculate_FDE(pred, GT):
    pred = tf.cast(pred, tf.float32)
    GT = tf.cast(GT, tf.float32)
    return tf.reduce_mean(tf.linalg.norm(pred[:, -1, :] - GT[:, -1, :], ord=2, axis=1))


import os
import pickle

import numpy as np

from astropy.io import fits
from astropy.table import Table
from astropy.wcs import WCS

from .config_file import get_values_from_config_file
from .gausspy_py3.gp_plus import get_fully_blended_gaussians
from .utils.fit_quality_checks import negative_residuals
from .utils.gaussian_functions import gaussian, combined_gaussian, area_of_gaussian
from .utils.output import say
from .utils.spectral_cube_functions import correct_header, change_header, update_header, save_fits, return_hdu_options
from .spatial_fitting import SpatialFitting

class Finalize(object):
    def __init__(self, path_to_pickle_file=None,
                 path_to_decomp_file=None, fin_filename=None,
                 config_file=''):
        """Class containing methods to finalize the GaussPy+ results.

        Parameters
        ----------
        path_to_pickle_file : type
            Description of parameter `path_to_pickle_file`.
        path_to_decomp_file : type
            Description of parameter `path_to_decomp_file`.
        fin_filename : type
            Description of parameter `fin_filename`.
        config_file : type
            Description of parameter `config_file`.

        """
        self.path_to_pickle_file = path_to_pickle_file
        self.path_to_decomp_file = path_to_decomp_file
        self.dirpath_gpy = None
        self.dirpath_table = None
        self.fin_filename = fin_filename

        self.dct_params = None
        self.config_file = config_file
        self.ncomps_max = None

        self.subcube_nr = None
        self.xpos_offset = 0
        self.ypos_offset = 0

        self.main_beam_efficiency = None

        self.initialized_state = False

        if config_file:
            get_values_from_config_file(
                self, config_file, config_key='DEFAULT')

    def check_settings(self):
        """Check user settings and raise error messages or apply corrections."""
        if self.path_to_pickle_file is None:
            raise Exception("Need to specify 'path_to_pickle_file'")
        if self.path_to_decomp_file is None:
            raise Exception("Need to specify 'path_to_decomp_file'")
        self.decomp_dirname = os.path.dirname(self.path_to_decomp_file)
        self.file = os.path.basename(self.path_to_decomp_file)
        self.filename, self.file_extension = os.path.splitext(self.file)

        if self.fin_filename is None:
            suffix = '_finalized'
            self.fin_filename = self.filename + suffix

        if self.dirpath_gpy is None:
            self.dirpath_gpy = os.path.dirname(self.decomp_dirname)

        if self.dirpath_table is None:
            self.dirpath_table = self.decomp_dirname

    def initialize(self):
        """Read in data files and initialize parameters."""
        with open(self.path_to_pickle_file, "rb") as pickle_file:
            self.pickled_data = pickle.load(pickle_file, encoding='latin1')

        with open(self.path_to_decomp_file, "rb") as pickle_file:
            self.decomposition = pickle.load(pickle_file, encoding='latin1')

        self.length = len(self.pickled_data['index'])

        if 'header' in self.pickled_data.keys():
            self.header = correct_header(self.pickled_data['header'])
            self.wcs = WCS(self.header)
            self.velocity_increment = (
                self.wcs.wcs.cdelt[2] * self.wcs.wcs.cunit[2]).to(
                    self.vel_unit).value
            self.to_unit = (self.wcs.wcs.cunit[2]).to(self.vel_unit)
        if 'location' in self.pickled_data.keys():
            self.location = self.pickled_data['location']
        if 'nan_mask' in self.pickled_data.keys():
            self.nan_mask = self.pickled_data['nan_mask']

        self.channels = self.pickled_data['x_values']

        self.initialized_state = True

    def finalize_dct(self):
        self.check_settings()
        self.initialize()
        sp = SpatialFitting(config_file=self.config_file)
        if self.dct_params is not None:
            for key, value in self.dct_params.items():
                try:
                    setattr(sp, key, value)
                except ValueError:
                    raise Exception('Could not parse parameter {} from dct_params'.format(key))

        sp.log_output = False
        sp.path_to_pickle_file = self.path_to_pickle_file
        sp.path_to_decomp_file = self.path_to_decomp_file

        results_list = sp.finalize()

        list_means_interval, list_n_centroids = (
            [{} for _ in range(self.length)] for _ in range(2))

        for i, item in enumerate(results_list):
            if not isinstance(item, list):
                say("Error for index {}: {}".format(i, item))
                continue

            index, means_interval, n_centroids = item
            list_means_interval[index] = means_interval
            list_n_centroids[index] = n_centroids

        self.decomposition['broad'] = sp.mask_broad_flagged
        self.decomposition['ncomps_wmedian'] = sp.ncomps_wmedian.astype('int')
        self.decomposition['ncomps_jumps'] = sp.ncomps_jumps.astype('int')
        self.decomposition['means_interval'] = list_means_interval
        self.decomposition['n_centroids'] = list_n_centroids

    def get_flag_blended(self, amps, fwhms, means):
        params_fit = amps + fwhms + means
        indices = get_fully_blended_gaussians(params_fit)
        flags = np.zeros(len(amps))
        flags[indices] = 1
        return flags.astype('int')

    def get_flag_broad(self, fwhms, broad):
        flags = np.zeros(len(fwhms))
        if broad:
            flags[np.argmax(fwhms)] = 1
        return flags.astype('int')

    def get_flag_centroid(self, means, means_interval, n_centroids):
        flag = 0
        for key in means_interval.keys():
            n_wanted = n_centroids[key]
            lower = means_interval[key][0]
            upper = means_interval[key][1]
            n_real = np.count_nonzero(
                np.logical_and(lower <= means, means <= upper))

            flag += abs(n_wanted - n_real)

        return flag

    def get_table_rows(self, idx, j):
        rows = []
        ncomps = self.decomposition['N_components'][idx]

        #  do not continue if spectrum was masked out, was not fitted,
        #  or was fitted by too many components
        if ncomps is None:
            return rows
        elif ncomps == 0:
            return rows
        elif self.ncomps_max is not None:
            if ncomps > self.ncomps_max:
                return rows

        yi, xi = self.pickled_data['location'][idx]
        spectrum = self.pickled_data['data_list'][idx]
        fit_amps = self.decomposition['amplitudes_fit'][idx]
        fit_fwhms = self.decomposition['fwhms_fit'][idx]
        fit_means = self.decomposition['means_fit'][idx]
        fit_e_amps = self.decomposition['amplitudes_fit_err'][idx]
        fit_e_fwhms = self.decomposition['fwhms_fit_err'][idx]
        fit_e_means = self.decomposition['means_fit_err'][idx]
        error = self.pickled_data['error'][idx][0]

        residual = spectrum - combined_gaussian(
            fit_amps, fit_fwhms, fit_means, self.pickled_data['x_values'])

        aicc = self.decomposition['best_fit_aicc'][idx]
        rchi2 = self.decomposition['best_fit_rchi2'][idx]
        pvalue = self.decomposition['pvalue'][idx]

        broad = self.decomposition['broad'][idx]

        ncomp_wmedian = self.decomposition['ncomps_wmedian'][idx]
        ncomp_jumps = self.decomposition['ncomps_jumps'][idx]

        means_interval = self.decomposition['means_interval'][idx]
        n_centroids = self.decomposition['n_centroids'][idx]

        flags_blended = self.get_flag_blended(fit_amps, fit_fwhms, fit_means)
        flags_neg_res_peak = negative_residuals(
            spectrum, residual, error, get_flags=True,
            fwhms=fit_fwhms, means=fit_means)
        flags_broad = self.get_flag_broad(fit_fwhms, broad)
        flag_centroid = self.get_flag_centroid(
            np.array(fit_means), means_interval, n_centroids)

        x_wcs, y_wcs, z_wcs = self.wcs.wcs_pix2world(
            xi, yi, np.array(fit_means), 0)

        velocities = z_wcs * self.to_unit
        e_velocities = np.array(fit_e_means) * self.velocity_increment
        vel_disps = (
            np.array(fit_fwhms) / 2.354820045) * self.velocity_increment
        e_vel_disps = (
            np.array(fit_e_fwhms) / 2.354820045) * self.velocity_increment

        amplitudes = np.array(fit_amps)
        e_amplitudes = np.array(fit_e_amps)

        if self.main_beam_efficiency is not None:
            amplitudes /= self.main_beam_efficiency
            e_amplitudes /= self.main_beam_efficiency
            error /= self.main_beam_efficiency

        integrated_intensity = area_of_gaussian(
            amplitudes, np.array(fit_fwhms) * self.velocity_increment)
        fit_fwhms_plus_error = np.array(fit_fwhms) + np.array(fit_e_fwhms)
        e_integrated_intensity = area_of_gaussian(
            amplitudes + e_amplitudes,
            fit_fwhms_plus_error * self.velocity_increment) -\
            integrated_intensity

        for i in range(ncomps):
            row = [
                xi + self.xpos_offset, yi + self.ypos_offset,
                x_wcs[0], y_wcs[0],
                amplitudes[i], e_amplitudes[i],
                velocities[i], e_velocities[i],
                vel_disps[i], e_vel_disps[i],
                integrated_intensity[i], e_integrated_intensity[i],
                error, pvalue, aicc, rchi2,
                ncomps, ncomp_wmedian, ncomp_jumps,
                flags_blended[i], flags_neg_res_peak[i],
                flags_broad[i], flag_centroid]

            if self.subcube_nr is not None:
                row.append(self.subcube_nr)

            rows.append(row)

        return rows

    def make_table(self, save=True):
        """Create table of the decomposition results.

        Parameters
        ----------
        save : bool
            Set to `True` if the table should be saved.

        Returns
        -------
        table_results : astropy.table.table.Table
            Table of the decomposition results.

        """
        import gausspyplus.parallel_processing
        gausspyplus.parallel_processing.init([self.decomposition['index_fit'], [self]])

        results_list = gausspyplus.parallel_processing.func(
            use_ncpus=self.use_ncpus, function='make_table')

        for i, item in enumerate(results_list):
            if not isinstance(item, list):
                say("Error for spectrum with index {}: {}".format(i, item))
                continue

        # results_list = [item for item in results_list if len(item) > 0]
        results_list = np.array([item for sublist in results_list
                                 for item in sublist])

        names = [
            'x_pos', 'y_pos', self.wcs.wcs.lngtyp, self.wcs.wcs.lattyp,
            'amp', 'e_amp', 'VLSR', 'e_VLSR', 'vel_disp', 'e_vel_disp',
            'int_tot', 'e_int_tot', 'rms', 'pvalue', 'aicc', 'rchi2',
            'ncomps', 'ncomp_wmedian', 'ncomp_jumps',
            'flag_blended', 'flag_neg_res_peak', 'flag_broad', 'flag_centroid']

        dtype = ['i4']*2 + ['f4']*14 + ['i4']*7

        if self.subcube_nr is not None:
            names.append('subcube_nr')
            dtype.append('i4')

        table_results = Table(data=results_list, names=names, dtype=dtype)

        for key in names[2:16]:
            table_results[key].format = "{0:.4f}"

        if save:
            filename = self.fin_filename + '.dat'
            path_to_table = os.path.join(self.dirpath_table, filename)
            table_results.write(path_to_table, format='ascii', overwrite=True)
            say("\033[92mSAVED FILE:\033[0m '{}' in '{}'".format(
                filename, self.decomp_dirname))

        return table_results

    def save_final_results(self):
        """Save the results of the spatially coherent refitting iterations."""
        filename = self.fin_filename + '.pickle'
        pathToFile = os.path.join(self.decomp_dirname, filename)
        pickle.dump(self.decomposition, open(pathToFile, 'wb'), protocol=2)
        say("\033[92mSAVED FILE:\033[0m '{}' in '{}'".format(
            filename, self.decomp_dirname))

    def produce_map(self, keyword='error', comments=[], suffix='', save=True,
                    dtype='float32'):
        """Generate an array/FITS map of a GaussPy+ quantity.

        Parameters
        ----------
        keyword : str ['error', 'best_fit_rchi2', 'N_components']
            Keyword contained in GaussPy+ dictionary whose values should be made into an array.
        comments : list
            List of strings that are added to the FITS header of the map.
        suffix : str
            Suffix that is added to the filename.
        save : bool
            Set to `True` if the map should be saved.
        dtype : str
            Valid dtype for the resulting map.

        Returns
        -------
        astropy.io.fits.HDUList
            FITS HDU object of the map.

        """
        #  TODO: routine in case pickled_data is missing the header key
        shape = (self.header['NAXIS2'], self.header['NAXIS1'])
        array = np.ones((shape[0], shape[1])) * np.nan

        if keyword in self.pickled_data.keys():
            data = self.pickled_data[keyword]
        elif keyword in self.decomposition.keys():
            data = self.decomposition[keyword]

        for (y, x), value in zip(self.pickled_data['location'], data):
            if value is None:
                continue

            try:
                array[y, x] = value
            except ValueError:
                array[y, x] = value[0]

        header = change_header(self.header.copy(), format='pp',
                               comments=comments)

        array = array.astype(dtype)

        if save:
            if keyword == 'error':
                filename, _ = os.path.splitext(
                    os.path.basename(self.path_to_pickle_file))
            else:
                filename = self.filename

            filename = "{}{}.fits".format(filename, suffix)
            path_to_file = os.path.join(
                self.dirpath_gpy, 'gpy_maps', filename)

            save_fits(array, header, path_to_file,
                      verbose=False)
            say("\n\033[92mSAVED FILE:\033[0m '{}' in '{}'".format(
                filename, os.path.dirname(path_to_file)), logger=self.logger)

        return fits.PrimaryHDU(array, header)

    def produce_noise_map(self, comments=['Rms noise values.'],
                          suffix='_noise_map', save=True, dtype='float32',
                          get_hdu=False, get_data=False, get_header=False):
        """Produce a map of the rms noise values.

        Parameters
        ----------
        comments : list
            List of strings that are added to the FITS header of the map.
        suffix : str
            Suffix that is added to the filename.
        save : bool
            Set to `True` if the map should be saved.
        dtype : str
            Valid dtype for the resulting map.
        get_hdu : bool
            Default is `False`. If set to `True`, an astropy.io.fits.HDUList is returned. Overrides 'get_data' and 'get_header'.
        get_data : bool
            Default is `False`. Returns a numpy.ndarray of the FITS array.
        get_header : bool
            Default is `False`. Returns a astropy.io.fits.Header of the FITS array.

        Returns
        -------
        tuple or None
            Result from `return_hdu_options`.

        """
        if not self.initialized_state:
            self.check_settings()
            self.initialize()
        hdu = self.produce_map(keyword='error', comments=comments,
                               suffix=suffix, save=save, dtype=dtype)

        return return_hdu_options(
            hdu, get_hdu=get_hdu, get_data=get_data, get_header=get_header)

    def produce_rchi2_map(self, suffix='_rchi2_map', save=True, dtype='float32',
                          get_hdu=False, get_data=False, get_header=False,
                          comments=['Reduced chi-square values.']):
        """Produce a map of the reduced chi-square values.

        Parameters
        ----------
        comments : list
            List of strings that are added to the FITS header of the map.
        suffix : str
            Suffix that is added to the filename.
        save : bool
            Set to `True` if the map should be saved.
        dtype : str
            Valid dtype for the resulting map.
        get_hdu : bool
            Default is `False`. If set to `True`, an astropy.io.fits.HDUList is returned. Overrides 'get_data' and 'get_header'.
        get_data : bool
            Default is `False`. Returns a numpy.ndarray of the FITS array.
        get_header : bool
            Default is `False`. Returns a astropy.io.fits.Header of the FITS array.

        Returns
        -------
        tuple or None
            Result from `return_hdu_options`.

        """
        if not self.initialized_state:
            self.check_settings()
            self.initialize()
        hdu = self.produce_map(keyword='best_fit_rchi2', comments=comments,
                               suffix=suffix, save=save, dtype=dtype)

        return return_hdu_options(
            hdu, get_hdu=get_hdu, get_data=get_data, get_header=get_header)

    def produce_component_map(self, suffix='_component_map', save=True,
                              dtype='float32', get_hdu=False, get_data=False, get_header=False, comments=['Number of fitted Gaussian components.']):
        """Produce a map of the number of fitted Gaussian components.

        Parameters
        ----------
        comments : list
            List of strings that are added to the FITS header of the map.
        suffix : str
            Suffix that is added to the filename.
        save : bool
            Set to `True` if the map should be saved.
        dtype : str
            Valid dtype for the resulting map.
        get_hdu : bool
            Default is `False`. If set to `True`, an astropy.io.fits.HDUList is returned. Overrides 'get_data' and 'get_header'.
        get_data : bool
            Default is `False`. Returns a numpy.ndarray of the FITS array.
        get_header : bool
            Default is `False`. Returns a astropy.io.fits.Header of the FITS array.

        Returns
        -------
        tuple or None
            Result from `return_hdu_options`.

        """
        if not self.initialized_state:
            self.check_settings()
            self.initialize()
        hdu = self.produce_map(keyword='N_components', comments=comments,
                               suffix=suffix, save=save, dtype=dtype)

        return return_hdu_options(
            hdu, get_hdu=get_hdu, get_data=get_data, get_header=get_header)

    def make_cube(self, mode='full_decomposition', save=True, get_hdu=False,
                  get_data=False, get_header=False, dtype='float32'):
        """Create FITS cube of the decomposition results.

        Parameters
        ----------
        mode : str
            'full_decomposition' recreates the whole FITS cube, 'integrated_intensity' creates a cube with the integrated intensity values of the Gaussian components placed at their mean positions, 'main_component' only retains the fitted component with the largest amplitude value
        save : bool
            Set to `True` if the map should be saved.
        dtype : str
            Valid dtype for the resulting map.
        get_hdu : bool
            Default is `False`. If set to `True`, an astropy.io.fits.HDUList is returned. Overrides 'get_data' and 'get_header'.
        get_data : bool
            Default is `False`. Returns a numpy.ndarray of the FITS array.
        get_header : bool
            Default is `False`. Returns a astropy.io.fits.Header of the FITS array.

        Returns
        -------
        tuple or None
            Result from `return_hdu_options`.

        """
        say('\ncreate {} cube...'.format(mode))

        x = self.header['NAXIS1']
        y = self.header['NAXIS2']
        z = self.header['NAXIS3']

        array = np.zeros([z, y, x], dtype=np.float32)
        nSpectra = len(self.decomposition['N_components'])

        for idx in range(nSpectra):
            ncomps = self.decomposition['N_components'][idx]
            if ncomps is None:
                continue

            yi = self.location[idx][0]
            xi = self.location[idx][1]

            amps = self.decomposition['amplitudes_fit'][idx]
            fwhms = self.decomposition['fwhms_fit'][idx]
            means = self.decomposition['means_fit'][idx]

            if self.main_beam_efficiency is not None:
                amps = [amp / self.main_beam_efficiency for amp in amps]

            if mode == 'main_component' and ncomps > 0:
                j = amps.index(max(amps))
                array[:, yi, xi] = gaussian(
                    amps[j], fwhms[j], means[j], self.channels)
            elif mode == 'integrated_intensity' and ncomps > 0:
                for j in range(ncomps):
                    integrated_intensity = area_of_gaussian(
                        amps[j], fwhms[j] * self.velocity_increment)
                    channel = int(round(means[j]))
                    if self.channels[0] <= channel <= self.channels[-1]:
                        array[channel, yi, xi] += integrated_intensity
            elif mode == 'full_decomposition':
                array[:, yi, xi] = combined_gaussian(
                    amps, fwhms, means, self.channels)

            nans = self.nan_mask[:, yi, xi]
            array[:, yi, xi][nans] = np.NAN

        array[self.nan_mask] = np.nan
        array = array.astype(dtype)

        if mode == 'main_component':
            comment = 'Fit component with highest amplitude per spectrum.'
            filename = "{}_main.fits".format(self.filename, self.suffix)
        elif mode == 'integrated_intensity':
            comment = 'Integrated intensity of fit component at VLSR position.'
            filename = "{}_int_tot.fits".format(self.filename, self.suffix)
        elif mode == 'full_decomposition':
            comment = 'Recreated dataset from fit components.'
            filename = "{}_decomp.fits".format(self.filename, self.suffix)

        comments = ['GaussPy+ decomposition results:']
        comments.append(comment)
        if self.main_beam_efficiency is not None:
            comments.append('Corrected for main beam efficiency of {}.'.format(
                self.main_beam_efficiency))

        header = update_header(
            self.header.copy(), comments=comments, write_meta=True)

        if save:
            pathToFile = os.path.join(self.decomp_dirname, 'FITS', filename)
            save_fits(array, header, pathToFile, verbose=False)
            say("\033[92mSAVED FILE:\033[0m '{}' in '{}'".format(
                filename, os.path.dirname(pathToFile)))

        hdu = fits.PrimaryHDU(array, header)

        return return_hdu_options(
            hdu, get_hdu=get_hdu, get_data=get_data, get_header=get_header)


#!/usr/bin/env python
"""

runsim.py by John Swoboda 3/16/2015
This script will run the SimISR code. The user can run a number of different
aspects including making data, fitting applying matrix formulation of the
space-time ambuiguty operator and calling inversion methods.

Inputs
             -f These will be the possible functions that can be used
                 spectrums :makespectrums,
                 radardata :makeradardata,
                 fitting :fitdata,'
                 fittingmat :fitdata,
                 fittinginv :fitdata,
                 applymat applymat
"""
from __future__ import print_function

#imported basic modules
import time
import sys
from datetime import datetime
import traceback
import argparse
import ipdb

# Imported scipy and matplotlib modules
import scipy as sp
# My modules
from SimISR import Path
from SimISR.IonoContainer import IonoContainer
from SimISR.radarData import RadarDataFile
import SimISR.specfunctions as specfuncs
from SimISR.fitterMethodGen import Fitterionoconainer
from SimISR.utilFunctions import readconfigfile, update_progress
from SimISR.operators import RadarSpaceTimeOperator

#%% Make spectrums
def makespectrums(basedir, configfile, printlines=True):
    """ This will make all of the spectra for a set of data and save it in a
    folder in basedir called Spectrums. It is assumed that the data in the Origparams
    is time tagged in the title with a string seperated by a white space and the
    rest of the file name. For example ~/DATA/Basedir/0 origdata.h5.
    Inputs:
        basedir: A string for the directory that will hold all of the data for the simulation.
        configfile: The configuration file for the simulation.
         """
    basedir = Path(basedir).expanduser()
    dirio = ('Origparams', 'Spectrums')
    inputdir = basedir/dirio[0]
    outputdir = basedir/dirio[1]
    # determine the list of h5 files in Origparams directory
    dirlist = sorted(inputdir.glob('*.h5'))
    # Make the lists of numbers and file names for the dictionary
    (listorder, _, _, timebeg, _) = IonoContainer.gettimes(dirlist)
    slist = [dirlist[ikey] for ikey in listorder]
    (sensdict, simparams) = readconfigfile(configfile)
    # Delete data
    outfiles = outputdir.glob('*.h5')
    for ifile in outfiles:
        ifile.unlink()

    for inum, curfile in zip(timebeg, slist):

        outfile = outputdir / (str(inum)+' spectrum.h5')
        update_progress(float(inum)/float(len(slist)),
                        'Processing file {} starting at {}'.format(curfile.name, datetime.now()))
        curiono = IonoContainer.readh5(str(curfile))

        curiono.makespectruminstanceopen(specfuncs.ISRSspecmake, sensdict,
                                         int(simparams['numpoints']), float(inum),
                                         float(len(slist)), printlines).saveh5(str(outfile))
        update_progress(float(inum+1)/float(len(slist)),
                        'Finished file {} starting at {}'.format(curfile.name, datetime.now()))

#%% Make Radar Data
def makeradardata(basedir,configfile,remakealldata):
    """ This function will make the radar data and create the acf estimates.
    Inputs:
        basedir: A string for the directory that will hold all of the data for the simulation.
        configfile: The configuration file for the simulation.
        remakealldata: A bool that determines if the radar data is remade. If false
            only the acfs will be estimated using the radar that is already made."""

    dirio = ('Spectrums', 'Radardata', 'ACF')
    inputdir = basedir/dirio[0]
    outputdir = basedir/dirio[1]
    outputdir2 = basedir/dirio[2]

    # determine the list of h5 files in Origparams directory
    dirlist = [str(i) for i in inputdir.glob('*.h5')]
    # Make the lists of numbers and file names for the dictionary
    if len(dirlist) > 0:
        (listorder, _, _, timebeg, _) = IonoContainer.gettimes(dirlist)

        Ionodict = {timebeg[itn]:dirlist[it] for itn, it in enumerate(listorder)}
    else:
        Ionodict = {0.:str(inputdir.joinpath('00.h5'))}
    # Find all of the raw data files
    radardatalist = outputdir.glob('*RawData.h5')
    if radardatalist and (not remakealldata):
        # XXX need to work on time stuff
        outlist2 = radardatalist
    else:
        outlist2 = None

    # create the radar data file class
    rdata = RadarDataFile(Ionodict, configfile, outputdir, outfilelist=outlist2)
    # From the ACFs and uncertainties
    (ionoout, ionosig) = rdata.processdataiono()
    # save the acfs and uncertianties in ionocontainer h5 files.
    ionoout.saveh5(str(outputdir2.joinpath('00lags.h5')))
    ionosig.saveh5(str(outputdir2.joinpath('00sigs.h5')))
    return ()
#%% Fit data
def fitdata(basedir,configfile,optinputs):
    """ This function will run the fitter on the estimated ACFs saved in h5 files.
        Inputs:
        basedir: A string for the directory that will hold all of the data for the simulation.
        configfile: The configuration file for the simulation.
        optinputs:A string that helps determine the what type of acfs will be fitted.
         """
    # determine the input folders which can be ACFs from the full simulation
    dirdict = {'fitting':('ACF', 'Fitted'), 'fittingmat':('ACFMat', 'FittedMat'),
               'fittinginv':('ACFInv', 'FittedInv'), 'fittingmatinv':('ACFMatInv', 'FittedMatInv')}
    dirio = dirdict[optinputs[0]]
    inputdir = basedir/dirio[0]
    outputdir = basedir/dirio[1]
    fitlist = optinputs[1]
    if len(optinputs) > 2:
        exstr = optinputs[2]
        printlines = optinputs[3]
    else:
        exstr = ''
    dirlist = [str(i) for i in inputdir.glob('*lags{0}.h5'.format(exstr))]
    dirlistsig = [str(i) for i in inputdir.glob('*sigs{0}.h5'.format(exstr))]

    Ionoin = IonoContainer.readh5(dirlist[0])
    if len(dirlistsig) == 0:
        Ionoinsig = None
    else:
        Ionoinsig = IonoContainer.readh5(dirlistsig[0])
    fitterone = Fitterionoconainer(Ionoin, Ionoinsig, configfile)

    fitoutput = fitterone.fitdata(specfuncs.ISRSfitfunction,
                                  fitterone.simparams['startfile'], fittimes=fitlist,
                                  printlines=printlines)

    #ipdb.set_trace()
    (fitteddata, fittederror, funcevals, fittedcov) = fitoutput
    if fitterone.simparams['Pulsetype'].lower() == 'barker':
        paramlist = fitteddata

        species = fitterone.simparams['species']
        paramnames = ['Ne']
        if not fittederror is None:
            fittederronly = sp.sqrt(fittederror)
            paramlist = sp.concatenate([fitteddata, fittederronly], axis=2)
            paramnamese = ['n'+ip for ip in paramnames]
            paranamsf = sp.array(paramnames+paramnamese)
        else:
            paranamsf = sp.array(paramnames)
    else:
        fittederronly = sp.sqrt(fittederror)
        paramnames = []
        species = fitterone.simparams['species']
        # Seperate Ti and put it in as an element of the ionocontainer.
        Ti = fitteddata[:, :, 1]

        nTi = fittederronly[:, :, 1]

        nTiTe = fittedcov[:, :, 0, 1]
        nTiNe = fittedcov[:, :, 0, 2]
        nTiVi = fittedcov[:, :, 0, 3]
        nTeNe = fittedcov[:, :, 1, 2]
        nTeVi = fittedcov[:, :, 1, 3]
        nNeVi = fittedcov[:, :, 2, 3]
        cov_list = [nTiTe[:, :, sp.newaxis], nTiNe[:, :, sp.newaxis],
                    nTiVi[:, :, sp.newaxis], nTeNe[:, :, sp.newaxis],
                    nTeVi[:, :, sp.newaxis], nNeVi[:, :, sp.newaxis]]
        cov_list_names = ['nTiTe', 'nTiNe', 'nTiVi', 'nTeNe', 'nTeVi','nNeVi']
        paramlist = sp.concatenate([fitteddata, Ti[:, :, sp.newaxis], fittederronly,
                                    nTi[:, :, sp.newaxis], funcevals[:, :, sp.newaxis]]
                                   + cov_list, axis=2)
        for isp in species[:-1]:
            paramnames.append('Ni_'+isp)
            paramnames.append('Ti_'+isp)
        paramnames = paramnames+['Ne', 'Te', 'Vi', 'Nepow', 'Ti']
        paramnamese = ['n'+ip for ip in paramnames]
        paranamsf = sp.array(paramnames+paramnamese+['FuncEvals']+cov_list_names)

    if fitlist is None:
        timevec = Ionoin.Time_Vector
    else:
        if len(fitlist) == 0:
            timevec = Ionoin.Time_Vector
        else:
            timevec = Ionoin.Time_Vector[fitlist]
    # This requires
    if set(Ionoin.Coord_Vecs) == {'x', 'y', 'z'}:
        newver = 0
        ionoout = IonoContainer(Ionoin.Cart_Coords, paramlist.real, timevec, ver=newver,
                                coordvecs=Ionoin.Coord_Vecs, paramnames=paranamsf,
                                species=species)
    elif set(Ionoin.Coord_Vecs) == {'r', 'theta', 'phi'}:
        newver = 1
        ionoout = IonoContainer(Ionoin.Sphere_Coords, paramlist.real, timevec, ver=newver,
                                coordvecs=Ionoin.Coord_Vecs, paramnames=paranamsf,
                                species=species)

    outfile = outputdir.joinpath('fitteddata{0}.h5'.format(exstr))
    ionoout.saveh5(str(outfile))
#%% apply the matrix for the data
def applymat(basedir, configfile, optinputs):
    """
        This function apply the matrix version of the space time ambiugty function
        to the ACFs and save the outcome in h5 files within the directory ACFMat.
        Inputs:
        basedir: A string for the directory that will hold all of the data for the simulation.
        configfile: The configuration file for the simulation.
    """
    dirio = ('Spectrums', 'Mat', 'ACFMat')
    basedir = Path(basedir)
    inputdir = basedir.joinpath(dirio[0])
    outputdir2 = basedir.joinpath(dirio[2])

    dirlist = [str(i) for i in inputdir.glob('*.h5')]
    (listorder, timevector, _, _, _) = IonoContainer.gettimes(dirlist)
    ionolist = [dirlist[ikey] for ikey in listorder]
    rsto = RadarSpaceTimeOperator(ionolist, configfile, timevector, mattype='matrix')
    ionoout = rsto.mult_iono(ionolist)
    outfile = outputdir2.joinpath('00lags.h5')
    ionoout.saveh5(str(outfile))

#%% For sorting
def ke(item):
    """
        Used for sorting names of files.
    """
    if item[0].isdigit():
        return int(item.partition(' ')[0])
    else:
        return float('inf')
#%% Main function
def main(funcnamelist,basedir,configfile,remakealldata,fitlist=None,invtype='',printlines=True):
    """ Main function for this module. The function will set up the directory
    structure, create or update a diary file and run the simulation depending
    on the user input.
    Inputs
        funcnamelist: A list of strings that coorespond to specific functions.
                The stirng and function they correspond to, that will be shown.

                 spectrums: makespectrums, This will create the ISR spectrums

                 radardata :makeradardata, This will make the radar data and
                     form the ACF estimates. If the raw radar data exists then
                     the user must use the -r option on the command line and set
                     it to y.

                 fitting :fitdata, This will apply the fitter to the data in
                 the ACF folder of the base directory.

                 fittingmat :fitdata,This will apply the fitter to the data in
                 the ACFMat folder of the base directory.

                 fittinginv :fitdata,This will apply the fitter to the data in
                 the ACFInv folder of the base directory.

                 applymat :applymat, This wil create and apply a matrix
                 formulation of thespace-time ambiguity function to ISR spectrums.
        basedir: The base directory that will contain all of the data. This directory
                must contain a directory called Origparams with h5 files to run
                the full simulation. The user can also start with a directory
                from a later stage of the simulation instead though.

        configfile: The configuration used for the simulation. Can be an ini file or
                a pickle file.

        remakealldata: A bool to determine if the raw radar data will be remade. If
                this is False the radar data will only be made if it does
                not exist in the file first.
        fitlist:  A list of time entries that will be fit.

        invtype

    """

    inputsep = '***************************************************************\n'

    funcdict = {'spectrums':makespectrums, 'radardata':makeradardata, 'fitting':fitdata,'fittingmat':fitdata,
                'fittinginv':fitdata,'applymat':applymat,'fittingmatinv':fitdata}
    #inout = {'spectrums':('Origparams','Spectrums'),'radardata':('Spectrums','Radardata'),'fitting':('ACF','Fitted')}
    #pdb.set_trace()

    # check for the directories
    dirnames = ['Origparams','Spectrums','Radardata','ACF','Fitted','ACFOrig','ACFMat','ACFInv','FittedMat','FittedInv','ACFMatInv','FittedMatInv']
    basedir=Path(basedir).expanduser()
    for idir in dirnames:
        curdir = basedir/idir
        curdir.mkdir(exist_ok=True,parents=True)

    if len(funcnamelist)==3:
        funcname='all'
    else:
        funcname=''.join(funcnamelist)

    dfilename = 'diary'+funcname+'.txt'
    dfullfilestr = basedir/dfilename

    with open(str(dfullfilestr),'a') as f:
        failure=False
        for curfuncn in funcnamelist:
            curfunc = funcdict[curfuncn]
            f.write(inputsep)
            f.write(curfunc.__name__+'\n')
            f.write(time.asctime()+'\n')
            if curfunc.__name__=='fitdata':
                ex_inputs=[curfuncn,fitlist,invtype,printlines]
            elif curfunc.__name__=='makeradardata':
                ex_inputs = remakealldata
            else:
                ex_inputs=printlines
            try:
                stime = datetime.now()
                curfunc(basedir,configfile,ex_inputs)
                ftime = datetime.now()
                ptime = ftime-stime
                f.write('Success!\n')
                f.write('Duration: {}\n'.format(ptime))
                f.write('Base directory: {}\n'.format(basedir))

            except Exception as e:
                f.write('Failed!\n')
                ftime = datetime.now()
                ptime = ftime-stime
                f.write('Duration: {}\n'.format(ptime))
                f.write('Base directory: {}\n'.format(basedir))
                traceback.print_exc(file=sys.stdout)
                traceback.print_exc(file=f)
                failure = True
                break

        f.write(inputsep)

    return failure

def parse_command_line(str_input=None):
    """
        This will parse through the command line arguments
    """
    # if str_input is None:
    parser = argparse.ArgumentParser()
    # else:
    #     parser = argparse.ArgumentParser(str_input)
    fstr = '''      These will be the possible strings for the argument and the
                    function they correspond to, that will be used ish shown.

                     spectrums: makespectrums, This will create the ISR spectrums

                     radardata :makeradardata, This will make the radar data and
                         form the ACF estimates. If the raw radar data exists then
                         the user must use the -r option on the command line and set
                         it to y.

                     fitting :fitdata, This will apply the fitter to the data in
                     the ACF folder of the base directory.

                     fittingmat :fitdata,This will apply the fitter to the data in
                     the ACFMat folder of the base directory.

                     fittinginv :fitdata,This will apply the fitter to the data in
                     the ACFInv folder of the base directory.

                     applymat :applymat, This wil create and apply a matrix
                     formulation of thespace-time ambiguity function to ISR spectrums.

                     all - This will run the commands from using the spectrums, radardata,
                         and fitting
                    '''
    parser.add_argument('-f', '--funclist', dest='funclist', default='all', help=fstr)

    parser.add_argument("-v", "--verbose", action="store_true",
                        dest="verbose", default=False,
                        help="prints debug output and additional detail.")
    parser.add_argument("-c", "--config", dest="config", default='default.ini',
                        help=" Config file used for the simulation, .ini or yaml file.")
    parser.add_argument('-p', "--path", dest='path', default=None,
                        help='Number of incoherent integrations in calculations.')
    parser.add_argument('-r', "--remake", action="store_true", dest='remake', default=False,
                        help='Remake data flag.')

    if str_input is None:
        return parser.parse_args()
    else:
        return parser.parse_args(str_input)

if __name__ == "__main__":

    outstr = '''
             Usage: python runsim.py -f <function: spectrums, radardata, fitting or all> -i <basedir> -c <config> -r <type y to remake data>

             or

             python runsim.py -h

             This script will run the SimISR code. The user
             can run a number of different aspects including making data, fitting
             applying matrix formulation sof the space-time operator and calling
             inversion methods.

             Manditory Arguments to run code, for help just use the -h option.

             -f These will be the possible strings for the argument and the
                function they correspond to, that will be used ish shown.

                 spectrums: makespectrums, This will create the ISR spectrums

                 radardata :makeradardata, This will make the radar data and
                     form the ACF estimates. If the raw radar data exists then
                     the user must use the -r option on the command line and set
                     it to y.

                 fitting :fitdata, This will apply the fitter to the data in
                 the ACF folder of the base directory.

                 fittingmat :fitdata,This will apply the fitter to the data in
                 the ACFMat folder of the base directory.

                 fittinginv :fitdata,This will apply the fitter to the data in
                 the ACFInv folder of the base directory.

                 applymat :applymat, This wil create and apply a matrix
                 formulation of thespace-time ambiguity function to ISR spectrums.

                 all - This will run the commands from using the spectrums, radardata,
                     and fitting

            -i The base directory that will contain all of the data. This directory
                must contain a directory called Origparams with h5 files to run
                the full simulation. The user can also start with a directory
                from a later stage of the simulation instead though.

            -c The configuration used for the simulation. Can be an ini file or
                a pickle file.

            Optional arguments

            -r If a y follows this then the raw radar data will be remade. If
                this is not used the radar data will only be made if it does
                not exist in the file first.

             Example:
             python runsim.py -f radardata -f fitting -i ~/DATA/ExampleLongPulse -c ~/DATA/Example -r y'''

    args_commd = parse_command_line()
    if args_commd.path is None:
        print("Please provide an input source with the -p option!")
        sys.exit(1)
    basedir =  str(Path(args_commd.path).expanduser())
    configfile = str(Path(args_commd.config).expanduser())
    funcname = args_commd.funclist
    remakealldata = args_commd.remake

    if funcname.lower() == 'all':
        funcnamelist = ['spectrums', 'radardata', 'fitting']
    else:
        funcnamelist = funcname.split()

    failflag = main(funcnamelist, basedir, configfile, remakealldata)


""" This module provides the basic classes for the pulse retrieval algorithms.
"""
import numpy as np
from types import SimpleNamespace
from .. import io
from ..mesh_data import MeshData
from ..pulse_error import pulse_error
from .. import lib
from ..pnps import BasePNPS
from math import isclose

# global dictionary that contains all PNPS classes
_RETRIEVER_CLASSES = {}

# =============================================================================
# Metaclass and factory
# =============================================================================
class MetaRetriever(type):
    """ Metaclass that registers Retriever classes in a global dictionary.
    """
    def __new__(cls, clsmethod, bases, attrs):
        global _RETRIEVER_CLASSES
        newclass = super().__new__(cls, clsmethod, bases, attrs)
        method = newclass.method
        if method is None:
            return newclass
        # register the Retriever method, e.g. 'copra'
        if method in _RETRIEVER_CLASSES:
            raise ValueError("Two retriever classes implement retriever '%s'."
                             % method)
        _RETRIEVER_CLASSES[method] = newclass
        return newclass

class MetaIORetriever(io.MetaIO, MetaRetriever):
    # to fix metaclass conflicts
    pass

# =============================================================================
# Retriever Base class
# =============================================================================
class BaseRetriever(io.IO, metaclass=MetaIORetriever):
    """ The abstract base class for pulse retrieval.

    This class implements common functionality for different retrieval
    algorithms.
    """
    method = None
    supported_schemes = None
    _io_store = ['pnps', 'options', 'logging', 'log',
                 '_retrieval_state', '_result']

    def __init__(self, pnps, logging=False, verbose=False, **kwargs):
        self.pnps = pnps
        self.ft = self.pnps.ft
        self.options = SimpleNamespace(**kwargs)
        self._result = None
        self.logging = logging
        self.verbose = verbose
        self.log = None
        rs = self._retrieval_state = SimpleNamespace()
        rs.running = False
        if (self.supported_schemes is not None and
                pnps.scheme not in self.supported_schemes):
            raise ValueError("Retriever '%s' does not support scheme '%s'. "
                             "It only supports %s." %
                             (self.method, pnps.scheme, self.supported_schemes)
                             )

    def retrieve(self, measurement, initial_guess, weights=None,
                 **kwargs):
        """ Retrieve pulse from ``measurement`` starting at ``initial_guess``.

        Parameters
        ----------
        measurement : MeshData
            A MeshData instance that contains the PNPS measurement. The first
            axis has to correspond to the PNPS parameter, the second to the
            frequency. The data has to be the measured _intensity_ over the
            frequency (not wavelength!). The second axis has to match exactly
            the frequency axis of the underlying PNPS instance. No
            interpolation is done.
        initial_guess : 1d-array
            The spectrum of the pulse that is used as initial guess in the
            iterative retrieval.
        weights : 1d-array
            Weights that are attributed to the measurement for retrieval.
            In the case of (assumed) Gaussian uncertainties with standard
            deviation sigma they should correspond to 1/sigma.
            Not all algorithms support using the weights.
        kwargs : dict
            Can override retrieval options specified in :func:`__init__`.

        Notes
        -----
        This function provides no interpolation or data processing. You have
        to write a retriever wrapper for that purpose.
        """
        self.options.__dict__.update(**kwargs)
        if not isinstance(measurement, MeshData):
            raise ValueError("measurement has to be a MeshData instance!")
        self._retrieve_begin(measurement, initial_guess, weights)
        self._retrieve()
        self._retrieve_end()

    def _retrieve_begin(self, measurement, initial_guess, weights):
        pnps = self.pnps
        if not np.allclose(pnps.process_w, measurement.axes[1], rtol=1e-6):
            raise ValueError("Measurement has to lie on simulation grid!")
        # Store measurement
        self.measurement = measurement
        self.parameter = measurement.axes[0]
        self.Tmn_meas = measurement.data

        self.initial_guess = initial_guess
        # set the size
        self.M, self.N = self.Tmn_meas.shape
        # Setup the weights
        if weights is None:
            self._weights = np.ones((self.M, self.N))
        else:
            self._weights = weights.copy()
        # Retrieval state
        rs = self._retrieval_state
        rs.approximate_error = False
        rs.running = True
        rs.steps_since_improvement = 0
        # Initialize result
        res = self._result = SimpleNamespace()
        res.trace_error = self.trace_error(self.initial_guess)
        res.approximate_error = False
        res.spectrum = self.initial_guess.copy()
        # Setup the logger
        if self.logging:
            log = self.log = SimpleNamespace()
            log.trace_error = []
            log.initial_guess = self.initial_guess.copy()
        else:
            self.log = None
        if self.verbose:
            print("Started retriever '%s'" % self.method)
            print("Options:")
            print(self.options)
            print("Initial trace error R = {:.10e}".format(res.trace_error))
            print("Starting retrieval...")
            print()

    def _retrieve_end(self):
        rs = self._retrieval_state
        rs.running = False
        res = self._result
        if res.approximate_error:
            res.trace_error = self.trace_error(res.spectrum)
            res.approximate_error = False

    def _project(self, measured, Smk):
        """ Performs the projection on the measured intensity.
        """
        # in frequency domain
        Smn = self.ft.forward(Smk)
        # project and specially handle values with zero amplitude
        absSmn = np.abs(Smn)
        f = (absSmn > 0.0)
        Smn[~f] = np.sqrt(measured[~f] + 0.0j)
        Smn[f] = Smn[f] / absSmn[f] * np.sqrt(measured[f] + 0.0j)
        # back in time domain
        Smk2 = self.ft.backward(Smn)
        return Smk2

    def _objective_function(self, spectrum):
        """ Calculates the minimization objective from the pulse spectrum.

        This is Eq. 11 in the paper:

            r = sum (Tmn^meas - mu * Tmn)
        """
        # calculate the PNPS trace
        Tmn = self.pnps.calculate(spectrum, self.parameter)
        return self._r(Tmn)

    def trace_error(self, spectrum, store=True):
        """ Calculates the trace error from the pulse spectrum.
        """
        Tmn = self.pnps.calculate(spectrum, self.parameter)
        return self._R(Tmn, store=store)

    def _r(self, Tmn, store=True):
        """ Calculates the minimization objective r from a simulated trace Tmn.
        """
        diff = self._error_vector(Tmn, store=store)
        return np.sum(diff * diff)

    def _error_vector(self, Tmn, store=True):
        """ Calculates the residual vector from measured to simulated
        intensity.
        """
        # rename
        rs = self._retrieval_state
        Tmn_meas = self.Tmn_meas
        # scaling factor
        w2 = self._weights * self._weights
        mu = np.sum(Tmn_meas * Tmn * w2) / np.sum(Tmn * Tmn * w2)
        # store intermediate results in current retrieval state
        if store:
            rs.mu = mu
            rs.Tmn = Tmn
            rs.Smk = self.pnps.Smk
        return np.ravel((Tmn_meas - mu * Tmn) * self._weights)

    def _R(self, Tmn, store=True):
        """ Calculates the trace error from a simulated trace Tmn.
        """
        r = self._r(Tmn, store=store)
        return self._Rr(r)

    def _Rr(self, r):
        """ Calculates the trace error from the minimization objective r.
        """
        return np.sqrt(r / (self.M * self.N *
                            (self.Tmn_meas * self._weights).max()**2))

    def result(self, pulse_original=None, full=True):
        """ Analyzes the retrieval results in one retrieval instance
            and processes it for plotting or storage.
        """
        rs = self._retrieval_state
        if self._result is None or self._retrieval_state.running:
            return None
        res = SimpleNamespace()
        # the meta data
        res.parameter = self.parameter
        res.options = self.options
        res.logging = self.logging
        res.measurement = self.measurement
        # store the retriever itself
        if full:
            res.pnps = self.pnps
        else:
            res.pnps = None

        # the pulse spectra
        # 1 - the retrieved pulse
        res.pulse_retrieved = self._result.spectrum
        # 2 - the original test pulse, optional
        res.pulse_original = pulse_original
        # 3 - the initial guess
        res.pulse_initial = self.initial_guess

        # the measurement traces
        # 1 - the original data used for retrieval
        res.trace_input = self.Tmn_meas
        # 2 - the trace error and the trace calculated from the retrieved pulse
        res.trace_error = self.trace_error(res.pulse_retrieved)
        res.trace_retrieved = rs.mu * rs.Tmn
        res.response_function = rs.mu
        # the weights
        res.weights = self._weights

        # this is set if the original spectrum is provided
        if res.pulse_original is not None:
            # the trace error of the test pulse (non-zero for noisy input)
            res.trace_error_optimal = self.trace_error(res.pulse_original)
            # 3 - the optimal trace calculated from the test pulse
            res.trace_original = rs.mu * rs.Tmn
            dot_ambiguity = False
            if self.pnps.method == "ifrog" or self.pnps.scheme == "shg-frog":
                dot_ambiguity = True
            # the pulse error to the test pulse
            res.pulse_error, res.pulse_retrieved = pulse_error(
                    res.pulse_retrieved, res.pulse_original, self.ft,
                    dot_ambiguity=dot_ambiguity)

        if res.logging:
            # the logged trace errors
            res.trace_errors = np.array(self.log.trace_error)
            # the running minimum of the trace errors (for plotting)
            res.rm_trace_errors = np.minimum.accumulate(res.trace_errors,
                                                        axis=-1)
        if self.verbose:
            lib.retrieval_report(res)
        return res

def Retriever(pnps: BasePNPS, method: str = "copra", maxiter=300, maxfev=None,
              logging=False, verbose=False, **kwargs) -> BaseRetriever:
    """ Creates a retriever instance.

    Parameters
    ----------
    pnps : PNPS
        A PNPS instance that is used to simulate a PNPS measurement.
    method : str, optional
        Type of solver.  Should be one of
            - 'copra'       :class:`(see here) <COPRARetriever>`
            - 'gpa'         :class:`(see here) <GPARetriever>`
            - 'gp-dscan'     :class:`(see here) <GPDSCANRetriever>`
            - 'pcgpa'       :class:`(see here) <PCGPARetriever>`
            - 'pie'         :class:`(see here) <PIERetriever>`
            - 'lm'          :class:`(see here) <LMRetriever>`
            - 'bfgs'        :class:`(see here) <BFGSRetriever>`
            - 'de'          :class:`(see here) <DERetriever>`
            - 'nelder-mead' :class:`(see here) <NMRetriever>`

        'copra' is the default choice.
    maxiter : int, optional
        The maximum number of algorithm iterations. The default is 300.
    maxfev : int, optional
        The maximum number of function evaluations. If given, the algorithms
        stop before this number is reached. Not all algorithms support this
        feature. Default is ``None``, in which case it is ignored.
    logging : bool, optional
        Stores trace errors and pulses over the iterations if supported
        by the retriever class. Default is `False`.
    verbose : bool, optional
        Prints out trace errors during the iteration if supported by the
        retriever class. Default is `False`.
    """
    method = method.lower()
    try:
        cls = _RETRIEVER_CLASSES[method]
    except KeyError:
        raise ValueError("Retriever '%s' is unknown!" % (method))
    return cls(pnps, maxiter=maxiter, maxfev=maxfev,
               logging=logging, verbose=verbose, **kwargs)