#!/usr/bin/env python
#
# Copyright 2009 Facebook
#
# 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.

"""Escaping/unescaping methods for HTML, JSON, URLs, and others.

Also includes a few other miscellaneous string manipulation functions that
have crept in over time.
"""

import htmlentitydefs
import re
import sys
import urllib

# Python3 compatibility:  On python2.5, introduce the bytes alias from 2.6
try: bytes
except Exception: bytes = str

try:
    from urlparse import parse_qs  # Python 2.6+
except ImportError:
    from cgi import parse_qs

# json module is in the standard library as of python 2.6; fall back to
# simplejson if present for older versions.
try:
    import json
    assert hasattr(json, "loads") and hasattr(json, "dumps")
    _json_decode = json.loads
    _json_encode = json.dumps
except Exception:
    try:
        import simplejson
        _json_decode = lambda s: simplejson.loads(_unicode(s))
        _json_encode = lambda v: simplejson.dumps(v)
    except ImportError:
        try:
            # For Google AppEngine
            from django.utils import simplejson
            _json_decode = lambda s: simplejson.loads(_unicode(s))
            _json_encode = lambda v: simplejson.dumps(v)
        except ImportError:
            def _json_decode(s):
                raise NotImplementedError(
                    "A JSON parser is required, e.g., simplejson at "
                    "http://pypi.python.org/pypi/simplejson/")
            _json_encode = _json_decode

_XHTML_ESCAPE_RE = re.compile('[&<>"]')
_XHTML_ESCAPE_DICT = {'&': '&', '<': '<', '>': '>', '"': '"'}
def xhtml_escape(value):
    """Escapes a string so it is valid within XML or XHTML."""
    return _XHTML_ESCAPE_RE.sub(lambda match: _XHTML_ESCAPE_DICT[match.group(0)],
                                to_basestring(value))

def xhtml_unescape(value):
    """Un-escapes an XML-escaped string."""
    return re.sub(r"&(#?)(\w+?);", _convert_entity, _unicode(value))

def json_encode(value):
    """JSON-encodes the given Python object."""
    # JSON permits but does not require forward slashes to be escaped.
    # This is useful when json data is emitted in a  tags from prematurely terminating
    # the javscript.  Some json libraries do this escaping by default,
    # although python's standard library does not, so we do it here.
    # http://stackoverflow.com/questions/1580647/json-why-are-forward-slashes-escaped
    return _json_encode(recursive_unicode(value)).replace("</", "<\\/")

def json_decode(value):
    """Returns Python objects for the given JSON string."""
    return _json_decode(to_basestring(value))

def squeeze(value):
    """Replace all sequences of whitespace chars with a single space."""
    return re.sub(r"[\x00-\x20]+", " ", value).strip()

def url_escape(value):
    """Returns a valid URL-encoded version of the given value."""
    return urllib.quote_plus(utf8(value))

# python 3 changed things around enough that we need two separate
# implementations of url_unescape.  We also need our own implementation
# of parse_qs since python 3's version insists on decoding everything.
if sys.version_info[0] < 3:
    def url_unescape(value, encoding='utf-8'):
        """Decodes the given value from a URL.

        The argument may be either a byte or unicode string.

        If encoding is None, the result will be a byte string.  Otherwise,
        the result is a unicode string in the specified encoding.
        """
        if encoding is None:
            return urllib.unquote_plus(utf8(value))
        else:
            return unicode(urllib.unquote_plus(utf8(value)), encoding)

    parse_qs_bytes = parse_qs
else:
    def url_unescape(value, encoding='utf-8'):
        """Decodes the given value from a URL.

        The argument may be either a byte or unicode string.

        If encoding is None, the result will be a byte string.  Otherwise,
        the result is a unicode string in the specified encoding.
        """
        if encoding is None:
            return urllib.parse.unquote_to_bytes(value)
        else:
            return urllib.unquote_plus(to_basestring(value), encoding=encoding)

    def parse_qs_bytes(qs, keep_blank_values=False, strict_parsing=False):
        """Parses a query string like urlparse.parse_qs, but returns the
        values as byte strings.

        Keys still become type str (interpreted as latin1 in python3!)
        because it's too painful to keep them as byte strings in
        python3 and in practice they're nearly always ascii anyway.
        """
        # This is gross, but python3 doesn't give us another way.
        # Latin1 is the universal donor of character encodings.
        result = parse_qs(qs, keep_blank_values, strict_parsing,
                          encoding='latin1', errors='strict')
        encoded = {}
        for k,v in result.iteritems():
            encoded[k] = [i.encode('latin1') for i in v]
        return encoded

_UTF8_TYPES = (bytes, type(None))
def utf8(value):
    """Converts a string argument to a byte string.

    If the argument is already a byte string or None, it is returned unchanged.
    Otherwise it must be a unicode string and is encoded as utf8.
    """
    if isinstance(value, _UTF8_TYPES):
        return value
    assert isinstance(value, unicode)
    return value.encode("utf-8")

_TO_UNICODE_TYPES = (unicode, type(None))
def to_unicode(value):
    """Converts a string argument to a unicode string.

    If the argument is already a unicode string or None, it is returned
    unchanged.  Otherwise it must be a byte string and is decoded as utf8.
    """
    if isinstance(value, _TO_UNICODE_TYPES):
        return value
    assert isinstance(value, bytes)
    return value.decode("utf-8")

# to_unicode was previously named _unicode not because it was private,
# but to avoid conflicts with the built-in unicode() function/type
_unicode = to_unicode

# When dealing with the standard library across python 2 and 3 it is
# sometimes useful to have a direct conversion to the native string type
if str is unicode:
    native_str = to_unicode
else:
    native_str = utf8

_BASESTRING_TYPES = (basestring, type(None))
def to_basestring(value):
    """Converts a string argument to a subclass of basestring.

    In python2, byte and unicode strings are mostly interchangeable,
    so functions that deal with a user-supplied argument in combination
    with ascii string constants can use either and should return the type
    the user supplied.  In python3, the two types are not interchangeable,
    so this method is needed to convert byte strings to unicode.
    """
    if isinstance(value, _BASESTRING_TYPES):
        return value
    assert isinstance(value, bytes)
    return value.decode("utf-8")

def recursive_unicode(obj):
    """Walks a simple data structure, converting byte strings to unicode.

    Supports lists, tuples, and dictionaries.
    """
    if isinstance(obj, dict):
        return dict((recursive_unicode(k), recursive_unicode(v)) for (k,v) in obj.iteritems())
    elif isinstance(obj, list):
        return list(recursive_unicode(i) for i in obj)
    elif isinstance(obj, tuple):
        return tuple(recursive_unicode(i) for i in obj)
    elif isinstance(obj, bytes):
        return to_unicode(obj)
    else:
        return obj

# I originally used the regex from
# http://daringfireball.net/2010/07/improved_regex_for_matching_urls
# but it gets all exponential on certain patterns (such as too many trailing
# dots), causing the regex matcher to never return.
# This regex should avoid those problems.
_URL_RE = re.compile(ur"""\b((?:([\w-]+):(/{1,3})|www[.])(?:(?:(?:[^\s&()]|&|")*(?:[^!"#$%&'()*+,.:;<=>?@\[\]^`{|}~\s]))|(?:\((?:[^\s&()]|&|")*\)))+)""")

def linkify(text, shorten=False, extra_params="",
            require_protocol=False, permitted_protocols=["http", "https"]):
    """Converts plain text into HTML with links.

    For example: ``linkify("Hello http://tornadoweb.org!")`` would return
    ``Hello http://tornadoweb.org!``

    Parameters:

    shorten: Long urls will be shortened for display.

    extra_params: Extra text to include in the link tag,
        e.g. linkify(text, extra_params='rel="nofollow" class="external"')

    require_protocol: Only linkify urls which include a protocol. If this is
        False, urls such as www.facebook.com will also be linkified.

    permitted_protocols: List (or set) of protocols which should be linkified,
        e.g. linkify(text, permitted_protocols=["http", "ftp", "mailto"]).
        It is very unsafe to include protocols such as "javascript".
    """
    if extra_params:
        extra_params = " " + extra_params.strip()

    def make_link(m):
        url = m.group(1)
        proto = m.group(2)
        if require_protocol and not proto:
            return url  # not protocol, no linkify

        if proto and proto not in permitted_protocols:
            return url  # bad protocol, no linkify

        href = m.group(1)
        if not proto:
            href = "http://" + href   # no proto specified, use http

        params = extra_params

        # clip long urls. max_len is just an approximation
        max_len = 30
        if shorten and len(url) > max_len:
            before_clip = url
            if proto:
                proto_len = len(proto) + 1 + len(m.group(3) or "")  # +1 for :
            else:
                proto_len = 0

            parts = url[proto_len:].split("/")
            if len(parts) > 1:
                # Grab the whole host part plus the first bit of the path
                # The path is usually not that interesting once shortened
                # (no more slug, etc), so it really just provides a little
                # extra indication of shortening.
                url = url[:proto_len] + parts[0] + "/" + \
                        parts[1][:8].split('?')[0].split('.')[0]

            if len(url) > max_len * 1.5:  # still too long
                url = url[:max_len]

            if url != before_clip:
                amp = url.rfind('&')
                # avoid splitting html char entities
                if amp > max_len - 5:
                    url = url[:amp]
                url += "..."

                if len(url) >= len(before_clip):
                    url = before_clip
                else:
                    # full url is visible on mouse-over (for those who don't
                    # have a status bar, such as Safari by default)
                    params += ' title="%s"' % href

        return u'%s' % (href, params, url)

    # First HTML-escape so that our strings are all safe.
    # The regex is modified to avoid character entites other than & so
    # that we won't pick up ", etc.
    text = _unicode(xhtml_escape(text))
    return _URL_RE.sub(make_link, text)

def _convert_entity(m):
    if m.group(1) == "#":
        try:
            return unichr(int(m.group(2)))
        except ValueError:
            return "&#%s;" % m.group(2)
    try:
        return _HTML_UNICODE_MAP[m.group(2)]
    except KeyError:
        return "&%s;" % m.group(2)

def _build_unicode_map():
    unicode_map = {}
    for name, value in htmlentitydefs.name2codepoint.iteritems():
        unicode_map[name] = unichr(value)
    return unicode_map

_HTML_UNICODE_MAP = _build_unicode_map()


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

class AuthPlugin(object):
    name = 'auth'

    def __init__(self, engine, keyword='auth'):
        """
        :param engine: Auth engine created function
        :param keyword: Keyword used to inject auth in a route
        """
        self.engine = engine
        self.keyword = keyword

    def setup(self, app):
        """ Make sure that other installed plugins don't affect the same
            keyword argument and check if metadata is available."""
        for other in app.plugins:
            if not isinstance(other, AuthPlugin):
                continue
            if other.keyword == self.keyword:
                raise bottle.PluginError("Found another auth plugin "
                                         "with conflicting settings ("
                                         "non-unique keyword).")

    def apply(self, callback, context):
        args = inspect.getargspec(context['callback'])[0]

        if self.keyword not in args:
            return callback

        def wrapper(*args, **kwargs):
            kwargs[self.keyword] = self.engine
            return callback(*args, **kwargs)

        return wrapper


"""
Routines for obtaining the class names
of an object and its parent classes.
"""

from __future__ import unicode_literals

def all_bases(c):
	"""
	return a tuple of all base classes the class c has as a parent.
	>>> object in all_bases(list)
	True
	"""
	return c.mro()[1:]

def all_classes(c):
	"""
	return a tuple of all classes to which c belongs
	>>> list in all_classes(list)
	True
	"""
	return c.mro()

