#encoding: utf-8
#
# <one line to give the library's name and an idea of what it does.>
# Copyright (C) 2015  <copyright holder> <email>
# 
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
# 

import entailment

__doc__ = '''
This application can be used in entailment task to summarize given 
pairs of graphs to one bigger. Each pair consists of two graphs G(f) and G(t).
G(f) is the graph which represents "from" sentence and G(t) "to" sentence.
As the output two graphs are produced G'(f) and G'(t) which contains summarized
graphs built from "from" and "to" sentences.
'''

import os, sys, networkx
try:
  import argcomplete
except ImportError:
  argcomplete = None

from grafon import grafon_base
from grafon.graphs.basegraph import BaseGraph
from grafon.utils.grafonconfig import GrafonConfig
from grafon.gbuilders.graphbuilder import GraphBuilder
from grafon.mergers.basegraphmerger import BaseGraphMerger
from grafon import constants

def make_parser(desc = ''):
  aparser = entailment.make_parser(__doc__)
  # w ilu grafach musi wystapic ten sam wezel
  aparser.add_argument(
    '-mn', '--min-count-node', dest='min_node_count',
    required=False, type=int, default=1,
    help=('Minimum number of node count into different graphs'))
  # how many tokens have to be represented by one node
  aparser.add_argument(
    '-mt', '--min-token-count', dest='min_token_count',
    required=False, type=int, default=1,
    help=('Minimum number of tokens belonging to the same node'))
  # removing nodes with zero degree
  aparser.add_argument(
    '--no-rem-zero-deg-nodes', dest='no_rem_zero_deg_nodes',
    default=False, action='store_true',
    help=('If set to true, then nodes with zero degree are not removed'))
  # output directory
  aparser.add_argument(
    '--out-dir', dest='out_dir', default='grafon_graphs',
    help=('Directory where all files will be stored'))
  # relation name
  aparser.add_argument(
    '--rel-name', dest='rel_name', required=True, 
    help=(
      'Name of processing relation. This name will '
      'be used to create any files: pickles and (optionally) pngs.'))

  if argcomplete:
    argcomplete.autocomplete(aparser)
  return aparser

def generalize_graphs(mn, mt, bg, graphs_dict, 
                      max_division=8, rem_zero_deg_nodes=True):
  """
  The main function to generalize graphs for relations

  @param mn minimum number of occurence node into graphs
  @type mn: Integer
  @param mt minimum number of tokens into node
  @type mt: Integer
  @param bg Base graph used to read document relations from it
  @type bg: BaseGraph
  @param graphs_dict Dictionary of graphs for sentences
  @type graphs_dict: dict{sid<str>: graph<networkx.MultiDiGraph>}
  @param max_division Number which will be used as normalization factor to 
         maximum occurence of nodes into graphs:
         if node_frequence < number_of_graphs/max_division then is accepted
         Default is set to 8
  @type max_division Integer
  @param rem_zero_deg_nodes If is set to true (as default) then nodes without
         any edges (with degree equal zero) will be removed
  @type rem_zero_deg_nodes: Integer

  @return Generalized graphs with experiment informations
  @rtype: triple: [networkx.MultiDiGraph, networkx.MultiDiGraph, str]
  """
  bmrg = BaseGraphMerger(bg)
  fg = tg = None

  # makes list of "from" and "to" graphs and merge graphs into two
  fgs = []
  tgs = []

  # maximum number of graphs merged into one
  max_g = 100
  gen_graphs = []
  rcount = len(bg.document.relations())
  for rnum, r in enumerate(bg.document.relations()):
    print >> sys.stderr, "\r\t[tmp] {}/{}".format(rnum, rcount),
    if rnum % max_g == 0:
      fg = networkx.MultiDiGraph()
      tg = networkx.MultiDiGraph()
    g_from = graphs_dict.get(r.rel_from().sentence_id(), None)
    g_to = graphs_dict.get(r.rel_to().sentence_id(), None)
    if not g_from or not g_to:
      continue

    fg = bmrg.merge(fg, g_from)
    tg = bmrg.merge(tg, g_to)
    fgs.append(g_from)
    tgs.append(g_to)
    if (rnum + 1) % max_g == 0:
      gen_graphs.append([fg, tg])

  if rcount and fg and tg:
    if rnum % max_g != 0:
      gen_graphs.append([fg, tg])

    gcount = len(gen_graphs)
    if gcount:
      print >> sys.stderr, ''
      fg, tg = gen_graphs[0]
      for gnum, [fgg, tgg] in enumerate(gen_graphs[1:]):
        print >> sys.stderr, "\r\t[gen] {}/{}".format(gnum + 1, gcount),
        fg = bmrg.merge(fg, fgg)
        tg = bmrg.merge(tg, tgg)
      print >> sys.stderr, ''

    assert(len(fgs) == len(tgs))

    # collect all "from" and "to" nodes into two separated dictionaries
    # dictionary maps node to number of occurece this node into all graphs
    print >> sys.stderr, '\t~>',
    f_nodes = [g.nodes() for g in fgs]
    print >> sys.stderr, '{}...'.format(1),
    t_nodes = [g.nodes() for g in tgs]
    print >> sys.stderr, '{}...'.format(2),
    f_nodes = [num for elem in f_nodes for num in elem]
    print >> sys.stderr, '{}...'.format(3),
    t_nodes = [num for elem in t_nodes for num in elem]
    print >> sys.stderr, '{}...'.format(4),
    fc = dict()
    for n in f_nodes:
      if not fc.has_key(n):
        fc[n] = 0
      else:
        fc[n] += 1
    print >> sys.stderr, '{}...'.format(5),
    tc = dict()
    for n in t_nodes:
      if not tc.has_key(n):
        tc[n] = 0
      else:
        tc[n] += 1
    print >> sys.stderr, '{}...'.format(6),
    # remove nodes if not occurs in required number of graphs
    fg.remove_nodes_from([n for n in fg.nodes() if fc[n] < mn])
    print >> sys.stderr, '{}...'.format(7),
    tg.remove_nodes_from([n for n in tg.nodes() if tc[n] < mn])
    print >> sys.stderr, '{}...'.format(8),
    # wywalame te wezly, ktore wystapily czesto w grafach.
    # Czesto = czesciej raz w kazdym
    ftgs_s = len(fgs) / max_division
    fg.remove_nodes_from([n for n in fg.nodes() if fc[n] > ftgs_s])
    print >> sys.stderr, '{}...'.format(9),
    tg.remove_nodes_from([n for n in tg.nodes() if tc[n] > ftgs_s])
    print >> sys.stderr, '{}...'.format(10),
    # remove nodes which represents to low number of tokens
    fg.remove_nodes_from(
      [n for n in fg.nodes() if fg.node[n][constants.NODE_COUNT] < mt])
    print >> sys.stderr, '{}...'.format(11),
    tg.remove_nodes_from(
      [n for n in tg.nodes() if tg.node[n][constants.NODE_COUNT] < mt])
    print >> sys.stderr, '{}...'.format(12),
    # remove nodes with 0 degree
    if rem_zero_deg_nodes:
      fg.remove_nodes_from(
        [n for n in fg.nodes() if fg.in_degree(n) + fg.out_degree(n) == 0])
      tg.remove_nodes_from(
        [n for n in tg.nodes() if tg.in_degree(n) + tg.out_degree(n) == 0])
    print >> sys.stderr, ''
    # remove token list from built graphs
    for n in fg.nodes():
      fg.node[n][constants.NODE_TOKENS] = []
    for n in tg.nodes():
      tg.node[n][constants.NODE_TOKENS] = []

  return [
    fg, tg, 
    'mintok-{}_minnode-{}_remzeronodes-{}_maxnode-div-{}'.format(
      mt, mn, rem_zero_deg_nodes, max_division)]

