# -*- coding: utf-8 -*-
"""
Color schemes for exception handling code in IPython.
"""

#*****************************************************************************
#       Copyright (C) 2005-2006 Fernando Perez <fperez@colorado.edu>
#
#  Distributed under the terms of the BSD License.  The full license is in
#  the file COPYING, distributed as part of this software.
#*****************************************************************************

from IPython.utils.coloransi import ColorSchemeTable, TermColors, ColorScheme

def exception_colors():
    """Return a color table with fields for exception reporting.

    The table is an instance of ColorSchemeTable with schemes added for
    'Linux', 'LightBG' and 'NoColor' and fields for exception handling filled
    in.

    Examples:

    >>> ec = exception_colors()
    >>> ec.active_scheme_name
    ''
    >>> print ec.active_colors
    None

    Now we activate a color scheme:
    >>> ec.set_active_scheme('NoColor')
    >>> ec.active_scheme_name
    'NoColor'
    >>> sorted(ec.active_colors.keys())
    ['Normal', 'caret', 'em', 'excName', 'filename', 'filenameEm', 'line',
    'lineno', 'linenoEm', 'name', 'nameEm', 'normalEm', 'topline', 'vName',
    'val', 'valEm']
    """

    ex_colors = ColorSchemeTable()

    # Populate it with color schemes
    C = TermColors # shorthand and local lookup
    ex_colors.add_scheme(ColorScheme(
        'NoColor',
        # The color to be used for the top line
        topline = C.NoColor,

        # The colors to be used in the traceback
        filename = C.NoColor,
        lineno = C.NoColor,
        name = C.NoColor,
        vName = C.NoColor,
        val = C.NoColor,
        em = C.NoColor,

        # Emphasized colors for the last frame of the traceback
        normalEm = C.NoColor,
        filenameEm = C.NoColor,
        linenoEm = C.NoColor,
        nameEm = C.NoColor,
        valEm = C.NoColor,

        # Colors for printing the exception
        excName = C.NoColor,
        line = C.NoColor,
        caret = C.NoColor,
        Normal = C.NoColor
        ))

    # make some schemes as instances so we can copy them for modification easily
    ex_colors.add_scheme(ColorScheme(
        'Linux',
        # The color to be used for the top line
        topline = C.LightRed,

        # The colors to be used in the traceback
        filename = C.Green,
        lineno = C.Green,
        name = C.Purple,
        vName = C.Cyan,
        val = C.Green,
        em = C.LightCyan,

        # Emphasized colors for the last frame of the traceback
        normalEm = C.LightCyan,
        filenameEm = C.LightGreen,
        linenoEm = C.LightGreen,
        nameEm = C.LightPurple,
        valEm = C.LightBlue,

        # Colors for printing the exception
        excName = C.LightRed,
        line = C.Yellow,
        caret = C.White,
        Normal = C.Normal
        ))

    # For light backgrounds, swap dark/light colors
    ex_colors.add_scheme(ColorScheme(
        'LightBG',
        # The color to be used for the top line
        topline = C.Red,

        # The colors to be used in the traceback
        filename = C.LightGreen,
        lineno = C.LightGreen,
        name = C.LightPurple,
        vName = C.Cyan,
        val = C.LightGreen,
        em = C.Cyan,

        # Emphasized colors for the last frame of the traceback
        normalEm = C.Cyan,
        filenameEm = C.Green,
        linenoEm = C.Green,
        nameEm = C.Purple,
        valEm = C.Blue,

        # Colors for printing the exception
        excName = C.Red,
        #line = C.Brown,  # brown often is displayed as yellow
        line = C.Red,
        caret = C.Normal,
        Normal = C.Normal,
        ))

    return ex_colors

# For backwards compatibility, keep around a single global object.  Note that
# this should NOT be used, the factory function should be used instead, since
# these objects are stateful and it's very easy to get strange bugs if any code
# modifies the module-level object's state.
ExceptionColors = exception_colors()


from django.core.exceptions import ImproperlyConfigured
try:
    # Django < 2.0
    from django.core.urlresolvers import RegexURLPattern as URLPattern, RegexURLResolver as URLResolver, get_callable
except:
    # Django >= 2.0
    from django.urls import URLPattern, URLResolver, get_callable

from django.utils import six
from django.views.generic import View

class CBVRegexURLPattern(URLPattern):
    _callback_processed = None

    @property
    def callback(self):
        if self._callback_processed is not None:
            return self._callback
        if getattr(self, '_callback_str', None) is not None:
            self._callback = get_callable(self._callback_str)
        if isinstance(self._callback, type) and issubclass(self._callback, View):
            self._callback = self._callback.as_view()
        else:
            self._callback = self._callback
        self._callback_processed = True
        return self._callback

def patterns(prefix, *args):
    """As patterns() in django."""
    pattern_list = []
    for t in args:
        if isinstance(t, (list, tuple)):
            t = url(prefix=prefix, *t)
        elif isinstance(t, RegexURLPattern):
            t.add_prefix(prefix)
        pattern_list.append(t)
    return pattern_list

def url(regex, view, kwargs=None, name=None, prefix=''):
    """As url() in Django."""
    if isinstance(view, (list, tuple)):
        # For include(...) processing.
        urlconf_module, app_name, namespace = view
        return URLResolver(regex, urlconf_module, kwargs, app_name=app_name, namespace=namespace)
    else:
        if isinstance(view, six.string_types):
            if not view:
                raise ImproperlyConfigured('Empty URL pattern view name not permitted (for pattern %r)' % regex)
            if prefix:
                view = prefix + '.' + view
            view = get_callable(view)
        return CBVRegexURLPattern(regex, view, kwargs, name)


#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
from itertools import izip

import polib
from odslib import ODS
from c3po.conf import settings
from c3po.converters.po_csv import _get_all_po_filenames
from c3po.converters.unicode import UnicodeReader