# borrowed from
# http://code.activestate.com/recipes/576949-find-all-subclasses-of-a-given-class/

def iter_subclasses(cls, _seen=None):
	"""
	Generator over all subclasses of a given class, in depth-first order.

	>>> bool in list(iter_subclasses(int))
	True
	>>> class A(object): pass
	>>> class B(A): pass
	>>> class C(A): pass
	>>> class D(B,C): pass
	>>> class E(D): pass
	>>>
	>>> for cls in iter_subclasses(A):
	...		print(cls.__name__)
	B
	D
	E
	C
	>>> # get ALL (new-style) classes currently defined
	>>> res = [cls.__name__ for cls in iter_subclasses(object)]
	>>> 'type' in res
	True
	>>> 'tuple' in res
	True
	>>> len(res) > 100
	True
	"""

	if not isinstance(cls, type):
		raise TypeError(
			'iter_subclasses must be called with '
			'new-style classes, not %.100r' % cls
		)
	if _seen is None:
		_seen = set()
	try:
		subs = cls.__subclasses__()
	except TypeError:  # fails only when cls is type
		subs = cls.__subclasses__(cls)
	for sub in subs:
		if sub in _seen:
			continue
		_seen.add(sub)
		yield sub
		for sub in iter_subclasses(sub, _seen):
			yield sub


# -*- encoding: utf-8 -*-
#
# This file is part of wsgicors
#
# wsgicors is a WSGI middleware that answers CORS preflight requests
#
# copyright 2014-2015 Norman Krämer
#
# 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.

import fnmatch
from functools import reduce
from collections import namedtuple
try:
    from functools import lru_cache
except ImportError:
    from backports.functools_lru_cache import lru_cache

class CORS(object):
    "WSGI middleware allowing CORS requests to succeed"

    @staticmethod
    def matchpattern(accu, pattern, host):
        return accu or fnmatch.fnmatch(host, pattern)

    @staticmethod
    def matchlist(origin, allowed_origins, case_sensitive=False):
        return reduce(lambda accu, x: CORS.matchpattern(accu, x, origin.lower() if not case_sensitive else origin), allowed_origins, False)

    def __init__(self, application, cfg=None, **kw):

        Policy = namedtuple("Policy", ["name", "origin", "methods", "headers", "expose_headers", "credentials", "maxage", "match"])

        self.policies = {}
        if kw and "policy" not in kw:  # direct config
            self.activepolicies = ["direct"]
            self.matchstrategy = "firstmatch"
            self.policies["direct"]=kw
        else:  # multiple policies programatically or via configfile (paster factory for instance)
            cfg = kw or cfg or {}
            self.activepolicies = list(map(lambda x: x.strip(), cfg.get("policy", "deny").split(",")))
            self.matchstrategy = cfg.get("matchstrategy", "firstmatch")

            for policy in self.activepolicies:
                kw = {}
                prefix = policy + "_"
                for k, v in filter(lambda key_value: key_value[0].startswith(prefix), cfg.items()):
                    kw[k.split(prefix)[-1]] = v
                self.policies[policy]=kw

        for policy in self.activepolicies:
            kw = self.policies[policy]
            # copy or * or a space separated list of hostnames, possibly with filename wildcards "*" and "?"
            pol_origin = kw.get("origin", "")
            if pol_origin not in ("copy", "*"):
                match = list(filter(lambda x: x != "*", map(lambda x: x.strip(), pol_origin.split(" "))))
            else:
                match = []

            pol_methods = kw.get("methods", "")
            methods = list(map(lambda x: x.strip(), pol_methods.split(",")))

            # copy or * or a space separated list of hostnames, possibly with filename wildcards "*" and "?"
            pol_headers = kw.get("headers", "")  # * or list of headers
            pol_expose_headers = kw.get("expose_headers", "")  # * or list of headers to expose to the client
            pol_credentials = kw.get("credentials", "false")  # true or false
            pol_maxage = kw.get("maxage", "")  # in seconds
            pol=Policy(name=policy,
                       origin=pol_origin,
                       methods=methods,
                       headers=pol_headers,
                       expose_headers=pol_expose_headers,
                       credentials=pol_credentials,
                       maxage=pol_maxage,
                       match=match)
            self.policies[policy] = pol

            # a little sanity check
            configkeys="origin,methods,headers,expose_headers,credentials,maxage".split(",")
            existingkeys=[k for k in configkeys if k in kw]

            if "origin" not in kw:
                if existingkeys:
                    print("The policy '%s' was referenced but has no value for 'origin' set. Nothing good can come from this." % policy)
                elif policy != "deny":
                    print("The policy '%s' was referenced but hasn't defined any keys. This might be an case sensitivity issue." % policy)

        self.application = application

    @lru_cache(maxsize=200)
    def selectPolicy(self, origin, request_method=None):
        "Based on the matching strategy and the origin and optionally the requested method a tuple of policyname and origin to pass back is returned."
        ret_origin = None
        policyname = None
        if self.matchstrategy in ("firstmatch", "verbmatch"):
            for pol in self.activepolicies:
                policy=self.policies[pol]
                ret_origin = None
                policyname = policy.name
                if policyname == "deny":
                    break
                if self.matchstrategy == "verbmatch":
                    if policy.methods != "*" and not CORS.matchlist(request_method, policy.methods, case_sensitive=True):
                        continue
                if origin and policy.match:
                    if CORS.matchlist(origin, policy.match):
                        ret_origin = origin
                elif policy.origin == "copy":
                    ret_origin = origin
                elif policy.origin:
                    ret_origin = policy.origin
                if ret_origin:
                    break
        return policyname, ret_origin

    def __call__(self, environ, start_response):

        # we handle the request ourself only if it is identified as a prefilght request
        if 'OPTIONS' == environ['REQUEST_METHOD'] and environ.get("HTTP_ACCESS_CONTROL_REQUEST_METHOD") is not None \
           and environ.get("HTTP_ORIGIN") is not None:
            resp = []

            orig = environ.get("HTTP_ORIGIN")
            ac_request_method = environ.get("HTTP_ACCESS_CONTROL_REQUEST_METHOD")
            policyname, origin = self.selectPolicy(orig, ac_request_method)

            if policyname == "deny":
                pass
            else:
                policy = self.policies[policyname]
                methods = None
                headers = None
                credentials = None
                maxage = None

                if "*" in policy.methods:
                    methods = environ.get("HTTP_ACCESS_CONTROL_REQUEST_METHOD", None)
                elif policy.methods:
                    methods = ", ".join(policy.methods)

                if policy.headers == "*":
                    headers = environ.get("HTTP_ACCESS_CONTROL_REQUEST_HEADERS", None)
                elif policy.headers:
                    headers = policy.headers

                if policy.credentials == "true":
                    credentials = "true"

                if policy.maxage:
                    maxage = policy.maxage

                if origin: resp.append(('Access-Control-Allow-Origin', origin))
                if methods: resp.append(('Access-Control-Allow-Methods', methods))
                if headers: resp.append(('Access-Control-Allow-Headers', headers))
                if credentials: resp.append(('Access-Control-Allow-Credentials', credentials))
                if maxage: resp.append(('Access-Control-Max-Age', maxage))

            status = '204 OK'
            start_response(status, resp)
            return []

        orig = environ.get("HTTP_ORIGIN", None)
        request_method = environ['REQUEST_METHOD']
        policyname, ret_origin = self.selectPolicy(orig, request_method)

        if orig and policyname != "deny":
            def custom_start_response(status, headers, exc_info=None):

                policyname, ret_origin = self.selectPolicy(orig, request_method)
                policy = self.policies[policyname]
                if policy.credentials == 'true' and policy.origin == "*":
                    # for credentialed access '*' are ignored in origin
                    ret_origin = orig

                if ret_origin:
                    headers.append(('Access-Control-Allow-Origin', ret_origin))

                    if policy.credentials == 'true':
                        headers.append(('Access-Control-Allow-Credentials', 'true'))

                    if policy.expose_headers:
                        headers.append(('Access-Control-Expose-Headers', policy.expose_headers))

                    if policy.origin != "*":
                        headers.append(('Vary', 'Origin'))

                return start_response(status, headers, exc_info)
        else:
            custom_start_response = start_response

        return self.application(environ, custom_start_response)

def make_middleware(app, cfg=None, **kw):
    cfg = (cfg or {}).copy()
    cfg.update(kw)
    app = CORS(app, cfg)
    return app


import struct
from collections import namedtuple

__all__ = ('load', 'loads', 'dump', 'dumps')

VDF_VERSIONS = [0x07564426, 0x07564427]
VDF_UNIVERSE = 0x00000001

LAST_SECTION = b'\x00'
LAST_APP = b'\x00\x00\x00\x00'
SECTION_END = b'\x08'

TYPE_SECTION = b'\x00'
TYPE_STRING = b'\x01'
TYPE_INT32 = b'\x02'
TYPE_INT64 = b'\x07'

# VDF has variable length integers (32-bit and 64-bit).
Integer = namedtuple('Integer', ('size', 'data'))

def load(fp, wrapper=dict):
    """
    Loads the contents of an Appinfo file into a Python object.
    :param fp: A file object.
    :param wrapper: A wrapping object for key-value pairs.
    :return: An Ordered Dictionary with Appinfo data.
    """
    return loads(fp.read(), wrapper=wrapper)

def loads(data, wrapper=dict):
    """
    Loads Appinfo content into a Python object.
    :param data: A byte-like object with the contents of an Appinfo file.
    :param wrapper: A wrapping object for key-value pairs.
    :return: An Ordered Dictionary with Appinfo data.
    """
    if not isinstance(data, (bytes, bytearray)):
        raise TypeError('can only load a bytes-like object as an Appinfo but got ' + type(data).__name__)

    return AppinfoDecoder(data, wrapper=wrapper).decode()

def dump(obj, fp):
    """
    Serializes a dictionary into Appinfo data and writes it to a file.
    :param obj: A dictionary to serialize.
    :param fp: A file object.
    """
    fp.write(dumps(obj))

def dumps(obj):
    """
    Serializes a dictionary into Appinfo data.
    :param obj: A dictionary to serialize.
    :return:
    """
    if not isinstance(obj, dict):
        raise TypeError('can only dump a dictionary as an Appinfo but got ' + type(obj).__name__)

    return b''.join(AppinfoEncoder(obj).iter_encode())