def save_graphs(fg, tg, out_dir, relname, ex_inform, store_to_png=False):
  out_dir = '{}__{}'.format(out_dir, ex_inform)
  if not os.path.exists(out_dir):
    os.makedirs(out_dir)

  networkx.write_gpickle(
    fg, "{}/{}_#_from_graph.pickle".format(out_dir, relname))
  networkx.write_gpickle(
    tg, "{}/{}_#_to___graph.pickle".format(out_dir, relname))

  if store_to_png:
    print >> sys.stderr, 'Saving both graphs to png files...'
    grafon_base.write_graph_to_file(
      fg, "{}/{}_#_f_g.png".format(out_dir, relname))
    grafon_base.write_graph_to_file(
      tg, "{}/{}_#_t_g.png".format(out_dir, relname))

def main(argv = None):
  p = make_parser(__doc__)
  args = p.parse_args(argv)
  cfg = GrafonConfig(args.config)

  tagset = cfg.tagset()
  node = cfg.node(tagset)
  filters = cfg.filters(tagset)

  bg = BaseGraph(node, filters)
  smerger = cfg.sumo_merger(bg)
  pmerger = cfg.plwn_merger(bg)

  sgraph    = grafon_base.read_sumo_graph(cfg) if smerger else None
  plwngraph = grafon_base.read_plwn_graph(cfg) if pmerger else None
  reducers  = cfg.reducers() if cfg.reducers() else None

  try:
    print >> sys.stderr, 'Building base graphs for relation {}...'.format(args.rel_name)
    relccl = args.relccl if os.path.exists(args.relccl) else args.ccl
    bgbuilder = GraphBuilder()
    bgbuilder.build_graph_for_rel_sentences(
      bg, args.ccl, relccl, tagset, cfg.accepted_relations())
  except Exception, e:
    print >> sys.stderr, "Exception:", e
    return 0

  for ebuilder in cfg.ebuilders():
    print >> sys.stderr, 'Adding edges into base graph using builder', ebuilder
    ebuilder.add_edges(bg)

  all_graphs = grafon_base.merge_graphs_to_dict(
    bg, smerger, sgraph, pmerger, plwngraph, reducers)

  print >> sys.stderr, 'Generalizing graphs for relation {}...'.format(args.rel_name)
  fg, tg, ex_infor = generalize_graphs(
    args.min_node_count, args.min_token_count, bg, all_graphs,
    rem_zero_deg_nodes=not args.no_rem_zero_deg_nodes)

  if fg and tg:
    print >> sys.stderr, \
      "\t\t>> from(nodes) =", len(fg.nodes()), \
      "from(edges) =", len(fg.edges()), \
      "\n\t\t>> to(nodes) =", len(tg.nodes()), \
      "to(edges) =", len(tg.edges())
  else:
    print >> sys.stderr, \
      "\t\t>> from(nodes) = 0 from(edges) = 0"\
      "\n\t\t>> to(nodes) = 0 to(edges) = 0"

  print >> sys.stderr, \
    'Saving graphs (relation {}) to files...'.format(args.rel_name), 
  save_graphs(fg, tg, args.out_dir, args.rel_name, ex_infor)
  print >> sys.stderr, 'done'
  
  print >> sys.stderr, 50 * '~'

if __name__ == '__main__':
  print >> sys.stderr, '[RUN] GENERALIZE GRAPHS'
  main()