def _escape_apostrophe(entry):
    return ("'" if entry.startswith("'") else "") + entry

def _prepare_ods_columns(ods, trans_title_row):
    """
    Prepare columns in new ods file, create new sheet for metadata,
    set columns color and width. Set formatting style info in your
    settings.py file in ~/.c3po/ folder.
    """
    ods.content.getSheet(0).setSheetName('Translations')
    ods.content.makeSheet('Meta options')
    ods.content.getColumn(0).setWidth('5.0in')
    ods.content.getCell(0, 0).stringValue('metadata')\
        .setCellColor(settings.TITLE_ROW_BG_COLOR) \
        .setBold(True).setFontColor(settings.TITLE_ROW_FONT_COLOR)

    ods.content.getSheet(0)
    ods.content.getColumn(0).setWidth('1.5in')
    ods.content.getCell(0, 0) \
        .setCellColor(settings.TITLE_ROW_BG_COLOR) \
        .setBold(True).setFontColor(settings.TITLE_ROW_FONT_COLOR)
    for i, title in enumerate(trans_title_row):
        ods.content.getColumn(i).setWidth(settings.MSGSTR_COLUMN_WIDTH)
        ods.content.getCell(i, 0).stringValue(title)\
            .setCellColor(settings.TITLE_ROW_BG_COLOR) \
            .setBold(True).setFontColor(settings.TITLE_ROW_FONT_COLOR)
    ods.content.getColumn(0).setWidth(settings.NOTES_COLUMN_WIDTH)

def _write_trans_into_ods(ods, languages, locale_root,
                          po_files_path, po_filename, start_row):
    """
    Write translations from po files into ods one file.
    Assumes a directory structure:
    <locale_root>/<lang>/<po_files_path>/<filename>.
    """
    ods.content.getSheet(0)
    for i, lang in enumerate(languages[1:]):
        lang_po_path = os.path.join(locale_root, lang,
                                    po_files_path, po_filename)
        if os.path.exists(lang_po_path):
            po_file = polib.pofile(lang_po_path)
            for j, entry in enumerate(po_file):
                # start from 4th column, 1st row
                row = j+start_row
                ods.content.getCell(i+4, row).stringValue(
                    _escape_apostrophe(entry.msgstr))
                if i % 2 == 1:
                    ods.content.getCell(i+4, row).setCellColor(
                        settings.ODD_COLUMN_BG_COLOR)
                else:
                    ods.content.getCell(i+4, row).setCellColor(
                        settings.EVEN_COLUMN_BG_COLOR)

def _write_row_into_ods(ods, sheet_no, row_no, row):
    """
    Write row with translations to ods file into specified sheet and row_no.
    """
    ods.content.getSheet(sheet_no)
    for j, col in enumerate(row):
        cell = ods.content.getCell(j, row_no+1)
        cell.stringValue(_escape_apostrophe(col))
        if j % 2 == 1:
            cell.setCellColor(settings.EVEN_COLUMN_BG_COLOR)
        else:
            cell.setCellColor(settings.ODD_COLUMN_BG_COLOR)

def po_to_ods(languages, locale_root, po_files_path, temp_file_path):
    """
    Converts po file to csv GDocs spreadsheet readable format.
    :param languages: list of language codes
    :param locale_root: path to locale root folder containing directories
                        with languages
    :param po_files_path: path from lang directory to po file
    :param temp_file_path: path where temporary files will be saved
    """
    title_row = ['file', 'comment', 'msgid']
    title_row += map(lambda s: s + ':msgstr', languages)

    ods = ODS()

    _prepare_ods_columns(ods, title_row)

    po_files = _get_all_po_filenames(locale_root, languages[0], po_files_path)

    i = 1
    for po_filename in po_files:
        po_file_path = os.path.join(locale_root, languages[0],
                                    po_files_path, po_filename)

        start_row = i

        po = polib.pofile(po_file_path)
        for entry in po:
            meta = dict(entry.__dict__)
            meta.pop('msgid', None)
            meta.pop('msgstr', None)
            meta.pop('tcomment', None)

            ods.content.getSheet(1)
            ods.content.getCell(0, i).stringValue(
                str(meta)).setCellColor(settings.EVEN_COLUMN_BG_COLOR)

            ods.content.getSheet(0)
            ods.content.getCell(0, i) \
                .stringValue(po_filename) \
                .setCellColor(settings.ODD_COLUMN_BG_COLOR)
            ods.content.getCell(1, i) \
                .stringValue(_escape_apostrophe(entry.tcomment)) \
                .setCellColor(settings.ODD_COLUMN_BG_COLOR)
            ods.content.getCell(2, i) \
                .stringValue(_escape_apostrophe(entry.msgid)) \
                .setCellColor(settings.EVEN_COLUMN_BG_COLOR)
            ods.content.getCell(3, i) \
                .stringValue(_escape_apostrophe(entry.msgstr))\
                .setCellColor(settings.ODD_COLUMN_BG_COLOR)

            i += 1

        _write_trans_into_ods(ods, languages, locale_root,
                              po_files_path, po_filename, start_row)

    ods.save(temp_file_path)

def csv_to_ods(trans_csv, meta_csv, local_ods):
    """
    Converts csv files to one ods file
    :param trans_csv: path to csv file with translations
    :param meta_csv: path to csv file with metadata
    :param local_ods: path to new ods file
    """
    trans_reader = UnicodeReader(trans_csv)
    meta_reader = UnicodeReader(meta_csv)

    ods = ODS()

    trans_title = trans_reader.next()
    meta_reader.next()

    _prepare_ods_columns(ods, trans_title)

    for i, (trans_row, meta_row) in enumerate(izip(trans_reader, meta_reader)):
        _write_row_into_ods(ods, 0, i, trans_row)
        _write_row_into_ods(ods, 1, i, meta_row)

    trans_reader.close()
    meta_reader.close()

    ods.save(local_ods)


