while True:
        for bundle in it:
            yield bundle

"""
Self-contained training loop for the project.
"""

import logging
import os
import pickle
from collections import defaultdict
from functools import partial
import gin
import numpy as np
import pandas as pd

import tensorflow as tf
from tensorflow.keras.metrics import Metric as keras_metric

from src.callbacks import ModelCheckpoint, LambdaCallback, History, DumpTensorboardSummaries, BaseLogger, Heartbeat, \
    MetaSaver, MeasureTime
from src.utils import save_model, restore_model_and_optimizer, restore_model, configure_neptune_exp

logger = logging.getLogger(__name__)

def _construct_default_callbacks(model, optimizer, H, save_path, checkpoint_monitor, save_freq, custom_callbacks,
                                 use_tb, save_history_every_k_examples):
    callbacks = []
    callbacks.append(MetaSaver())
    callbacks.append(BaseLogger())
    callbacks.append(MeasureTime())
    callbacks.append(Heartbeat())
    callbacks.append(LambdaCallback(on_epoch_end=partial(_append_to_history_csv, H=H)))
    callbacks.append(LambdaCallback(on_epoch_end=partial(_save_history_csv, save_path=save_path, H=H)))
    callbacks.append(History(save_every_k_examples=save_history_every_k_examples))

    if len(checkpoint_monitor):
        callbacks.append(ModelCheckpoint(monitor=checkpoint_monitor,
                                         save_best_only=True,
                                         mode='max',
                                         filepath=os.path.join(save_path, "model_best_val")))
    if save_freq > 0:
        def save_weights_fnc(epoch, logs):
            if epoch % save_freq == 0:
                logger.info("Saving model from epoch " + str(epoch))
                save_model(model, optimizer, os.path.join(save_path, "model_last_epoch"))

        callbacks.append(LambdaCallback(on_epoch_end=save_weights_fnc))

    if use_tb:
        callbacks.append(DumpTensorboardSummaries())
    callbacks.append(LambdaCallback(on_epoch_end=partial(_save_loop_state, save_callbacks=custom_callbacks,
                                                         save_path=save_path)))
    return callbacks

def _save_loop_state(epoch, logs, save_path, save_callbacks):
    loop_state = {"epochs_done": epoch, "callbacks": save_callbacks}
    pickle.dump(loop_state, open(os.path.join(save_path, "loop_state.pkl"), "wb"))

def _save_history_csv(epoch, logs, save_path, H):
    out = ""
    keys_to_save = []
    for key, value in logs.items():
        if isinstance(value, (int, float, complex, np.float32, np.float64)):
            out += "{key}={value}\t".format(key=key, value=value)
            keys_to_save.append(key)
    logger.info(out)
    logger.info("Saving history to " + os.path.join(save_path, "history.csv"))
    try:
        pd.DataFrame(H)[keys_to_save].to_csv(os.path.join(save_path, "history.csv"), index=False)
    except ValueError:
        logger.warning("Couldnt' save history.csv. Please read in history.pkl")

def _append_to_history_csv(epoch, logs, H):
    for key, value in logs.items():
        if isinstance(value, (int, float, complex, np.float32, np.float64)):
            if key not in H:
                H[key] = [value]
            else:
                H[key].append(value)

            assert len(H[key]) == epoch + 1, "Len H[{}] is {}, expected {} ".format(key, len(H[key]), epoch + 1)
        else:
            pass

def _reload(model, optimizer, save_path, callbacks):
    model_last_epoch_path = os.path.join(save_path, "model_last_epoch")
    loop_state_path = os.path.join(save_path, "loop_state.pkl")
    history_csv_path = os.path.join(save_path, "history.csv")

    if not os.path.exists(model_last_epoch_path) or not os.path.exists(loop_state_path):
        logger.warning("Failed to find last epoch model or loop state")
        return {}, 0

    # Reload everything (model, optimizer, loop state)
    logger.warning("Reloading weights!")
    model, optimizer = restore_model_and_optimizer(model, optimizer, model_last_epoch_path)
    logger.info("Reloading loop state!")
    loop_state = pickle.load(open(loop_state_path, 'rb'))
    logger.info("Reloading history!")
    H = pd.read_csv(history_csv_path)
    H = {col: list(H[col].values) for col in H.columns}
    print(H)
    logger.info("Done reloading!")

    # Small back-up
    os.system("cp " + os.path.join(save_path, "history.pkl") + " " + os.path.join(save_path, "history.pkl.bckp"))

    # Setup the rest
    epoch_start = loop_state['epochs_done'] + 1
    if not len(H[next(iter(H))]) == loop_state['epochs_done'] + 1:
        raise IOError("Mismatch between saved history and epochs recorded. "
                      "Found len(H)={0} and epoch_start={1} "
                      "Run was likely interrupted incorrectly and cannot be rerun.".format(len(H[next(iter(H))]),
                                                                                epoch_start))

    # Load all callbacks from the loop_state
    for e, e_loaded in zip(callbacks, loop_state['callbacks']):
        assert type(e) == type(e_loaded)
        if hasattr(e, "__setstate__"):
            e.__setstate__(e_loaded.__dict__)
        else:
            e.__dict__.update(e_loaded.__dict__)

    # Some diagnostics
    logger.info(loop_state)
    for k in H:
        logger.info((k, len(H)))
        break

    logger.info("epoch_start={}".format(epoch_start))

    return model, optimizer, H, epoch_start

def _to_infinite_iterator(it):

def evaluate(model, data_generators, loss_function, metrics):
    results_all = defaultdict(float)
    for dataset_id in range(len(data_generators)):
        results = defaultdict(float)
        seen = 0

        for m in metrics:
            if isinstance(m, keras_metric):
                m.reset_states()

        for x, y in data_generators[dataset_id]:
            outputs = model.predict(x)

            if isinstance(x, dict):
                for k in x:
                    seen += len(x[k])
                    assert len(outputs) == len(x[k])
                    break
            else:
                seen += len(x)
                assert len(outputs) == len(x)

            results[f'loss:{dataset_id}'] += float(np.sum(loss_function(y, outputs)))

            for m in metrics:
                results[m.__name__ + ":" + str(dataset_id)] += float(np.sum(m(y, outputs)))

        for m in metrics:
            results[m.__name__ + ":" + str(dataset_id)] /= seen
        results_all.update(results)

    # Just to make sure
    for m in metrics:
        if isinstance(m, keras_metric):
            m.reset_states()

    return results_all

def _train_on_batch(model, optimizer, x_trains, y_trains, metrics, loss_function,
                    weight_decay=0):
    batch_logs = {}

    with tf.GradientTape(persistent=True) as tape:  # TODO: Fix persistency
        loss = 0
        all_outputs = []
        all_inputs = []
        for dataset_id, (x_train, y_train) in enumerate(zip(x_trains, y_trains)):
            outputs = model(x_train, training=True)

            all_outputs.append(outputs)
            all_inputs.append(x_train)

            loss += tf.math.reduce_mean(loss_function(y_train, outputs))

            # Update logs
            for m in metrics:
                if isinstance(m, str):
                    continue
                batch_logs[m.__name__ + ":" + str(dataset_id)] = tf.math.reduce_mean(
                    m(outputs, y_train))  # If not scalar, reduce it

            batch_logs[f'loss:{dataset_id}'] = loss

    grads = tape.gradient(loss, model.trainable_weights)

    if weight_decay > 0:
        grads = [g + weight_decay * p for (g, p) in zip(grads, model.trainable_weights)]

    optimizer.apply_gradients(zip(grads, model.trainable_variables))

    return batch_logs

@tf.function
def _train_on_batch_optimized(model, optimizer, x_trains, y_trains, metrics,
                              loss_function, weight_decay=0):
    return _train_on_batch(model=model,
                           optimizer=optimizer,
                           x_trains=x_trains, y_trains=y_trains,
                           metrics=metrics,
                           loss_function=loss_function, weight_decay=weight_decay)

def _training_loop(model, datasets, optimizer, loss_function, initial_epoch, epochs, callbacks,
                   steps_per_epoch, train_on_batch, evaluate_model, metrics=[], weight_decay=0,
                   evaluation_freq=1):
    """
    Internal implementation of the training loop.
    """
    tf.keras.backend.set_learning_phase(1)

    train_generators = [_to_infinite_iterator(d[0]) for d in datasets]
    valid_generators = [d[1] for d in datasets]

    for c in callbacks:
        c.on_train_begin(model)

    cumulative_batch_id = 0

    for epoch in range(initial_epoch, epochs):
        logger.info(f"Start epoch {epoch}")

        epoch_logs = {}

        for c in callbacks:
            c.on_epoch_begin(epoch, epoch_logs)
        for batch_id in range(steps_per_epoch):
            cumulative_batch_id += 1
            batch_logs = {}

            x_trains, y_trains, y_trains_a, y_trains_b, lams = [], [], [], [], []

            for dataset_id in range(len(datasets)):
                x_train, y_train = next(train_generators[dataset_id])

                x_trains.append(x_train)
                y_trains.append(y_train)

                if isinstance(x_train, dict):
                    batch_logs.update({"size:" + str(dataset_id): len(list(x_train.values())[0])})
                else:
                    batch_logs.update({"size:" + str(dataset_id): len(x_train)})

            for c in callbacks:
                c.on_batch_begin(batch=batch_id, logs=batch_logs)

            batch_logs_step = train_on_batch(model, optimizer, x_trains, y_trains,
                                             metrics,
                                             loss_function,
                                             weight_decay=weight_decay)

            batch_logs.update(batch_logs_step)
            for k in batch_logs:
                if hasattr(batch_logs[k], "numpy"):
                    batch_logs[k] = batch_logs[k].numpy()

                if (hasattr(batch_logs[k], 'ndim') and batch_logs[k].ndim > 0) or isinstance(batch_logs[k], list):
                    batch_logs[k] = batch_logs[k]  # .numpy()
                    if isinstance(batch_logs[k], list):
                        batch_logs[k] = np.array(batch_logs[k])
                else:
                    batch_logs[k] = float(batch_logs[k])

            for c in callbacks:
                c.on_batch_end(batch=batch_id, logs=batch_logs)

        if evaluation_freq > 0 and (epoch % evaluation_freq == 0 or epoch == epochs - 1):
            tf.keras.backend.set_learning_phase(0)
            val_results = evaluate_model(model, valid_generators, loss_function, metrics)
            tf.keras.backend.set_learning_phase(1)
            for k, v in val_results.items():
                epoch_logs[f'val_{k}'] = v
        else:
            if evaluation_freq > 0:
                for k in previous_epoch_logs:
                    if k not in epoch_logs:
                        epoch_logs[k] = np.nan

        for c in callbacks:
            c.on_epoch_end(epoch, epoch_logs)

        logger.info('End of epoch {}, loss={}'.format(epoch, epoch_logs['loss:0']))

        previous_epoch_logs = dict(epoch_logs)

    for c in callbacks:
        c.on_train_end(model)

