# -*- coding: utf-8 -*-
# Python 2.7.6
# tensorflow (0.7.1)

import tensorflow as tf
import numpy as np
import random
from param import *

# =================================================================
class Model:
    '''機械学習のモデル定義を管理するクラス 
        TensorFlowのSessionや学習済みモデルの保存・読込機能も持たせる
    '''

    def __init__(self):
        ''' 初期化処理（準コンストラクタ） '''
        self.__setup_placeholder()
        self.__setup_model()
        self.__setup_ops()

    def __enter__(self):
        ''' コンテキストマネージャ（With開始時に呼ばれる） '''
        self.session = tf.Session()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        ''' コンテキストマネージャ（With終了時に呼ばれる） '''
        self.session.close()
        return False

    def save(self, str_path):
        """TensorFlowのモデルを保存する
        Args:
            str_path: モデル保存先のパス
        Returns:
            -
        """
        saver = tf.train.Saver()
        saver.save(self.session, str_path)

    def restore(self, str_path):
        """TensorFlowのモデルを読み込むする
        Args:
            str_path: モデル保存先のパス
        Returns:
            -
        """
        saver = tf.train.Saver()
        saver.restore(self.session, str_path)

    def __setup_placeholder(self):
        ''' プレイスホルダーの初期化 '''
        self.ph_trade_data   = tf.placeholder(tf.float32, [None, Param.IN_COLUMN_SIZE, 1], name="input")
        self.ph_actual_class = tf.placeholder(tf.float32, [None, Param.OUT_NODES_SIZE])
        self.ph_label        = tf.placeholder(tf.string)
        self.ph_forget_bias  = tf.placeholder(tf.float32)
        self.ph_state_size   = tf.placeholder(tf.int32, shape=[])
        self.ph_dropout_keep_prob = tf.placeholder(tf.float32)

    def __setup_model(self):
        ''' ニューラルネットワークのモデル定義 '''

        with tf.name_scope("inference") as scope:
             
            var_weight     = tf.Variable(tf.truncated_normal([Param.HIDDEN_UNIT_SIZE, Param.OUT_NODES_SIZE], stddev=0.1), name="weight2")
            var_bias       = tf.Variable(tf.truncated_normal([Param.OUT_NODES_SIZE], stddev=0.1), name="bias2")
            var_weight_512 = tf.Variable(tf.truncated_normal([Param.HIDDEN_UNIT_SIZE, Param.HIDDEN_UNIT_SIZE], stddev=0.1), name="weight")
            var_bias_512   = tf.Variable(tf.truncated_normal([Param.HIDDEN_UNIT_SIZE], stddev=0.1), name="bias")
            var_weight_h1  = tf.Variable(tf.truncated_normal([Param.HIDDEN_UNIT_SIZE, 512], stddev=0.1), name="weight")
            var_bias_h1    = tf.Variable(tf.truncated_normal([512], stddev=0.1), name="bias")
            var_weight_h2  = tf.Variable(tf.truncated_normal([512, 256], stddev=0.1), name="weight")
            var_bias_h2    = tf.Variable(tf.truncated_normal([256], stddev=0.1), name="bias")
            var_weight_h3  = tf.Variable(tf.truncated_normal([256, 128], stddev=0.1), name="weight")
            var_bias_h3    = tf.Variable(tf.truncated_normal([128], stddev=0.1), name="bias")
            # var_weight_o   = tf.Variable(tf.truncated_normal([128, Param.OUT_NODES_SIZE], stddev=0.1), name="weight")
            var_weight_o   = tf.Variable(tf.truncated_normal([256, Param.OUT_NODES_SIZE], stddev=0.1), name="weight")
            var_bias_o     = tf.Variable(tf.truncated_normal([Param.OUT_NODES_SIZE], stddev=0.1), name="bias")

            ## Cell定義
            # cell = tf.nn.rnn_cell.BasicLSTMCell(Param.HIDDEN_UNIT_SIZE, forget_bias=Param.FORGET_BIAS, activation=tf.nn.relu, state_is_tuple=False)
            cell = tf.nn.rnn_cell.BasicLSTMCell(Param.HIDDEN_UNIT_SIZE, forget_bias=Param.FORGET_BIAS)
            
            initial_state = cell.zero_state(self.ph_state_size, tf.float32)

            gr_state = initial_state
            ary_outputs = []  # 過去の隠れ層の出力を保存
            with tf.variable_scope('LSTM'):
                for t in range(Param.IN_COLUMN_SIZE):
                    if t > 0:
                        tf.get_variable_scope().reuse_variables()
                    (cell_output, gr_state) = cell(self.ph_trade_data[:, t, :], gr_state)
                    ary_outputs.append(cell_output)

            gr_lstm_output = ary_outputs[-1]
            gr_lstm_output = tf.nn.dropout(gr_lstm_output, self.ph_dropout_keep_prob)
            
            h1 = tf.nn.relu(tf.matmul(gr_lstm_output, var_weight_h1) + var_bias_h1)
            h1 = tf.nn.dropout(h1, self.ph_dropout_keep_prob)

            h2 = tf.nn.relu(tf.matmul(h1, var_weight_h2) + var_bias_h2)
            h2 = tf.nn.dropout(h2, self.ph_dropout_keep_prob)
            self.output = tf.matmul(h2, var_weight_o) + var_bias_o

            # self.output = tf.nn.softmax(tf.matmul(gr_lstm_output, var_weight) + var_bias)
            #self.output = tf.matmul(gr_lstm_output, var_weight) + var_bias
            #self.output = tf.nn.relu(tf.matmul(gr_lstm_output, var_weight) + var_bias)
    
    def __setup_ops(self):
        ''' 各種オペレーションの定義 '''

        #reduce_sum : テンソルの次元間の要素の合計を計算します。
        #cross_entropy : コスト関数。目標をどれだけ達成できたのかを図る。
        #self.loss = tf.reduce_mean(tf.square(self.output - self.ph_actual_class))
        # self.loss = tf.reduce_mean(-tf.reduce_sum(
        #               self.ph_actual_class * tf.log(tf.clip_by_value(self.output, 1e-10, 1.0)),
        #               reduction_indices=[1]))
        #self.loss = -tf.reduce_sum(self.ph_actual_class * tf.log(tf.clip_by_value(self.output, 1e-10, 1.0)))
        self.loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=self.output, labels=self.ph_actual_class))
        self.out_softmax = tf.nn.softmax(self.output)

        #クロスエントロピーの最適化。0.0001は学習レート。
        #self.train_op = tf.train.AdamOptimizer(0.0001).minimize(cross_entropy)
        #self.train_op = tf.train.GradientDescentOptimizer(0.5).minimize(cross_entropy)
        #self.train_op = tf.train.AdamOptimizer(learning_rate=0.0005,beta1=0.5).minimize(self.loss)
        # self.train_op = tf.train.AdamOptimizer(learning_rate=0.00005).minimize(self.loss)
        self.train_op = tf.train.RMSPropOptimizer(learning_rate=0.001).minimize(self.loss)
        #self.train_op = tf.train.AdamOptimizer(learning_rate=0.001, beta1=0.9, beta2=0.999).minimize(self.loss)