""" Utilities for accessing the platform's clipboard.
"""

import subprocess
import sys

from IPython.core.error import TryNext

def win32_clipboard_get():
    """ Get the current clipboard's text on Windows.

    Requires Mark Hammond's pywin32 extensions.
    """
    try:
        import win32clipboard
    except ImportError:
        raise TryNext("Getting text from the clipboard requires the pywin32 "
                      "extensions: http://sourceforge.net/projects/pywin32/")
    win32clipboard.OpenClipboard()
    text = win32clipboard.GetClipboardData(win32clipboard.CF_TEXT)
    # FIXME: convert \r\n to \n?
    win32clipboard.CloseClipboard()
    return text

def osx_clipboard_get():
    """ Get the clipboard's text on OS X.
    """
    p = subprocess.Popen(['pbpaste', '-Prefer', 'ascii'],
        stdout=subprocess.PIPE)
    text, stderr = p.communicate()
    # Text comes in with old Mac \r line endings. Change them to \n.
    text = text.replace('\r', '\n')
    return text

def tkinter_clipboard_get():
    """ Get the clipboard's text using Tkinter.

    This is the default on systems that are not Windows or OS X. It may
    interfere with other UI toolkits and should be replaced with an
    implementation that uses that toolkit.
    """
    try:
        import Tkinter
    except ImportError:
        raise TryNext("Getting text from the clipboard on this platform "
                      "requires Tkinter.")
    root = Tkinter.Tk()
    root.withdraw()
    text = root.clipboard_get()
    root.destroy()
    return text


"""Locations where we look for configs, install stuff, etc"""
from __future__ import absolute_import

import getpass
import os
import os.path
import site
import sys
import tempfile

from distutils import sysconfig
from distutils.command.install import install, SCHEME_KEYS

from pip.compat import get_path_uid, WINDOWS
from pip.utils import appdirs
from pip import exceptions

# Hack for flake8
install

# CA Bundle Locations
CA_BUNDLE_PATHS = [
    # Debian/Ubuntu/Gentoo etc.
    "/etc/ssl/certs/ca-certificates.crt",

    # Fedora/RHEL
    "/etc/pki/tls/certs/ca-bundle.crt",

    # OpenSUSE
    "/etc/ssl/ca-bundle.pem",

    # OpenBSD
    "/etc/ssl/cert.pem",

    # FreeBSD/DragonFly
    "/usr/local/share/certs/ca-root-nss.crt",

    # Homebrew on OSX
    "/usr/local/etc/openssl/cert.pem",
]

# Attempt to locate a CA Bundle that we can pass into requests, we have a list
# of possible ones from various systems. If we cannot find one then we'll set
# this to None so that we default to whatever requests is setup to handle.
#
# Note to Downstream: If you wish to disable this autodetection and simply use
#                     whatever requests does (likely you've already patched
#                     requests.certs.where()) then simply edit this line so
#                     that it reads ``CA_BUNDLE_PATH = None``.
CA_BUNDLE_PATH = next((x for x in CA_BUNDLE_PATHS if os.path.exists(x)), None)

# Application Directories
USER_CACHE_DIR = appdirs.user_cache_dir("pip")

DELETE_MARKER_MESSAGE = '''\
This file is placed here by pip to indicate the source was put
here by pip.

Once this package is successfully installed this source code will be
deleted (unless you remove this file).
'''
PIP_DELETE_MARKER_FILENAME = 'pip-delete-this-directory.txt'

def write_delete_marker_file(directory):
    """
    Write the pip delete marker file into this directory.
    """
    filepath = os.path.join(directory, PIP_DELETE_MARKER_FILENAME)
    with open(filepath, 'w') as marker_fp:
        marker_fp.write(DELETE_MARKER_MESSAGE)

def running_under_virtualenv():
    """
    Return True if we're running inside a virtualenv, False otherwise.

    """
    if hasattr(sys, 'real_prefix'):
        return True
    elif sys.prefix != getattr(sys, "base_prefix", sys.prefix):
        return True

    return False

def virtualenv_no_global():
    """
    Return True if in a venv and no system site packages.
    """
    # this mirrors the logic in virtualenv.py for locating the
    # no-global-site-packages.txt file
    site_mod_dir = os.path.dirname(os.path.abspath(site.__file__))
    no_global_file = os.path.join(site_mod_dir, 'no-global-site-packages.txt')
    if running_under_virtualenv() and os.path.isfile(no_global_file):
        return True

def __get_username():
    """ Returns the effective username of the current process. """
    if WINDOWS:
        return getpass.getuser()
    import pwd
    return pwd.getpwuid(os.geteuid()).pw_name

def _get_build_prefix():
    """ Returns a safe build_prefix """
    path = os.path.join(
        tempfile.gettempdir(),
        'pip_build_%s' % __get_username().replace(' ', '_')
    )
    if WINDOWS:
        """ on windows(tested on 7) temp dirs are isolated """
        return path
    try:
        os.mkdir(path)
        write_delete_marker_file(path)
    except OSError:
        file_uid = None
        try:
            # raises OSError for symlinks
            # https://github.com/pypa/pip/pull/935#discussion_r5307003
            file_uid = get_path_uid(path)
        except OSError:
            file_uid = None

        if file_uid != os.geteuid():
            msg = (
                "The temporary folder for building (%s) is either not owned by"
                " you, or is a symlink." % path
            )
            print(msg)
            print(
                "pip will not work until the temporary folder is either "
                "deleted or is a real directory owned by your user account."
            )
            raise exceptions.InstallationError(msg)
    return path

if running_under_virtualenv():
    build_prefix = os.path.join(sys.prefix, 'build')
    src_prefix = os.path.join(sys.prefix, 'src')