class AppinfoDecoder:

    def __init__(self, data, wrapper=dict):
        self.wrapper = wrapper        # Wrapping container
        self.data = memoryview(data)  # Incoming data (bytes)
        self.offset = 0               # Parsing offset

        # Commonly used structs
        self._read_int32 = self.make_custom_reader('<I', single_value=True)
        self._read_int64 = self.make_custom_reader('<Q', single_value=True)
        self.read_vdf_header = self.make_custom_reader('<2I')
        self.read_game_header = self.make_custom_reader('<3IQ20sI')

        # Functions to parse different data structures.
        self.value_parsers = {
            0x00: self.parse_subsections,
            0x01: self.read_string,
            0x02: self.read_int32,
            0x07: self.read_int64,
        }

    def decode(self):
        parsed = self.wrapper()

        # These should always be present.
        header_fields = ('version', 'universe')
        header = self.wrapper((zip(header_fields, self.read_vdf_header())))
        if len(header) != len(header_fields):
            raise ValueError('Not all VDF headers are present, only found {num}: {header!r}'.format(
                num=len(header),
                header=header,
            ))

        # Currently these are the only possible values for
        # a valid appinfo.vdf
        if header['version'] not in VDF_VERSIONS:
            raise ValueError('Unknown VDF_VERSION: 0x{0:08x}'.format(header['version']))

        if header['universe'] != VDF_UNIVERSE:
            raise ValueError('Unknown VDF_UNIVERSE: 0x{0:08x}'.format(header['version']))

        # Store VDF_VERSION and VDF_UNIVERSE internally, as it's needed for proper encoding.
        parsed[b'__vdf_version'], parsed[b'__vdf_universe'] = header['version'], header['universe']

        # Parsing applications
        app_fields = ('size', 'state', 'last_update', 'access_token', 'checksum', 'change_number')
        while True:
            app_id = self._read_int32()

            # AppID = 0 marks the last application in the Appinfo
            if not app_id:
                break

            # All fields are required.
            app = self.wrapper((zip(app_fields, self.read_game_header())))
            if len(app) != len(app_fields):
                raise ValueError('Not all App headers are present, only found {num}: {header!r}'.format(
                    num=len(app),
                    header=app,
                ))

            # The newest VDF format is a bit simpler to parse.
            if header['version'] == 0x07564427:
                app['sections'] = self.parse_subsections()
            else:
                app['sections'] = self.wrapper()
                while True:
                    section_id = self.read_byte()
                    if not section_id:
                        break

                    # Skip the 0x00 byte before section name.
                    self.offset += 1

                    section_name = self.read_string()
                    app['sections'][section_name] = self.parse_subsections(root_section=True)

                    # New Section ID's could be added in the future, or changes could be made to
                    # existing ones, so instead of maintaining a table of section names and their
                    # corresponding IDs, we are going to store the IDs with all the data.
                    app['sections'][section_name][b'__steamfiles_section_id'] = section_id

            parsed[app_id] = app

        return parsed

    def parse_subsections(self, root_section=False):
        subsection = self.wrapper()

        while True:
            value_type = self.read_byte()
            if value_type == 0x08:
                if root_section:
                    # There's one additional 0x08 byte at the end of
                    # the root subsection.
                    self.offset += 1
                break

            key = self.read_string()
            value = self.value_parsers.get(value_type, self._unknown_value_type)()

            subsection[key] = value

        return subsection

    def make_custom_reader(self, fmt, single_value=False):
        custom_struct = struct.Struct(fmt)

        def return_many():
            result = custom_struct.unpack_from(self.data, self.offset)
            self.offset += custom_struct.size
            return result

        def return_one():
            result = custom_struct.unpack_from(self.data, self.offset)
            self.offset += custom_struct.size
            return result[0]

        if single_value:
            return return_one
        else:
            return return_many

    def read_int32(self):
        number = self._read_int32()
        return Integer(data=number, size=32)

    def read_int64(self):
        number = self._read_int64()
        return Integer(data=number, size=64)

    def read_byte(self):
        byte = self.data[self.offset]
        self.offset += 1
        return byte

    def read_string(self):
        # This method is pretty fast, provided we iterate over a memoryview.
        # It's also easier to read then the most performant ones, which is more important.
        for index, value in enumerate(self.data[self.offset:]):
            # NUL-byte – a string's end
            if value != 0:
                continue

            string = slice(self.offset, self.offset + index)
            self.offset += index + 1
            return self.data[string].tobytes()

    @staticmethod
    def _unknown_value_type():
        raise ValueError("Cannot parse the provided data type.")

class AppinfoEncoder:

    def __init__(self, data):
        self.data = data
        self.version = self.data[b'__vdf_version']

    def iter_encode(self):
        # VDF Header
        yield struct.pack('<2I', self.data[b'__vdf_version'], self.data[b'__vdf_universe'])

        for app_id, app_data in self.data.items():
            # Don't encode internal variables.
            if app_id in (b'__vdf_version', b'__vdf_universe'):
                continue

            # Game Header
            yield struct.pack('<I', app_id)
            yield struct.pack('<3IQ20sI', app_data['size'], app_data['state'],
                              app_data['last_update'], app_data['access_token'],
                              app_data['checksum'], app_data['change_number'])

            # Once again, new VDF format is much simpler.
            if self.version == 0x07564427:
                yield from self.iter_encode_section(app_data['sections'])
            else:
                for section_name, section_data in app_data['sections'].items():
                    # Delete '_section_id' from the dictionary, as it was placed there by
                    # the decoding class only to preserve the section id number.
                    section_id = section_data[b'__steamfiles_section_id']

                    yield struct.pack('<H', section_id)
                    yield self.encode_string(section_name)
                    yield from self.iter_encode_section(section_data, root_section=True)

                yield LAST_SECTION

        yield LAST_APP

    def iter_encode_section(self, section_data, root_section=False):
        for key, value in section_data.items():
            if key == b'__steamfiles_section_id':
                continue

            # Encode different types using their corresponding generators.
            if isinstance(value, dict):
                yield TYPE_SECTION
                yield self.encode_string(key)
                yield from self.iter_encode_section(value)
            elif isinstance(value, bytes):
                yield TYPE_STRING
                yield self.encode_string(key)
                yield self.encode_string(value)
            elif isinstance(value, Integer):
                yield from self.encode_integer(key, value)
            else:
                raise TypeError('Unknown value type ' + type(value).__name__)

        yield SECTION_END
        if root_section:
            # There's one additional 0x08 byte at the end of
            # the root subsection.
            yield SECTION_END

    def encode_integer(self, key, value):
        if value.size == 32:
            yield TYPE_INT32
            yield self.encode_string(key)
            yield struct.pack('<I', value.data)
        elif value.size == 64:
            yield TYPE_INT64
            yield self.encode_string(key)
            yield struct.pack('<Q', value.data)
        else:
            raise TypeError('Unknown type of an Integer')

    @staticmethod
    def encode_string(string):
        # A string with a NUL-byte at the end.
        # Example format for 'gameid': "7s".
        # The bytes packed with above format: b'gameid\x00'.
        fmt = str(len(string) + 1) + 's'
        return struct.pack(fmt, string)


import struct
from protobuf_to_dict import protobuf_to_dict, dict_to_protobuf

__all__ = ('load', 'loads', 'dump', 'dumps')

try:
    from .manifest_pb2 import Payload, Metadata, Signature
except ImportError:
    print('Run with protoc>=3.0.0: protoc --python_out=. manifest.proto')
    raise

MSG_PAYLOAD = 0x71F617D0
MSG_METADATA = 0x1F4812BE
MSG_SIGNATURE = 0x1B81B817
MSG_EOF = 0x32C415AB

MSG_NAMES = {
    MSG_PAYLOAD: 'payload',
    MSG_METADATA: 'metadata',
    MSG_SIGNATURE: 'signature',
}

MSG_IDS = {
    'payload': MSG_PAYLOAD,
    'metadata': MSG_METADATA,
    'signature': MSG_SIGNATURE,
}

MessageClass = {
    MSG_PAYLOAD: Payload,
    MSG_METADATA: Metadata,
    MSG_SIGNATURE: Signature
}

def load(fp, wrapper=dict):
    """
    Loads the contents of a Manifest file into a Python object.
    :param fp: A file object.
    :param wrapper: A wrapping object for key-value pairs.
    :return: A dictionary with Manifest data.
    """
    return loads(fp.read(), wrapper=wrapper)

def loads(data, wrapper=dict):
    """
    Loads Manifest content into a Python object.
    :param data: A byte-like object with the contents of an Appinfo file.
    :param wrapper: A wrapping object for key-value pairs.
    :return: A dictionary with Manifest data.
    """
    if not isinstance(data, (bytes, bytearray)):
        raise TypeError('can only load a bytes-like object as a Manifest but got ' + type(data).__name__)

    offset = 0
    parsed = wrapper()
    int32 = struct.Struct('<I')

    while True:
        msg_id, = int32.unpack_from(data, offset)
        offset += int32.size

        if msg_id == MSG_EOF:
            break

        msg_size, = int32.unpack_from(data, offset)
        offset += int32.size

        msg_data = data[offset:offset + msg_size]
        offset += msg_size

        message = MessageClass[msg_id]()
        message.ParseFromString(msg_data)

        parsed[MSG_NAMES[msg_id]] = wrapper(protobuf_to_dict(message))

    return parsed

def dump(obj, fp):
    """
    Serializes a dictionary into Manifest data and writes it to a file.
    :param obj: A dictionary to serialize.
    :param fp: A file object.
    """
    fp.write(dumps(obj))

def dumps(obj):
    """
    Serializes a dictionary into Manifest data.
    :param obj: A dictionary to serialize.
    :return: A file object.
    """
    if not isinstance(obj, dict):
        raise TypeError('can only dump a dictionary as a Manifest but got ' + type(obj).__name__)

    data = []
    int32 = struct.Struct('<I')

    for message_name in ('payload', 'metadata', 'signature'):
        message_data = obj[message_name]
        message_id = MSG_IDS[message_name]
        message_class = MessageClass[message_id]
        message = dict_to_protobuf(message_class, message_data)
        message_bytes = message.SerializeToString()
        message_size = len(message_bytes)

        data.append(int32.pack(message_id))
        data.append(int32.pack(message_size))
        data.append(message_bytes)

    # MSG_EOF marks the end of messages.
    data.append(int32.pack(MSG_EOF))
    return b''.join(data)


__all__ = ('load', 'loads', 'dump', 'dumps')

SECTION_START = '{'
SECTION_END = '}'

def loads(data, wrapper=dict):
    """
    Loads ACF content into a Python object.
    :param data: An UTF-8 encoded content of an ACF file.
    :param wrapper: A wrapping object for key-value pairs.
    :return: An Ordered Dictionary with ACF data.
    """
    if not isinstance(data, str):
        raise TypeError('can only load a str as an ACF but got ' + type(data).__name__)

    parsed = wrapper()
    current_section = parsed
    sections = []

    lines = (line.strip() for line in data.splitlines())

    for line in lines:
        try:
            key, value = line.split(None, 1)
            key = key.replace('"', '').lstrip()
            value = value.replace('"', '').rstrip()
        except ValueError:
            if line == SECTION_START:
                # Initialize the last added section.
                current_section = _prepare_subsection(parsed, sections, wrapper)
            elif line == SECTION_END:
                # Remove the last section from the queue.
                sections.pop()
            else:
                # Add a new section to the queue.
                sections.append(line.replace('"', ''))
            continue

        current_section[key] = value

    return parsed

def load(fp, wrapper=dict):
    """
    Loads the contents of an ACF file into a Python object.
    :param fp: A file object.
    :param wrapper: A wrapping object for key-value pairs.
    :return: An Ordered Dictionary with ACF data.
    """
    return loads(fp.read(), wrapper=wrapper)

def dumps(obj):
    """
    Serializes a dictionary into ACF data.
    :param obj: A dictionary to serialize.
    :return: ACF data.
    """
    if not isinstance(obj, dict):
        raise TypeError('can only dump a dictionary as an ACF but got ' + type(obj).__name__)

    return '\n'.join(_dumps(obj, level=0)) + '\n'

def dump(obj, fp):
    """
    Serializes a dictionary into ACF data and writes it to a file.
    :param obj: A dictionary to serialize.
    :param fp: A file object.
    """
    fp.write(dumps(obj))

def _dumps(obj, level):
    """
    Does the actual serializing of data into an ACF format.
    :param obj: A dictionary to serialize.
    :param level: Nesting level.
    :return: A List of strings.
    """
    lines = []
    indent = '\t' * level

    for key, value in obj.items():
        if isinstance(value, dict):
            # [INDENT]"KEY"
            # [INDENT]{
            line = indent + '"{}"\n'.format(key) + indent + '{'
            lines.append(line)
            # Increase intendation of the nested dict
            lines.extend(_dumps(value, level+1))
            # [INDENT]}
            lines.append(indent + '}')
        else:
            # [INDENT]"KEY"[TAB][TAB]"VALUE"
            lines.append(indent + '"{}"'.format(key) + '\t\t' + '"{}"'.format(value))

    return lines