# =================================================================
class Trainer(Model):
    '''
    学習用クラス
    jijiトレードデータをもとに、TensorFlowでの学習を行わせる
    '''

    def train(self, steps, data):
        '''
        学習の実施
        起動スクリプトからは本メソッドがキックされる
        '''
        self.__prepare_train(self.session)
        f_actual_rate_train = 0

        #=== エポック分のループ処理 ========================
        for i in range(steps):
            
            ary_trade_data = None
            ary_actual_data = None
            i_actual_cnt = 0
            k = 0

            #=== 学習処理の実行 ===================
            
            # データ取得    
            (ary_trade_data, ary_actual_data, df_jiji_data) = data.train_data(Param.MINI_BATCH_SIZE)

            # 学習の実行
            self.session.run(self.train_op, feed_dict=self.train_feed_dict(ary_trade_data, ary_actual_data))

            #=== 途中経過（サマリ）の表示 ===================
            # 現在のモデルの出力とコスト関数の結果を取得する
            output = self.session.run(self.out_softmax, feed_dict=self.train_feed_dict(ary_trade_data, ary_actual_data)) 
            loss   = self.session.run(self.loss,        feed_dict=self.train_feed_dict(ary_trade_data, ary_actual_data)) 
            #import pdb; pdb.set_trace()
            cnt_correct = 0
            for j in range(len(output)):
                actual_val = 'up' if ary_actual_data[j][0] == 1.0 else 'down'
                output_val = 'up' if output[j][0] > output[j][1] else 'down'

                if actual_val == output_val:
                    cnt_correct += 1
            rate_correct = float(cnt_correct) / len(output)

            print 'step:{} rate_correct:{} loss:{}'.format(i, rate_correct, loss)

        #=== テストデータで検証 ============================
        print '==========================================='
        print 'テストデータでの検証'
        i_data_cnt = 1000
        i_actual_cnt = 0
        (ary_trade_data, ary_actual_data, df_jiji_data) = data.test_data(i_data_cnt)

        for i in range(i_data_cnt):

            output = self.session.run(self.out_softmax, feed_dict=self.test_feed_dict(np.array([ary_trade_data[i]]), np.array([ary_actual_data[i]]))) 
            p = '✕'
            actual_val = 'up' if ary_actual_data[i][0] == 1.0 else 'down'
            output_val = 'up' if output[0][0] > output[0][1] else 'down'

            if actual_val == output_val:
                p = '○'
                i_actual_cnt += 1

            print '{} {} {} output={} actual={}'.format(
                i, p, df_jiji_data[i:]['n01_target_date'].values[0], output, ary_actual_data[i])

        # 教師データに対する正解率を表示
        f_actual_rate_test = float(i_actual_cnt) / i_data_cnt
        print '### f_actual_rate_test={}'.format(f_actual_rate_test)
        print '==========================================='

    def __prepare_train(self, session):
        ''' 初期処理 '''
        #グラフの初期化
        self.summary_writer = tf.summary.FileWriter('logs', graph=session.graph)
        session.run(tf.global_variables_initializer())

    def __print_status(self, session, k, ary_trade_data, ary_actual_data, df_jiji_data, bl_detail):
        # 現在のモデルの出力とコスト関数の結果を取得する
        output = session.run(self.out_softmax, feed_dict=self.train_feed_dict(ary_trade_data, ary_actual_data)) 
        loss   = session.run(self.loss,        feed_dict=self.train_feed_dict(ary_trade_data, ary_actual_data)) 

        # 教師データとモデル出力の比較
        result = False
        p = '✕'
        for i in range(len(output)):
            actual_val = 'up' if ary_actual_data[i][0] == 1.0 else 'down'
            output_val = 'up' if output[i][0] > output[i][1] else 'down'

            if actual_val == output_val:
                result = True
                p = '○'
        #import pdb; pdb.set_trace()
        if bl_detail == True:
            print df_jiji_data
        print '{} {} {} output={} actual={} loss={}'.format(
            k, p, df_jiji_data[-1:]['n01_target_date'].values[0], output, ary_actual_data, loss)
        
        return result

    def train_feed_dict(self, ary_trade_data, ary_actual_data):
        return {
            self.ph_trade_data: ary_trade_data,      # 為替データ
            self.ph_actual_class: ary_actual_data,   # 上がったか下がったか？
            self.ph_label: "train",
            self.ph_state_size: Param.MINI_BATCH_SIZE,
            self.ph_dropout_keep_prob: Param.DROPOUT_KEEP_PROB,
        }

    def test_feed_dict(self, ary_trade_data, ary_actual_data):
        return {
            self.ph_trade_data: ary_trade_data,
            self.ph_actual_class: ary_actual_data,
            self.ph_label: "test",
            self.ph_state_size: 1,
            self.ph_dropout_keep_prob: 1.0,
        }


# =================================================================
class Estimator(Model):

    def estimate( self, data ):
        output = self.session.run(self.out_softmax, feed_dict=self.estimate_feed_dict(data))
        #loss   = self.session.run(self.loss,   feed_dict=self.estimate_feed_dict(data)) 
        #print output
        return 'up' if output[0][0] > output[0][1] else 'down'
        #return self.session.run(tf.argmax(self.output,1), feed_dict=self.estimate_feed_dict(data))

    def estimate_feed_dict(self, data):
        return {
            self.ph_trade_data: data,
            self.ph_state_size: 1,
            self.ph_dropout_keep_prob: 1.0,
        }