else:
    # Note: intentionally NOT using mkdtemp
    # See https://github.com/pypa/pip/issues/906 for plan to move to mkdtemp
    build_prefix = _get_build_prefix()

    # FIXME: keep src in cwd for now (it is not a temporary folder)
    try:
        src_prefix = os.path.join(os.getcwd(), 'src')
    except OSError:
        # In case the current working directory has been renamed or deleted
        sys.exit(
            "The folder you are executing pip from can no longer be found."
        )

# under Mac OS X + virtualenv sys.prefix is not properly resolved
# it is something like /path/to/python/bin/..
# Note: using realpath due to tmp dirs on OSX being symlinks
build_prefix = os.path.abspath(os.path.realpath(build_prefix))
src_prefix = os.path.abspath(src_prefix)

# FIXME doesn't account for venv linked to global site-packages

site_packages = sysconfig.get_python_lib()
user_site = site.USER_SITE
user_dir = os.path.expanduser('~')
if WINDOWS:
    bin_py = os.path.join(sys.prefix, 'Scripts')
    bin_user = os.path.join(user_site, 'Scripts')
    # buildout uses 'bin' on Windows too?
    if not os.path.exists(bin_py):
        bin_py = os.path.join(sys.prefix, 'bin')
        bin_user = os.path.join(user_site, 'bin')

    config_basename = 'pip.ini'

    legacy_storage_dir = os.path.join(user_dir, 'pip')
    legacy_config_file = os.path.join(
        legacy_storage_dir,
        config_basename,
    )
else:
    bin_py = os.path.join(sys.prefix, 'bin')
    bin_user = os.path.join(user_site, 'bin')

    config_basename = 'pip.conf'

    legacy_storage_dir = os.path.join(user_dir, '.pip')
    legacy_config_file = os.path.join(
        legacy_storage_dir,
        config_basename,
    )

    # Forcing to use /usr/local/bin for standard Mac OS X framework installs
    # Also log to ~/Library/Logs/ for use with the Console.app log viewer
    if sys.platform[:6] == 'darwin' and sys.prefix[:16] == '/System/Library/':
        bin_py = '/usr/local/bin'

site_config_files = [
    os.path.join(path, config_basename)
    for path in appdirs.site_config_dirs('pip')
]

def distutils_scheme(dist_name, user=False, home=None, root=None,
                     isolated=False):
    """
    Return a distutils install scheme
    """
    from distutils.dist import Distribution

    scheme = {}

    if isolated:
        extra_dist_args = {"script_args": ["--no-user-cfg"]}
    else:
        extra_dist_args = {}
    dist_args = {'name': dist_name}
    dist_args.update(extra_dist_args)

    d = Distribution(dist_args)
    d.parse_config_files()
    i = d.get_command_obj('install', create=True)
    # NOTE: setting user or home has the side-effect of creating the home dir
    # or user base for installations during finalize_options()
    # ideally, we'd prefer a scheme class that has no side-effects.
    i.user = user or i.user
    i.home = home or i.home
    i.root = root or i.root
    i.finalize_options()
    for key in SCHEME_KEYS:
        scheme[key] = getattr(i, 'install_' + key)

    if i.install_lib is not None:
        # install_lib takes precedence over purelib and platlib
        scheme.update(dict(purelib=i.install_lib, platlib=i.install_lib))

    if running_under_virtualenv():
        scheme['headers'] = os.path.join(
            sys.prefix,
            'include',
            'site',
            'python' + sys.version[:3],
            dist_name,
        )

        if root is not None:
            scheme["headers"] = os.path.join(
                root,
                os.path.abspath(scheme["headers"])[1:],
            )

    return scheme


#!/usr/bin/env python
"""A rectangular domain partitioner and associated communication
functionality for solving PDEs in (1D,2D) using FDM
written in the python language

The global solution domain is assumed to be of rectangular shape,
where the number of cells in each direction is stored in nx, ny, nz

The numerical scheme is fully explicit

Authors
-------

 * Xing Cai
 * Min Ragan-Kelley

"""
from __future__ import print_function
import time

from numpy import zeros, ascontiguousarray, frombuffer
try:
    from mpi4py import MPI
except ImportError:
    pass
else:
    mpi = MPI.COMM_WORLD

class RectPartitioner:
    """
    Responsible for a rectangular partitioning of a global domain,
    which is expressed as the numbers of cells in the different
    spatial directions. The partitioning info is expressed as an
    array of integers, each indicating the number of subdomains in
    one spatial direction.
    """

    def __init__(self, my_id=-1, num_procs=-1, \
                 global_num_cells=[], num_parts=[]):
        self.nsd = 0
        self.my_id = my_id
        self.num_procs = num_procs
        self.redim (global_num_cells, num_parts)

    def redim (self, global_num_cells, num_parts):
        nsd_ = len(global_num_cells)