@gin.configurable
def training_loop(model, loss_function, metrics, optimizer,
                  config, save_path, datasets, steps_per_epoch, seed,
                  custom_callbacks=[], checkpoint_monitor="val_categorical_accuracy:0",
                  use_tb=False, reload=False, evaluation_freq=1,
                  n_epochs=100, save_freq=1, save_history_every_k_examples=1,
                  load_weights_from="", load_weights_and_optimizer_from="",
                  weight_decay=0, load_classifier=True):
    if load_weights_and_optimizer_from != "":
        assert load_weights_from == ""
        assert load_weights_and_optimizer_from.endswith("h5")
        logger.info(f"load_weights_and_optimizer_from={load_weights_and_optimizer_from}")
        _, optimizer = restore_model_and_optimizer(model, optimizer, load_weights_and_optimizer_from)
        logger.info("Loaded optimizer")
        model.load_weights(load_weights_and_optimizer_from, by_name=True)
        model.optimizer = optimizer
    elif load_weights_from != "":
        assert load_weights_and_optimizer_from == ""
        model.load_weights(load_weights_from)

    if reload:
        assert load_classifier is True
        logger.warning("Only custom_callbacks can be reloaded for now")
        previous_model = model
        model, optimizer, H, epoch_start = _reload(model, optimizer, save_path, custom_callbacks)
        del previous_model
        logger.warning("Changed model reference internally in the training loop!")
    else:
        if hasattr(model, "compile"):
            model.compile(optimizer=optimizer,
                          loss=loss_function,
                          metrics=[m for m in metrics if not isinstance(m, str)])  # FIXME

        save_model(model, optimizer, os.path.join(save_path, "init_weights"))

        history_csv_path, history_pkl_path = os.path.join(save_path, "history.csv"), os.path.join(save_path,
                                                                                "history.pkl")
        logger.info("Removing {} and {}".format(history_pkl_path, history_csv_path))
        os.system("rm " + history_pkl_path)
        os.system("rm " + history_csv_path)
        H, epoch_start = {}, 0

    callbacks = list(custom_callbacks)
    callbacks += _construct_default_callbacks(model, optimizer, H, save_path, checkpoint_monitor,
                                              save_freq, custom_callbacks, use_tb,
                                              save_history_every_k_examples)

    # Configure callbacks
    for clbk in callbacks:
        clbk.set_save_path(save_path)
        clbk.set_optimizer(optimizer)
        clbk.set_model(model)
        clbk.set_seed(seed)
        clbk.set_datasets(datasets)
        clbk.set_config(config)
        clbk.set_callbacks(callbacks)

    _training_loop(model, datasets, optimizer, loss_function, epoch_start, n_epochs, callbacks,
                   steps_per_epoch, train_on_batch=_train_on_batch_optimized,
                   evaluate_model=evaluate,
                   metrics=metrics, evaluation_freq=evaluation_freq,
                   weight_decay=weight_decay)

    if save_freq != -1:
        save_model(model, optimizer, os.path.join(save_path, "model_last_epoch"))


try:
        pd.DataFrame(H)[keys_to_save].to_csv(os.path.join(save_path, "history.csv"), index=False)
    except ValueError:
        logger.warning("Couldnt' save history.csv. Please read in history.pkl")

"""
Self-contained training loop for the project.
"""

import logging
import os
import pickle
from collections import defaultdict
from functools import partial
import gin
import numpy as np
import pandas as pd

import tensorflow as tf
from tensorflow.keras.metrics import Metric as keras_metric

from src.callbacks import ModelCheckpoint, LambdaCallback, History, DumpTensorboardSummaries, BaseLogger, Heartbeat, \
    MetaSaver, MeasureTime
from src.utils import save_model, restore_model_and_optimizer, restore_model, configure_neptune_exp

logger = logging.getLogger(__name__)

def _construct_default_callbacks(model, optimizer, H, save_path, checkpoint_monitor, save_freq, custom_callbacks,
                                 use_tb, save_history_every_k_examples):
    callbacks = []
    callbacks.append(MetaSaver())
    callbacks.append(BaseLogger())
    callbacks.append(MeasureTime())
    callbacks.append(Heartbeat())
    callbacks.append(LambdaCallback(on_epoch_end=partial(_append_to_history_csv, H=H)))
    callbacks.append(LambdaCallback(on_epoch_end=partial(_save_history_csv, save_path=save_path, H=H)))
    callbacks.append(History(save_every_k_examples=save_history_every_k_examples))

    if len(checkpoint_monitor):
        callbacks.append(ModelCheckpoint(monitor=checkpoint_monitor,
                                         save_best_only=True,
                                         mode='max',
                                         filepath=os.path.join(save_path, "model_best_val")))
    if save_freq > 0:
        def save_weights_fnc(epoch, logs):
            if epoch % save_freq == 0:
                logger.info("Saving model from epoch " + str(epoch))
                save_model(model, optimizer, os.path.join(save_path, "model_last_epoch"))

        callbacks.append(LambdaCallback(on_epoch_end=save_weights_fnc))

    if use_tb:
        callbacks.append(DumpTensorboardSummaries())
    callbacks.append(LambdaCallback(on_epoch_end=partial(_save_loop_state, save_callbacks=custom_callbacks,
                                                         save_path=save_path)))
    return callbacks

def _save_loop_state(epoch, logs, save_path, save_callbacks):
    loop_state = {"epochs_done": epoch, "callbacks": save_callbacks}
    pickle.dump(loop_state, open(os.path.join(save_path, "loop_state.pkl"), "wb"))

def _save_history_csv(epoch, logs, save_path, H):
    out = ""
    keys_to_save = []
    for key, value in logs.items():
        if isinstance(value, (int, float, complex, np.float32, np.float64)):
            out += "{key}={value}\t".format(key=key, value=value)
            keys_to_save.append(key)
    logger.info(out)
    logger.info("Saving history to " + os.path.join(save_path, "history.csv"))

def _append_to_history_csv(epoch, logs, H):
    for key, value in logs.items():
        if isinstance(value, (int, float, complex, np.float32, np.float64)):
            if key not in H:
                H[key] = [value]
            else:
                H[key].append(value)

            assert len(H[key]) == epoch + 1, "Len H[{}] is {}, expected {} ".format(key, len(H[key]), epoch + 1)
        else:
            pass

def _reload(model, optimizer, save_path, callbacks):
    model_last_epoch_path = os.path.join(save_path, "model_last_epoch")
    loop_state_path = os.path.join(save_path, "loop_state.pkl")
    history_csv_path = os.path.join(save_path, "history.csv")

    if not os.path.exists(model_last_epoch_path) or not os.path.exists(loop_state_path):
        logger.warning("Failed to find last epoch model or loop state")
        return {}, 0

    # Reload everything (model, optimizer, loop state)
    logger.warning("Reloading weights!")
    model, optimizer = restore_model_and_optimizer(model, optimizer, model_last_epoch_path)
    logger.info("Reloading loop state!")
    loop_state = pickle.load(open(loop_state_path, 'rb'))
    logger.info("Reloading history!")
    H = pd.read_csv(history_csv_path)
    H = {col: list(H[col].values) for col in H.columns}
    print(H)
    logger.info("Done reloading!")

    # Small back-up
    os.system("cp " + os.path.join(save_path, "history.pkl") + " " + os.path.join(save_path, "history.pkl.bckp"))

    # Setup the rest
    epoch_start = loop_state['epochs_done'] + 1
    if not len(H[next(iter(H))]) == loop_state['epochs_done'] + 1:
        raise IOError("Mismatch between saved history and epochs recorded. "
                      "Found len(H)={0} and epoch_start={1} "
                      "Run was likely interrupted incorrectly and cannot be rerun.".format(len(H[next(iter(H))]),
                                                                                epoch_start))

    # Load all callbacks from the loop_state
    for e, e_loaded in zip(callbacks, loop_state['callbacks']):
        assert type(e) == type(e_loaded)
        if hasattr(e, "__setstate__"):
            e.__setstate__(e_loaded.__dict__)
        else:
            e.__dict__.update(e_loaded.__dict__)

    # Some diagnostics
    logger.info(loop_state)
    for k in H:
        logger.info((k, len(H)))
        break

    logger.info("epoch_start={}".format(epoch_start))

    return model, optimizer, H, epoch_start

def _to_infinite_iterator(it):
    while True:
        for bundle in it:
            yield bundle

def evaluate(model, data_generators, loss_function, metrics):
    results_all = defaultdict(float)
    for dataset_id in range(len(data_generators)):
        results = defaultdict(float)
        seen = 0

        for m in metrics:
            if isinstance(m, keras_metric):
                m.reset_states()

        for x, y in data_generators[dataset_id]:
            outputs = model.predict(x)

            if isinstance(x, dict):
                for k in x:
                    seen += len(x[k])
                    assert len(outputs) == len(x[k])
                    break
            else:
                seen += len(x)
                assert len(outputs) == len(x)

            results[f'loss:{dataset_id}'] += float(np.sum(loss_function(y, outputs)))

            for m in metrics:
                results[m.__name__ + ":" + str(dataset_id)] += float(np.sum(m(y, outputs)))

        for m in metrics:
            results[m.__name__ + ":" + str(dataset_id)] /= seen
        results_all.update(results)

    # Just to make sure
    for m in metrics:
        if isinstance(m, keras_metric):
            m.reset_states()

    return results_all

def _train_on_batch(model, optimizer, x_trains, y_trains, metrics, loss_function,
                    weight_decay=0):
    batch_logs = {}

    with tf.GradientTape(persistent=True) as tape:  # TODO: Fix persistency
        loss = 0
        all_outputs = []
        all_inputs = []
        for dataset_id, (x_train, y_train) in enumerate(zip(x_trains, y_trains)):
            outputs = model(x_train, training=True)

            all_outputs.append(outputs)
            all_inputs.append(x_train)

            loss += tf.math.reduce_mean(loss_function(y_train, outputs))

            # Update logs
            for m in metrics:
                if isinstance(m, str):
                    continue
                batch_logs[m.__name__ + ":" + str(dataset_id)] = tf.math.reduce_mean(
                    m(outputs, y_train))  # If not scalar, reduce it

            batch_logs[f'loss:{dataset_id}'] = loss

    grads = tape.gradient(loss, model.trainable_weights)

    if weight_decay > 0:
        grads = [g + weight_decay * p for (g, p) in zip(grads, model.trainable_weights)]

    optimizer.apply_gradients(zip(grads, model.trainable_variables))

    return batch_logs

@tf.function
def _train_on_batch_optimized(model, optimizer, x_trains, y_trains, metrics,
                              loss_function, weight_decay=0):
    return _train_on_batch(model=model,
                           optimizer=optimizer,
                           x_trains=x_trains, y_trains=y_trains,
                           metrics=metrics,
                           loss_function=loss_function, weight_decay=weight_decay)

def _training_loop(model, datasets, optimizer, loss_function, initial_epoch, epochs, callbacks,
                   steps_per_epoch, train_on_batch, evaluate_model, metrics=[], weight_decay=0,
                   evaluation_freq=1):
    """
    Internal implementation of the training loop.
    """
    tf.keras.backend.set_learning_phase(1)

    train_generators = [_to_infinite_iterator(d[0]) for d in datasets]
    valid_generators = [d[1] for d in datasets]

    for c in callbacks:
        c.on_train_begin(model)

    cumulative_batch_id = 0

    for epoch in range(initial_epoch, epochs):
        logger.info(f"Start epoch {epoch}")

        epoch_logs = {}

        for c in callbacks:
            c.on_epoch_begin(epoch, epoch_logs)
        for batch_id in range(steps_per_epoch):
            cumulative_batch_id += 1
            batch_logs = {}

            x_trains, y_trains, y_trains_a, y_trains_b, lams = [], [], [], [], []

            for dataset_id in range(len(datasets)):
                x_train, y_train = next(train_generators[dataset_id])

                x_trains.append(x_train)
                y_trains.append(y_train)

                if isinstance(x_train, dict):
                    batch_logs.update({"size:" + str(dataset_id): len(list(x_train.values())[0])})
                else:
                    batch_logs.update({"size:" + str(dataset_id): len(x_train)})

            for c in callbacks:
                c.on_batch_begin(batch=batch_id, logs=batch_logs)

            batch_logs_step = train_on_batch(model, optimizer, x_trains, y_trains,
                                             metrics,
                                             loss_function,
                                             weight_decay=weight_decay)

            batch_logs.update(batch_logs_step)
            for k in batch_logs:
                if hasattr(batch_logs[k], "numpy"):
                    batch_logs[k] = batch_logs[k].numpy()

                if (hasattr(batch_logs[k], 'ndim') and batch_logs[k].ndim > 0) or isinstance(batch_logs[k], list):
                    batch_logs[k] = batch_logs[k]  # .numpy()
                    if isinstance(batch_logs[k], list):
                        batch_logs[k] = np.array(batch_logs[k])
                else:
                    batch_logs[k] = float(batch_logs[k])

            for c in callbacks:
                c.on_batch_end(batch=batch_id, logs=batch_logs)

        if evaluation_freq > 0 and (epoch % evaluation_freq == 0 or epoch == epochs - 1):
            tf.keras.backend.set_learning_phase(0)
            val_results = evaluate_model(model, valid_generators, loss_function, metrics)
            tf.keras.backend.set_learning_phase(1)
            for k, v in val_results.items():
                epoch_logs[f'val_{k}'] = v
        else:
            if evaluation_freq > 0:
                for k in previous_epoch_logs:
                    if k not in epoch_logs:
                        epoch_logs[k] = np.nan

        for c in callbacks:
            c.on_epoch_end(epoch, epoch_logs)

        logger.info('End of epoch {}, loss={}'.format(epoch, epoch_logs['loss:0']))

        previous_epoch_logs = dict(epoch_logs)

    for c in callbacks:
        c.on_train_end(model)