def _prepare_subsection(data, sections, wrapper):
    """
    Creates a subsection ready to be filled.
    :param data: Semi-parsed dictionary.
    :param sections: A list of sections.
    :param wrapper: A wrapping object for key-value pairs.
    :return: A newly created subsection.
    """
    current = data
    for i in sections[:-1]:
        current = current[i]

    current[sections[-1]] = wrapper()
    return current[sections[-1]]


from __future__ import print_function

import numpy as np

from .converters import *
from ._data import SOLVENTS, CHARGES, APOLARS

def occupancy(grid, points, spacing=0.01):
    """Return a vector with the occupancy of each grid point for
    given array of points"""
    distances = ((grid[:,None,:] - points[None,:,:])**2).sum(axis=2)
    occupied = (distances < spacing).sum(axis=1)
    return occupied

def isPDBAtom(l):
    return l.startswith("ATOM") or l.startswith("HETATM")

def pdbAtom(a):
    ##01234567890123456789012345678901234567890123456789012345678901234567890123456789
    ##ATOM   2155 HH11 ARG C 203     116.140  48.800   6.280  1.00  0.00
    ## ===>   atom name,   res name,     res id, chain,       x,            y,             z
    return (a[12:16], a[17:20], int(a[22:26]), a[21], float(a[30:38])/10, float(a[38:46])/10, float(a[46:54])/10)

# Reformatting of lines in structure file
pdbBoxLine  = "CRYST1%9.3f%9.3f%9.3f%7.2f%7.2f%7.2f P 1           1"
pdbline = "ATOM  %5i %-4s %4s%1s%4i%1s   %8.3f%8.3f%8.3f%6.2f%6.2f           %1s  "

def angle(a, b):
    p = (a*b).sum()
    q = np.sqrt((a**2).sum()*(b**2).sum())
    return np.arccos(np.clip(p/q, -1.0, 1.0))*180/np.pi

def pdbBoxString(box):
    box = np.array(box)

    # Box vectors
    u, v, w  = box

    # Box vector lengths
    nu, nv, nw = np.sqrt((box**2).sum(axis=1))

    # Box vector angles
    alpha = nv*nw == 0 and 90 or angle(v, w)
    beta  = nu*nw == 0 and 90 or angle(u, w)
    gamma = nu*nv == 0 and 90 or angle(u, v)

    # Conversion from nm to A
    return pdbBoxLine % (10*nu, 10*nv, 10*nw, alpha, beta, gamma)

def groAtom(a):
    #012345678901234567890123456789012345678901234567890
    #    1PRN      N    1   4.168  11.132   5.291
    ## ===>   atom name,   res name,     res id, chain,       x,          y,          z
    return (a[10:15], a[5:10],   int(a[:5]), " ", float(a[20:28]), float(a[28:36]), float(a[36:44]))

def groBoxRead(a):
    b = [float(i) for i in a.split()] + 6*[0] # Padding for rectangular boxes
    return b[0], b[3], b[4], b[5], b[1], b[6], b[7], b[8], b[2]

class Structure(object):
    def __init__(self, filename=None, options=None):
        self.title   = ""
        self.atoms   = []
        self._coord  = None
        self.rest    = []
        self.box     = []
        self._center = None

        if filename:
            lines = open(filename).readlines()
            # Try extracting PDB atom/hetatm definitions
            self.rest   = []
            self.atoms  = [pdbAtom(i) for i in lines if isPDBAtom(i) or self.rest.append(i)]
            if self.atoms:
                # This must be a PDB file
                self.title = "THIS IS INSANE!\n"
                for i in self.rest:
                    if i.startswith("TITLE"):
                        self.title = i
                self.box   = [0, 0, 0, 0, 0, 0, 0, 0, 0]
                for i in self.rest:
                    if i.startswith("CRYST1"):
                        self.box = pdbBoxRead(i)
            else:
                # This should be a GRO file
                self.atoms = [groAtom(i) for i in lines[2:-1]]
                self.rest  = [lines[0], lines[1], lines[-1]]
                self.box   = groBoxRead(lines[-1])
                self.title = lines[0]

        if options:
            self.setup(**options)

    def __nonzero__(self):
        return bool(self.atoms)

    def __len__(self):
        return len(self.atoms)

    def __iadd__(self, s):
        if self.coord.shape[0]:
            self.coord += s ###
        return self

    def __add__(self, other):
        if hasattr(other, 'atoms') and hasattr(other, 'coord'):
            result = self.__class__()
            result.atoms.extend(self.atoms)
            result.atoms.extend(other.atoms)
            result.coord = np.concatenate((
                self.coord.reshape((-1,3)),
                other.coord.reshape((-1,3)))) ###
            return result
        raise TypeError('Cannot add {} to {}'
                        .format(self.__class__, other.__class__))

    def __iter__(self):
        atom_enumeration = enumerate(zip(self.atoms, self.coord), start=1)
        for idx, (atom, (x, y, z)) in atom_enumeration:
            atname, resname, resid = atom[:3]
            if resname.endswith('.o'):
                resname = resname[:-2]
            yield idx, atname, resname, resid, x, y, z

    @property
    def coord(self):
        if self._coord is None:
            self._coord = np.array([i[4:7] for i in self.atoms]).reshape((-1,3))
        return self._coord

    @coord.setter
    def coord(self, other):
        self._coord = np.array(other).reshape((-1,3))

    @property
    def charge(self):
        last = None
        charge = 0
        for j in self.atoms:
            if not j[0].strip().startswith('v') and j[1:3] != last:
                charge += CHARGES.get(j[1].strip(), 0)
            last = j[1:3]
        return charge

    @property
    def center(self):
        if self._center is None:
            self._center = self.coord.mean(axis=0)
        return self._center

    @center.setter
    def center(self, other):
        s = other - self.coord.mean(axis=0)
        self.coord += s ###

    def diam(self):
        if np.any(self._center):
            self.center = (0, 0, 0)
        return 2*np.sqrt(max([i*i+j*j+k*k for i, j, k in self.coord]))

    def diamxy(self):
        if np.any(self._center):
            self.center = (0, 0, 0)
        return 2*np.sqrt(max([i*i+j*j for i, j, k in self.coord]))

    def areaxy(self, lowerbound=-np.inf, upperbound=np.inf, spacing=0.1):
        mask = (self.coord[:,2] > lowerbound) & (self.coord[:,2] < upperbound)
        if not mask.sum():
            # No cross section with membrane
            return 0
        points = self.coord[mask, :2]
        # The magic number factor 1.1 is not critical at all
        # Just a number to set a margin to the bounding box and
        # have all points fall within the boundaries
        bbmin, bbmax = 1.1*points.min(axis=0), 1.1*points.max(axis=0)
        size = bbmax - bbmin
        cells = (size / spacing + 0.5).astype('int')
        # Grid points over bounding box with specified spacing
        grid = np.mgrid[bbmin[0]:bbmax[0]:(cells[0]*1j),
                        bbmin[1]:bbmax[1]:(cells[1]*1j)].reshape((2,-1)).T
        # Occupied cells is approximately equal to grid points within
        # gridspacing distance of points
        occupied = occupancy(grid, points, spacing)
        # The occupied area follows from the fraction of occupied
        # cells times the area spanned by the bounding box
        return size[0]*size[1]*sum(occupied > 0)/occupied.size

    def fun(self, fn):
        return [fn(i) for i in zip(*self.coord)]

    def orient(self, d, pw):
        # Determine grid size
        m = self.coord.min(axis=0)
        r = self.coord.max(axis=0)-self.coord.min(axis=0)+1e-8

        # Number of grid cells
        n = (r/d + 0.5).astype('int')
        nx, ny, nz = n

        # Initialize grids
        atom = np.zeros(n+2)
        phobic = np.zeros(n+2)
        binned = (n * (self.coord - m) / r).astype('int')
        notdummy = np.array([ i[1] != "DUM" for i in self.atoms ])
        apolar = np.array([ i[1].strip() in APOLARS for i in self.atoms ])
        for i,j,k in binned[notdummy]:
            atom[i,j,k] += 1
        for i,j,k in binned[apolar]:
            phobic[i,j,k] += 1
        #atom = np.histogramdd(binned[notdummy], n+2)[0]

        with np.errstate(divide='ignore', invalid='ignore'):
            ratio = (2 * phobic / atom) ** pw

        # Determine average density
        occupd = atom.astype('bool').sum()
        avdens = float(atom.sum())/occupd
        threshold = 0.1*avdens
        above = atom > threshold
        surface  = []
        #cgofile  = open('density.cgo', "w")
        #cgofile.write('[\n')
        for i,j,k in zip(*np.where(atom > threshold)):
            # Check the neighbouring cells;
            # if one of them is not occupied, count cell as surface
            if not (atom[i-1,j,k] and atom[i+1,j,k] and
                    atom[i,j-1,k] and atom[i,j+1,k] and
                    atom[i,j,k-1] and atom[i,j,k+1]):
                sx, sy, sz = m + (r*(i,j,k)+0.5*r)/n
                sw = ratio[i,j,k]
                surface.append((sx, sy, sz, sw))
                #cgofile.write("    7.0, %f, %f, %f, %f, \n"%(10*sx, 10*sy, 10*sz, 0.25*sw))
        #cgofile.write(']\n')
        #cgofile.close()

        surface = np.array(surface)
        sx, sy, sz, w = zip(*surface)
        W             = 1.0/sum(w)

        # Weighted center of apolar region; has to go to (0, 0, 0)
        sxm, sym, szm = np.average(surface[:,:-1], axis=0, weights=surface[:,-1])
        apolar_center = np.average(surface[:,:-1], axis=0, weights=surface[:,-1])

        # Place apolar center at origin
        self.center = -apolar_center

        # Determine weighted deviations from centers
        dev = (surface[:,:-1] - apolar_center) * surface[:, -1, None]

        # Covariance matrix for surface
        val, vec = np.linalg.eig(np.dot(dev.T, dev))
        vec = vec[:,val.argsort()[::-1]]

        # Rotate the coordinates
        self.coord = np.dot(self.coord, vec)

    def rotate(self, what):
            if what == "princ":
                self.rotate_princ()
            ## ii. Randomly
            elif what == "random":
                self.rotate_random()
            ## iii. Specifically
            elif what:
                self.rotate_degrees(float(what))

    def rotate_princ(self):
        R = np.linalg.eig(np.dot(self.coord[:,:2].T,self.coord[:,:2]))
        self.coord[:,:2] = np.dot(self.coord[:,:2], R[1][:,np.argsort(R[0])[::-1]])
        return

    def rotate_random(self):
        ux   = np.cos(random.random()*2*np.pi)
        uy   = np.sqrt(1-ux*ux)
        self.coord[:,:2] = np.dot(self.coord[:,:2],[[ux,-uy],[uy,ux]])

    def rotate_degrees(self, angle):
        ux   = np.cos(angle*np.pi/180.)
        uy   = np.sin(angle*np.pi/180.)
        self.coord[:,:2] = np.dot(self.coord[:,:2], [[ux, -uy],[uy, ux]])

    def setup(self, **kwargs):
        # Center the protein and store the shift
        shift = self.center
        self.center = (0, 0, 0)

        ## 1. Orient with respect to membrane
        # Orient the protein according to the TM region, if requested
        # This doesn't actually work very well...
        if kwargs["orient"]:
            self.orient(kwargs["origriddist"], kwargs["oripower"])

        ## 4. Orient the protein in the xy-plane
        ## i. According to principal axes and unit cell
        self.rotate(kwargs["rotate"])

        ## 5. Determine the minimum and maximum x and y of the protein
        pmin, pmax = self.coord.min(axis=0), self.coord.max(axis=0)

        # At this point we should shift the subsequent proteins such
        # that they end up at the specified distance, in case we have
        # a number of them to do
        # y-shift is always -ycenter
        # x-shift is -xmin+distance+xmax(current)
        # xshifts.append(xshifts[-1]+pmax[0]+(options["distance"] or 0))

        ## 2. Shift of protein relative to the membrane center
        zshift = kwargs["memshift"]
        if not kwargs["center"]:
            zshift += shift[2]

        # Now we center the system in the rectangular
        # brick corresponding to the unit cell
        # If -center is given, also center z in plane
        self += (0, 0, zshift)

        # The z position is now correct with respect to the membrane
        # at z = 0. The x/y need to be set still

