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

import numpy as np
import random
from param import *
import matplotlib
matplotlib.use('Agg') # GUI Off設定
import matplotlib.pyplot as plt

from keras.models import Sequential
from keras.layers.normalization import BatchNormalization
from keras.layers.core import Dense, Activation, Dropout
from keras.layers.recurrent import LSTM
from keras.optimizers import Adam
from keras.callbacks import EarlyStopping

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

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

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

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

    def save(self, str_path):
        """TensorFlowのモデルを保存する
        Args:
            str_path: モデル保存先のパス
        Returns:
            -
        """

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

    def __setup_model(self):
        ''' NNモデル定義 '''

        def weight_variable(shape, name=None):
            return np.random.normal(scale=.01, size=shape)

        model = Sequential()
        model.add(BatchNormalization(input_shape=(Param.IN_CYCLE_SIZE, Param.IN_COLUMN_SIZE)))
        model.add(
            LSTM(Param.HIDDEN_UNIT_SIZE,
                 kernel_initializer=weight_variable,
                 input_shape=(Param.IN_CYCLE_SIZE, Param.IN_COLUMN_SIZE), # seq_length, dim
                 dropout=Param.DROPOUT_KEEP_PROB,
                 recurrent_dropout=Param.DROPOUT_KEEP_PROB))
        model.add(Dropout(Param.DROPOUT_KEEP_PROB))
        model.add(Dense(Param.OUT_NODES_SIZE, kernel_initializer=weight_variable))
        
        # 分類
        # model.add(Activation('softmax'))
        # optimizer = Adam(lr=0.001, beta_1=0.9, beta_2=0.999)
        # model.compile(loss='categorical_crossentropy', optimizer=optimizer, metrics=["accuracy"])

        # 回帰
        model.add(Activation('linear'))
        model.compile(loss='mean_squared_error', optimizer="rmsprop")

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

    def train(self, steps, data):
        '''
        モデルの学習
        '''
        (ary_in, ary_actual, df_jiji) = data.train_data()
        hist = self.model.fit(
            ary_in,
            ary_actual,
            batch_size=Param.IN_BATCH_SIZE,
            epochs=Param.TRANING_EPOCH_SIZE,
        )

        '''
        学習結果の可視化
        '''
        loss = hist.history['loss']

        plt.rc('font', family='serif')
        fig = plt.figure()
        plt.plot(range(len(loss)), loss, label='loss', color='black')
        plt.xlabel('epochs')
        plt.show()
        plt.savefig(__file__ + '.eps')

        '''
        テストデータでの確認
        '''
        (ary_in, ary_actual, df_jiji) = data.test_data()

        acc_count = 0
        for i in range(1000):
            k = random.randrange(len(ary_in))
            
            x = ary_in[k]
            y = self.model.predict(np.array([x]), batch_size=1)
            l = ary_actual[k]

            close = df_jiji.loc[k]['n03_close']

            sub_y = close - y
            sub_l = close - l

            if sub_y > 0 and sub_l > 0:
                acc_count = acc_count + 1
            if sub_y < 0 and sub_l < 0:
                acc_count = acc_count + 1
            
            print 'y : {} , l : {}, sub_y:{}, sub_l:{}'.format(y, l, sub_y ,sub_l)

        print '=============================================='
        print 'acc_count:{}'.format(acc_count)

        # # score = self.model.evaluate(ary_in, ary_actual, batch_size=Param.IN_BATCH_SIZE)
        # # print '¥n evaluate : {} '.format(score)

        # score = self.model.predict(ary_in, batch_size=Param.IN_BATCH_SIZE)
        # print '\n predict : {} '.format(score)

        # score = self.model.predict_classes(ary_in, batch_size=Param.IN_BATCH_SIZE)
        # print '\n predict_classes : {} '.format(score)

        # score = self.model.predict_proba(ary_in, batch_size=Param.IN_BATCH_SIZE)
        # print '\n predict_proba : {} '.format(score)

# =================================================================
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,
        }