import time
import math
import sys
import pickle
import copy
import os
import re

import numpy as np
from chainer import cuda, Variable, FunctionSet, optimizers
import chainer.functions as F

"""
Explanation: Introduction
Chainer とはニューラルネットの実装を簡単にしたフレームワークです。

今回は言語の分野でニューラルネットを適用してみました。



今回は言語モデルを作成していただきます。

言語モデルとはある単語が来たときに次の単語に何が来やすいかを予測するものです。
言語モデルにはいくつか種類があるのでここでも紹介しておきます。

n-グラム言語モデル
単語の数を単純に数え挙げて作成されるモデル。考え方としてはデータにおけるある単語の頻度に近い
ニューラル言語モデル

単語の辞書ベクトルを潜在空間ベクトルに落とし込み、ニューラルネットで次の文字を学習させる手法


リカレントニューラル言語モデル

基本的なアルゴリズムはニューラル言語モデルと同一だが過去に使用した単語を入力に加えることによって文脈を考慮した言語モデルの学習が可能となる。ニューラル言語モデルとは異なり、より古い情報も取得可能

以下では、このChainerを利用しデータを準備するところから実際に言語モデルを構築し学習・評価を行うまでの手順を解説します。

各種ライブラリ導入 
初期設定 
データ入力
[リカレントニューラル言語モデル設定](#リカレントニューラル言語モデル設定(ハンズオン）) 
学習を始める前の設定
パラメータ更新方法
言語の予測

もしGPUを使用したい方は、以下にまとめてあるのでご参考ください。
Chainer を用いてリカレントニューラル言語モデル作成のサンプルコードを解説してみた
1.各種ライブラリ導入
Chainerの言語処理では多数のライブラリを導入します。
End of explanation
"""


#-------------Explain7 in the Qiita-------------
n_epochs    = 30
n_units     = 625
batchsize   = 100
bprop_len   = 10
grad_clip   = 0.5
data_dir = "data_hands_on"
checkpoint_dir = "cv"
#-------------Explain7 in the Qiita-------------

"""
Explanation: `導入するライブラリの代表例は下記です。

numpy: 行列計算などの複雑な計算を行なうライブラリ
chainer: Chainerの導入

2.初期設定
下記を設定しています。
* 学習回数：n_epochs
* ニューラルネットのユニット数：n_units
* 確率的勾配法に使用するデータの数：batchsize
* 学習に使用する文字列の長さ：bprop_len
* 勾配法で使用する敷居値：grad_clip
* 学習データの格納場所：data_dir
* モデルの出力場所：checkpoint_dir
End of explanation
"""


# input data
#-------------Explain1 in the Qiita-------------
def source_to_words(source):
    line = source.replace("¥n", " ").replace("¥t", " ")
    for spacer in ["(", ")", "{", "}", "[", "]", ",", ";", ":", "++", "!", "$", '"', "'"]:
        line = line.replace(spacer, " " + spacer + " ")
    
    words = [w.strip() for w in line.split()]
    return words

def load_data():
    vocab = {}
    print ('%s/angular.js'% data_dir)
    source = open('%s/angular_full_remake.js' % data_dir, 'r').read()
    words = source_to_words(source)
    freq = {}
    dataset = np.ndarray((len(words),), dtype=np.int32)
    for i, word in enumerate(words):
        if word not in vocab:
            vocab[word] = len(vocab)
            freq[word] = 0
        dataset[i] = vocab[word]
        freq[word] += 1

    print('corpus length:', len(words))
    print('vocab size:', len(vocab))
    return dataset, words, vocab, freq
#-------------Explain1 in the Qiita-------------

if not os.path.exists(checkpoint_dir):
    os.mkdir(checkpoint_dir)

train_data, words, vocab, freq = load_data()

for f in ["frequent", "rarely"]:
    print("{0} words".format(f))
    print(sorted(freq.items(), key=lambda i: i[1], reverse=True if f == "frequent" else False)[:50])

"""
Explanation: 3.データ入力
学習用にダウンロードしたファイルをプログラムに読ませる処理を関数化しています
文字列の場合は通常のデータと異なり、数字ベクトル化する必要があります。

学習データをテキスト形式で読み込んでいます。
ソースコードを単語で扱えるようにsource_to_wordsのメソッドで分割しています。
文字データを確保するための行列を定義しています。
データは単語をキー、語彙数の連番idを値とした辞書データにして行列データセットに登録しています。

学習データ、単語の長さ、語彙数を取得しています。
上記をそれぞれ行列データとして保持しています。
End of explanation
"""