def write_gro(outfile, title, atoms, box):
    """
    Write a GRO file.

    Parameters
    ----------
    outfile
        The stream to write in.
    title
        The title of the GRO file. Must be a single line.
    atoms
        An instance of Structure containing the atoms to write.
    box
        The periodic box as a 3x3 matrix.
    """
    # Print the title
    print(title, file=outfile)

    # Print the number of atoms
    print("{:5d}".format(len(atoms)), file=outfile)

    # Print the atoms
    atom_template = "{:5d}{:<5s}{:>5s}{:5d}{:8.3f}{:8.3f}{:8.3f}"
    for idx, atname, resname, resid, x, y, z in atoms:
        print(atom_template
              .format(int(resid % 1e5), resname, atname, int(idx % 1e5),
                      x, y, z),
              file=outfile)

    # Print the box
    grobox = (box[0][0], box[1][1], box[2][2],
              box[0][1], box[0][2], box[1][0],
              box[1][2], box[2][0], box[2][1])
    box_template = '{:10.5f}' * 9
    print(box_template.format(*grobox), file=outfile)

def write_pdb(outfile, title, atoms, box):
    """
    Write a PDB file.

    Parameters
    ----------
    outfile
        The stream to write in.
    title
        The title of the GRO file. Must be a single line.
    atoms
        An instance of Structure containing the atoms to write.
    box
        The periodic box as a 3x3 matrix.
    """
    # Print the title
    print('TITLE ' + title, file=outfile)

    # Print the box
    print(pdbBoxString(box), file=outfile)

    # Print the atoms
    for idx, atname, resname, resid, x, y, z in atoms:
        print(pdbline % (idx % 1e5, atname[:4], resname[:3], "",
                         resid % 1e4, '', 10*x, 10*y, 10*z, 0, 0, ''),
              file=outfile)

def write_structure(output, title, atoms, box):
    oStream = output and open(output, "w") or sys.stdout
    # The title is always truncated to 80 characters to avoid
    # Gromacs (or other programs) choking on them.
    with oStream:
        if output.endswith(".gro"):
            write_gro(oStream, title[:80], atoms, box.tolist())
        else:
            write_pdb(oStream, title[:80], atoms, box.tolist())


#!/usr/bin/env python

# INSert membrANE
# A simple, versatile tool for building coarse-grained simulation systems
# Copyright (C) 2017  Tsjerk A. Wassenaar and contributors
#
# 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 2
# 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, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

"""
INSANE: A versatile tool for building membranes and/or solvent with proteins.

... Someone ought to write a more extensive docstring here...
"""

from __future__ import print_function

import os
import sys
import random
import collections
import numpy as np

from simopt import opt_func

from . import lipids
from .pbc import PBC
from .structure import *
from .converters import *
from .constants import d2r
from ._data import SOLVENTS, CHARGES, APOLARS
from .options import OPTIONS

try:
    range = xrange
except NameError:
    pass

RTOL = 1e-8

# Set the random seed.
# The seed is set to an arbitary value set in the INSANE_SEED environment
# variable. If the environment variable is not set, then the system time is
# used to set the seed.
random.seed(os.environ.get('INSANE_SEED', None))

def _point(y, phi):
    r = np.sqrt(1-y*y)
    return np.cos(phi)*r, y, np.sin(phi)*r

def pointsOnSphere(n):
    return np.array([_point((2.*k+1)/n-1, k*2.3999632297286531) for k in range(n)])

def determine_molecule_numbers(total, molecules, absolute, relative):
    """Determine molecule numbers for given total,
    absolute and relative numbers"""
    weight = sum(relative)
    if not any(absolute):
        # Only relative numbers
        numbers = [int(total*i/weight) for i in relative]
    elif any(relative):
        # Absolute numbers and fill the rest with relative numbers
        rest = total - sum(absolute)
        numbers = [int(rest*i/weight) if i else j
                   for i,j in zip(relative, absolute)]
    else:
        # Only absolute numbers
        numbers = absolute
    return list(zip(molecules, numbers))

def resize_pbc_for_lipids(pbc, relL, relU, absL, absU,
                          uparea, area, hole, proteins):
    """
    Adapt the size of the box to accomodate the lipids.

    The PBC is changed **in place**.
    """
    if any(relL) and any(relU):
        # Determine box from size
        # If one leaflet is defined with an absolute number of lipids
        # then the other leaflet (with relative numbers) will follow
        # from that.
        # box/d/x/y/z needed to define unit cell

        # box is set up already...
        # yet it may be set to 0, then there is nothing we can do.
        if 0 in (pbc.x, pbc.y, pbc.z):
            raise PBCException('Not enough information to set the box size.')
    elif any(absL) or any(absU):
        # All numbers are absolute.. determine size from number of lipids
        # Box x/y will be set, d/dz/z is needed to set third box vector.
        # The area is needed to determine the size of the x/y plane OR
        # the area will be SET if a box is given.

        if pbc.z == 0:
            raise PBCException('Not enough information to set the box size.')

        if 0 in (pbc.x, pbc.y):
            # We do not know what size the box should be.
            # Let X and Y be the same.
            #T This does not agree with the default pbc being hexagonal...
            pbc.x = pbc.y = 1
        # A scaling factor is needed for the box
        # This is the area for the given number of lipids
        upsize = sum(absU) * uparea
        losize = sum(absL) * area
        # This is the size of the hole, going through both leaflets
        holesize = np.pi * hole ** 2
        # This is the area of the PBC xy plane
        xysize = pbc.x * pbc.y
        # This is the total area of the proteins per leaflet (IMPLEMENT!)
        psize_up = sum([p.areaxy(0, 2.4) for p in proteins])
        psize_lo = sum([p.areaxy(-2.4, 0) for p in proteins])
        # So this is unavailable:
        unavail_up = holesize + psize_up
        unavail_lo = holesize + psize_lo
        # This is the current area marked for lipids
        # xysize_up = xysize - unavail_up
        # xysize_lo = xysize - unavail_lo
        # This is how much the unit cell xy needs to be scaled
        # to accomodate the fixed amount of lipids with given area.
        upscale = (upsize + unavail_up)/xysize
        loscale = (losize + unavail_lo)/xysize
        area_scale = max(upscale, loscale)
        aspect_ratio = pbc.x / pbc.y
        scale_x = np.sqrt(area_scale / aspect_ratio)
        scale_y = np.sqrt(area_scale / aspect_ratio)
        pbc.box[:2,:] *= math.sqrt(area_scale)

def setup_solvent(pbc, protein, membrane, options):
    # Charge of the system so far

    if not options["solvent"]:
        return Structure(), []

    solv = options["solvent"]
    charge = membrane.charge + protein.charge

    # Set up a grid
    d        = 1/options["soldiam"]

    nx, ny, nz = int(1+d*pbc.x), int(1+d*pbc.y), int(1+d*pbc.z)
    dx, dy, dz = pbc.x/nx, pbc.y/ny, pbc.z/nz

    # Exclusion range for solvent
    excl = int(nz*options["solexcl"]/pbc.z)

    # Layer index at middle of box
    hz = int(0.5*nz)

    zshift   = 0
    if membrane:
        memz   = [i[2] for i in membrane.coord]
        midz   = (max(memz)+min(memz))/2
        hz     = int(nz*midz/pbc.z)  # Grid layer in which the membrane is located
        zshift = (hz+0.5)*nz - midz # Shift of membrane middle to center of grid layer

    # Initialize a grid of solvent, spanning the whole cell
    # Exclude all cells within specified distance from membrane center
    grid   = [[[i < hz-excl or i > hz+excl for i in range(nz)] for j in range(ny)] for i in range(nx)]

    # Flag all cells occupied by protein or membrane
    for coord in (protein+membrane).coord:
        for x, y, z in coord + options["lipradius"] * pointsOnSphere(options["lipdensity"]):
            if z >= pbc.z:
                x -= pbc.box[2,0]
                y -= pbc.box[2,1]
                z -= pbc.box[2,2]
            if z < 0:
                x += pbc.box[2,0]
                y += pbc.box[2,1]
                z += pbc.box[2,2]
            if y >= pbc.y:
                x -= pbc.box[1,0]
                y -= pbc.box[1,1]
            if y < 0:
                x += pbc.box[1,0]
                y += pbc.box[1,1]
            if x >= pbc.x:
                x -= pbc.box[0,0]
            if x < 0:
                x += pbc.box[0,0]
            grid[int(nx*x/pbc.rx)][int(ny*y/pbc.ry)][int(nz*z/pbc.rz)] = False

    ##-T grid should be a wrapper around a numpy.ndarray
    # Set the center for each solvent molecule
    kick = options["solrandom"]
    grid = [ (random.random(), (i+0.5+random.random()*kick)*dx, (j+0.5+random.random()*kick)*dy, (k+0.5+random.random()*kick)*dz)
             for i in range(nx) for j in range(ny) for k in range(nz) if grid[i][j][k] ]

    # Sort on the random number
    grid.sort()

    # 'grid' contains all positions on which a solvent molecule can be placed.
    # The number of positions is taken as the basis for determining the salt concentration.
    # This is fine for simple salt solutions, but may not be optimal for complex mixtures
    # (like when mixing a 1M solution of this with a 1M solution of that

    # First get names and relative numbers for each solvent
    solnames, solabs, solnums = zip(*solv)
    solnames, solnums = list(solnames), list(solnums)
    totS       = float(sum(solnums))

    # Set the number of ions to add
    nna, ncl = 0, 0
    if options["salt"]:

        # If the concentration is set negative, set the charge to zero
        if options["salt"].startswith("-"):
            charge = 0
            options["salt"] = -float(options["salt"])
        else:
            options["salt"] = float(options["salt"])

        # Determine charge to use, either determined or given on command line
        if options["charge"] != "0":
            charge = (options["charge"] != "auto") and int(options["charge"]) or charge
        else:
            charge = 0

        # Determine number of sodium and chloride to add
        concentration = options["salt"]
        nsol = ("SPC" in solnames and 1 or 4)*len(grid)
        ncl  = max(max(0, charge), int(.5+.5*(concentration*nsol/(27.7+concentration)+charge)))
        nna  = ncl - charge

    # Correct number of grid cells for placement of solvent
    ngrid   = len(grid) - nna - ncl
    num_sol = [int(ngrid*i/totS) for i in solnums]

    # Add salt to solnames and num_sol
    if nna:
        solnames.append("NA+")
        num_sol.append(nna)
        solv.append("NA+")
    if ncl:
        solnames.append("CL-")
        num_sol.append(ncl)
        solv.append("CL-")

    # Names and grid positions for solvent molecules
    solvent    = list(zip([s for i, s in zip(num_sol, solnames) for j in range(i)], grid))

    # Build the solvent
    resi = 0
    if protein:
        resi = protein.atoms[-1][2]
    if membrane:
        resi = membrane.atoms[-1][2]
    sol = Structure()
    solcoord = []
    for resn, (rndm, x, y, z) in solvent:
        resi += 1
        solmol = SOLVENTS.get(resn)
        if solmol and len(solmol) > 1:
            # Random rotation (quaternion)
            u,  v,  w       = random.random(), 2*np.pi*random.random(), 2*np.pi*random.random()
            s,  t           = np.sqrt(1-u), np.sqrt(u)
            qw, qx, qy, qz  = s*np.sin(v), s*np.cos(v), t*np.sin(w), t*np.cos(w)
            qq              = qw*qw-qx*qx-qy*qy-qz*qz
            for atnm, (px, py, pz) in solmol:
                qp = 2*(qx*px + qy*py + qz*pz)
                rx = x + qp*qx + qq*px + qw*(qy*pz-qz*py)
                ry = y + qp*qy + qq*py + qw*(qz*px-qx*pz)
                rz = z + qp*qz + qq*pz + qw*(qx*py-qy*px)
                sol.atoms.append((atnm, resn, resi, 0, 0, 0))
                solcoord.append((rx, ry, rz))
        else:
            sol.atoms.append((solmol and solmol[0][0] or resn,
                              resn, resi,
                              0, 0, 0))
            solcoord.append((x, y, z))
    sol.coord = solcoord

    return sol, zip(solnames, num_sol)