#        print("Inside the redim function, nsd=%d" %nsd_)

        if nsd_<1 | nsd_>3 | nsd_!=len(num_parts):
            print('The input global_num_cells is not ok!')
            return

        self.nsd = nsd_
        self.global_num_cells = global_num_cells
        self.num_parts = num_parts

    def prepare_communication (self):
        """
        Find the subdomain rank (tuple) for each processor and
        determine the neighbor info.
        """

        nsd_ = self.nsd
        if nsd_<1:
            print('Number of space dimensions is %d, nothing to do' %nsd_)
            return

        self.subd_rank = [-1,-1,-1]
        self.subd_lo_ix = [-1,-1,-1]
        self.subd_hi_ix = [-1,-1,-1]
        self.lower_neighbors = [-1,-1,-1]
        self.upper_neighbors = [-1,-1,-1]

        num_procs = self.num_procs
        my_id = self.my_id

        num_subds = 1
        for i in range(nsd_):
            num_subds = num_subds*self.num_parts[i]
        if my_id==0:
            print("# subds=", num_subds)
            # should check num_subds againt num_procs

        offsets = [1, 0, 0]

        # find the subdomain rank
        self.subd_rank[0] = my_id%self.num_parts[0]
        if nsd_>=2:
            offsets[1] = self.num_parts[0]
            self.subd_rank[1] = my_id/offsets[1]
        if nsd_==3:
            offsets[1] = self.num_parts[0]
            offsets[2] = self.num_parts[0]*self.num_parts[1]
            self.subd_rank[1] = (my_id%offsets[2])/self.num_parts[0]
            self.subd_rank[2] = my_id/offsets[2]

        print("my_id=%d, subd_rank: "%my_id, self.subd_rank)
        if my_id==0:
            print("offsets=", offsets)

        # find the neighbor ids
        for i in range(nsd_):
            rank = self.subd_rank[i]
            if rank>0:
                self.lower_neighbors[i] = my_id-offsets[i]
            if rank<self.num_parts[i]-1:
                self.upper_neighbors[i] = my_id+offsets[i]

            k = self.global_num_cells[i]/self.num_parts[i]
            m = self.global_num_cells[i]%self.num_parts[i]

            ix = rank*k+max(0,rank+m-self.num_parts[i])
            self.subd_lo_ix[i] = ix

            ix = ix+k
            if rank>=(self.num_parts[i]-m):
                ix = ix+1  # load balancing
            if rank<self.num_parts[i]-1:
                ix = ix+1  # one cell of overlap
            self.subd_hi_ix[i] = ix

        print("subd_rank:",self.subd_rank,\
              "lower_neig:", self.lower_neighbors, \
              "upper_neig:", self.upper_neighbors)
        print("subd_rank:",self.subd_rank,"subd_lo_ix:", self.subd_lo_ix, \
              "subd_hi_ix:", self.subd_hi_ix)

class RectPartitioner1D(RectPartitioner):
    """
    Subclass of RectPartitioner, for 1D problems
    """
    def prepare_communication (self):
        """
        Prepare the buffers to be used for later communications
        """

        RectPartitioner.prepare_communication (self)

        if self.lower_neighbors[0]>=0:
            self.in_lower_buffers = [zeros(1, float)]
            self.out_lower_buffers = [zeros(1, float)]
        if self.upper_neighbors[0]>=0:
            self.in_upper_buffers = [zeros(1, float)]
            self.out_upper_buffers = [zeros(1, float)]

    def get_num_loc_cells(self):
        return [self.subd_hi_ix[0]-self.subd_lo_ix[0]]

class RectPartitioner2D(RectPartitioner):
    """
    Subclass of RectPartitioner, for 2D problems
    """
    def prepare_communication (self):
        """
        Prepare the buffers to be used for later communications
        """

        RectPartitioner.prepare_communication (self)

        self.in_lower_buffers = [[], []]
        self.out_lower_buffers = [[], []]
        self.in_upper_buffers = [[], []]
        self.out_upper_buffers = [[], []]

        size1 = self.subd_hi_ix[1]-self.subd_lo_ix[1]+1
        if self.lower_neighbors[0]>=0:
            self.in_lower_buffers[0] = zeros(size1, float)
            self.out_lower_buffers[0] = zeros(size1, float)
        if self.upper_neighbors[0]>=0:
            self.in_upper_buffers[0] = zeros(size1, float)
            self.out_upper_buffers[0] = zeros(size1, float)

        size0 = self.subd_hi_ix[0]-self.subd_lo_ix[0]+1
        if self.lower_neighbors[1]>=0:
            self.in_lower_buffers[1] = zeros(size0, float)
            self.out_lower_buffers[1] = zeros(size0, float)
        if self.upper_neighbors[1]>=0:
            self.in_upper_buffers[1] = zeros(size0, float)
            self.out_upper_buffers[1] = zeros(size0, float)

    def get_num_loc_cells(self):
        return [self.subd_hi_ix[0]-self.subd_lo_ix[0],\
                self.subd_hi_ix[1]-self.subd_lo_ix[1]]

class MPIRectPartitioner2D(RectPartitioner2D):
    """
    Subclass of RectPartitioner2D, which uses MPI via mpi4py for communication
    """

    def __init__(self, my_id=-1, num_procs=-1,
                 global_num_cells=[], num_parts=[],
                 slice_copy=True):
        RectPartitioner.__init__(self, my_id, num_procs,
                                 global_num_cells, num_parts)
        self.slice_copy = slice_copy

    def update_internal_boundary (self, solution_array):
        nsd_ = self.nsd
        if nsd_!=len(self.in_lower_buffers) | nsd_!=len(self.out_lower_buffers):
            print("Buffers for communicating with lower neighbors not ready")
            return
        if nsd_!=len(self.in_upper_buffers) | nsd_!=len(self.out_upper_buffers):
            print("Buffers for communicating with upper neighbors not ready")
            return

        loc_nx = self.subd_hi_ix[0]-self.subd_lo_ix[0]
        loc_ny = self.subd_hi_ix[1]-self.subd_lo_ix[1]

        lower_x_neigh = self.lower_neighbors[0]
        upper_x_neigh = self.upper_neighbors[0]
        lower_y_neigh = self.lower_neighbors[1]
        upper_y_neigh = self.upper_neighbors[1]

        # communicate in the x-direction first
        if lower_x_neigh>-1:
            if self.slice_copy:
                self.out_lower_buffers[0] = ascontiguousarray(solution_array[1,:])
            else:
                for i in xrange(0,loc_ny+1):
                    self.out_lower_buffers[0][i] = solution_array[1,i]
            mpi.Isend(self.out_lower_buffers[0], lower_x_neigh)

        if upper_x_neigh>-1:
            mpi.Recv(self.in_upper_buffers[0], upper_x_neigh)
            if self.slice_copy:
                solution_array[loc_nx,:] = self.in_upper_buffers[0]
                self.out_upper_buffers[0] = ascontiguousarray(solution_array[loc_nx-1,:])
            else:
                for i in xrange(0,loc_ny+1):
                    solution_array[loc_nx,i] = self.in_upper_buffers[0][i]
                    self.out_upper_buffers[0][i] = solution_array[loc_nx-1,i]
            mpi.Isend(self.out_upper_buffers[0], upper_x_neigh)

        if lower_x_neigh>-1:
            mpi.Recv(self.in_lower_buffers[0], lower_x_neigh)
            if self.slice_copy:
                solution_array[0,:] = self.in_lower_buffers[0]
            else:
                for i in xrange(0,loc_ny+1):
                    solution_array[0,i] = self.in_lower_buffers[0][i]

        # communicate in the y-direction afterwards
        if lower_y_neigh>-1:
            if self.slice_copy:
                self.out_lower_buffers[1] = ascontiguousarray(solution_array[:,1])
            else:
                for i in xrange(0,loc_nx+1):
                    self.out_lower_buffers[1][i] = solution_array[i,1]
            mpi.Isend(self.out_lower_buffers[1], lower_y_neigh)

        if upper_y_neigh>-1:
            mpi.Recv(self.in_upper_buffers[1], upper_y_neigh)
            if self.slice_copy:
                solution_array[:,loc_ny] = self.in_upper_buffers[1]
                self.out_upper_buffers[1] = ascontiguousarray(solution_array[:,loc_ny-1])
            else:
                for i in xrange(0,loc_nx+1):
                    solution_array[i,loc_ny] = self.in_upper_buffers[1][i]
                    self.out_upper_buffers[1][i] = solution_array[i,loc_ny-1]
            mpi.Isend(self.out_upper_buffers[1], upper_y_neigh)

        if lower_y_neigh>-1:
            mpi.Recv(self.in_lower_buffers[1], lower_y_neigh)
            if self.slice_copy:
                solution_array[:,0] = self.in_lower_buffers[1]
            else:
                for i in xrange(0,loc_nx+1):
                    solution_array[i,0] = self.in_lower_buffers[1][i]