#-------------Explain2 in the Qiita-------------
class CharRNN(FunctionSet):

    """
    ニューラルネットワークを定義している部分です。
    上から順に入力された辞書ベクトル空間を隠れ層のユニット数に変換し、次に隠れ層の入
    力と隠れ層を設定しています。
    同様の処理を2層にも行い、出力層では語彙数に修正して出力しています。
    なお最初に設定するパラメータは-0.08から0.08の間でランダムに設定しています。
    """

    def __init__(self, n_vocab, n_units):


    """
    順伝搬の記述です。
    順伝搬の入力をVariableで定義し、入力と答えを渡しています。
    入力層を先ほど定義したembedを用います。
    隠れ層の入力には、先ほど定義したl1_xを用いて、引数にdropout、隠れ層の状態を渡して
    います。
    lstmに隠れ層第1層の状態とh1_inを渡します。
    2層目も同様に記述し、出力層は状態を渡さずに定義します。
    次回以降の入力に使用するため各状態は保持しています。
    出力されたラベルと答えのラベル比較し、損失を返すのと状態を返しています。
    """

    def forward_one_step(self, x_data, y_data, state, train=True, dropout_ratio=0.5):


    """
    dropoutの記述を外して予測用のメソッドとして記述しています。
    dropoutにはtrainという引数が存在し、trainの引数をfalseにしておくと動作しない
    ので、予測の時は渡す引数を変えて学習と予測を変えても良いですが、今回は明示的に分る
    ように分けて記述しました。
    """

    def predict(self, x_data, state):


"""
状態の初期化です。
"""

def make_initial_state(n_units, batchsize=100, train=True):

#-------------Explain2 in the Qiita-------------

"""
Explanation: 4.リカレントニューラル言語モデル設定(ハンズオン）
RNNLM(リカレントニューラル言語モデルの設定を行っています）
この部分で自由にモデルを変えることが出来ます。
この部分でリカレントニューラル言語モデル独特の特徴を把握してもらうことが目的です。

EmbedIDで行列変換を行い、疎なベクトルを密なベクトルに変換しています。辞書データを、入力ユニット数分のデータに変換する処理（潜在ベクトル空間への変換）を行っています。
出力が4倍の理由は入力層、入力制限層、出力制限層、忘却層をLSTMでは入力に使用するためです。LSTMの魔法の工夫について知りたい方は下記をご覧下さい。
http://www.slideshare.net/nishio/long-shortterm-memory
h1_in   = self.l1_x(F.dropout(h0, ratio=dropout_ratio, train=train)) + self.l1_h(state['h1'])は隠れ層に前回保持した隠れ層の状態を入力することによってLSTMを実現しています。
F.dropoutは過去の情報を保持しながらどれだけのdropoutでユニットを削るかを表しています。これにより過学習するのを抑えています。
 Drop outについては下記をご覧下さい。

http://olanleed.hatenablog.com/entry/2013/12/03/010945

c1, h1  = F.lstm(state['c1'], h1_in)はlstmと呼ばれる魔法の工夫によってリカレントニューラルネットがメモリ破綻を起こさずにいい感じで学習するための工夫です。詳しく知りたい人は下記をご覧下さい。
return state, F.softmax_cross_entropy(y, t)は予測した文字と実際の文字を比較して損失関数を更新している所になります。ソフトマックス関数を使用している理由は出力層の一つ前の層の全入力を考慮して出力を決定できるので一般的に出力層の計算にはソフトマックス関数が使用されます。
予測を行なうメソッドも実装しており、入力されたデータ、状態を元に次の文字列と状態を返すような関数になっています。
モデルの初期化を行なう関数もここで定義しています。

下記をコーディングして下さい！！！！
End of explanation
"""


# Prepare RNNLM model
model = CharRNN(len(vocab), n_units)

optimizer = optimizers.RMSprop(lr=2e-3, alpha=0.95, eps=1e-8)
optimizer.setup(model.collect_parameters())

"""
Explanation: RNNLM(リカレントニューラル言語モデルの設定を行っています）

作成したリカレントニューラル言語モデルを導入しています。
最適化の手法はRMSpropを使用
http://qiita.com/skitaoka/items/e6afbe238cd69c899b2a
RMSpropは勾配がマイナスであれば重みに加算、正であれば重みを減算する手法です。勾配の加算、減算の度合いを表しています。基本的に勾配が急になれなるほど緩やかに演算が行なわれるように工夫がされています。alphaは過去の勾配による影響を減衰させるパラメータで、lrは勾配の影響を減衰させるパラメータです。epsは０割を防ぐために導入されています。
初期のパラメータを-0.1〜0.1の間で与えています。
End of explanation
"""


whole_len    = train_data.shape[0]
jump         = whole_len // batchsize
epoch        = 0
start_at     = time.time()
cur_at       = start_at
state        = make_initial_state(n_units, batchsize=batchsize)
accum_loss   = Variable(np.zeros((), dtype=np.float32))
cur_log_perp = np.zeros(())

"""
Explanation: 5.学習を始める前の設定

学習データのサイズを取得 
ジャンプの幅を設定（順次学習しない）
パープレキシティを0で初期化 
最初の時間情報を取得 
初期状態を現在の状態に付与 
状態の初期化 
損失を0で初期化
End of explanation
"""