def setup_membrane(pbc, protein, lipid, options):
    membrane = Structure()
    molecules = []

    lower, upper = lipid
    lipL, absL, relL = lower
    lipU, absU, relU = upper
    nabsL, nabsU = sum(absL), sum(absU)

    if not any((absL, relL, absU, relU)):
        return membrane, molecules

    lo_lipd = np.sqrt(options["area"])
    up_lipd = np.sqrt(options["uparea"])

    #print('lipd', lo_lipd, up_lipd)

    num_up, num_lo = [], []

    # Lipids are added on grid positions, using the prototypes defined above.
    # If a grid position is already occupied by protein, the position is untagged.

    lipd = lo_lipd

    # Number of lipids in x and y in lower leaflet if there were no solute
    lo_lipids_x = int(pbc.x/lo_lipd+0.5)
    lo_lipids_y = int(pbc.y/lo_lipd+0.5)
    q = False
    while lo_lipids_x*lo_lipids_y < nabsL:
        if q:
            lo_lipids_x += 1
        else:
            lo_lipids_y += 1
        q = not q

    lo_lipdx    = pbc.x/lo_lipids_x
    lo_rlipx    = list(range(lo_lipids_x))
    lo_lipdy    = pbc.y/lo_lipids_y
    lo_rlipy    = list(range(lo_lipids_y))

    if options["uparea"]:
        lipd = up_lipd

    # Number of lipids in x and y in upper leaflet if there were no solute
    up_lipids_x = int(pbc.x/up_lipd+0.5)
    up_lipids_y = int(pbc.y/up_lipd+0.5)
    while up_lipids_x*up_lipids_y < nabsU:
        if q:
            up_lipids_x += 1
        else:
            up_lipids_y += 1
        q = not q

    up_lipdx    = pbc.x/up_lipids_x
    up_rlipx    = list(range(up_lipids_x))
    up_lipdy    = pbc.y/up_lipids_y
    up_rlipy    = list(range(up_lipids_y))

    # Set up grids to check where to place the lipids
    grid_lo = [[0 for j in lo_rlipy] for i in lo_rlipx]
    grid_up = [[0 for j in up_rlipy] for i in up_rlipx]

    ## NEW GRID
    lo_gx = slice(0, pbc.x, int(pbc.x/lo_lipd + 0.5)*1j)
    lo_gy = slice(0, pbc.x, int(pbc.x/lo_lipd + 0.5)*1j)
    up_gx = slice(0, pbc.x, int(pbc.x/up_lipd + 0.5)*1j)
    up_gy = slice(0, pbc.x, int(pbc.x/up_lipd + 0.5)*1j)
    grid_l = np.mgrid[lo_gx, lo_gy][:,:-1,:-1].reshape((2,-1)).T
    grid_u = np.mgrid[up_gx, up_gy][:,:-1,:-1].reshape((2,-1)).T
    # If there is a protein, mark the corresponding POINTS as occupied
    occupied_lo = np.zeros(grid_l.shape[0])
    occupied_up = np.zeros(grid_u.shape[0])
    if protein:
        upmask = (protein.coord[:,2] <  2.4) & (protein.coord[:,2] > 0)
        lomask = (protein.coord[:,2] > -2.4) & (protein.coord[:,2] < 0)
        occupied_lo += occupancy(grid_l, protein.coord[lomask,:2], lo_lipd)
        occupied_up += occupancy(grid_u, protein.coord[upmask,:2], up_lipd)
        maxd = max(occupied_lo.max(), occupied_up.max())
        if maxd:
            occupied_up = (occupied_up/maxd) > options["fudge"]
            occupied_lo = (occupied_lo/maxd) > options["fudge"]

    #print(absU, sum(~occupied_up.astype('bool')),
    #      absL, sum(~occupied_lo.astype('bool')))
    #print('->', len(grid_lo), len(grid_lo[0]))
    #print('->', len(grid_up), len(grid_up[0]))

    ## OLD GRID

    maxd = 1

    # If there is a protein, mark the corresponding cells as occupied
    if protein:
        # Extract the parts of the protein that are in either leaflet
        sphere = options["protradius"] * pointsOnSphere(options["protdensity"])

        # Calculate number density per cell
        mem_mask_lo = (0 > protein.coord[:,2]) & (protein.coord[:,2] > -2.4)
        prot_lo = protein.coord[mem_mask_lo, :]
        for i in prot_lo:
            for x, y, z in i + sphere:
                grid_lo[ int(lo_lipids_x*x/pbc.rx)%lo_lipids_x ][ int(lo_lipids_y*y/pbc.ry)%lo_lipids_y ] += 1

        mem_mask_up = (0 < protein.coord[:,2]) & (protein.coord[:,2] < 2.4)
        prot_up = protein.coord[mem_mask_up, :]
        for i in prot_up:
            for x, y, z in i + sphere:
                grid_up[ int(up_lipids_x*x/pbc.rx)%up_lipids_x ][ int(up_lipids_y*y/pbc.ry)%up_lipids_y ] += 1

        # Determine which cells to consider occupied, given the fudge factor
        # The array is changed to boolean type here
        maxd = float(max([max(i) for i in grid_up+grid_lo]))
        if  maxd == 0:
            print("; The protein seems not to be inside the membrane.", file=sys.stderr)
            print("; Run with -orient to put it in.", file=sys.stderr)
            maxd = 1

    fudge   = options["fudge"]
    grid_up = [[(j/maxd) <= fudge for j in i] for i in grid_up]
    grid_lo = [[(j/maxd) <= fudge for j in i] for i in grid_lo]

    # If we don't want lipids inside of the protein
    # we also mark everything from the center up to the first cell filled
    if not options["inside"]:

        # Upper leaflet
        marked = [(i, j) for i in up_rlipx for j in up_rlipy if not grid_up[i][j]]
        if marked:
            # Find the center
            cx, cy  = [float(sum(i))/len(marked) for i in zip(*marked)]
            for i, j in marked:
                md = int(abs(i-cx)+abs(j-cy)) # Manhattan length/distance
                for f in range(md):
                    ii = int(cx+f*(i-cx)/md)
                    jj = int(cy+f*(j-cy)/md)
                    grid_up[ii][jj] = False

        # Lower leaflet
        marked = [(i, j) for i in lo_rlipx for j in lo_rlipy if not grid_lo[i][j]]
        if marked:
            # Find the center
            cx, cy  = [float(sum(i))/len(marked) for i in zip(*marked)]
            for i, j in marked:
                md = int(abs(i-cx)+abs(j-cy)) # Manhattan length
                for f in range(md):
                    ii = int(cx+f*(i-cx)/md)
                    jj = int(cy+f*(j-cy)/md)
                    grid_lo[ii][jj] = False

    # If we make a circular patch, we flag the cells further from the
    # protein or box center than the given radius as occupied.
    if options["disc"]:
        if protein:
            cx, cy = protein.center[:2]
        else:
            cx, cy = 0.5*pbc.x, 0.5*pbc.y
        for i in range(len(grid_lo)):
            for j in range(len(grid_lo[i])):
                if (i*pbc.x/lo_lipids_x - cx)**2 + (j*pbc.y/lo_lipids_y - cy)**2 > options["disc"]**2:
                    grid_lo[i][j] = False
        for i in range(len(grid_up)):
            for j in range(len(grid_up[i])):
                if (i*pbc.x/up_lipids_x - cx)**2 + (j*pbc.y/up_lipids_y - cy)**2 > options["disc"]**2:
                    grid_up[i][j] = False

    # If we need to add a hole, we simply flag the corresponding cells
    # as occupied. The position of the hole depends on the type of PBC,
    # to ensure an optimal arrangement of holes around the protein. If
    # there is no protein, the hole is just put in the center.
    if options["hole"]:
        # Lower leaflet
        if protein:
            if ("square".startswith(options["pbc"]) or
                "rectangular".startswith(options["pbc"])):
                hx, hy = (0, 0)
            else:
                hx, hy = (0, int(lo_lipids_y*np.cos(np.pi/6)/9+0.5))
        else:
            hx, hy = (int(0.5*lo_lipids_x), int(0.5*lo_lipids_y))
        hr = int(options["hole"]/min(lo_lipdx,  lo_lipdy)+0.5)
        ys = int(lo_lipids_x*pbc.box[1,0]/pbc.box[0,0]+0.5)
        print("; Making a hole with radius %f nm centered at grid cell (%d,%d)"%(options["hole"], hx, hy), hr, file=sys.stderr)
        hr -= 1
        for ii in range(hx-hr-1, hx+hr+1):
            for jj in range(hx-hr-1, hx+hr+1):
                xi, yj = ii, jj
                if (ii-hx)**2+(jj-hy)**2 < hr**2:
                    if jj < 0:
                        xi += ys
                        yj += lo_lipids_y
                    if jj >= lo_lipids_y:
                        xi -= ys
                        yj -= lo_lipids_y
                    if xi < 0:
                        xi += lo_lipids_x
                    if xi >= lo_lipids_x:
                        xi -= lo_lipids_x
                    grid_lo[xi][yj] = False
                    grid_up[xi][yj] = False
        # Upper leaflet
        if protein:
            if ("square".startswith(options["pbc"]) or
                "rectangular".startswith(options["pbc"])):
                hx, hy = (0, 0)
            else:
                hx, hy = (0, int(up_lipids_y*np.cos(np.pi/6)/9+0.5))
        else:
            hx, hy = (int(0.5*up_lipids_x), int(0.5*up_lipids_y))
        hr = int(options["hole"]/min(up_lipdx, up_lipdy)+0.5)
        ys = int(up_lipids_x*pbc.box[1,0]/pbc.box[0,0]+0.5)
        print("; Making a hole with radius %f nm centered at grid cell (%d,%d)"%(options["hole"], hx, hy), hr, file=sys.stderr)
        hr -= 1
        for ii in range(hx-hr-1, hx+hr+1):
            for jj in range(hx-hr-1, hx+hr+1):
                xi, yj = ii, jj
                if (ii-hx)**2+(jj-hy)**2 < hr**2:
                    if jj < 0:
                        xi += ys
                        yj += up_lipids_y
                    if jj >= up_lipids_y:
                        xi -= ys
                        yj -= up_lipids_y
                    if xi < 0:
                        xi += up_lipids_x
                    if xi >= up_lipids_x:
                        xi -= up_lipids_x
                    grid_up[xi][yj] = False

    # Set the XY coordinates
    # To randomize the lipids we add a random number which is used for sorting
    upper, lower = [], []
    for i in range(up_lipids_x):
        for j in range(up_lipids_y):
            if grid_up[i][j]:
                upper.append((random.random(), i*pbc.x/up_lipids_x, j*pbc.y/up_lipids_y))
    for i in range(lo_lipids_x):
        for j in range(lo_lipids_y):
            if grid_lo[i][j]:
                lower.append((random.random(), i*pbc.x/lo_lipids_x, j*pbc.y/lo_lipids_y))

    # Sort on the random number
    upper.sort()
    lower.sort()

    # Extract coordinates, taking asymmetry in account
    asym  = options["asymmetry"] or 0
    upper = [i[1:] for i in upper[max(0, asym):]]
    lower = [i[1:] for i in lower[max(0, -asym):]]

    print("; X: %.3f (%d bins) Y: %.3f (%d bins) in upper leaflet"%(pbc.x, up_lipids_x, pbc.y, up_lipids_y), file=sys.stderr)
    print("; X: %.3f (%d bins) Y: %.3f (%d bins) in lower leaflet"%(pbc.x, lo_lipids_x, pbc.y, lo_lipids_y), file=sys.stderr)
    print("; %d lipids in upper leaflet, %d lipids in lower leaflet"%(len(upper), len(lower)), file=sys.stderr)

    ##> Types of lipids, relative numbers, fractions and numbers

    # Upper leaflet (+1)
    numbers = determine_molecule_numbers(len(upper), lipU, absU, relU)
    lip_up     = [l for l, i in numbers for j in range(i)]
    leaf_up    = ( 1, zip(lip_up, upper), up_lipd, up_lipdx, up_lipdy)
    molecules.extend(numbers)

    # Lower leaflet (-1)
    numbers = determine_molecule_numbers(len(lower), lipL, absL, relL)
    lip_lo = [l for l, i in numbers for j in range(i)]
    leaf_lo = (-1, zip(lip_lo, lower), lo_lipd, lo_lipdx, lo_lipdy)
    molecules.extend(numbers)

    ##< Done determining numbers per lipid

    ##> Building lipids

    kick       = options["randkick"]

    mematoms = []
    memcoords = []

    # Build the membrane

    ## ==> LIPID  BOOKKEEPING:
    # Read lipids defined in insane
    liplist = lipids.get_lipids()
    # Then add lipids from file
    liplist.add_from_files(options["molfile"])
    # Last, add lipids from command line
    liplist.add_from_def(options["lipnames"], options["lipheads"], options["liplinks"],
                         options["liptails"], options["lipcharge"])

    if protein:
        resi = protein.atoms[-1][2]
    else:
        resi = 0

    inshift = options["indist"] / 2
    for leaflet, leaf_lip, lipd, lipdx, lipdy in [leaf_up, leaf_lo]:
        for lipid, pos in leaf_lip:
            # Increase the residue number by one
            resi += 1

            # Fetch the atom list with x, y, z coordinates
            at, ax, ay, az = zip(*liplist[lipid].build(diam=lipd))
            # The z-coordinates are spaced at 0.3 nm,
            # starting with the first bead at 0.15 nm
            az = [ leaflet*(inshift + (i-min(az)))*options["beaddist"] for i in az ]
            xx = np.array((ax, ay)).T

            # Set the random rotation for this lipid
            rangle   = 2*random.random()*np.pi
            if options["norotate"]:
                nx = pos[0] + lipdx/2 + [ kick*random.random() for i in az ]
                ny = pos[1] + lipdy/2 + [ kick*random.random() for i in az ]
            else:
                rcos     = np.cos(rangle)
                rsin     = np.sin(rangle)
                rcosx    = rcos*lipdx*2/3
                rcosy    = rcos*lipdy*2/3
                rsinx    = rsin*lipdx*2/3
                rsiny    = rsin*lipdy*2/3
                nx = np.dot(xx,(rcosx, -rsiny)) + pos[0] + lipdx/2 + [ kick*random.random() for i in az ]
                ny = np.dot(xx,(rsinx, rcosy)) + pos[1] + lipdy/2 + [ kick*random.random() for i in az ]

            # Add the atoms to the list
            memcoords.extend([(nx[i], ny[i], az[i]) for i in range(len(at))])
            mematoms.extend([(at[i], lipid, resi, 0, 0, 0) for i in range(len(at))])

    ##< Done building lipids

    membrane.coord = memcoords
    membrane.atoms = mematoms

    return membrane, molecules