class ZMQRectPartitioner2D(RectPartitioner2D):
    """
    Subclass of RectPartitioner2D, which uses 0MQ via pyzmq for communication
    The first two arguments must be `comm`, an EngineCommunicator object,
    and `addrs`, a dict of connection information for other EngineCommunicator
    objects.
    """

    def __init__(self, comm, addrs, my_id=-1, num_procs=-1,
                 global_num_cells=[], num_parts=[],
                 slice_copy=True):
        RectPartitioner.__init__(self, my_id, num_procs,
                                 global_num_cells, num_parts)
        self.slice_copy = slice_copy
        self.comm = comm # an Engine
        self.addrs = addrs

    def prepare_communication(self):
        RectPartitioner2D.prepare_communication(self)
        # connect west/south to east/north
        west_id,south_id = self.lower_neighbors[:2]
        west = self.addrs.get(west_id, None)
        south = self.addrs.get(south_id, None)
        self.comm.connect(south, west)

    def update_internal_boundary_x_y (self, solution_array):
        """update the inner boundary with the same send/recv pattern as the MPIPartitioner"""
        nsd_ = self.nsd
        dtype = solution_array.dtype
        if nsd_!=len(self.in_lower_buffers) | nsd_!=len(self.out_lower_buffers):
            print("Buffers for communicating with lower neighbors not ready")
            return
        if nsd_!=len(self.in_upper_buffers) | nsd_!=len(self.out_upper_buffers):
            print("Buffers for communicating with upper neighbors not ready")
            return

        loc_nx = self.subd_hi_ix[0]-self.subd_lo_ix[0]
        loc_ny = self.subd_hi_ix[1]-self.subd_lo_ix[1]

        lower_x_neigh = self.lower_neighbors[0]
        upper_x_neigh = self.upper_neighbors[0]
        lower_y_neigh = self.lower_neighbors[1]
        upper_y_neigh = self.upper_neighbors[1]
        trackers = []
        flags = dict(copy=False, track=False)
        # communicate in the x-direction first
        if lower_x_neigh>-1:
            if self.slice_copy:
                self.out_lower_buffers[0] = ascontiguousarray(solution_array[1,:])
            else:
                for i in xrange(0,loc_ny+1):
                    self.out_lower_buffers[0][i] = solution_array[1,i]
            t = self.comm.west.send(self.out_lower_buffers[0], **flags)
            trackers.append(t)

        if upper_x_neigh>-1:
            msg = self.comm.east.recv(copy=False)
            self.in_upper_buffers[0] = frombuffer(msg, dtype=dtype)
            if self.slice_copy:
                solution_array[loc_nx,:] = self.in_upper_buffers[0]
                self.out_upper_buffers[0] = ascontiguousarray(solution_array[loc_nx-1,:])
            else:
                for i in xrange(0,loc_ny+1):
                    solution_array[loc_nx,i] = self.in_upper_buffers[0][i]
                    self.out_upper_buffers[0][i] = solution_array[loc_nx-1,i]
            t = self.comm.east.send(self.out_upper_buffers[0], **flags)
            trackers.append(t)

        if lower_x_neigh>-1:
            msg = self.comm.west.recv(copy=False)
            self.in_lower_buffers[0] = frombuffer(msg, dtype=dtype)
            if self.slice_copy:
                solution_array[0,:] = self.in_lower_buffers[0]
            else:
                for i in xrange(0,loc_ny+1):
                    solution_array[0,i] = self.in_lower_buffers[0][i]

        # communicate in the y-direction afterwards
        if lower_y_neigh>-1:
            if self.slice_copy:
                self.out_lower_buffers[1] = ascontiguousarray(solution_array[:,1])
            else:
                for i in xrange(0,loc_nx+1):
                    self.out_lower_buffers[1][i] = solution_array[i,1]
            t = self.comm.south.send(self.out_lower_buffers[1], **flags)
            trackers.append(t)

        if upper_y_neigh>-1:
            msg = self.comm.north.recv(copy=False)
            self.in_upper_buffers[1] = frombuffer(msg, dtype=dtype)
            if self.slice_copy:
                solution_array[:,loc_ny] = self.in_upper_buffers[1]
                self.out_upper_buffers[1] = ascontiguousarray(solution_array[:,loc_ny-1])
            else:
                for i in xrange(0,loc_nx+1):
                    solution_array[i,loc_ny] = self.in_upper_buffers[1][i]
                    self.out_upper_buffers[1][i] = solution_array[i,loc_ny-1]
            t = self.comm.north.send(self.out_upper_buffers[1], **flags)
            trackers.append(t)

        if lower_y_neigh>-1:
            msg = self.comm.south.recv(copy=False)
            self.in_lower_buffers[1] = frombuffer(msg, dtype=dtype)
            if self.slice_copy:
                solution_array[:,0] = self.in_lower_buffers[1]
            else:
                for i in xrange(0,loc_nx+1):
                    solution_array[i,0] = self.in_lower_buffers[1][i]

        # wait for sends to complete:
        if flags['track']:
            for t in trackers:
                t.wait()

    def update_internal_boundary_send_recv (self, solution_array):
        """update the inner boundary, sending first, then recving"""
        nsd_ = self.nsd
        dtype = solution_array.dtype
        if nsd_!=len(self.in_lower_buffers) | nsd_!=len(self.out_lower_buffers):
            print("Buffers for communicating with lower neighbors not ready")
            return
        if nsd_!=len(self.in_upper_buffers) | nsd_!=len(self.out_upper_buffers):
            print("Buffers for communicating with upper neighbors not ready")
            return

        loc_nx = self.subd_hi_ix[0]-self.subd_lo_ix[0]
        loc_ny = self.subd_hi_ix[1]-self.subd_lo_ix[1]

        lower_x_neigh = self.lower_neighbors[0]
        upper_x_neigh = self.upper_neighbors[0]
        lower_y_neigh = self.lower_neighbors[1]
        upper_y_neigh = self.upper_neighbors[1]
        trackers = []
        flags = dict(copy=False, track=False)

        # send in all directions first
        if lower_x_neigh>-1:
            if self.slice_copy:
                self.out_lower_buffers[0] = ascontiguousarray(solution_array[1,:])
            else:
                for i in xrange(0,loc_ny+1):
                    self.out_lower_buffers[0][i] = solution_array[1,i]
            t = self.comm.west.send(self.out_lower_buffers[0], **flags)
            trackers.append(t)

        if lower_y_neigh>-1:
            if self.slice_copy:
                self.out_lower_buffers[1] = ascontiguousarray(solution_array[:,1])
            else:
                for i in xrange(0,loc_nx+1):
                    self.out_lower_buffers[1][i] = solution_array[i,1]
            t = self.comm.south.send(self.out_lower_buffers[1], **flags)
            trackers.append(t)

        if upper_x_neigh>-1:
            if self.slice_copy:
                self.out_upper_buffers[0] = ascontiguousarray(solution_array[loc_nx-1,:])
            else:
                for i in xrange(0,loc_ny+1):
                    self.out_upper_buffers[0][i] = solution_array[loc_nx-1,i]
            t = self.comm.east.send(self.out_upper_buffers[0], **flags)
            trackers.append(t)

        if upper_y_neigh>-1:
            if self.slice_copy:
                self.out_upper_buffers[1] = ascontiguousarray(solution_array[:,loc_ny-1])
            else:
                for i in xrange(0,loc_nx+1):
                    self.out_upper_buffers[1][i] = solution_array[i,loc_ny-1]
            t = self.comm.north.send(self.out_upper_buffers[1], **flags)
            trackers.append(t)

        # now start receiving
        if upper_x_neigh>-1:
            msg = self.comm.east.recv(copy=False)
            self.in_upper_buffers[0] = frombuffer(msg, dtype=dtype)
            if self.slice_copy:
                solution_array[loc_nx,:] = self.in_upper_buffers[0]
            else:
                for i in xrange(0,loc_ny+1):
                    solution_array[loc_nx,i] = self.in_upper_buffers[0][i]

        if lower_x_neigh>-1:
            msg = self.comm.west.recv(copy=False)
            self.in_lower_buffers[0] = frombuffer(msg, dtype=dtype)
            if self.slice_copy:
                solution_array[0,:] = self.in_lower_buffers[0]
            else:
                for i in xrange(0,loc_ny+1):
                    solution_array[0,i] = self.in_lower_buffers[0][i]

        if upper_y_neigh>-1:
            msg = self.comm.north.recv(copy=False)
            self.in_upper_buffers[1] = frombuffer(msg, dtype=dtype)
            if self.slice_copy:
                solution_array[:,loc_ny] = self.in_upper_buffers[1]
            else:
                for i in xrange(0,loc_nx+1):
                    solution_array[i,loc_ny] = self.in_upper_buffers[1][i]

        if lower_y_neigh>-1:
            msg = self.comm.south.recv(copy=False)
            self.in_lower_buffers[1] = frombuffer(msg, dtype=dtype)
            if self.slice_copy:
                solution_array[:,0] = self.in_lower_buffers[1]
            else:
                for i in xrange(0,loc_nx+1):
                    solution_array[i,0] = self.in_lower_buffers[1][i]

        # wait for sends to complete:
        if flags['track']:
            for t in trackers:
                t.wait()

    # use send/recv pattern instead of x/y sweeps
    update_internal_boundary = update_internal_boundary_send_recv


