# -*- coding:utf8 -*-
import logging
import os

from utensor_cgen.backend.base import BackendPart
from utensor_cgen.backend.utensor.snippets.composer import Composer
from utensor_cgen.backend.utensor.snippets.legacy import (
    CommentSnippet, ContextGlobalArrayContainer, ContextHeaderSnippet,
    ContextSnippetsContainer)
from utensor_cgen.transformer.optimizer import RefCntOptimizer
from utensor_cgen.utils import (Configuration, LazyLoader,
                                NamescopedKWArgsParser, class_property)

from ._operators import OperatorFactory

__all__ = ["uTensorLegacyCodeGenerator"]
_logger = logging.getLogger('utensor-cli')
tf = LazyLoader('tensorflow')

class uTensorLegacyCodeGenerator(BackendPart):

  TARGET = 'utensor'
  PART = 'legacy_code_generator'

  def __init__(self, config):
    final_config = Configuration(self.default_config, config)
    self.src_fname = final_config['src_fname']
    self.params_dir = final_config['params_dir'].rstrip('/')
    if not os.path.exists(self.params_dir):
      os.makedirs(self.params_dir)
    self.embed_data_dir = final_config['embed_params_dir'].rstrip('/')
    self.model_dir = final_config['model_dir'].rstrip('/')
    self.debug_cmt = final_config['debug_cmt']

  @classmethod
  def from_kwargs(cls, **kwargs):
    return cls(config=kwargs)

  def apply(self, ugraph):
    """Generate source and header files
    """
    src_fname = self.src_fname
    if src_fname == 'None':
      src_fname = '{}.cpp'.format(ugraph.name)
    header_snippet = ContextHeaderSnippet(
      '_MODELS_{}'.format(ugraph.name),
      ugraph.name
    )
    weight_container = ContextGlobalArrayContainer()
    composer = Composer()
    header_fname = '{}.hpp'.format(ugraph.name)
    weight_header_fname = '{}_weight.hpp'.format(ugraph.name)
    container = ContextSnippetsContainer(ugraph.name, header_fname, weight_header_fname)

    opFactory = OperatorFactory()

    if not os.path.exists(os.path.join(self.params_dir, ugraph.name)):
      os.makedirs(os.path.join(self.params_dir, ugraph.name))
    for op_id, op_name in enumerate(ugraph.topo_order):
      op_info = ugraph.ops_info[op_name]
      op_type = op_info.op_type
      # TODO: better abstraction for snippet
      if op_type == "Placeholder":
        parser = NamescopedKWArgsParser(RefCntOptimizer.KWARGS_NAMESCOPE, 
                                        op_info.op_attr)
        out_tname = op_info.output_tensors[0].name
        ref_count = parser.get('ref_counts', [0])[0]
        container.template_vars["placeholders"].append(out_tname)
        container.template_vars["ref_counts"].append(ref_count)
        header_snippet.template_vars["placeholders"].append(out_tname)
      else:
        # TODO: the operator may correspond to multiple snippets (such as InlinTensor)
        # weight_container is passed to function for workaround
        snippet = opFactory.createOperatorSnippet(
          op_info,
          idx_dir=os.path.join(self.params_dir, ugraph.name),
          embed_data_dir=self.embed_data_dir,
          weight_container=weight_container,
        )
        container.add_snippet(snippet)

      if self.debug_cmt:
        comments = ["<<< Operation id {}: {}".format(op_id, op_name),
                    ">>> Operation id {}: {}".format(op_id + 1, op_name)]
        cmt_snippet = CommentSnippet(comments)
        container.add_snippet(cmt_snippet)
    composer.add_snippet(container)

    # generate cpp/hpp files
    if not os.path.exists(self.model_dir):
      os.makedirs(self.model_dir)
    if weight_container.snippets:  
      _logger.info("Generate weight file: %s", weight_header_fname)
      with open(os.path.join(self.model_dir, weight_header_fname), "w") as wf:
        wf.write('// Auto generated by utensor-cli\n\n')
        wf.write(weight_container.render())
    else:
      container.remove_header('"{}"'.format(weight_header_fname))
      
    _logger.info("Generate header file: %s", header_fname)
    with open(os.path.join(self.model_dir, header_fname), "w") as wf:
      wf.write('// Auto generated by utensor-cli\n\n')
      wf.write(header_snippet.render())
    _logger.info("Generate source file: %s", src_fname)
    with open(os.path.join(self.model_dir, src_fname), "w") as wf:
      wf.write('// Auto generated by utensor-cli\n\n')
      wf.write(composer.compose())

  @class_property
  def default_config(cls):
    config = {}
    config['src_fname'] = 'None'
    config['params_dir'] = 'constants'
    config['embed_params_dir'] = '/fs/constants'
    config['model_dir'] = 'models'
    config['debug_cmt'] = False
    return config