for i in range(int(jump * n_epochs)):
    #-------------Explain4 in the Qiita-------------
    x_batch = np.array([train_data[(jump * j + i) % whole_len]
                        for j in range(batchsize)])
    y_batch = np.array([train_data[(jump * j + i + 1) % whole_len]
                        for j in range(batchsize)])

    state, loss_i = model.forward_one_step(x_batch, y_batch, state, dropout_ratio=0.7)
    accum_loss   += loss_i
    cur_log_perp += loss_i.data.reshape(())

    if (i + 1) % bprop_len == 0:  # Run truncated BPTT
        now = time.time()
        cur_at = now
        #print('{}/{}, train_loss = {}, time = {:.2f}'.format((i + 1)/bprop_len, jump, accum_loss.data / bprop_len, now-cur_at))

        optimizer.zero_grads()
        accum_loss.backward()
        accum_loss.unchain_backward()  # truncate
        accum_loss = Variable(np.zeros((), dtype=np.float32))

        optimizer.clip_grads(grad_clip)
        optimizer.update()
        
    if (i + 1) % 1000 == 0:
        perp = math.exp(cuda.to_cpu(cur_log_perp) / 1000)
        print('iter {} training perplexity: {:.2f} '.format(i + 1, perp))
        fn = ('%s/charrnn_epoch_%i.chainermodel' % (checkpoint_dir, epoch))
        pickle.dump(copy.deepcopy(model).to_cpu(), open(fn, 'wb'))
        cur_log_perp.fill(0)

    if (i + 1) % jump == 0:
        epoch += 1

    #-------------Explain4 in the Qiita-------------

    sys.stdout.flush()


"""
Explanation: 6.パラメータ更新方法（ミニバッチ学習）

ミニバッチを用いて学習している。
x_batch = np.array([train_data[(jump * j + i) % whole_len] for j in range(batchsize)])はややこしいので下記の図を用いて説明します。
下図のように縦に文字が入っている配列があるとします。
jのインデックスはjump(全データのサイズをバッチサイズで割った数)を掛けている各バッチサイズ分移動させる役目を持ち、iのインデックスはバッチサイズ内で移動することを表しています。
whole_lenで余りを出しているのは(jump * j + i)がデータのサイズを超えるので、最初の位置に戻すために行なっている。


* y_batch = np.array([train_data[(jump * j + i + 1) % whole_len] for j in range(batchsize)])はxの一つ先の文字を与えて学習させて
* state, loss_i = model.forward_one_step(x_batch, y_batch, state, dropout_ratio=0.5)は損失と状態を計算しています。ここで過学習を防ぐdropアウトの率も設定可能です。
* if (i + 1) % bprop_len == 0はどれだけ過去の文字を保持するかを表しています。bprop_lenが大きければ大きいほど過去の文字を保持できますが、メモリ破綻を起こす可能性があるのでタスクによって適切な数値に設定する必要があります。
* bprop_lenの詳細についてtruncate
* optimizer.clip_grads(grad_clip)は勾配（重みの更新幅）の大きさに上限を設けており、重みが爆発するのを防いでいます。
End of explanation
"""


# load model
#-------------Explain6 in the Qiita-------------
model = pickle.load(open("cv/charrnn_epoch_1.chainermodel", 'rb'))
#-------------Explain6 in the Qiita-------------
n_units = model.embed.W.shape[1]


"""
Explanation: 7.言語の予測
予測では作成されたモデル変更と文字列予測を行ないます。

モデルを変更する。
文字列を予測する。

予測するモデルの変更はここではiPython notebook内の下記のコードを変更します。
作成されたモデルはcvフォルダの中にあるので
あまり数は出来ていませんが、確認して見て下さい。
End of explanation
"""


# initialize generator
state = make_initial_state(n_units, batchsize=1, train=False)
index = np.random.randint(0, len(vocab), 1)[0]
ivocab = {v:k for k, v in vocab.items()}
sampling_range = 5

for i in range(1000):
    if ivocab[index] in ["}", ";"]:
        sys.stdout.write(ivocab[index] + "\n")
    else:
        sys.stdout.write(ivocab[index] + " ")
    
    #-------------Explain7 in the Qiita-------------
    state, prob = model.predict(np.array([index], dtype=np.int32), state)
    #index = np.argmax(prob.data)
    index = np.random.choice(prob.data.argsort()[0,-sampling_range:][::-1], 1)[0]
    #-------------Explain7 in the Qiita-------------

print

"""
Explanation: state, prob = model.predict(prev_char, state)で予測された確率と状態を取得しています。次の予測にも使用するため状態も取得しています。
index = np.argmax(cuda.to_cpu(prob.data))はcuda.to_cpu(prob.data)部分で各単語の重み確率を取得できるため、その中で一番確率が高いものが予測された文字なのでその文字のインデックスを返すようにしています。
sys.stdout.write(ivocab[index] + " ")で予測した文字を出力するための準備です。
prev_char = np.array([index], dtype=np.int32)は次の予測に使用するために過去の文字を保持するのに使用しています。
End of explanation
"""