@opt_func(OPTIONS)
def old_main(**options):

    molecules = []

    #########################################
    ## I. PROTEIN and other macromolecules ##
    #########################################

    # Read in the structures (if any)
    tm = [ Structure(i, options) for i in options["solute"] ]

    if tm:
        molecules.append(('Protein', len(tm)))

    xshifts  = [0] # Shift in x direction per protein

    #############
    ## II. PBC ##
    #############

    # Periodic boundary conditions
    if options["pbc"] == 'keep' and tm:
        box = tm[0].box
    else:
        box = options.get("box")

    zdist = options["zdistance"]
    if zdist == None:
        zdist = options["distance"]

    # Set up base PBC
    # Override where needed to accomodate additional components
    # box/shape are final - if these are given and a solute does
    # not fit in it raises an exception
    pbc = PBC(shape=options["pbc"], box=box,
              distance=(options["distance"], zdist),
              xyz=(options["xvector"], options["yvector"], options["zvector"]),
              disc=options["disc"], hole=options["hole"],
              membrane=options["lower"], protein=tm)

    #################
    ## III. LIPIDS ##
    #################

    lipL      = options["lower"]
    lipU      = options["upper"]
    relU, relL, absU, absL = [], [], [], []
    if lipL:
        lipU = lipU or lipL
        lipL, absL, relL = zip(*lipL)
        totL       = float(sum(relL))
        lipU, absU, relU = zip(*lipU)
        totU       = float(sum(relU))
    else:
        options["solexcl"] = -1

    if options["uparea"] is None:
        options["uparea"] = options["area"]

    #print(pbc.box)
    resize_pbc_for_lipids(pbc=pbc, relL=relL, relU=relU, absL=absL, absU=absU,
                          uparea=options["uparea"], area=options["area"],
                          hole=options["hole"], proteins=tm)
    #print(pbc.box)

    ##################
    ## IV. MEMBRANE ##
    ##################

    ## 1. Proteins

    # Now that PBC is set, we can shift the proteins
    for xshft, prot in zip(xshifts,tm):
        # Half the distance should be added to xshft
        # to center the whole lot.
        #xshft += options["distance"]/2
        xshft = pbc.x/2
        prot.coord += (xshft, pbc.y/2, (not lipL)*pbc.z/2)

    prot_coord = []
    protein  = Structure()
    for prot in tm:
        # And we collect the atoms
        protein.atoms.extend(prot.atoms)
        prot_coord.append(prot.coord)
    if tm:
        protein.coord = np.concatenate(prot_coord)

    # Current residue ID is set to that of the last atom
    resi = 0
    if protein:
        resi = protein.atoms[-1][2]

    ## 2. Lipids

    lipid = ((lipL, absL, relL), (lipU, absU, relU))
    membrane, added = setup_membrane(pbc, protein, lipid, options)
    molecules.extend(added)

    if added:
        # Now move everything to the center of the box before adding solvent
        mz  = pbc.z/2
        z   = (protein+membrane).coord[:,2]
        mz -= (max(z)+min(z))/2
        protein.coord += (0, 0, mz)
        membrane.coord += (0, 0, mz)

    if membrane:
        resi = membrane.atoms[-1][2]

    ################
    ## 3. SOLVENT ##
    ################

    solvent, added = setup_solvent(pbc, protein, membrane, options)
    molecules.extend(added)

    return (molecules, protein, membrane, solvent, lipid, pbc.box)

def write_summary(protein, membrane, solvent):
    charge  = protein.charge + membrane.charge
    plen = len(protein)
    print("; NDX Solute %d %d" % (1, protein and plen or 0), file=sys.stderr)
    print("; Charge of protein: %f" % protein.charge, file=sys.stderr)

    mlen = len(membrane)
    print("; NDX Membrane %d %d" % (1 + plen, membrane and plen + mlen or 0),
          file=sys.stderr)
    print("; Charge of membrane: %f" % membrane.charge, file=sys.stderr)
    print("; Total charge: %f" % charge, file=sys.stderr)

    slen = len(solvent)
    print("; NDX Solvent %d %d" % (1+plen+mlen, solvent and plen+mlen+slen or 0), file=sys.stderr)
    print("; NDX System %d %d" % (1, plen+mlen+slen), file=sys.stderr)
    print("; \"I mean, the good stuff is just INSANE\" --Julia Ormond",
          file=sys.stderr)

def write_top(outpath, molecules, title):
    """
    Write a basic TOP file.

    The topology is written in *outpath*. If *outpath* is en empty string, or
    anything for which ``bool(outpath) == False``, the topology is written on
    the standard error, and the header is omitted, and only what has been buit
    by Insane id displayed (e.g. Proteins are excluded).

    Parameters
    ----------
    outpath
        The path to the file to write. If empty, a simplify topology is
        written on stderr.
    molecules
        List of molecules with the number of them.
    title
        Title of the system.
    """
    topmolecules = []
    for i in molecules:
        if i[0].endswith('.o'):
            topmolecules.append(tuple([i[0][:-2]]+list(i[1:])))
        else:
            topmolecules.append(i)

    if outpath:
        # Write a rudimentary topology file
        with open(outpath, "w") as top:
            print('#include "martini.itp"\n', file=top)
            print('[ system ]', file=top)
            print('; name', file=top)
            print(title, file=top)
            print('\n', file=top)
            print('[ molecules ]', file=top)
            print('; name  number', file=top)
            print("\n".join("%-10s %7d"%i for i in topmolecules), file=top)
    else:
        # Here we only include molecules that have beed added by insane.
        # This is usually concatenated at the end of an existint top file.
        # As the existing file usually contain the proteins already, we do not
        # include them here.
        added_molecules = (molecule for molecule in topmolecules
                           if molecule[0] != 'Protein')
        print("\n".join("%-10s %7d"%i for i in added_molecules), file=sys.stderr)

def system_title(membrane, protein, lipids):
    (lipL, absL, relL), (lipU, absU, relU) = lipids
    if membrane.atoms:
        title  = "INSANE! Membrane UpperLeaflet>"+":".join(lipU)+"="+":".join([str(i) for i in relU])
        title += " LowerLeaflet>"+":".join(lipL)+"="+":".join([str(i) for i in relL])

        if protein:
            title = "Protein in " + title
    else:
        title = "Insanely solvated protein."
    return title

def insane(**options):

    ## PROTEINS
    # Collect and build input structures (proteins)
    # Determine protein/hole layout
    # Collect and build lipid/solvent dictionaries

    ## PBC
    # Determine size of system

    # MEMBRANE
    # Make grid - two leaflets
    # Mask grid occupied cells
    # Determine lipid compositions per leaflet
    # Build membrane

    ## SOLVENT
    # Determine solvent composition
    # Build solvent

    return 0


