# nlogoxml.py
# Copyright 2008 Richard Lawrence

#  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/>.
#  The GPL is available at: http://www.gnu.org/licenses/gpl.html

import math
from xml.dom import minidom

class SteppedVariable(object):
    """
    Helper class to represent a steppedValueSet node.
    """
    def __init__(self, node, chunk_size):
        "Initializes a SteppedVariable by introspecting a steppedValueSet node."
        # A little type checking
        if node.nodeName != "steppedValueSet":
            raise ValueError(
                "Tried to process a %s node as a SteppedVariable" % \
                 node.nodeName)

        # Someone has to tell us the chunk_size for this variable
        # ahead of time, since it's not in the XML.  Non-positive
        # chunk sizes indicate that this variable should not be
        # chunked.
        if (chunk_size > 0):
            self._chunk_size = chunk_size
        else:
            self._chunk_size = None

        # Introspect the node to get the variable name and the range
        # of values it takes on.  We convert numeric values to floats;
        # type conversion errors will propagate correctly.
        for key, val in node.attributes.items():
            if key == 'variable':
                self._var_name = val
            elif key == 'first':
                self._first_val = float(val)
            elif key == 'last':
                self._last_val = float(val)
            elif key == 'step':
                self._step = float(val)
            else:
                raise ValueError(
                    "Unexpected attribute '%s' on steppedValueSet node" % key)

    def _get_width_of_interval(self):
        "Returns the width of the variable's entire interval"
        return self._last_val - self._first_val

    def _get_number_of_chunks(self):
        "Returns the number of chunks in the variable's interval"
        if self._chunk_size:
            return int(math.ceil(
                    self._get_width_of_interval() / self._chunk_size))
        else:
            return 1

    def _get_endpoints_for_chunk(self, i):
        "Returns the start and end values for a chunk"
        if self._chunk_size:
            start = self._first_val + (i * self._chunk_size)
            end = min(start + self._chunk_size, self._last_val)
        else:
            start = self._first_val
            end = self._last_val
        return start, end


    def get_nodes(self):
        "Returns a list of steppedValueSet nodes, one for each chunk."
        nodes = []
        chunk_num = self._get_number_of_chunks()

        for i in range(chunk_num):
            new_node = minidom.Element(u"steppedValueSet")
            start, end = self._get_endpoints_for_chunk(i)

            new_node.setAttribute(u"variable", self._var_name)
            new_node.setAttribute(u"first", str(start))
            new_node.setAttribute(u"last", str(end))
            new_node.setAttribute(u"step", str(self._step))

            nodes.append(new_node)

        return nodes
            
def cartesian_product(L, *lists):
    """
    Returns the a generator for the Cartesian product of {L, *lists}.

    Usage: 
      rows = [(1, 2, 3), ('a', 'b', 'c'), (4, 5, 6, 7)]
      for tup in cartesian_product(*rows):
          do_something_with(tup)

    Courtesy of Kay Schluehr, http://bytes.com/forum/thread576730.html
    """
    if not lists:
        for x in L:
            yield (x,)
    else:
        for x in L:
            for y in cartesian_product(lists[0],*lists[1:]):
                yield (x,)+y


def split_experiments(doc, variable_chunk_sizes):
    """
    Returns a new minidom.Document object with experiments split into chunks.

    doc should be a minidom.Document object representing a collection
    of NetLogo experiments.

    variable_chunk_sizes should be a dictionary with steppedValueSet
    variable names (as they appear in the XML) for keys, and positive
    numbers for values, e.g., {'community_size' : 10}.

    The new Document returned will have an <experiment> node for each
    item in the Cartesian product of chunked steppedValueSet nodes for
    each variable.  e.g., if there are 3 variables with 2 chunks each,
    the new Document will have 2*2*2 = 8 <experiment> nodes.
    """
    # To avoid mutating the doc, copy it first
    new_doc = doc.cloneNode(True)

    experiment_nodes = new_doc.getElementsByTagName("experiment")
    

    for e in experiment_nodes:
        stepped_val_nodes = e.getElementsByTagName("steppedValueSet")
        chunked_stepped_nodes = []

        for s in stepped_val_nodes:
            # Take steppedValueNodes off the original experiment node;
            # this is so we can copy the experiment node and then
            # insert new steppedValueSet nodes
            e.removeChild(s)

            # Get the chunk size if available; if not, use -1 as an
            # indicator for SteppedVariable that there should only be
            # one chunk
            var_name = s.attributes['variable'].value
            chunk_size = variable_chunk_sizes.get(var_name, -1)

            # Get the chunked steppedValueSet nodes for the variable
            chunky = SteppedVariable(s, chunk_size)
            chunked_stepped_nodes.append(chunky.get_nodes())

            # Free memory associated with s
            s.unlink()

        # Ok, so now chunked_stepped_nodes is a list of lists of
        # parentless steppedValueSet nodes...let's put them back in a
        # new experiment
        count = 0
        for node_set in cartesian_product(*chunked_stepped_nodes):
            new_experiment_node = e.cloneNode(True) # Deep copy
            for node in node_set:
                new_experiment_node.appendChild(node)

            # put a count in the new experiment name before appending it
            count += 1
            experiment_name = "%s-%d" % \
                (e.attributes['name'].value, count)
            new_experiment_node.setAttribute('name', experiment_name)
            e.parentNode.appendChild(new_experiment_node)

        # Before we go, remove the (now empty-of-steppedValSets)
        # experiment node
        e.parentNode.removeChild(e)
        e.unlink()

    return new_doc