# -*- coding: utf-8 -*-
import os
import math
import time
import signal
import threading

import tensorflow as tf

import constants
from game_ac_network import GameACFFNetwork, GameACLSTMNetwork
from a3c_training_thread import A3CTrainingThread
from rmsprop_applier import RMSPropApplier


def log_uniform(lo, hi, rate):
    log_lo = math.log(lo)
    log_hi = math.log(hi)
    v = log_lo * (1-rate) + log_hi * rate
    return math.exp(v)


def _main():
    device = "/cpu:0"
    if constants.USE_GPU:
        device = "/gpu:0"

    initial_learning_rate = log_uniform(
        constants.INITIAL_ALPHA_LOW,
        constants.INITIAL_ALPHA_HIGH,
        constants.INITIAL_ALPHA_LOG_RATE
    )

    global_t = 0

    stop_requested = False

    Network = GameACLSTMNetwork if constants.USE_LSTM else GameACFFNetwork
    global_network = Network(constants.ACTION_SIZE, -1, device)

    training_threads = []

    learning_rate_input = tf.placeholder("float")

    grad_applier = RMSPropApplier(
        learning_rate=learning_rate_input,
        decay=constants.RMSP_ALPHA,
        momentum=0.0,
        epsilon=constants.RMSP_EPSILON,
        clip_norm=constants.GRAD_NORM_CLIP,
        device=device)

    for i in range(constants.PARALLEL_SIZE):
        training_thread = A3CTrainingThread(
            i, global_network, initial_learning_rate,
            learning_rate_input, grad_applier, constants.MAX_TIME_STEP,
            device=device
        )
        training_threads.append(training_thread)

    # prepare session
    sess = tf.Session(
        config=tf.ConfigProto(
            log_device_placement=False, allow_soft_placement=True
        )
    )

    init = tf.global_variables_initializer()
    sess.run(init)

    # summary for tensorboard
    score_input = tf.placeholder(tf.int32)
    tf.summary.scalar("score", score_input)

    summary_op = tf.summary.merge_all()
    summary_writer = tf.summary.FileWriter(constants.LOG_FILE, sess.graph)

    # init or load checkpoint with saver
    saver = tf.train.Saver()
    checkpoint = tf.train.get_checkpoint_state(constants.CHECKPOINT_DIR)
    if checkpoint and checkpoint.model_checkpoint_path:
        saver.restore(sess, checkpoint.model_checkpoint_path)
        print("checkpoint loaded:", checkpoint.model_checkpoint_path)
        tokens = checkpoint.model_checkpoint_path.split("-")
        # set global step
        global_t = int(tokens[1])
        print(">>> global step set: ", global_t)
        # set wall time
        wall_t_fname = constants.CHECKPOINT_DIR + '/' + 'wall_t.' + str(global_t)
        with open(wall_t_fname, 'r') as f:
            wall_t = float(f.read())
    else:
        print("Could not find old checkpoint")
        # set wall time
        wall_t = 0.0

    def train_function(parallel_index):
        global global_t

        training_thread = training_threads[parallel_index]
        # set start_time
        start_time = time.time() - wall_t
        training_thread.set_start_time(start_time)

        while True:
            if stop_requested:
                break
            if global_t > constants.MAX_TIME_STEP:
                break

            diff_global_t = training_thread.process(
                sess, global_t, summary_writer,
                summary_op, score_input)
            global_t += diff_global_t

    def signal_handler(signal, frame):
        global stop_requested
        print('You pressed Ctrl+C!')
        stop_requested = True

    train_threads = []
    for i in range(constants.PARALLEL_SIZE):
        train_threads.append(
            threading.Thread(target=train_function, args=(i,))
        )

    signal.signal(signal.SIGINT, signal_handler)

    # set start time
    start_time = time.time() - wall_t

    for t in train_threads:
        t.start()

    print('Press Ctrl+C to stop')
    signal.pause()

    print('Now saving data. Please wait')

    for t in train_threads:
        t.join()

    if not os.path.exists(constants.CHECKPOINT_DIR):
        os.mkdir(constants.CHECKPOINT_DIR)

    # write wall time
    wall_t = time.time() - start_time
    wall_t_fname = constants.CHECKPOINT_DIR + '/' + 'wall_t.' + str(global_t)
    with open(wall_t_fname, 'w') as f:
        f.write(str(wall_t))

    saver.save(
        sess, constants.CHECKPOINT_DIR + '/' + 'checkpoint',
        global_step=global_t
    )


if __name__ == '__main__':
    _main()
