#!/usr/bin/env python

# Copyright 2014 Open Connectome Project (http://openconnecto.me)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

# graphml_adapter.py
# Created by Disa Mhembere on 2014-01-06.
# Email: disa@jhu.edu

# File performs conversion to and from scipy.sparse matrices to graphml and back

import argparse
import os
import re

__weight__=True # If the graph is weighted this will be set

def csc_to_graphml(g, is_weighted=True, desikan=False, is_directed=False, save_fn="default_name.graphml", is_tri=False, test=False):
  """
  Convert a csc graph to graphml format for writing to disk

  Positional arguments:
  ====================
  g - the csc graph

  Optional arguments:
  ===================
  is_weighted - is the graph weighted. Type: boolean.
  desikan - use the desikan mapping to label nodes. Type: boolean
  is_directed - is g symmetric ? Type: boolean
  save_fn - file name to use when saving. Type: boolean
  is_tri - is the adjacency mat upper or lower triangular. Type: boolean
  test - are we running a test. Type: boolean
  """

  print "Beginning graphml construction .."
  if test: test_str = ""

  tabs = 2 # How many tabs on affix to the front

  src = """<?xml version="1.0" encoding="UTF-8"?>
  <graphml xmlns="http://graphml.graphdrawing.org/xmlns"
    xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
    xsi:schemaLocation="http://graphml.graphdrawing.org/xmlns
    http://graphml.graphdrawing.org/xmlns/1.0/graphml.xsd">
    <!-- Created by script: %s -->\n""" % __file__

  # Do we have desikan labels ?
  if desikan:
    from mrcap import desikan

    src += "  "*2+"<key id=\"v_region\" for=\"node\" attr.name=\"region\" attr.type=\"string\"/>\n" # Desikan vertex attr called v_region
    tabs = 3

  # Is our graph weighted ?
  if is_weighted:
    src += "  "*2+"<key id=\"e_weight\" for=\"edge\" attr.name=\"weight\" attr.type=\"double\"/>\n" # Desikan vertex attr called v_region
    tabs = 3

  # Directed graph ?
  if is_directed:
    src += "\n    <graph id=\"G\" edgedefault=\"undirected\">"

  # Undirected graph?
  else:  # not directed so just use upper tri
    if not is_tri:
      print "Converting to upper triangular ..."
      from scipy.sparse.csc import csc_matrix
      from scipy.sparse import triu

      g = g = csc_matrix(triu(g, k=0))
      src += "\n    <graph id=\"G\" edgedefault=\"undirected\">\n"

  NUM_NODES = g.shape[0]

  if not test: f = open(save_fn if os.path.splitext(save_fn)[1] == ".graphml" else save_fn+".graphml", "wb")

  # Can be #pragma for
  for node in xrange(NUM_NODES): # Cycle through all nodes
    s = "<node id=\"n%d\">\n" % node

    if desikan:
      s += "  "*(tabs+1)+"<data key=\"v_region\">\"%s\"</data>\n" % (desikan.des_map.get(node, "Undefined"))

    s += "  "*tabs+"</node>\n"
    src += "  "*tabs+s

    if node % 50000 == 0:
      print "Processing node %d / %d ..." % (node, NUM_NODES)
      if test: test_str += src
      else: f.write(src)
      src = ""

  del s # free mem

  print "Adding edges to graph ..."
  # Get all edge data
  nodes_from, nodes_to = g.nonzero()
  data = g.data
  del g # free some mem

  # Can be #pragma for
  NUM_EDGES = nodes_from.shape[0]
  for idx in xrange(NUM_EDGES): # Only the edges that exist
    src += "  "*tabs+"<edge source=\"n%d\" target=\"n%d\">\n" % (nodes_from[idx], nodes_to[idx])
    if is_weighted:
      src += "  "*(tabs+1)+"<data key=\"e_weight\">%d</data>\n" % data[idx]
    src += "  "*tabs+"</edge>\n"

    if idx % 100000 == 0:
      print "Processing edge %d / %d ..." % (idx, NUM_EDGES)
      if test: test_str += src
      else: f.write(src)
      src = ""

  src += "  </graph>\n</graphml>"

  if test:
    test_str += src
    return test_str

  f.write(src)
  f.close

def graphml_to_csc(fh):
  """
  Take the filehandle of a graphml formatted graph written to disk and convert to
  scipy.sparse.csc.matrix

  *Cannot account for any node attributes*
  *CAN ONLY account for edge weight attributes*
  *All other attributes are ignored if any*

  Positional arguments:
  =====================
  fh - the file handle to the graphml file
  """

  from scipy.sparse.lil import lil_matrix
  from scipy.sparse.csc import csc_matrix

  print "Processing nodes ..."
  # Assume all header stuff is ok
  while True: # Infinite loop
    pos = fh.tell()
    line = fh.readline().replace(" ", "").strip() # remove if inefficient

    if line.startswith("<node"):
      node_pos = pos # May be last node so take position

    if line.startswith("<edge"): # Wait for edges to begin
      break

  fh.seek(node_pos)