#!/usr/bin/env python3
"""
Creates files for a given attribute present in the database

Usage:
  fastsubtrees-attributes-construct [options] <outfile> <tree>
                                              <attrmod> [<attrmod_data>...]

Arguments:
  outfile       desired name for the output file
  tree          fastsubtrees tree representation file,
                e.g. output by fastsubtrees-construct
  attribute     name of the attribute
  attrmod       Python module defining a function attribute_values()
                which may take arguments (<attrmod_data>) and returns pairs
                (element_id, attribute_value) for each node to which an
                attribute value exists.
  attrmod_data  [optional] arguments to be passed to the attribute_values()
                function of the module specified as <attrmod>; to pass keyword
                arguments, use the syntax "key=value" and the option --keyargs

Options:
  --keyargs      split the arguments specified in <idsmod_data> into
                 keywords and values by splitting on the first instance of '=';
                 arguments which do not contain '=' are passed as positional,
                 before any keyword argument
  --quiet        disable log messages
  --debug        print debug information
  --help         show this help message and exit
  --version      show program's version number and exit
  --datatype DT  function to apply to the attribute values before writing
                 them to the output file; <DT> must be a valid Python function;
                 the function can optionally be defined in the module specified
                 by <attrmod> [default: str]
"""

from docopt import docopt
from fastsubtrees import _scripts_support, logger, Tree
import importlib
import json
from collections import defaultdict
from pathlib import Path

def get_datatype_cast_function(args, m):
  if not args["--datatype"]:
    logger.debug("No --datatype option specified, using str")
    cast = str
  else:
    try:
      cast = eval(args["--datatype"])
      logger.debug("Using function '{}' as datatype cast".\
          format(args["--datatype"]))
    except NameError:
      if m.__dict__.get(args["--datatype"]):
        cast = m.__dict__[args["--datatype"]]
        logger.debug("Using function {} from module {} as datatype cast".format(
          args["--datatype"], args["<attrmod>"]))
      else:
        raise ValueError("Invalid attribute datatype casting function '{}'".\
            format(args["--datatype"]))
    return cast

def get_function_arguments(args):
  if args["--keyargs"]:
    keyargs = {k: v for k, v in \
               [a.split("=", 1) for a in args["<attrmod_data>"] if "=" in a]}
    posargs = [a for a in args["<attrmod_data>"] if "=" not in a]
  else:
    keyargs = {}
    posargs = args["<attrmod_data>"]
  if posargs:
    logger.debug(f"Positional arguments passed to the generator: {posargs}")
  if keyargs:
    logger.debug(f"Keyword arguments passed to the generator: {keyargs}")
  return posargs, keyargs

def get_module(args):
  logger.debug("Loading Python module '{}'".format(args['<attrmod>']))
  modulename = Path(args["<attrmod>"]).stem
  spec = importlib.util.spec_from_file_location(modulename, args["<attrmod>"])
  m = importlib.util.module_from_spec(spec)
  spec.loader.exec_module(m)
  if not m.__dict__.get("attribute_values"):
    raise ValueError("The specified Python module {} does not define a "
                     "function attribute_values()".format(args["<attrmod>"]))
  logger.success("Attributes module loaded, found generator attribute_values()")
  return m

def read_attribute_values(m, posargs, keyargs, cast):
  result = defaultdict(list)
  for k, v in m.attribute_values(*posargs, **keyargs):
    result[int(k)].append(cast(v))
  return result

def write_attribute_file(args, attrvalues):
  logger.debug("Loading tree from file '{}'".format(args['<tree>']))
  tree = Tree.from_file(args["<tree>"])
  logger.debug("Writing attribute values to file '{}'".\
      format(args['<outfile>']))
  with open(args["<outfile>"], "w") as outfile:
    for element_id in tree.subtree_ids(tree.root_id):
      attribute = attrvalues.get(element_id, None)
      outfile.write(json.dumps(attribute) + "\n")
  logger.success("Attribute values successfully written to file '{}'".\
      format(args['<outfile>']))

def main(args):
  m = get_module(args)
  posargs, keyargs = get_function_arguments(args)
  cast = get_datatype_cast_function(args, m)
  attrvalues = read_attribute_values(m, posargs, keyargs, cast)
  write_attribute_file(args, attrvalues)

if __name__ == "__main__":
  args = docopt(__doc__, version="0.1")
  _scripts_support.setup_verbosity(args)
  main(args)