@gin.configurable
def training_loop(model, loss_function, metrics, optimizer,
                  config, save_path, datasets, steps_per_epoch, seed,
                  custom_callbacks=[], checkpoint_monitor="val_categorical_accuracy:0",
                  use_tb=False, reload=False, evaluation_freq=1,
                  n_epochs=100, save_freq=1, save_history_every_k_examples=1,
                  load_weights_from="", load_weights_and_optimizer_from="",
                  weight_decay=0, load_classifier=True):
    if load_weights_and_optimizer_from != "":
        assert load_weights_from == ""
        assert load_weights_and_optimizer_from.endswith("h5")
        logger.info(f"load_weights_and_optimizer_from={load_weights_and_optimizer_from}")
        _, optimizer = restore_model_and_optimizer(model, optimizer, load_weights_and_optimizer_from)
        logger.info("Loaded optimizer")
        model.load_weights(load_weights_and_optimizer_from, by_name=True)
        model.optimizer = optimizer
    elif load_weights_from != "":
        assert load_weights_and_optimizer_from == ""
        model.load_weights(load_weights_from)

    if reload:
        assert load_classifier is True
        logger.warning("Only custom_callbacks can be reloaded for now")
        previous_model = model
        model, optimizer, H, epoch_start = _reload(model, optimizer, save_path, custom_callbacks)
        del previous_model
        logger.warning("Changed model reference internally in the training loop!")
    else:
        if hasattr(model, "compile"):
            model.compile(optimizer=optimizer,
                          loss=loss_function,
                          metrics=[m for m in metrics if not isinstance(m, str)])  # FIXME

        save_model(model, optimizer, os.path.join(save_path, "init_weights"))

        history_csv_path, history_pkl_path = os.path.join(save_path, "history.csv"), os.path.join(save_path,
                                                                                "history.pkl")
        logger.info("Removing {} and {}".format(history_pkl_path, history_csv_path))
        os.system("rm " + history_pkl_path)
        os.system("rm " + history_csv_path)
        H, epoch_start = {}, 0

    callbacks = list(custom_callbacks)
    callbacks += _construct_default_callbacks(model, optimizer, H, save_path, checkpoint_monitor,
                                              save_freq, custom_callbacks, use_tb,
                                              save_history_every_k_examples)

    # Configure callbacks
    for clbk in callbacks:
        clbk.set_save_path(save_path)
        clbk.set_optimizer(optimizer)
        clbk.set_model(model)
        clbk.set_seed(seed)
        clbk.set_datasets(datasets)
        clbk.set_config(config)
        clbk.set_callbacks(callbacks)

    _training_loop(model, datasets, optimizer, loss_function, epoch_start, n_epochs, callbacks,
                   steps_per_epoch, train_on_batch=_train_on_batch_optimized,
                   evaluate_model=evaluate,
                   metrics=metrics, evaluation_freq=evaluation_freq,
                   weight_decay=weight_decay)

    if save_freq != -1:
        save_model(model, optimizer, os.path.join(save_path, "model_last_epoch"))


# pylmm is a python-based linear mixed-model solver with applications to GWAS

import pdb
import time
import sys

def printOutHead(): out.write("\t".join(["SNP_ID","BETA","BETA_SD","F_STAT","P_VALUE"]) + "\n")

out.write("\t".join([str(x) for x in [id,beta,betaSD,ts,ps]]) + "\n")

from optparse import OptionParser,OptionGroup
usage = """usage: %prog [options] --kfile kinshipFile --[tfile | bfile] plinkFileBase outfileBase

This program provides basic genome-wide association (GWAS) functionality.  You provide a phenotype and genotype file as well as a pre-computed (use pylmmKinship.py) kinship matrix and the program outputs a result file with information about each SNP, including the association p-value.
The input file are all standard plink formatted with the first two columns specifiying the individual and family ID.  For the phenotype file, we accept either NA or -9 to denote missing values.

Basic usage:

      python pylmmGWAS.py -v --bfile plinkFile --kfile preComputedKinship.kin --phenofile plinkFormattedPhenotypeFile resultFile

	    """
parser = OptionParser(usage=usage)

basicGroup = OptionGroup(parser, "Basic Options")
advancedGroup = OptionGroup(parser, "Advanced Options")
experimentalGroup = OptionGroup(parser, "Experimental Options")

#                  help="The base for a PLINK ped file")
basicGroup.add_option("--tfile", dest="tfile",
                  help="The base for a PLINK tped file")
basicGroup.add_option("--bfile", dest="bfile",
                  help="The base for a PLINK binary bed file")
basicGroup.add_option("--phenofile", dest="phenoFile", default=None,
                  help="Without this argument the program will look for a file with .pheno that has the plinkFileBase root.  If you want to specify an alternative phenotype file, then use this argument.  This file should be in plink format. ")

# EMMA Options
basicGroup.add_option("--emmaSNP", dest="emmaFile", default=None,
                  help="For backwards compatibility with emma, we allow for \"EMMA\" file formats.  This is just a text file with individuals on the columns and snps on the rows.")
basicGroup.add_option("--emmaPHENO", dest="emmaPheno", default=None,
                  help="For backwards compatibility with emma, we allow for \"EMMA\" file formats.  This is just a text file with each phenotype as one row.")
basicGroup.add_option("--emmaCOV", dest="emmaCov", default=None,
                  help="For backwards compatibility with emma, we allow for \"EMMA\" file formats.  This is just a text file with each covariate as one row.")

basicGroup.add_option("--kfile", dest="kfile",
                  help="The location of a kinship file.  This is an nxn plain text file and can be computed with the pylmmKinship program.")
basicGroup.add_option("--covfile", dest="covfile",
                  help="The location of a covariate file file.  This is a plink formatted covariate file.")
basicGroup.add_option("-p", type="int", dest="pheno", help="The phenotype index to be used in association.", default=0)

advancedGroup.add_option("--removeMissingGenotypes",
                  action="store_false", dest="normalizeGenotype", default=True,
                  help="By default the program replaces missing genotypes with the minor allele frequency.  This option overrides that behavior making the program remove missing individuals.  NOTE: This can increase running time due to the need to recompute the eigendecomposition for each SNP with missing values.")
advancedGroup.add_option("--refit",
                  action="store_true", dest="refit", default=False,
                  help="Refit the variance components at each SNP (default is to lock in the variance components under the null).")

advancedGroup.add_option("--REML",
                  action="store_true", dest="REML", default=False,
                  help="Use restricted maximum-likelihood (REML) (default is maximum-likelihood).")
#advancedGroup.add_option("-e", "--efile", dest="saveEig", help="Save eigendecomposition to this file.")
advancedGroup.add_option("--eigen", dest="eigenfile",
                  help="The location of the precomputed eigendecomposition for the kinship file.  These can be computed with pylmmKinship.py.")
advancedGroup.add_option("--noMean", dest="noMean", default=False,action="store_true",
                  help="This option only applies when --cofile is used.  When covfile is provided, the program will automatically add a global mean covariate to the model unless this option is specified.")

advancedGroup.add_option("-v", "--verbose",
                  action="store_true", dest="verbose", default=False,
                  help="Print extra info")

# Experimental Group Options
experimentalGroup.add_option("--kfile2", dest="kfile2",
                  help="The location of a second kinship file.  This file has the same format as the first kinship.  This might be used if you want to correct for another form of confounding.")

parser.add_option_group(basicGroup)
parser.add_option_group(advancedGroup)
parser.add_option_group(experimentalGroup)

(options, args) = parser.parse_args()

import sys
import os
import numpy as np
from scipy import linalg
from pylmm.lmm import LMM
from pylmm import input

if len(args) != 1:
   parser.print_help()
   sys.exit()

outFile = args[0]

if not options.tfile and not options.bfile and not options.emmaFile:
   parser.error("You must provide at least one PLINK input file base (--tfile or --bfile) or an EMMA formatted file (--emmaSNP).")
if not options.kfile:
   parser.error("Please provide a pre-computed kinship file")

# READING PLINK input
if options.verbose: sys.stderr.write("Reading SNP input...\n")
if options.bfile: IN = input.plink(options.bfile,type='b', phenoFile=options.phenoFile,normGenotype=options.normalizeGenotype)
elif options.tfile: IN = input.plink(options.tfile,type='t', phenoFile=options.phenoFile,normGenotype=options.normalizeGenotype)
elif options.emmaFile: IN = input.plink(options.emmaFile,type='emma', phenoFile=options.phenoFile,normGenotype=options.normalizeGenotype)
else: parser.error("You must provide at least one PLINK input file base")

if not os.path.isfile(options.phenoFile or IN.fbase + '.phenos') and not os.path.isfile(options.emmaPheno):
   parser.error("No .pheno file exist for %s.  Please provide a phenotype file using the --phenofile or --emmaPHENO argument." % (options.phenoFile or IN.fbase + '.phenos'))

# Read the emma phenotype file if provided.
# Format should be rows are phenotypes and columns are individuals.
if options.emmaPheno:
   f = open(options.emmaPheno,'r')
   P = []
   for line in f:
      v = line.strip().split()
      p = []
      for x in v:
	 try:
	    p.append(float(x))
	 except: p.append(np.nan)
      P.append(p)
   f.close()
   IN.phenos = np.array(P).T

# READING Covariate File
if options.covfile:
   if options.verbose:
      sys.stderr.write("Reading covariate file...\n")
   P = IN.getCovariates(options.covfile)
   if options.noMean:
      X0 = P
   else:
      X0 = np.hstack([np.ones((IN.phenos.shape[0],1)),P])
elif options.emmaCov:
   if options.verbose:
      sys.stderr.write("Reading covariate file...\n")
   P = IN.getCovariatesEMMA(options.emmaCov)
   if options.noMean:
      X0 = P
   else:
      X0 = np.hstack([np.ones((IN.phenos.shape[0],1)),P])
else:
   X0 = np.ones((IN.phenos.shape[0],1))

if np.isnan(X0).sum():
      parser.error("The covariate file %s contains missing values. At this time we are not dealing with this case.  Either remove those individuals with missing values or replace them in some way.")

# READING Kinship - major bottleneck for large datasets
if options.verbose: sys.stderr.write("Reading kinship...\n")
begin = time.time()
# This method seems to be the fastest and works if you already know the size of the matrix
if options.kfile[-3:] == '.gz':
   import gzip
   f = gzip.open(options.kfile,'r')
   F = f.read() # might exhaust mem if the file is huge
   K = np.fromstring(F,sep=' ') # Assume that space separated
   f.close()
else:
   K = np.fromfile(open(options.kfile,'r'),sep=" ")