# INSert membrANE
# A simple, versatile tool for building coarse-grained simulation systems
# Copyright (C) 2017  Tsjerk A. Wassenaar and contributors
#
# 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 2
# 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, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

"""
Utility functions
"""

import pkg_resources

def iter_resource(filename):
    """
    Return a stream for a given resource file in the module.

    The resource file has to be part of the module and its filenane given
    relative to the module.
    """
    with pkg_resources.resource_stream(__name__, filename) as resource:
        for line in resource:
            yield line.decode('utf-8')


# INSert membrANE
# A simple, versatile tool for building coarse-grained simulation systems
# Copyright (C) 2017  Tsjerk A. Wassenaar and contributors
#
# 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 2
# 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, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

import collections
import math
import os

from . import utils

__all__ = ['Lipid', 'Lipid_List', 'get_lipids']

# Lipid data file
LIPID_FILE = 'lipids.dat'

# Define supported lipid head beads. One letter name mapped to atom name
HEADBEADS = {
"C":  "NC3", # NC3 = Choline
    "E":  "NH3", # NH3 = Ethanolamine
    "G":  "GL0", # GL0 = Glycerol
    "S":  "CNO", # CNO = Serine
    "P":  "PO4", # PO4 = Phosphate
    }

# Define supported lipid link beads. One letter name mapped to atom name
LINKBEADS = {
    "G":  "GL",  # Glycerol
    "A":  "AM",  # Amide (Ceramide/Sphingomyelin)
}

class Lipid:
    """Lipid structure"""

    def __init__(self, **kwargs):
        self.name      = kwargs.get("name")
        self.head      = kwargs.get("head")
        self.link      = kwargs.get("link")
        self.tail      = kwargs.get("tail")
        self.beads     = kwargs.get("beads")
        if type(self.beads) == str:
            self.beads = self.beads.split()
        self.charge    = kwargs.get("charge")
        self.template  = kwargs.get("template")
        self.area      = kwargs.get("area")
        self.diam      = kwargs.get("diam", math.sqrt(kwargs.get("area", 0)))
        self.coords    = None
        if kwargs.get("string"):
            self.parse(kwargs["string"])

    def parse(self, string):
        """
        Parse lipid definition from string:

            alhead=C P, allink=A A, altail=TCC CCCC, alname=DPSM, charge=0.0
        """
        fields = [i.split("=") for i in string.split(', ')]
        for what, val in fields:
            what = what.strip()
            val  = val.split()
            if what.endswith("head"):
                self.head = val
            elif what.endswith("link"):
                self.link = val
            elif what.endswith("tail"):
                self.tail = val
            elif what == "charge":
                self.charge = float(val[0])
            elif what.endswith("name") and not self.name:
                self.name = val[0]
        if self.charge is None:
            # Infer charge from head groups
            self.charge = sum([headgroup_charges[bead] for bead in self.head])

    def build(self, **kwargs):
        """Build/return a list of [(bead, x, y, z), ...]"""

        if not self.coords:
            if self.beads and self.template:
                stuff = zip(self.beads, self.template)
                self.coords = [[i, x, y, z] for i, (x, y, z) in stuff if i != "-"]
            else:
                # Set beads/structure from head/link/tail
                # Set bead names
                if self.beads:
                    beads = list(self.beads)
                else:
                    beads = [HEADBEADS[i] for i in self.head]
                    beads.extend([LINKBEADS[n] + str(i + 1)
                                  for i, n in enumerate(self.link)])
                    for i, t in enumerate(self.tail):
                        beads.extend([n + chr(65 + i) + str(j + 1)
                                      for j, n in enumerate(t)])

                taillength = max([0]+[len(i) for i in self.tail])
                length = len(self.head)+taillength

                # Add the pseudocoordinates for the head
                rl     = range(len(self.head))
                struc  = [(0, 0, length-i) for i in rl]

                # Add the linkers
                rl     = range(len(self.link))
                struc.extend([(i%2, i//2, taillength) for i in rl ])

                # Add the tails
                for j, tail in enumerate(self.tail):
                    rl = range(len(tail))
                    struc.extend([(j%2, j//2, taillength-1-i) for i in rl])

                mx, my, mz = [ (max(i)+min(i))/2 for i in zip(*struc) ]
                self.coords = [
                    [i, 0.25 * (x - mx), 0.25 * (y - my), z]
                    for i, (x, y, z) in zip(beads, struc)
                ]

        # Scale the x/y based on the lipid's APL - diameter is less than sqrt(APL)
        diam   = kwargs.get("diam", self.diam)
        radius = diam*0.45
        minmax = [ (min(i), max(i)) for i in list(zip(*self.coords))[1:] ]
        mx, my, mz = [ sum(i)/2. for i in minmax ]
        scale  = radius/math.sqrt((minmax[0][0]-mx)**2 + (minmax[1][0]-my)**2)

        for i in self.coords:
            i[1] = scale*(i[1]-mx)
            i[2] = scale*(i[2]-my)
            i[3] -= minmax[2][0]

        return self.coords

    def h(self, head):
        self.head = head.replace(".", " ").split()

    def l(self, link):
        self.link = link.replace(".", " ").split()

    def t(self, tail):
        self.tail = tail.replace(".", " ").split()

    def c(self, charge):
        self.charge = float(charge)

class Lipid_List(collections.MutableMapping):
    """Container class for lipid definitions"""

    def __init__(self):
        self.store = dict()
        self.last  = None

    def __getitem__(self, key):
        if key == -1:
            return self.last
        return self.store[key]

    def __setitem__(self, key, value):
        self.store[key] = value

    def __delitem__(self, key):
        del self.store[key]

    def __iter__(self):
        return iter(self.store)

    def __len__(self):
        return len(self.store)

    def add(self, name=None, string=None):
        lip = Lipid(name=name, string=string)
        self.store[lip.name] = lip
        self.last = lip.name

    def add_from_stream(self, stream):
        current_lipid = None
        for line in stream:
            if line[0] == ';':
                line = line[1:]
            if line.startswith("@INSANE"):
                current_lipid = Lipid(string=line[len('@INSANE'):].strip())
                self[current_lipid.name] = current_lipid
            elif line.startswith("@BEADS"):
                beads = line.split()[1:]
                current_lipid.beads = beads

    def add_from_file(self, path):
        with open(path) as infile:
            self.add_from_stream(infile)

    def add_from_files(self, multi_path):
        for path in multi_path:
            self.add_from_file(path)

    def add_from_def(self, usrnames, usrheads, usrlinks, usrtails, usrcharges):
        if not usrcharges:
            usrcharges = [None] * len(usrnames)
        lipids = zip(usrnames, usrheads, usrlinks, usrtails, usrcharges)
        for name, head, link, tail, charge in lipids:
            heads = head.replace(".", " ").split()
            linkers = link.replace(".", " ").split()
            tails = tail.replace(".", " ").split()
            charge = None if charge is None else float(charge)
            self[name] = Lipid(name=name,
                               head=heads,
                               link=linkers,
                               tail=tails,
                               charge=charge)

def read_lipids(lipfile):
    lipids = Lipid_List()
    x, y, z = None, None, None
    for line in lipfile:
        stripped = line.strip()
        if (not stripped) or stripped.startswith((';','#')):
            # Comment or empty line
            continue
        elif stripped.startswith('['):
            # Moleculetype tag
            x, y, z = None, None, None
            # moltype = stripped[2:stripped.find(']')].strip()
        elif x is None:
            x = [float(i) for i in line.split() ]
        elif y is None:
            y = [float(i) for i in line.split() ]
        elif z is None:
            z = [float(i) for i in line.split() ]
        else:
            splitted = line.split()
            name = splitted.pop(0)
            lipids[name] = Lipid(
                name=name,
                beads=splitted,
                template=zip(x, y, z)
            )
    return lipids

def get_lipids():
    lipid_stream = utils.iter_resource(LIPID_FILE)
    lipids = read_lipids(lipid_stream)
    return lipids


# INSert membrANE
# A simple, versatile tool for building coarse-grained simulation systems
# Copyright (C) 2017  Tsjerk A. Wassenaar and contributors
#
# 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 2
# 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, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

import math

from .constants import d2r

# These functions are not to stay here... -TAW
def vector(v):
    if type(v) == str and "," in v:
        return [float(i) for i in v.split(",")]
    return float(v)

def pdbBoxRead(a):
    # Convert a PDB CRYST1 entry to a lattice definition.
    # Convert from Angstrom to nanometer
    fa, fb, fc, aa, ab, ac = [float(i) for i in a.split()[1:7]]
    ca, cb, cg, sg         = math.cos(d2r*aa), math.cos(d2r*ab), math.cos(d2r*ac) , math.sin(d2r*ac)
    wx, wy                 = 0.1*fc*cb, 0.1*fc*(ca-cb*cg)/sg
    wz                     = math.sqrt(0.01*fc*fc - wx*wx - wy*wy)
    return [0.1*fa, 0, 0, 0.1*fb*cg, 0.1*fb*sg, 0, wx, wy, wz]

def box3d(a):
    x = [ float(i) for i in a.split(",") ] + 6*[0]
    if len(x) == 12: # PDB format
        return pdbBoxRead("CRYST1 "+" ".join([str(i) for i in x]))
    elif len(x) == 7:
        return x[0], x[0], x[0], 0, 0, 0, 0, 0, 0
    else:            # GRO format
        return x[0], x[3], x[4], x[5], x[1], x[6], x[7], x[8], x[2]

def molspec(x):
    """
    Parse a string for a lipid or a solvent as given on the command line
    (MOLECULE[=NUMBER|:NUMBER]); where `=NUMBER` sets an absolute number of the
    molecule, and `:NUMBER` sets a relative number of it.
    If both absolute and relative number are set False, then relative count is 1.
    """
    lip = x.split(":")
    abn = lip[0].split("=")
    names = abn[0]
    if len(abn) > 1:
        nrel = 0
        nabs = int(abn[1])
    else:
        nabs = 0
        if len(lip) > 1:
            nrel = float(lip[1])
        else:
            nrel = 1
    return abn[0], nabs, nrel


from django.core.cache import cache
from django.contrib.messages import constants

def message_user(user, message, level=constants.INFO):
    """
    Send a message to a particular user.

    :param user: User instance
    :param message: Message to show
    :param level: Message level
    """
    # We store a list of messages in the cache so we can have multiple messages
    # queued up for a user.
    user_key = _user_key(user)
    messages = cache.get(user_key) or []
    messages.append((message, level))
    cache.set(user_key, messages)

def message_users(users, message, level=constants.INFO):
    """
    Send a message to a group of users.

    :param users: Users queryset
    :param message: Message to show
    :param level: Message level
    """
    for user in users:
        message_user(user, message, level)

def get_messages(user):
    """
    Fetch messages for given user.  Returns None if no such message exists.

    :param user: User instance
    """
    key = _user_key(user)
    result = cache.get(key)
    if result:
        cache.delete(key)
        return result
    return None

def _user_key(user):
    return '_async_message_%d' % user.pk


from django.contrib import messages

from async_messages import get_messages

class AsyncMiddleware(object):

    def process_response(self, request, response):
        """
        Check for messages for this user and, if it exists,
        call the messages API with it
        """
        if hasattr(request, "session") and hasattr(request, "user") and request.user.is_authenticated():
            msgs = get_messages(request.user)
            if msgs:
                for msg, level in msgs:
                    messages.add_message(request, level, msg)
        return response