#!/usr/bin/env python
#
# Pucktada Treeratpituk (https://pucktada.github.io/)
# License: MIT
# 2017-05-01
#
# A recurrent neural network model (LSTM) for thai word segmentation
import logging
import re
import numpy as np
import tensorflow as tf
import tensorflow.contrib.layers as layers
import tensorflow.contrib.rnn as rnn

#from . import char_dictionary

def load_settings(sess):

    model_settings = dict()        
    model_vars = dict()

    graph = tf.get_default_graph()
    #for v in sess.graph.get_operations():
    #    print('P:', v.name)

    configs = ['cell_sizes', 'look_ahead', 'num_layers', 'input_classes', 'label_classes', 'learning_rate', 'l2_regularization', 'cell_type'] #, 'direction']
    for c in configs:
        name = 'prefix/%s:0' % c
        model_settings[c] = sess.run(graph.get_tensor_by_name(name))

    model_vars['inputs']  = graph.get_tensor_by_name('prefix/placeholder/inputs:0')
    model_vars['fw_state'] = graph.get_tensor_by_name('prefix/placeholder/fw_state:0')    
    model_vars['seq_lengths'] = graph.get_tensor_by_name('prefix/placeholder/seq_lengths:0')
    model_vars['keep_prob'] = graph.get_tensor_by_name('prefix/placeholder/keep_prob:0')

    model_vars['probs'] = graph.get_tensor_by_name('prefix/probs:0')

    return model_settings, model_vars

def load_graph(model_file):
    """ loading necessary configuration of the network from the meta file & 
        the checkpoint file together with variables that are needed for the inferences
    """

    # We load the protobuf file from the disk and parse it to retrieve the 
    # unserialized graph_def
    with tf.gfile.GFile(model_file, "rb") as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())

    # Then, we can use again a convenient built-in function to import a graph_def into the 
    # current default Graph
    with tf.Graph().as_default() as graph:
        tf.import_graph_def(
            graph_def,
            input_map=None, 
            return_elements=None, 
            name='prefix',
            op_dict=None,
            producer_op_list=None)

    return graph

def load_model2(sess, meta_file, checkpoint_file):
    """ loading necessary configuration of the network from the meta file & 
        the checkpoint file together with variables that are needed for the inferences
    """
    saver = tf.train.import_meta_graph(meta_file, clear_devices=True)
    saver.restore(sess, checkpoint_file)    

    configs = tf.get_collection('configs')
    pvars   = tf.get_collection('placeholders')

    model_settings = dict()
    for c in configs:
        name = c.name.split(':')[0]
        model_settings[name] = sess.run(c)
        
    model_vars = dict()
    for p in pvars:
        scope, name, _ = re.split('[:/]', p.name)
        model_vars[name] = p
    model_vars['probs'] = tf.get_collection('probs')[0]

    return model_settings, model_vars

class CkModel:
    """ cutkum model: LSTM recurrent neural network model """
    
    def __init__(self, model_settings):
        logging.info('...init WordSegmentor')

        self.num_layers = model_settings["num_layers"]        
      
        self.cell_sizes  = model_settings["cell_sizes"] # list of cell_size, same length as num_layers
        self.total_cells = sum(self.cell_sizes)
        self.cell_start  = [sum(self.cell_sizes [:i]) for i in range(self.num_layers)]

        # keep number of look_ahead (not used in the training, but so that people know how to use the model)
        self.look_ahead = model_settings['look_ahead']

        #self.num_unroll = model_settings["num_unroll"]
        self.input_classes = model_settings['input_classes']
        self.label_classes = model_settings['label_classes']
        self.learning_rate = model_settings['learning_rate']
        self.l2_regularization = model_settings['l2_regularization'] # 0.1