K.resize((len(IN.indivs),len(IN.indivs)))
end = time.time()
# Other slower ways
if options.verbose: sys.stderr.write("Read the %d x %d kinship matrix in %0.3fs \n" % (K.shape[0],K.shape[1],end-begin))

if options.kfile2:
   if options.verbose: sys.stderr.write("Reading second kinship...\n")
   begin = time.time()
   # This method seems to be the fastest and works if you already know the size of the matrix
   if options.kfile2[-3:] == '.gz':
      import gzip
      f = gzip.open(options.kfile2,'r')
      F = f.read() # might exhaust mem if the file is huge
      K2 = np.fromstring(F,sep=' ') # Assume that space separated
      f.close()
   else: K2 = np.fromfile(open(options.kfile2,'r'),sep=" ")
   K2.resize((len(IN.indivs),len(IN.indivs)))
   end = time.time()
   if options.verbose: sys.stderr.write("Read the %d x %d second kinship matrix in %0.3fs \n" % (K2.shape[0],K2.shape[1],end-begin))

# PROCESS the phenotype data -- Remove missing phenotype values
# Keep will now index into the "full" data to select what we keep (either everything or a subset of non missing data
Y = IN.phenos[:,options.pheno]
v = np.isnan(Y)
keep = True - v
if v.sum():
   if options.verbose: sys.stderr.write("Cleaning the phenotype vector by removing %d individuals...\n" % (v.sum()))
   Y = Y[keep]
   X0 = X0[keep,:]
   K = K[keep,:][:,keep]
   if options.kfile2: K2 = K2[keep,:][:,keep]
   Kva = []
   Kve = []

# Only load the decomposition if we did not remove individuals.
# Otherwise it would not be correct and we would have to compute it again.
if not v.sum() and options.eigenfile:
   if options.verbose: sys.stderr.write("Loading pre-computed eigendecomposition...\n")
   Kva = np.load(options.eigenfile + ".Kva")
   Kve = np.load(options.eigenfile + ".Kve")
else:
   Kva = []
   Kve = []

# CREATE LMM object for association
n = K.shape[0]
if not options.kfile2:  L = LMM(Y,K,Kva,Kve,X0,verbose=options.verbose)
else:  L = LMM_withK2(Y,K,Kva,Kve,X0,verbose=options.verbose,K2=K2)

# Fit the null model -- if refit is true we will refit for each SNP, so no reason to run here
if not options.refit:
   if options.verbose: sys.stderr.write("Computing fit for null model\n")
   L.fit()
   if options.verbose and not options.kfile2: sys.stderr.write("\t heritability=%0.3f, sigma=%0.3f\n" % (L.optH,L.optSigma))
   if options.verbose and options.kfile2: sys.stderr.write("\t heritability=%0.3f, sigma=%0.3f, w=%0.3f\n" % (L.optH,L.optSigma,L.optW))

# Buffers for pvalues and t-stats
PS = []
TS = []
count = 0
out = open(outFile,'w')
printOutHead()

for snp,id in IN:
   count += 1
   if options.verbose and count % 1000 == 0:
      sys.stderr.write("At SNP %d\n" % count)

   x = snp[keep].reshape((n,1))
   v = np.isnan(x).reshape((-1,))
   # Check SNPs for missing values
   if v.sum():
      keeps = True - v
      xs = x[keeps,:]
      if keeps.sum() <= 1 or xs.var() <= 1e-6:
	 PS.append(np.nan)
	 TS.append(np.nan)
	 outputResult(id,np.nan,np.nan,np.nan,np.nan)
	 continue

      # Its ok to center the genotype -  I used options.normalizeGenotype to
      # force the removal of missing genotypes as opposed to replacing them with MAF.
      if not options.normalizeGenotype: xs = (xs - xs.mean()) / np.sqrt(xs.var())
      Ys = Y[keeps]
      X0s = X0[keeps,:]
      Ks = K[keeps,:][:,keeps]
      if options.kfile2: K2s = K2[keeps,:][:,keeps]
      if options.kfile2: Ls = LMM_withK2(Ys,Ks,X0=X0s,verbose=options.verbose,K2=K2s)
      else: Ls = LMM(Ys,Ks,X0=X0s,verbose=options.verbose)
      if options.refit: Ls.fit(X=xs,REML=options.REML)
      else:
	 Ls.fit(REML=options.REML)
      ts,ps,beta,betaVar = Ls.association(xs,REML=options.REML,returnBeta=True)
   else:
      if x.var() == 0:
	 PS.append(np.nan)
	 TS.append(np.nan)
	 outputResult(id,np.nan,np.nan,np.nan,np.nan)
	 continue

      if options.refit: L.fit(X=x,REML=options.REML)
      ts,ps,beta,betaVar = L.association(x,REML=options.REML,returnBeta=True)

   outputResult(id,beta,np.sqrt(betaVar).sum(),ts,ps)
   PS.append(ps)
   TS.append(ts)


out.write("\t".join([str(x) for x in [id,beta,betaSD,ts,ps]]) + "\n")

# pylmm is a python-based linear mixed-model solver with applications to GWAS

import pdb
import time
import sys

def printOutHead(): out.write("\t".join(["SNP_ID","BETA","BETA_SD","F_STAT","P_VALUE"]) + "\n")
def outputResult(id,beta,betaSD,ts,ps):

from optparse import OptionParser,OptionGroup
usage = """usage: %prog [options] --kfile kinshipFile --[tfile | bfile] plinkFileBase outfileBase

This program provides basic genome-wide association (GWAS) functionality.  You provide a phenotype and genotype file as well as a pre-computed (use pylmmKinship.py) kinship matrix and the program outputs a result file with information about each SNP, including the association p-value.
The input file are all standard plink formatted with the first two columns specifiying the individual and family ID.  For the phenotype file, we accept either NA or -9 to denote missing values.

Basic usage:

      python pylmmGWAS.py -v --bfile plinkFile --kfile preComputedKinship.kin --phenofile plinkFormattedPhenotypeFile resultFile

	    """
parser = OptionParser(usage=usage)

basicGroup = OptionGroup(parser, "Basic Options")
advancedGroup = OptionGroup(parser, "Advanced Options")
experimentalGroup = OptionGroup(parser, "Experimental Options")

#                  help="The base for a PLINK ped file")
basicGroup.add_option("--tfile", dest="tfile",
                  help="The base for a PLINK tped file")
basicGroup.add_option("--bfile", dest="bfile",
                  help="The base for a PLINK binary bed file")
basicGroup.add_option("--phenofile", dest="phenoFile", default=None,
                  help="Without this argument the program will look for a file with .pheno that has the plinkFileBase root.  If you want to specify an alternative phenotype file, then use this argument.  This file should be in plink format. ")

# EMMA Options
basicGroup.add_option("--emmaSNP", dest="emmaFile", default=None,
                  help="For backwards compatibility with emma, we allow for \"EMMA\" file formats.  This is just a text file with individuals on the columns and snps on the rows.")
basicGroup.add_option("--emmaPHENO", dest="emmaPheno", default=None,
                  help="For backwards compatibility with emma, we allow for \"EMMA\" file formats.  This is just a text file with each phenotype as one row.")
basicGroup.add_option("--emmaCOV", dest="emmaCov", default=None,
                  help="For backwards compatibility with emma, we allow for \"EMMA\" file formats.  This is just a text file with each covariate as one row.")

basicGroup.add_option("--kfile", dest="kfile",
                  help="The location of a kinship file.  This is an nxn plain text file and can be computed with the pylmmKinship program.")
basicGroup.add_option("--covfile", dest="covfile",
                  help="The location of a covariate file file.  This is a plink formatted covariate file.")
basicGroup.add_option("-p", type="int", dest="pheno", help="The phenotype index to be used in association.", default=0)

advancedGroup.add_option("--removeMissingGenotypes",
                  action="store_false", dest="normalizeGenotype", default=True,
                  help="By default the program replaces missing genotypes with the minor allele frequency.  This option overrides that behavior making the program remove missing individuals.  NOTE: This can increase running time due to the need to recompute the eigendecomposition for each SNP with missing values.")
advancedGroup.add_option("--refit",
                  action="store_true", dest="refit", default=False,
                  help="Refit the variance components at each SNP (default is to lock in the variance components under the null).")

advancedGroup.add_option("--REML",
                  action="store_true", dest="REML", default=False,
                  help="Use restricted maximum-likelihood (REML) (default is maximum-likelihood).")
#advancedGroup.add_option("-e", "--efile", dest="saveEig", help="Save eigendecomposition to this file.")
advancedGroup.add_option("--eigen", dest="eigenfile",
                  help="The location of the precomputed eigendecomposition for the kinship file.  These can be computed with pylmmKinship.py.")
advancedGroup.add_option("--noMean", dest="noMean", default=False,action="store_true",
                  help="This option only applies when --cofile is used.  When covfile is provided, the program will automatically add a global mean covariate to the model unless this option is specified.")

advancedGroup.add_option("-v", "--verbose",
                  action="store_true", dest="verbose", default=False,
                  help="Print extra info")

# Experimental Group Options
experimentalGroup.add_option("--kfile2", dest="kfile2",
                  help="The location of a second kinship file.  This file has the same format as the first kinship.  This might be used if you want to correct for another form of confounding.")

parser.add_option_group(basicGroup)
parser.add_option_group(advancedGroup)
parser.add_option_group(experimentalGroup)

(options, args) = parser.parse_args()

import sys
import os
import numpy as np
from scipy import linalg
from pylmm.lmm import LMM
from pylmm import input

if len(args) != 1:
   parser.print_help()
   sys.exit()

outFile = args[0]

if not options.tfile and not options.bfile and not options.emmaFile:
   parser.error("You must provide at least one PLINK input file base (--tfile or --bfile) or an EMMA formatted file (--emmaSNP).")
if not options.kfile:
   parser.error("Please provide a pre-computed kinship file")

# READING PLINK input
if options.verbose: sys.stderr.write("Reading SNP input...\n")
if options.bfile: IN = input.plink(options.bfile,type='b', phenoFile=options.phenoFile,normGenotype=options.normalizeGenotype)
elif options.tfile: IN = input.plink(options.tfile,type='t', phenoFile=options.phenoFile,normGenotype=options.normalizeGenotype)
elif options.emmaFile: IN = input.plink(options.emmaFile,type='emma', phenoFile=options.phenoFile,normGenotype=options.normalizeGenotype)
else: parser.error("You must provide at least one PLINK input file base")

if not os.path.isfile(options.phenoFile or IN.fbase + '.phenos') and not os.path.isfile(options.emmaPheno):
   parser.error("No .pheno file exist for %s.  Please provide a phenotype file using the --phenofile or --emmaPHENO argument." % (options.phenoFile or IN.fbase + '.phenos'))

# Read the emma phenotype file if provided.
# Format should be rows are phenotypes and columns are individuals.
if options.emmaPheno:
   f = open(options.emmaPheno,'r')
   P = []
   for line in f:
      v = line.strip().split()
      p = []
      for x in v:
	 try:
	    p.append(float(x))
	 except: p.append(np.nan)
      P.append(p)
   f.close()
   IN.phenos = np.array(P).T

# READING Covariate File
if options.covfile:
   if options.verbose:
      sys.stderr.write("Reading covariate file...\n")
   P = IN.getCovariates(options.covfile)
   if options.noMean:
      X0 = P
   else:
      X0 = np.hstack([np.ones((IN.phenos.shape[0],1)),P])
elif options.emmaCov:
   if options.verbose:
      sys.stderr.write("Reading covariate file...\n")
   P = IN.getCovariatesEMMA(options.emmaCov)
   if options.noMean:
      X0 = P
   else:
      X0 = np.hstack([np.ones((IN.phenos.shape[0],1)),P])