"""Utilities to manipulate JSON objects.
"""
#-----------------------------------------------------------------------------
#  Copyright (C) 2010-2011  The IPython Development Team
#
#  Distributed under the terms of the BSD License.  The full license is in
#  the file COPYING.txt, distributed as part of this software.
#-----------------------------------------------------------------------------

#-----------------------------------------------------------------------------
# Imports
#-----------------------------------------------------------------------------
# stdlib
import math
import re
import sys
import types
from base64 import encodestring
from datetime import datetime

from IPython.utils import py3compat
from IPython.utils.encoding import DEFAULT_ENCODING
from IPython.utils import text
next_attr_name = '__next__' if py3compat.PY3 else 'next'

#-----------------------------------------------------------------------------
# Globals and constants
#-----------------------------------------------------------------------------

# timestamp formats
ISO8601="%Y-%m-%dT%H:%M:%S.%f"
ISO8601_PAT=re.compile(r"^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\.\d+$")

#-----------------------------------------------------------------------------
# Classes and functions
#-----------------------------------------------------------------------------

def rekey(dikt):
    """Rekey a dict that has been forced to use str keys where there should be
    ints by json."""
    for k in dikt.iterkeys():
        if isinstance(k, basestring):
            ik=fk=None
            try:
                ik = int(k)
            except ValueError:
                try:
                    fk = float(k)
                except ValueError:
                    continue
            if ik is not None:
                nk = ik
            else:
                nk = fk
            if nk in dikt:
                raise KeyError("already have key %r"%nk)
            dikt[nk] = dikt.pop(k)
    return dikt