else:
   X0 = np.ones((IN.phenos.shape[0],1))

if np.isnan(X0).sum():
      parser.error("The covariate file %s contains missing values. At this time we are not dealing with this case.  Either remove those individuals with missing values or replace them in some way.")

# READING Kinship - major bottleneck for large datasets
if options.verbose: sys.stderr.write("Reading kinship...\n")
begin = time.time()
# This method seems to be the fastest and works if you already know the size of the matrix
if options.kfile[-3:] == '.gz':
   import gzip
   f = gzip.open(options.kfile,'r')
   F = f.read() # might exhaust mem if the file is huge
   K = np.fromstring(F,sep=' ') # Assume that space separated
   f.close()
else:
   K = np.fromfile(open(options.kfile,'r'),sep=" ")
K.resize((len(IN.indivs),len(IN.indivs)))
end = time.time()
# Other slower ways
if options.verbose: sys.stderr.write("Read the %d x %d kinship matrix in %0.3fs \n" % (K.shape[0],K.shape[1],end-begin))

if options.kfile2:
   if options.verbose: sys.stderr.write("Reading second kinship...\n")
   begin = time.time()
   # This method seems to be the fastest and works if you already know the size of the matrix
   if options.kfile2[-3:] == '.gz':
      import gzip
      f = gzip.open(options.kfile2,'r')
      F = f.read() # might exhaust mem if the file is huge
      K2 = np.fromstring(F,sep=' ') # Assume that space separated
      f.close()
   else: K2 = np.fromfile(open(options.kfile2,'r'),sep=" ")
   K2.resize((len(IN.indivs),len(IN.indivs)))
   end = time.time()
   if options.verbose: sys.stderr.write("Read the %d x %d second kinship matrix in %0.3fs \n" % (K2.shape[0],K2.shape[1],end-begin))

# PROCESS the phenotype data -- Remove missing phenotype values
# Keep will now index into the "full" data to select what we keep (either everything or a subset of non missing data
Y = IN.phenos[:,options.pheno]
v = np.isnan(Y)
keep = True - v
if v.sum():
   if options.verbose: sys.stderr.write("Cleaning the phenotype vector by removing %d individuals...\n" % (v.sum()))
   Y = Y[keep]
   X0 = X0[keep,:]
   K = K[keep,:][:,keep]
   if options.kfile2: K2 = K2[keep,:][:,keep]
   Kva = []
   Kve = []

# Only load the decomposition if we did not remove individuals.
# Otherwise it would not be correct and we would have to compute it again.
if not v.sum() and options.eigenfile:
   if options.verbose: sys.stderr.write("Loading pre-computed eigendecomposition...\n")
   Kva = np.load(options.eigenfile + ".Kva")
   Kve = np.load(options.eigenfile + ".Kve")
else:
   Kva = []
   Kve = []

# CREATE LMM object for association
n = K.shape[0]
if not options.kfile2:  L = LMM(Y,K,Kva,Kve,X0,verbose=options.verbose)
else:  L = LMM_withK2(Y,K,Kva,Kve,X0,verbose=options.verbose,K2=K2)

# Fit the null model -- if refit is true we will refit for each SNP, so no reason to run here
if not options.refit:
   if options.verbose: sys.stderr.write("Computing fit for null model\n")
   L.fit()
   if options.verbose and not options.kfile2: sys.stderr.write("\t heritability=%0.3f, sigma=%0.3f\n" % (L.optH,L.optSigma))
   if options.verbose and options.kfile2: sys.stderr.write("\t heritability=%0.3f, sigma=%0.3f, w=%0.3f\n" % (L.optH,L.optSigma,L.optW))

# Buffers for pvalues and t-stats
PS = []
TS = []
count = 0
out = open(outFile,'w')
printOutHead()

for snp,id in IN:
   count += 1
   if options.verbose and count % 1000 == 0:
      sys.stderr.write("At SNP %d\n" % count)

   x = snp[keep].reshape((n,1))
   v = np.isnan(x).reshape((-1,))
   # Check SNPs for missing values
   if v.sum():
      keeps = True - v
      xs = x[keeps,:]
      if keeps.sum() <= 1 or xs.var() <= 1e-6:
	 PS.append(np.nan)
	 TS.append(np.nan)
	 outputResult(id,np.nan,np.nan,np.nan,np.nan)
	 continue

      # Its ok to center the genotype -  I used options.normalizeGenotype to
      # force the removal of missing genotypes as opposed to replacing them with MAF.
      if not options.normalizeGenotype: xs = (xs - xs.mean()) / np.sqrt(xs.var())
      Ys = Y[keeps]
      X0s = X0[keeps,:]
      Ks = K[keeps,:][:,keeps]
      if options.kfile2: K2s = K2[keeps,:][:,keeps]
      if options.kfile2: Ls = LMM_withK2(Ys,Ks,X0=X0s,verbose=options.verbose,K2=K2s)
      else: Ls = LMM(Ys,Ks,X0=X0s,verbose=options.verbose)
      if options.refit: Ls.fit(X=xs,REML=options.REML)
      else:
	 Ls.fit(REML=options.REML)
      ts,ps,beta,betaVar = Ls.association(xs,REML=options.REML,returnBeta=True)
   else:
      if x.var() == 0:
	 PS.append(np.nan)
	 TS.append(np.nan)
	 outputResult(id,np.nan,np.nan,np.nan,np.nan)
	 continue

      if options.refit: L.fit(X=x,REML=options.REML)
      ts,ps,beta,betaVar = L.association(x,REML=options.REML,returnBeta=True)

   outputResult(id,beta,np.sqrt(betaVar).sum(),ts,ps)
   PS.append(ps)
   TS.append(ts)


# pylmm is a python-based linear mixed-model solver with applications to GWAS

import pdb
import time
import sys