def extract_dates(obj):
    """extract ISO8601 dates from unpacked JSON"""
    if isinstance(obj, dict):
        obj = dict(obj) # don't clobber
        for k,v in obj.iteritems():
            obj[k] = extract_dates(v)
    elif isinstance(obj, (list, tuple)):
        obj = [ extract_dates(o) for o in obj ]
    elif isinstance(obj, basestring):
        if ISO8601_PAT.match(obj):
            obj = datetime.strptime(obj, ISO8601)
    return obj

def squash_dates(obj):
    """squash datetime objects into ISO8601 strings"""
    if isinstance(obj, dict):
        obj = dict(obj) # don't clobber
        for k,v in obj.iteritems():
            obj[k] = squash_dates(v)
    elif isinstance(obj, (list, tuple)):
        obj = [ squash_dates(o) for o in obj ]
    elif isinstance(obj, datetime):
        obj = obj.strftime(ISO8601)
    return obj

def date_default(obj):
    """default function for packing datetime objects in JSON."""
    if isinstance(obj, datetime):
        return obj.strftime(ISO8601)
    else:
        raise TypeError("%r is not JSON serializable"%obj)

# constants for identifying png/jpeg data
PNG = b'\x89PNG\r\n\x1a\n'
JPEG = b'\xff\xd8'

def encode_images(format_dict):
    """b64-encodes images in a displaypub format dict

    Perhaps this should be handled in json_clean itself?

    Parameters
    ----------

    format_dict : dict
        A dictionary of display data keyed by mime-type

    Returns
    -------

    format_dict : dict
        A copy of the same dictionary,
        but binary image data ('image/png' or 'image/jpeg')
        is base64-encoded.

    """
    encoded = format_dict.copy()
    pngdata = format_dict.get('image/png')
    if isinstance(pngdata, bytes) and pngdata[:8] == PNG:
        encoded['image/png'] = encodestring(pngdata).decode('ascii')
    jpegdata = format_dict.get('image/jpeg')
    if isinstance(jpegdata, bytes) and jpegdata[:2] == JPEG:
        encoded['image/jpeg'] = encodestring(jpegdata).decode('ascii')
    return encoded

def json_clean(obj):
    """Clean an object to ensure it's safe to encode in JSON.

    Atomic, immutable objects are returned unmodified.  Sets and tuples are
    converted to lists, lists are copied and dicts are also copied.

    Note: dicts whose keys could cause collisions upon encoding (such as a dict
    with both the number 1 and the string '1' as keys) will cause a ValueError
    to be raised.

    Parameters
    ----------
    obj : any python object

    Returns
    -------
    out : object

      A version of the input which will not cause an encoding error when
      encoded as JSON.  Note that this function does not *encode* its inputs,
      it simply sanitizes it so that there will be no encoding errors later.

    Examples
    --------
    >>> json_clean(4)
    4
    >>> json_clean(range(10))
    [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
    >>> sorted(json_clean(dict(x=1, y=2)).items())
    [('x', 1), ('y', 2)]
    >>> sorted(json_clean(dict(x=1, y=2, z=[1,2,3])).items())
    [('x', 1), ('y', 2), ('z', [1, 2, 3])]
    >>> json_clean(True)
    True
    """
    # types that are 'atomic' and ok in json as-is.  bool doesn't need to be
    # listed explicitly because bools pass as int instances
    atomic_ok = (unicode, int, types.NoneType)

    # containers that we need to convert into lists
    container_to_list = (tuple, set, types.GeneratorType)

    if isinstance(obj, float):
        # cast out-of-range floats to their reprs
        if math.isnan(obj) or math.isinf(obj):
            return repr(obj)
        return obj

    if isinstance(obj, atomic_ok):
        return obj

    if isinstance(obj, bytes):
        return obj.decode(DEFAULT_ENCODING, 'replace')

    if isinstance(obj, container_to_list) or (
        hasattr(obj, '__iter__') and hasattr(obj, next_attr_name)):
        obj = list(obj)

    if isinstance(obj, list):
        return [json_clean(x) for x in obj]

    if isinstance(obj, dict):
        # First, validate that the dict won't lose data in conversion due to
        # key collisions after stringification.  This can happen with keys like
        # True and 'true' or 1 and '1', which collide in JSON.
        nkeys = len(obj)
        nkeys_collapsed = len(set(map(str, obj)))
        if nkeys != nkeys_collapsed:
            raise ValueError('dict can not be safely converted to JSON: '
                             'key collision would lead to dropped values')
        # If all OK, proceed by making the new dict that will be json-safe
        out = {}
        for k,v in obj.iteritems():
            out[str(k)] = json_clean(v)
        return out

    # If we get here, we don't know how to handle the object, so we just get
    # its repr and return that.  This will catch lambdas, open sockets, class
    # objects, and any other complicated contraption that json can't encode
    return repr(obj)