def printOutHead(): out.write("\t".join(["SNP_ID","BETA","BETA_SD","F_STAT","P_VALUE"]) + "\n")
def outputResult(id,beta,betaSD,ts,ps):
   out.write("\t".join([str(x) for x in [


for x in v:
	 try:
	    p.append(float(x))
	 except: p.append(np.nan)

# pylmm is a python-based linear mixed-model solver with applications to GWAS

import pdb
import time
import sys

def printOutHead(): out.write("\t".join(["SNP_ID","BETA","BETA_SD","F_STAT","P_VALUE"]) + "\n")
def outputResult(id,beta,betaSD,ts,ps):
   out.write("\t".join([str(x) for x in [id,beta,betaSD,ts,ps]]) + "\n")

from optparse import OptionParser,OptionGroup
usage = """usage: %prog [options] --kfile kinshipFile --[tfile | bfile] plinkFileBase outfileBase

This program provides basic genome-wide association (GWAS) functionality.  You provide a phenotype and genotype file as well as a pre-computed (use pylmmKinship.py) kinship matrix and the program outputs a result file with information about each SNP, including the association p-value.
The input file are all standard plink formatted with the first two columns specifiying the individual and family ID.  For the phenotype file, we accept either NA or -9 to denote missing values.

Basic usage:

      python pylmmGWAS.py -v --bfile plinkFile --kfile preComputedKinship.kin --phenofile plinkFormattedPhenotypeFile resultFile

	    """
parser = OptionParser(usage=usage)

basicGroup = OptionGroup(parser, "Basic Options")
advancedGroup = OptionGroup(parser, "Advanced Options")
experimentalGroup = OptionGroup(parser, "Experimental Options")

#                  help="The base for a PLINK ped file")
basicGroup.add_option("--tfile", dest="tfile",
                  help="The base for a PLINK tped file")
basicGroup.add_option("--bfile", dest="bfile",
                  help="The base for a PLINK binary bed file")
basicGroup.add_option("--phenofile", dest="phenoFile", default=None,
                  help="Without this argument the program will look for a file with .pheno that has the plinkFileBase root.  If you want to specify an alternative phenotype file, then use this argument.  This file should be in plink format. ")

# EMMA Options
basicGroup.add_option("--emmaSNP", dest="emmaFile", default=None,
                  help="For backwards compatibility with emma, we allow for \"EMMA\" file formats.  This is just a text file with individuals on the columns and snps on the rows.")
basicGroup.add_option("--emmaPHENO", dest="emmaPheno", default=None,
                  help="For backwards compatibility with emma, we allow for \"EMMA\" file formats.  This is just a text file with each phenotype as one row.")
basicGroup.add_option("--emmaCOV", dest="emmaCov", default=None,
                  help="For backwards compatibility with emma, we allow for \"EMMA\" file formats.  This is just a text file with each covariate as one row.")

basicGroup.add_option("--kfile", dest="kfile",
                  help="The location of a kinship file.  This is an nxn plain text file and can be computed with the pylmmKinship program.")
basicGroup.add_option("--covfile", dest="covfile",
                  help="The location of a covariate file file.  This is a plink formatted covariate file.")
basicGroup.add_option("-p", type="int", dest="pheno", help="The phenotype index to be used in association.", default=0)

advancedGroup.add_option("--removeMissingGenotypes",
                  action="store_false", dest="normalizeGenotype", default=True,
                  help="By default the program replaces missing genotypes with the minor allele frequency.  This option overrides that behavior making the program remove missing individuals.  NOTE: This can increase running time due to the need to recompute the eigendecomposition for each SNP with missing values.")
advancedGroup.add_option("--refit",
                  action="store_true", dest="refit", default=False,
                  help="Refit the variance components at each SNP (default is to lock in the variance components under the null).")

advancedGroup.add_option("--REML",
                  action="store_true", dest="REML", default=False,
                  help="Use restricted maximum-likelihood (REML) (default is maximum-likelihood).")
#advancedGroup.add_option("-e", "--efile", dest="saveEig", help="Save eigendecomposition to this file.")
advancedGroup.add_option("--eigen", dest="eigenfile",
                  help="The location of the precomputed eigendecomposition for the kinship file.  These can be computed with pylmmKinship.py.")
advancedGroup.add_option("--noMean", dest="noMean", default=False,action="store_true",
                  help="This option only applies when --cofile is used.  When covfile is provided, the program will automatically add a global mean covariate to the model unless this option is specified.")

advancedGroup.add_option("-v", "--verbose",
                  action="store_true", dest="verbose", default=False,
                  help="Print extra info")

# Experimental Group Options
experimentalGroup.add_option("--kfile2", dest="kfile2",
                  help="The location of a second kinship file.  This file has the same format as the first kinship.  This might be used if you want to correct for another form of confounding.")

parser.add_option_group(basicGroup)
parser.add_option_group(advancedGroup)
parser.add_option_group(experimentalGroup)

(options, args) = parser.parse_args()

import sys
import os
import numpy as np
from scipy import linalg
from pylmm.lmm import LMM
from pylmm import input

if len(args) != 1:
   parser.print_help()
   sys.exit()

outFile = args[0]

if not options.tfile and not options.bfile and not options.emmaFile:
   parser.error("You must provide at least one PLINK input file base (--tfile or --bfile) or an EMMA formatted file (--emmaSNP).")
if not options.kfile:
   parser.error("Please provide a pre-computed kinship file")

# READING PLINK input
if options.verbose: sys.stderr.write("Reading SNP input...\n")
if options.bfile: IN = input.plink(options.bfile,type='b', phenoFile=options.phenoFile,normGenotype=options.normalizeGenotype)
elif options.tfile: IN = input.plink(options.tfile,type='t', phenoFile=options.phenoFile,normGenotype=options.normalizeGenotype)
elif options.emmaFile: IN = input.plink(options.emmaFile,type='emma', phenoFile=options.phenoFile,normGenotype=options.normalizeGenotype)
else: parser.error("You must provide at least one PLINK input file base")

if not os.path.isfile(options.phenoFile or IN.fbase + '.phenos') and not os.path.isfile(options.emmaPheno):
   parser.error("No .pheno file exist for %s.  Please provide a phenotype file using the --phenofile or --emmaPHENO argument." % (options.phenoFile or IN.fbase + '.phenos'))

# Read the emma phenotype file if provided.
# Format should be rows are phenotypes and columns are individuals.
if options.emmaPheno:
   f = open(options.emmaPheno,'r')
   P = []
   for line in f:
      v = line.strip().split()
      p = []

      P.append(p)
   f.close()
   IN.phenos = np.array(P).T

# READING Covariate File
if options.covfile:
   if options.verbose:
      sys.stderr.write("Reading covariate file...\n")
   P = IN.getCovariates(options.covfile)
   if options.noMean:
      X0 = P
   else:
      X0 = np.hstack([np.ones((IN.phenos.shape[0],1)),P])
elif options.emmaCov:
   if options.verbose:
      sys.stderr.write("Reading covariate file...\n")
   P = IN.getCovariatesEMMA(options.emmaCov)
   if options.noMean:
      X0 = P
   else:
      X0 = np.hstack([np.ones((IN.phenos.shape[0],1)),P])
else:
   X0 = np.ones((IN.phenos.shape[0],1))

if np.isnan(X0).sum():
      parser.error("The covariate file %s contains missing values. At this time we are not dealing with this case.  Either remove those individuals with missing values or replace them in some way.")

# READING Kinship - major bottleneck for large datasets
if options.verbose: sys.stderr.write("Reading kinship...\n")
begin = time.time()
# This method seems to be the fastest and works if you already know the size of the matrix
if options.kfile[-3:] == '.gz':
   import gzip
   f = gzip.open(options.kfile,'r')
   F = f.read() # might exhaust mem if the file is huge
   K = np.fromstring(F,sep=' ') # Assume that space separated
   f.close()
else:
   K = np.fromfile(open(options.kfile,'r'),sep=" ")
K.resize((len(IN.indivs),len(IN.indivs)))
end = time.time()
# Other slower ways
if options.verbose: sys.stderr.write("Read the %d x %d kinship matrix in %0.3fs \n" % (K.shape[0],K.shape[1],end-begin))

if options.kfile2:
   if options.verbose: sys.stderr.write("Reading second kinship...\n")
   begin = time.time()
   # This method seems to be the fastest and works if you already know the size of the matrix
   if options.kfile2[-3:] == '.gz':
      import gzip
      f = gzip.open(options.kfile2,'r')
      F = f.read() # might exhaust mem if the file is huge
      K2 = np.fromstring(F,sep=' ') # Assume that space separated
      f.close()
   else: K2 = np.fromfile(open(options.kfile2,'r'),sep=" ")
   K2.resize((len(IN.indivs),len(IN.indivs)))
   end = time.time()
   if options.verbose: sys.stderr.write("Read the %d x %d second kinship matrix in %0.3fs \n" % (K2.shape[0],K2.shape[1],end-begin))

# PROCESS the phenotype data -- Remove missing phenotype values
# Keep will now index into the "full" data to select what we keep (either everything or a subset of non missing data
Y = IN.phenos[:,options.pheno]
v = np.isnan(Y)
keep = True - v
if v.sum():
   if options.verbose: sys.stderr.write("Cleaning the phenotype vector by removing %d individuals...\n" % (v.sum()))
   Y = Y[keep]
   X0 = X0[keep,:]
   K = K[keep,:][:,keep]
   if options.kfile2: K2 = K2[keep,:][:,keep]
   Kva = []
   Kve = []

# Only load the decomposition if we did not remove individuals.
# Otherwise it would not be correct and we would have to compute it again.
if not v.sum() and options.eigenfile:
   if options.verbose: sys.stderr.write("Loading pre-computed eigendecomposition...\n")
   Kva = np.load(options.eigenfile + ".Kva")
   Kve = np.load(options.eigenfile + ".Kve")
else:
   Kva = []
   Kve = []

# CREATE LMM object for association
n = K.shape[0]
if not options.kfile2:  L = LMM(Y,K,Kva,Kve,X0,verbose=options.verbose)
else:  L = LMM_withK2(Y,K,Kva,Kve,X0,verbose=options.verbose,K2=K2)

# Fit the null model -- if refit is true we will refit for each SNP, so no reason to run here
if not options.refit:
   if options.verbose: sys.stderr.write("Computing fit for null model\n")
   L.fit()
   if options.verbose and not options.kfile2: sys.stderr.write("\t heritability=%0.3f, sigma=%0.3f\n" % (L.optH,L.optSigma))
   if options.verbose and options.kfile2: sys.stderr.write("\t heritability=%0.3f, sigma=%0.3f, w=%0.3f\n" % (L.optH,L.optSigma,L.optW))

# Buffers for pvalues and t-stats
PS = []
TS = []
count = 0
out = open(outFile,'w')
printOutHead()

for snp,id in IN:
   count += 1
   if options.verbose and count % 1000 == 0:
      sys.stderr.write("At SNP %d\n" % count)

   x = snp[keep].reshape((n,1))
   v = np.isnan(x).reshape((-1,))
   # Check SNPs for missing values
   if v.sum():
      keeps = True - v
      xs = x[keeps,:]
      if keeps.sum() <= 1 or xs.var() <= 1e-6:
	 PS.append(np.nan)
	 TS.append(np.nan)
	 outputResult(id,np.nan,np.nan,np.nan,np.nan)
	 continue

      # Its ok to center the genotype -  I used options.normalizeGenotype to
      # force the removal of missing genotypes as opposed to replacing them with MAF.
      if not options.normalizeGenotype: xs = (xs - xs.mean()) / np.sqrt(xs.var())
      Ys = Y[keeps]
      X0s = X0[keeps,:]
      Ks = K[keeps,:][:,keeps]
      if options.kfile2: K2s = K2[keeps,:][:,keeps]
      if options.kfile2: Ls = LMM_withK2(Ys,Ks,X0=X0s,verbose=options.verbose,K2=K2s)
      else: Ls = LMM(Ys,Ks,X0=X0s,verbose=options.verbose)
      if options.refit: Ls.fit(X=xs,REML=options.REML)
      else:
	 Ls.fit(REML=options.REML)
      ts,ps,beta,betaVar = Ls.association(xs,REML=options.REML,returnBeta=True)
   else:
      if x.var() == 0:
	 PS.append(np.nan)
	 TS.append(np.nan)
	 outputResult(id,np.nan,np.nan,np.nan,np.nan)
	 continue

      if options.refit: L.fit(X=x,REML=options.REML)
      ts,ps,beta,betaVar = L.association(x,REML=options.REML,returnBeta=True)

   outputResult(id,beta,np.sqrt(betaVar).sum(),ts,ps)
   PS.append(ps)
   TS.append(ts)


try:
	    p.append(float(x))
	 except: p.append(np.nan)

# pylmm is a python-based linear mixed-model solver with applications to GWAS

import pdb
import time
import sys

def printOutHead(): out.write("\t".join(["SNP_ID","BETA","BETA_SD","F_STAT","P_VALUE"]) + "\n")
def outputResult(id,beta,betaSD,ts,ps):
   out.write("\t".join([str(x) for x in [id,beta,betaSD,ts,ps]]) + "\n")

from optparse import OptionParser,OptionGroup
usage = """usage: %prog [options] --kfile kinshipFile --[tfile | bfile] plinkFileBase outfileBase

This program provides basic genome-wide association (GWAS) functionality.  You provide a phenotype and genotype file as well as a pre-computed (use pylmmKinship.py) kinship matrix and the program outputs a result file with information about each SNP, including the association p-value.
The input file are all standard plink formatted with the first two columns specifiying the individual and family ID.  For the phenotype file, we accept either NA or -9 to denote missing values.

Basic usage:

      python pylmmGWAS.py -v --bfile plinkFile --kfile preComputedKinship.kin --phenofile plinkFormattedPhenotypeFile resultFile

	    """
parser = OptionParser(usage=usage)

basicGroup = OptionGroup(parser, "Basic Options")
advancedGroup = OptionGroup(parser, "Advanced Options")
experimentalGroup = OptionGroup(parser, "Experimental Options")

#                  help="The base for a PLINK ped file")
basicGroup.add_option("--tfile", dest="tfile",
                  help="The base for a PLINK tped file")
basicGroup.add_option("--bfile", dest="bfile",
                  help="The base for a PLINK binary bed file")
basicGroup.add_option("--phenofile", dest="phenoFile", default=None,
                  help="Without this argument the program will look for a file with .pheno that has the plinkFileBase root.  If you want to specify an alternative phenotype file, then use this argument.  This file should be in plink format. ")

# EMMA Options
basicGroup.add_option("--emmaSNP", dest="emmaFile", default=None,
                  help="For backwards compatibility with emma, we allow for \"EMMA\" file formats.  This is just a text file with individuals on the columns and snps on the rows.")
basicGroup.add_option("--emmaPHENO", dest="emmaPheno", default=None,
                  help="For backwards compatibility with emma, we allow for \"EMMA\" file formats.  This is just a text file with each phenotype as one row.")
basicGroup.add_option("--emmaCOV", dest="emmaCov", default=None,
                  help="For backwards compatibility with emma, we allow for \"EMMA\" file formats.  This is just a text file with each covariate as one row.")

basicGroup.add_option("--kfile", dest="kfile",
                  help="The location of a kinship file.  This is an nxn plain text file and can be computed with the pylmmKinship program.")
basicGroup.add_option("--covfile", dest="covfile",
                  help="The location of a covariate file file.  This is a plink formatted covariate file.")
basicGroup.add_option("-p", type="int", dest="pheno", help="The phenotype index to be used in association.", default=0)

advancedGroup.add_option("--removeMissingGenotypes",
                  action="store_false", dest="normalizeGenotype", default=True,
                  help="By default the program replaces missing genotypes with the minor allele frequency.  This option overrides that behavior making the program remove missing individuals.  NOTE: This can increase running time due to the need to recompute the eigendecomposition for each SNP with missing values.")
advancedGroup.add_option("--refit",
                  action="store_true", dest="refit", default=False,
                  help="Refit the variance components at each SNP (default is to lock in the variance components under the null).")

advancedGroup.add_option("--REML",
                  action="store_true", dest="REML", default=False,
                  help="Use restricted maximum-likelihood (REML) (default is maximum-likelihood).")
#advancedGroup.add_option("-e", "--efile", dest="saveEig", help="Save eigendecomposition to this file.")
advancedGroup.add_option("--eigen", dest="eigenfile",
                  help="The location of the precomputed eigendecomposition for the kinship file.  These can be computed with pylmmKinship.py.")
advancedGroup.add_option("--noMean", dest="noMean", default=False,action="store_true",
                  help="This option only applies when --cofile is used.  When covfile is provided, the program will automatically add a global mean covariate to the model unless this option is specified.")

advancedGroup.add_option("-v", "--verbose",
                  action="store_true", dest="verbose", default=False,
                  help="Print extra info")

# Experimental Group Options
experimentalGroup.add_option("--kfile2", dest="kfile2",
                  help="The location of a second kinship file.  This file has the same format as the first kinship.  This might be used if you want to correct for another form of confounding.")

parser.add_option_group(basicGroup)
parser.add_option_group(advancedGroup)
parser.add_option_group(experimentalGroup)

(options, args) = parser.parse_args()

import sys
import os
import numpy as np
from scipy import linalg
from pylmm.lmm import LMM
from pylmm import input

if len(args) != 1:
   parser.print_help()
   sys.exit()

outFile = args[0]

if not options.tfile and not options.bfile and not options.emmaFile:
   parser.error("You must provide at least one PLINK input file base (--tfile or --bfile) or an EMMA formatted file (--emmaSNP).")
if not options.kfile:
   parser.error("Please provide a pre-computed kinship file")

# READING PLINK input
if options.verbose: sys.stderr.write("Reading SNP input...\n")
if options.bfile: IN = input.plink(options.bfile,type='b', phenoFile=options.phenoFile,normGenotype=options.normalizeGenotype)
elif options.tfile: IN = input.plink(options.tfile,type='t', phenoFile=options.phenoFile,normGenotype=options.normalizeGenotype)
elif options.emmaFile: IN = input.plink(options.emmaFile,type='emma', phenoFile=options.phenoFile,normGenotype=options.normalizeGenotype)
else: parser.error("You must provide at least one PLINK input file base")

if not os.path.isfile(options.phenoFile or IN.fbase + '.phenos') and not os.path.isfile(options.emmaPheno):
   parser.error("No .pheno file exist for %s.  Please provide a phenotype file using the --phenofile or --emmaPHENO argument." % (options.phenoFile or IN.fbase + '.phenos'))

# Read the emma phenotype file if provided.
# Format should be rows are phenotypes and columns are individuals.
if options.emmaPheno:
   f = open(options.emmaPheno,'r')
   P = []
   for line in f:
      v = line.strip().split()
      p = []
      for x in v:


"""
    Zernike polynomial

       +-m
      R    as in Born and Wolf, p. 465, sixth edition, Pergamon
         n

    The implementation have not been optimized for speed.

    Parameters
    ----------
    n : int
        DESCRIPTION.
    m : int
        DESCRIPTION.
    rho : rho
        DESCRIPTION.
    phi : phi
        DESCRIPTION.

    Returns
    -------
    double.

    """
    mabs = _np.abs(m)
    prod = 1.0
    sign = 1
    summ = 0.0
    for s in range(int((n-mabs)/2) + 1):
        if n-2*s != 0:
            prod = _np.power(rho, n-2*s)
        else:
            prod = 1.0
        prod *= _np.math.factorial(n-s)*sign
        prod /= (_np.math.factorial(s)
                * _np.math.factorial(((n+mabs)/2)-s)
                * _np.math.factorial(((n-mabs)/2)-s))
        summ += prod
        sign = -sign
    if m>=0:
        return summ*_np.cos(m*phi)
    else:
        return (-1)*summ*_np.sin(m*phi)


import numpy as _np

def zernike(n,m,rho,phi):
    """
    Zernike polynomial

       +-m
      R    as in Born and Wolf, p. 465, sixth edition, Pergamon
         n

    The implementation have not been optimized for speed.

    Parameters
    ----------
    n : int
        DESCRIPTION.
    m : int
        DESCRIPTION.
    rho : rho
        DESCRIPTION.
    phi : phi
        DESCRIPTION.

    Returns
    -------
    double.

    """
    mabs = _np.abs(m)
    prod = 1.0
    sign = 1
    summ = 0.

    for s in range(int((n-mabs)/2) + 1):
        if n-2*s != 0:
            prod = _np.power(rho, n-2*s)
        else:
            prod = 1.0
        prod *= _np.math.factorial(n-s)*sign
        prod /= (_np.math.factorial(s)
                * _np.math.factorial(((n+mabs)/2)-s)
                * _np.math.factorial(((n-mabs)/2)-s))
        summ += prod
        sign = -sign
    if m>=0:
        return summ*_np.cos(m*phi)
    else:
        return (-1)*summ*_np.sin(m*phi)


for s in range(int((n-mabs)/2) + 1):
        if n-2*s != 0:
            prod = _np.power(rho, n-2*s)
        else:
            prod = 1.0
        prod *= _np.math.factorial(n-s)*sign
        prod /= (_np.math.factorial(s)
                * _np.math.factorial(((n+mabs)/2)-s)
                * _np.math.factorial(((n-mabs)/2)-s))
        summ += prod
        sign = -sign

import numpy as _np

def zernike(n,m,rho,phi):
    """
    Zernike polynomial

       +-m
      R    as in Born and Wolf, p. 465, sixth edition, Pergamon
         n

    The implementation have not been optimized for speed.

    Parameters
    ----------
    n : int
        DESCRIPTION.
    m : int
        DESCRIPTION.
    rho : rho
        DESCRIPTION.
    phi : phi
        DESCRIPTION.

    Returns
    -------
    double.

    """
    mabs = _np.abs(m)
    prod = 1.0
    sign = 1
    summ = 0.0

    if m>=0:
        return summ*_np.cos(m*phi)
    else:
        return (-1)*summ*_np.sin(m*phi)


"""
@Author: vega13
"""
import platform
import time
import tkinter
import requests
from tkinter import *
from tkinter import filedialog
from tkinter import messagebox
from .request import json_request
from . import overwatch as OW
from . import game
from . import pool

def log(*args):
    print(args)

class Gui(object):
    def __init__(self):

        self.root = tkinter.Tk()
        self.root.title('Overwatch Replay Analyzer v0.1.2')
        self.root.geometry('550x350+400+200')

        self.read_path = None
        self.save_path = None
        self.left_frame = None
        self.right_frame = None
        # path
        self.create_path()
        self.create_time_inputs()
        # player
        self.create_player()
        # run
        self.run_btn = tkinter.Button(self.root, text="Analyze", command=self.run)
        self.run_btn.pack()
        self.create_text()
        # check for update
        self.root.after(500, self.check_update)

    def create_path(self):
        width_msg = 100
        width_path = 400

        path_frame = tkinter.Frame(self.root)
        path_frame.pack(fill=X)
        # read
        read_frame = tkinter.Frame(path_frame)
        read_frame.pack(fill=X)
        read_msg = tkinter.Message(read_frame, width=width_msg, text='Video path:')
        read_msg.pack(side=LEFT)
        read_path = tkinter.Message(read_frame, width=width_path, text='file1')
        read_path.pack(side=LEFT)
        read_btn = tkinter.Button(read_frame, text="Choose file", command=self.click_read)
        read_btn.pack(side=RIGHT)
        # save
        save_frame = tkinter.Frame(path_frame)
        save_frame.pack(fill=X)
        save_msg = tkinter.Message(save_frame, width=width_msg, text='Save to:')
        save_msg.pack(side=LEFT)
        save_path = tkinter.Message(save_frame, width=width_path, text='file2')
        save_path.pack(side=LEFT)
        save_btn = tkinter.Button(save_frame, text="Choose file", command=self.click_save)
        save_btn.pack(side=RIGHT)

        self.read_path = read_path
        self.save_path = save_path

    def create_player(self):
        player = tkinter.Frame(self.root)
        player.pack(fill=X, expand=1)
        # left
        left_frame = tkinter.Frame(player)
        left_frame.pack(side=LEFT)
        left_team_name = tkinter.Entry(left_frame, bg='pink', fg='black')
        left_team_name.insert(0, 'Team A')
        left_team_name.pack()
        for i in range(1, 7):
            e = tkinter.Entry(left_frame, bg='red', fg='white')
            name = 'Player' + str(i)
            e.insert(0, name)
            e.pack()
        # right
        right_frame = tkinter.Frame(player)
        right_frame.pack(side=RIGHT)
        right_team_name = tkinter.Entry(right_frame, bg='lightBlue', fg='black')
        right_team_name.insert(0, 'Team B')
        right_team_name.pack()
        for i in range(7, 13):
            e = tkinter.Entry(right_frame, bg='blue', fg='white')
            name = 'Player' + str(i)
            e.insert(0, name)
            e.pack()

        self.left_frame = left_frame
        self.right_frame = right_frame

    def create_time_inputs(self):
        time_inputs_frame = tkinter.Frame(self.root)
        time_inputs_frame.pack(fill=X, expand=1)

        left_frame = tkinter.Frame(time_inputs_frame)
        left_frame.pack(side=LEFT)

        label_start_time = tkinter.Label(left_frame, text="Start time in seconds (0 = start from beginning):")
        label_end_time = tkinter.Label(left_frame, text="End time in seconds (0 = analyze till the end):")
        label_fps = tkinter.Label(left_frame, text='FPS of analyzer:')
        label_game_type = tkinter.Label(left_frame, text='Game type (0 = OWL, 1 = Custom game):')
        label_game_version = tkinter.Label(left_frame, text='OWL stage number(0=preseason, 1, 2, 3, 4):')
        label_start_time.pack()
        label_end_time.pack()
        label_fps.pack()
        label_game_type.pack()
        label_game_version.pack()

        right_frame = tkinter.Frame(time_inputs_frame)
        right_frame.pack(side=RIGHT)

        start_time = tkinter.Entry(right_frame, bg='lightBlue', fg='black')
        start_time.insert(0, '0')
        start_time.pack()
        end_time = tkinter.Entry(right_frame, bg='lightBlue', fg='black')
        end_time.insert(0, '0')
        end_time.pack()
        fps = tkinter.Entry(right_frame, bg='lightBlue', fg='black')
        fps.insert(0, '2')
        fps.pack()
        game_type = tkinter.Entry(right_frame, bg='lightBlue', fg='black')
        game_type.insert(0, '0')
        game_type.pack()
        game_version = tkinter.Entry(right_frame, bg='lightBlue', fg='black')
        game_version.insert(0, '0')
        game_version.pack()

        self.time_inputs_frame = right_frame
    def create_text(self):
        self.notice_window = tkinter.Toplevel(self.root)
        self.notice_window.title('Notice')
        self.notice_window.geometry('400x400+300+100')
        self.notice = tkinter.Text(self.notice_window)
        self.notice.insert(tkinter.INSERT,
            """Overwatch Replay Analyzer, a data extractor of Overwatch game replays

You can contact the author or report issues by: http://example.com""")
        self.notice.pack()

    def click_save(self):
        filename = filedialog.askdirectory(initialdir='~/')
        self.save_path.config(text=filename)

    def click_read(self):
        filename = filedialog.askopenfilename(initialdir='~/OW/OverwatchDataAnalysis/videos')
        self.read_path.config(text=filename)

    def check_update(self):
        import configparser
        config = configparser.ConfigParser()
        config.read('ora/etc/config.ini')

        if platform.system() == "Windows":
            osname = "win"
        else:
            osname = "osx"

        current_version = config['version']['client_version']
        r = requests.get(
                config['api']['url'] + config['api']['check_update']
                + osname + '/' + current_version
            ).json()

        if not r['is_latest']:
            tkinter.messagebox.showinfo('Update Available',
                'There is a newer version of ORA available, please download at \n{}'.format(r['url']))

    def info(self):
        valid = True
        info = {
            "name_team_left": "left",
            "name_team_right": "right",
            "name_players_team_left": [],
            "name_players_team_right": [],
            "video_path": "/",
            "output_path": "/",
            "start_time": 0,
            "end_time": 0,
            "fps": 0,
            "game_type": 0,
            "game_version": 1
        }
        frame_left = self.left_frame.pack_slaves()
        frame_right = self.right_frame.pack_slaves()
        time_inputs_frame = self.time_inputs_frame.pack_slaves()
        info['name_team_left'] = frame_left[0].get()
        info['name_team_right'] = frame_right[0].get()
        try:
            info['start_time'] = int(time_inputs_frame[0].get())
        except ValueError:
            tkinter.messagebox.showinfo('Error', 'Invalid video start time!')
            valid = False
            return [info, valid]

        try:
            info['end_time'] = int(time_inputs_frame[1].get())
        except ValueError:
            tkinter.messagebox.showinfo('Error', 'Invalid video end time!')
            valid = False
            return [info, valid]

        try:
            info['fps'] = int(time_inputs_frame[2].get())
        except ValueError:
            tkinter.messagebox.showinfo('Error', 'Invalid analysis fps!')
            valid = False
            return [info, valid]

        if not (info['end_time'] >= 0 and info['start_time'] >= 0 \
            and info['end_time'] >= info['start_time']):
            tkinter.messagebox.showinfo('Error', 'Invalid video end time!')
            valid = False
            return [info, valid]

        if not (info['fps'] > 0):
            tkinter.messagebox.showinfo('Error', 'Invalid analysis fps!')
            valid = False
            return [info, valid]

        try:
            info['game_type'] = int(time_inputs_frame[3].get())
        except ValueError:
            tkinter.messagebox.showinfo('Error', 'Invalid game type!')
            valid = False
            return [info, valid]

        try:
            info['game_version'] = int(time_inputs_frame[4].get())
        except ValueError:
            tkinter.messagebox.showinfo('Error', 'Invalid owl stage number!')
            valid = False
            return [info, valid]

        if not (info['game_type'] == 0 or info['game_type'] == 1):
            tkinter.messagebox.showinfo('Error', 'Invalid game type!')
            valid = False
            return [info, valid]

        if not (info['game_version'] >= 0 or info['game_version'] < 5):
            tkinter.messagebox.showinfo('Error', 'Invalid OWL stage number!')
            valid = False
            return [info, valid]

        team_left = []
        team_right = []
        for i in range(1, 7):
            player_left = frame_left[i].get()
            team_left.append(player_left)
            player_right = frame_right[i].get()
            team_right.append(player_right)

        info['name_players_team_left'] = team_left
        info['name_players_team_right'] = team_right
        info['video_path'] = self.read_path['text']
        info['output_path'] = self.save_path['text']

        return [info, valid]

self.root.mainloop()

    def show_finish_msg(self):
        tkinter.messagebox.showinfo('An alert box', 'Output file successfully saved!')

    def show_progress(self, progress):
        self.notice.insert(Tkinter.INSERT, str(progress))

    def run(self):
        info, valid = self.info()
        game_type = OW.GAMETYPE_OWL if info['game_type'] == 0 else OW.GAMETYPE_CUSTOM
        self.game_instance = game.Game(game_type)
        if valid is True:
            self.game_instance.set_game_info(info)
            pool.initPool()

            try:
                self.game_instance.analyze(info['start_time'], info['end_time'], is_test=False)
            except Exception as err:
                print(err)
            else:
                self.game_instance.output()
                self.show_finish_msg()

            pool.PROCESS_POOL.close()
            pool.PROCESS_POOL.join()

gui_instance = Gui()


tkinter.messagebox.showinfo('An alert box', 'Output file successfully saved!')

"""
@Author: vega13
"""
import platform
import time
import tkinter
import requests
from tkinter import *
from tkinter import filedialog
from tkinter import messagebox
from .request import json_request
from . import overwatch as OW
from . import game
from . import pool

def log(*args):
    print(args)

class Gui(object):
    def __init__(self):

        self.root = tkinter.Tk()
        self.root.title('Overwatch Replay Analyzer v0.1.2')
        self.root.geometry('550x350+400+200')

        self.read_path = None
        self.save_path = None
        self.left_frame = None
        self.right_frame = None
        # path
        self.create_path()
        self.create_time_inputs()
        # player
        self.create_player()
        # run
        self.run_btn = tkinter.Button(self.root, text="Analyze", command=self.run)
        self.run_btn.pack()
        self.create_text()
        # check for update
        self.root.after(500, self.check_update)

    def create_path(self):
        width_msg = 100
        width_path = 400

        path_frame = tkinter.Frame(self.root)
        path_frame.pack(fill=X)
        # read
        read_frame = tkinter.Frame(path_frame)
        read_frame.pack(fill=X)
        read_msg = tkinter.Message(read_frame, width=width_msg, text='Video path:')
        read_msg.pack(side=LEFT)
        read_path = tkinter.Message(read_frame, width=width_path, text='file1')
        read_path.pack(side=LEFT)
        read_btn = tkinter.Button(read_frame, text="Choose file", command=self.click_read)
        read_btn.pack(side=RIGHT)
        # save
        save_frame = tkinter.Frame(path_frame)
        save_frame.pack(fill=X)
        save_msg = tkinter.Message(save_frame, width=width_msg, text='Save to:')
        save_msg.pack(side=LEFT)
        save_path = tkinter.Message(save_frame, width=width_path, text='file2')
        save_path.pack(side=LEFT)
        save_btn = tkinter.Button(save_frame, text="Choose file", command=self.click_save)
        save_btn.pack(side=RIGHT)

        self.read_path = read_path
        self.save_path = save_path

    def create_player(self):
        player = tkinter.Frame(self.root)
        player.pack(fill=X, expand=1)
        # left
        left_frame = tkinter.Frame(player)
        left_frame.pack(side=LEFT)
        left_team_name = tkinter.Entry(left_frame, bg='pink', fg='black')
        left_team_name.insert(0, 'Team A')
        left_team_name.pack()
        for i in range(1, 7):
            e = tkinter.Entry(left_frame, bg='red', fg='white')
            name = 'Player' + str(i)
            e.insert(0, name)
            e.pack()
        # right
        right_frame = tkinter.Frame(player)
        right_frame.pack(side=RIGHT)
        right_team_name = tkinter.Entry(right_frame, bg='lightBlue', fg='black')
        right_team_name.insert(0, 'Team B')
        right_team_name.pack()
        for i in range(7, 13):
            e = tkinter.Entry(right_frame, bg='blue', fg='white')
            name = 'Player' + str(i)
            e.insert(0, name)
            e.pack()

        self.left_frame = left_frame
        self.right_frame = right_frame

    def create_time_inputs(self):
        time_inputs_frame = tkinter.Frame(self.root)
        time_inputs_frame.pack(fill=X, expand=1)

        left_frame = tkinter.Frame(time_inputs_frame)
        left_frame.pack(side=LEFT)

        label_start_time = tkinter.Label(left_frame, text="Start time in seconds (0 = start from beginning):")
        label_end_time = tkinter.Label(left_frame, text="End time in seconds (0 = analyze till the end):")
        label_fps = tkinter.Label(left_frame, text='FPS of analyzer:')
        label_game_type = tkinter.Label(left_frame, text='Game type (0 = OWL, 1 = Custom game):')
        label_game_version = tkinter.Label(left_frame, text='OWL stage number(0=preseason, 1, 2, 3, 4):')
        label_start_time.pack()
        label_end_time.pack()
        label_fps.pack()
        label_game_type.pack()
        label_game_version.pack()

        right_frame = tkinter.Frame(time_inputs_frame)
        right_frame.pack(side=RIGHT)

        start_time = tkinter.Entry(right_frame, bg='lightBlue', fg='black')
        start_time.insert(0, '0')
        start_time.pack()
        end_time = tkinter.Entry(right_frame, bg='lightBlue', fg='black')
        end_time.insert(0, '0')
        end_time.pack()
        fps = tkinter.Entry(right_frame, bg='lightBlue', fg='black')
        fps.insert(0, '2')
        fps.pack()
        game_type = tkinter.Entry(right_frame, bg='lightBlue', fg='black')
        game_type.insert(0, '0')
        game_type.pack()
        game_version = tkinter.Entry(right_frame, bg='lightBlue', fg='black')
        game_version.insert(0, '0')
        game_version.pack()

        self.time_inputs_frame = right_frame
    def create_text(self):
        self.notice_window = tkinter.Toplevel(self.root)
        self.notice_window.title('Notice')
        self.notice_window.geometry('400x400+300+100')
        self.notice = tkinter.Text(self.notice_window)
        self.notice.insert(tkinter.INSERT,
            """Overwatch Replay Analyzer, a data extractor of Overwatch game replays

You can contact the author or report issues by: http://example.com""")
        self.notice.pack()

    def click_save(self):
        filename = filedialog.askdirectory(initialdir='~/')
        self.save_path.config(text=filename)

    def click_read(self):
        filename = filedialog.askopenfilename(initialdir='~/OW/OverwatchDataAnalysis/videos')
        self.read_path.config(text=filename)

    def check_update(self):
        import configparser
        config = configparser.ConfigParser()
        config.read('ora/etc/config.ini')

        if platform.system() == "Windows":
            osname = "win"
        else:
            osname = "osx"

        current_version = config['version']['client_version']
        r = requests.get(
                config['api']['url'] + config['api']['check_update']
                + osname + '/' + current_version
            ).json()

        if not r['is_latest']:
            tkinter.messagebox.showinfo('Update Available',
                'There is a newer version of ORA available, please download at \n{}'.format(r['url']))

    def info(self):
        valid = True
        info = {
            "name_team_left": "left",
            "name_team_right": "right",
            "name_players_team_left": [],
            "name_players_team_right": [],
            "video_path": "/",
            "output_path": "/",
            "start_time": 0,
            "end_time": 0,
            "fps": 0,
            "game_type": 0,
            "game_version": 1
        }
        frame_left = self.left_frame.pack_slaves()
        frame_right = self.right_frame.pack_slaves()
        time_inputs_frame = self.time_inputs_frame.pack_slaves()
        info['name_team_left'] = frame_left[0].get()
        info['name_team_right'] = frame_right[0].get()
        try:
            info['start_time'] = int(time_inputs_frame[0].get())
        except ValueError:
            tkinter.messagebox.showinfo('Error', 'Invalid video start time!')
            valid = False
            return [info, valid]

        try:
            info['end_time'] = int(time_inputs_frame[1].get())
        except ValueError:
            tkinter.messagebox.showinfo('Error', 'Invalid video end time!')
            valid = False
            return [info, valid]

        try:
            info['fps'] = int(time_inputs_frame[2].get())
        except ValueError:
            tkinter.messagebox.showinfo('Error', 'Invalid analysis fps!')
            valid = False
            return [info, valid]

        if not (info['end_time'] >= 0 and info['start_time'] >= 0 \
            and info['end_time'] >= info['start_time']):
            tkinter.messagebox.showinfo('Error', 'Invalid video end time!')
            valid = False
            return [info, valid]

        if not (info['fps'] > 0):
            tkinter.messagebox.showinfo('Error', 'Invalid analysis fps!')
            valid = False
            return [info, valid]

        try:
            info['game_type'] = int(time_inputs_frame[3].get())
        except ValueError:
            tkinter.messagebox.showinfo('Error', 'Invalid game type!')
            valid = False
            return [info, valid]

        try:
            info['game_version'] = int(time_inputs_frame[4].get())
        except ValueError:
            tkinter.messagebox.showinfo('Error', 'Invalid owl stage number!')
            valid = False
            return [info, valid]

        if not (info['game_type'] == 0 or info['game_type'] == 1):
            tkinter.messagebox.showinfo('Error', 'Invalid game type!')
            valid = False
            return [info, valid]

        if not (info['game_version'] >= 0 or info['game_version'] < 5):
            tkinter.messagebox.showinfo('Error', 'Invalid OWL stage number!')
            valid = False
            return [info, valid]

        team_left = []
        team_right = []
        for i in range(1, 7):
            player_left = frame_left[i].get()
            team_left.append(player_left)
            player_right = frame_right[i].get()
            team_right.append(player_right)

        info['name_players_team_left'] = team_left
        info['name_players_team_right'] = team_right
        info['video_path'] = self.read_path['text']
        info['output_path'] = self.save_path['text']

        return [info, valid]

    def show(self):
        self.root.mainloop()

    def show_finish_msg(self):

    def show_progress(self, progress):
        self.notice.insert(Tkinter.INSERT, str(progress))

    def run(self):
        info, valid = self.info()
        game_type = OW.GAMETYPE_OWL if info['game_type'] == 0 else OW.GAMETYPE_CUSTOM
        self.game_instance = game.Game(game_type)
        if valid is True:
            self.game_instance.set_game_info(info)
            pool.initPool()

            try:
                self.game_instance.analyze(info['start_time'], info['end_time'], is_test=False)
            except Exception as err:
                print(err)
            else:
                self.game_instance.output()
                self.show_finish_msg()

            pool.PROCESS_POOL.close()
            pool.PROCESS_POOL.join()

gui_instance = Gui()