# Copyright (c) 2013 Spotify AB
#
# 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 datetime

import core
import copy
from core import ParseError, Field, auto_store, PySchema
import binascii
try:
    from collections import OrderedDict
except ImportError:
    from ordereddict import OrderedDict

def ordereddict_push_front(dct, key, value):
    """Set a value at the front of an OrderedDict

    The original dict isn't modified, instead a copy is returned
    """
    d = OrderedDict()
    d[key] = value
    d.update(dct)
    return d

class Text(Field):
    def load(self, obj):
        if not isinstance(obj, (unicode, type(None))):
            raise ParseError("%r not a unicode object" % obj)
        return obj

    def dump(self, obj):
        if isinstance(obj, (unicode, type(None))):
            return obj
        else:
            try:
                return obj.decode('utf8')
            except:
                raise ValueError(
                    "%r is not a valid UTF-8 string" % obj
                )

class Bytes(Field):
    """Binary data"""

    def __init__(self, custom_encoding=False, **kwargs):
        super(Bytes, self).__init__(**kwargs)
        self.custom_encoding = custom_encoding

    def _load_utf8_codepoints(self, obj):
        return obj.encode("iso-8859-1")

    def _dump_utf8_codepoints(self, binary_data):
        return binary_data.decode("iso-8859-1")

    def _load_b64(self, obj):
        return binascii.a2b_base64(obj.encode("ascii"))

    def _dump_b64(self, binary_data):
        return binascii.b2a_base64(binary_data).rstrip('\n')

    def load(self, obj):
        if not self.custom_encoding:
            return self._load_utf8_codepoints(obj)
        return self._load_b64(obj)

    def dump(self, binary_data):
        if isinstance(binary_data, unicode):
            raise ValueError(
                "Unicode objects are not accepted values for Bytes (%r)"
                % (binary_data,)
            )
        if not self.custom_encoding:
            return self._dump_utf8_codepoints(binary_data)
        return self._dump_b64(binary_data)

    def is_similar_to(self, other):
        return super(Bytes, self).is_similar_to(other) and self.custom_encoding == other.custom_encoding

class List(Field):
    """List of one other Field type

    Differs from other fields in that it is not nullable
    and defaults to empty array instead of null
    """
    def __init__(self, field_type=Text(), nullable=False, default=[], **kwargs):
        super(List, self).__init__(nullable=nullable, default=default, **kwargs)
        self.field_type = field_type

    def load(self, obj):
        if not isinstance(obj, list):
            raise ParseError("%r is not a list object" % obj)
        return [self.field_type.load(o) for o in obj]

    def dump(self, obj):
        if not isinstance(obj, (tuple, list)):
            raise ValueError("%r is not a list object" % obj)
        return [self.field_type.dump(o) for o in obj]

    def set_parent(self, schema):
        self.field_type.set_parent(schema)

    def default_value(self):
        #  avoid default-sharing between records
        return copy.deepcopy(self.default)

    def is_similar_to(self, other):
        return super(List, self).is_similar_to(other) and self.field_type.is_similar_to(other.field_type)

    def repr_vars(self):
        return ordereddict_push_front(
            super(List, self).repr_vars(),
            "field_type",
            repr(self.field_type)
        )

class Enum(Field):
    _field_type = Text()  # don't change

    def __init__(self, values, name=None, **kwargs):
        super(Enum, self).__init__(**kwargs)
        self.values = set(values)
        self.name = name

        if name is not None and PySchema.auto_register:
            auto_store.add_enum(self)

    def dump(self, obj):
        if obj not in self.values:
            raise ValueError(
                "%r is not an allowed value of Enum%r"
                % (obj, tuple(self.values)))
        return self._field_type.dump(obj)

    def load(self, obj):
        parsed = self._field_type.load(obj)
        if parsed not in self.values and parsed is not None:
            raise ParseError(
                "Parsed value %r not in allowed value of Enum(%r)"
                % (parsed, tuple(self.values)))
        return parsed

    def is_similar_to(self, other):
        return super(Enum, self).is_similar_to(other) and self.values == other.values

    def repr_vars(self):
        return OrderedDict([
            ("values", self.values),
            ("name", repr(self.name))
        ] + super(Enum, self).repr_vars().items()
        )

class Integer(Field):
    def __init__(self, size=8, **kwargs):
        super(Integer, self).__init__(**kwargs)
        self.size = size

    def dump(self, obj):
        if not isinstance(obj, (int, long, type(None))) or isinstance(obj, bool):
            raise ValueError("%r is not a valid Integer" % (obj,))
        return obj

    def load(self, obj):
        if not isinstance(obj, (int, long, type(None))) or isinstance(obj, bool):
            raise ParseError("%r is not a valid Integer" % (obj,))
        return obj

    def is_similar_to(self, other):
        return super(Integer, self).is_similar_to(other) and self.size == other.size

    def repr_vars(self):
        return ordereddict_push_front(
            super(Integer, self).repr_vars(),
            "size",
            self.size
        )

class Boolean(Field):
    VALUE_MAP = {True: '1', 1: '1',
                 False: '0', 0: '0'}

    def dump(self, obj):
        if obj not in self.VALUE_MAP:
            raise ValueError(
                "Invalid value for Boolean field: %r" % obj)
        return bool(obj)

    def load(self, obj):
        if obj not in self.VALUE_MAP:
            raise ParseError(
                "Invalid value for Boolean field: %r" % obj)
        return bool(obj)

class Float(Field):
    def __init__(self, size=8, **kwargs):
        super(Float, self).__init__(**kwargs)
        self.size = size

    def dump(self, obj):
        if not isinstance(obj, float):
            raise ValueError("Invalid value for Float field: %r" % obj)
        return float(obj)

    def load(self, obj):
        if not isinstance(obj, (float, int, long)):
            raise ParseError("Invalid value for Float field: %r" % obj)
        return float(obj)

    def is_similar_to(self, other):
        return super(Float, self).is_similar_to(other) and self.size == other.size

class Date(Text):
    def dump(self, obj):
        if not isinstance(obj, datetime.date):
            raise ValueError("Invalid value for Date field: %r" % obj)
        return str(obj)

    def load(self, obj):
        try:
            # This is much faster than calling strptime
            (year, month, day) = obj.split('-')
            return datetime.date(int(year), int(month), int(day))
        except ValueError:
            raise ValueError("Invalid value for Date field: %r" % obj)

class DateTime(Text):
    def dump(self, obj):
        if not isinstance(obj, datetime.datetime):
            raise ValueError("Invalid value for DateTime field: %r" % obj)
        return str(obj)

    def load(self, obj):
        try:
            if '.' in obj:
                return datetime.datetime.strptime(obj, "%Y-%m-%d %H:%M:%S.%f")
            return datetime.datetime.strptime(obj, "%Y-%m-%d %H:%M:%S")
        except ValueError:
            raise ValueError("Invalid value for DateField field: %r" % obj)

# special value for SubRecord's schema parameter
# that signifies a SubRecord accepting records of the
# same type as the container/parent Record.
SELF = object()

class SubRecord(Field):
    """"Field for storing other :class:`record.Record`s"""

    def __init__(self, schema, **kwargs):
        super(SubRecord, self).__init__(**kwargs)
        self._schema = schema

    def dump(self, obj):
        if not isinstance(obj, self._schema):
            raise ValueError("%r is not a %r"
                             % (obj, self._schema))
        return core.to_json_compatible(obj)

    def load(self, obj):
        return core.from_json_compatible(self._schema, obj)

    def set_parent(self, schema):
        """This method gets called by the metaclass
        once the container class has been created
        to let the field store a reference to its
        parent if needed. Its needed for SubRecords
        in case it refers to the container record.
        """
        if self._schema == SELF:
            self._schema = schema

    def default_value(self):
        #  avoid default-sharing between records
        return copy.deepcopy(self.default)

    def is_similar_to(self, other):
        return super(SubRecord, self).is_similar_to(other) and self._schema == other._schema

    def repr_vars(self):
        return ordereddict_push_front(
            super(SubRecord, self).repr_vars(),
            "schema",
            self._schema._schema_name
        )

class Map(Field):
    """List of one other Field type

    Differs from other fields in that it is not nullable
    and defaults to empty array instead of null
    """
    def __init__(self, value_type, nullable=False, default={}, **kwargs):
        super(Map, self).__init__(nullable=nullable, default=default, **kwargs)
        self.value_type = value_type
        self.key_type = Text()

    def load(self, obj):
        return dict([
            (self.key_type.load(k),
             self.value_type.load(v))
            for k, v in obj.iteritems()
        ])

    def dump(self, obj):
        if not isinstance(obj, dict):
            raise ValueError("%r is not a dict" % (obj,))

        return dict([
            (self.key_type.dump(k),
             self.value_type.dump(v))
            for k, v in obj.iteritems()
        ])

    def set_parent(self, schema):
        self.value_type.set_parent(schema)

    def default_value(self):
        #  avoid default-sharing between records
        return copy.deepcopy(self.default)

    def is_similar_to(self, other):
        return super(Map, self).is_similar_to(other) and self.value_type.is_similar_to(other.value_type)

    def repr_vars(self):
        return ordereddict_push_front(
            super(Map, self).repr_vars(),
            "value_type",
            repr(self.value_type)
        )


# -*- coding: utf-8 -*-
from .utils import escape_filter

OPERATORS = {u'=', u'!=', u'<', u'<=', u'>=', u'>'}

def gen_filter(name, op, value, is_or=False):
    """Generates a single filter expression for ``filter[]``."""
    if op not in OPERATORS:
        raise ValueError('Unknown operator {}'.format(op))
    result = u'{} {} {}'.format(name, op, escape_filter(value))
    if is_or:
        result = u'or ' + result
    return result

class Q(object):
    """A Django-like query composition class.

    Note:
        Only supports simple chaining using ``|`` and ``&`` (as ``filter[]`` supports)

    Args:
        name: Name of the value to be compared
        op: An operator
        value: Value to be matched.

    .. code-block:: python

        api.collections.vms.filter(Q('name', '=', 'foo') | Q('name', '=', 'bar'))

    Et cetera ... You can use ``|`` and ``&`` operators and they only work in chain (no parentheses)
    because that is what filter supports.

    """
    @classmethod
    def from_dict(cls, d):
        """Creates a query (AND and =) from a dictionary."""
        if not d:
            raise ValueError('Empty dictionary!')
        items = list(d.items())
        key, value = items.pop(0)
        q = cls(key, u'=', value)
        for key, value in items:
            q = q & cls(key, u'=', value)
        return q

    def __init__(self, name, op, value):
        self.name = name
        if op not in OPERATORS:
            raise ValueError('Invalid operator {}'.format(op))
        self.op = op
        self.value = value
        self.preceeding = None
        self.is_or = False

    def _set_preceeding(self, preceeding):
        self.preceeding = preceeding

    @property
    def as_filters(self):
        if self.preceeding is not None:
            filters = self.preceeding.as_filters
        else:
            filters = []

        filters.append(gen_filter(self.name, self.op, self.value, self.is_or))

        return filters

    def __or__(self, other_q):
        other_q._set_preceeding(self)
        other_q.is_or = True
        return other_q

    def __and__(self, other_q):
        other_q._set_preceeding(self)
        return other_q


# -*- coding: utf-8 -*-
# TODO: WIP WIP WIP WIP!
# I got another stage of this library aside, but this is perfectly usable with some restrictions :)
import json
import logging
import re
import requests
import simplejson
import six
import iso8601
from copy import copy
from distutils.version import LooseVersion
from functools import partial
from wait_for import wait_for

from .filters import Q

class APIException(Exception):
    pass

class ManageIQClient(object):
    def __init__(self, entry_point, auth, logger=None, verify_ssl=True, ca_bundle_path=None):
        """ If ca_bundle_path is specified it replaces the system's trusted root CAs"""
        self._entry_point = entry_point
        self._auth = auth
        self._verify_ssl = verify_ssl
        self._ca_bundle_path = ca_bundle_path
        self._session = requests.Session()
        self._session.headers.update({'Content-Type': 'application/json; charset=utf-8'})
        self._build_auth(auth)
        if not verify_ssl:
            self._session.verify = False
        elif ca_bundle_path:
            self._session.verify = ca_bundle_path
        self.logger = logger or logging.getLogger(__name__)
        self.response = None
        self._load_data()

    def _build_auth(self, auth):
        valid = False
        if isinstance(auth, dict):
            if set(("user", "password")) <= set(auth):
                self._session.auth = (auth["user"], auth["password"])
                valid = True
            if "token" in auth:
                self._session.headers.update({'X-Auth-Token': auth['token']})
                valid = True
        elif isinstance(auth, (tuple, list)):  # for backward compatibility
            self._session.auth = tuple(auth[:2])
            valid = True

        if not valid:
            raise ValueError("Unknown values provided for auth")

    def _load_data(self):
        data = self.get(self._entry_point)
        self.collections = CollectionsIndex(self, data.pop("collections", []))
        self._version = data.pop("version", None)
        self._versions = {}
        for version in data.pop("versions", []):
            self._versions[version["name"]] = version["href"]
        for key, value in data.items():
            setattr(self, key, value)

    @property
    def version(self):
        return self._version

    def _result_processor(self, result):
        # Save last full response
        self.response = result

        result_json = None
        try:
            result_json = result.json()
        except simplejson.scanner.JSONDecodeError:
            result_text = result.text.strip()
            # HTTP methods other than GET and OPTIONS are allowed to return empty result
            if result.request.method in ("GET", "OPTIONS") or result_text:
                raise APIException("JSONDecodeError: {}".format(result_text or "empty result"))

        if result_json and "error" in result_json:
            if isinstance(result_json["error"], dict):
                raise APIException(
                    "{}: {}".format(result_json["error"]["klass"], result_json["error"]["message"]))
            else:
                raise APIException(
                    "{}: {}".format(result_json.get("status", None), result_json["error"]))

        # Check HTTP response status
        if not result:
            raise APIException("The request failed with HTTP status {}: {}".format(
                result.status_code, result.reason))

        return result_json

    def _sending_request(self, func, retries=2):
        while retries:
            try:
                return func()
            except requests.ConnectionError as e:
                last_connection_exception = e
                retries -= 1
        raise last_connection_exception

    def get(self, api_endpoint_url=None, **get_params):
        if not api_endpoint_url:
            if "url" not in get_params:
                raise ValueError(
                    "url must be specified, as positional parameter, "
                    "or (deprecated) 'url' keyword parameter"
                )
            api_endpoint_url = get_params.pop("url")

        self.logger.info("[RESTAPI] GET %s %r", api_endpoint_url, get_params)
        data = self._sending_request(
            partial(self._session.get, api_endpoint_url, params=get_params))
        return self._result_processor(data)

    def post(self, api_endpoint_url=None, **payload):
        if not api_endpoint_url:
            if "url" not in payload:
                raise ValueError(
                    "url must be specified, as positional parameter, "
                    "or (deprecated) 'url' keyword parameter"
                )
            api_endpoint_url = payload.pop("url")

        self.logger.info("[RESTAPI] POST %s %r", api_endpoint_url, payload)
        data = self._sending_request(
            partial(self._session.post, api_endpoint_url, data=json.dumps(payload)))
        self.logger.info("[RESTAPI] RESPONSE %s", data)
        return self._result_processor(data)

    def put(self, api_endpoint_url=None, **payload):
        if not api_endpoint_url:
            if "url" not in payload:
                raise ValueError(
                    "url must be specified, as positional parameter, "
                    "or (deprecated) 'url' keyword parameter"
                )
            api_endpoint_url = payload.pop("url")

        self.logger.info("[RESTAPI] PUT %s %r", api_endpoint_url, payload)
        data = self._sending_request(
            partial(self._session.put, api_endpoint_url, data=json.dumps(payload)))
        self.logger.info("[RESTAPI] RESPONSE %s", data)
        return self._result_processor(data)

    def patch(self, url, *payload):
        self.logger.info("[RESTAPI] PATCH %s %r", url, payload)
        data = self._sending_request(
            partial(self._session.patch, url, data=json.dumps(payload)))
        self.logger.info("[RESTAPI] RESPONSE %s", data)
        return self._result_processor(data)

    def delete(self, api_endpoint_url=None, **payload):
        if not api_endpoint_url:
            if "url" not in payload:
                raise ValueError(
                    "url must be specified, as positional parameter, "
                    "or (deprecated) 'url' keyword parameter"
                )
            api_endpoint_url = payload.pop("url")

        self.logger.info("[RESTAPI] DELETE %s %r", api_endpoint_url, payload)
        data = self._sending_request(
            partial(self._session.delete, api_endpoint_url, data=json.dumps(payload)))
        self.logger.info("[RESTAPI] RESPONSE %s", data)
        return self._result_processor(data)

    def options(self, api_endpoint_url=None, **opt_params):
        if not api_endpoint_url:
            if "url" not in opt_params:
                raise ValueError(
                    "url must be specified, as positional parameter, "
                    "or (deprecated) 'url' keyword parameter"
                )
            api_endpoint_url = opt_params.pop("url")

        self.logger.info("[RESTAPI] OPTIONS %s %r", api_endpoint_url, opt_params)
        data = self._sending_request(
            partial(self._session.options, api_endpoint_url, params=opt_params))
        return self._result_processor(data)

    def get_entity(self, collection_or_name, entity_id, attributes=None):
        if not isinstance(collection_or_name, Collection):
            collection = Collection(
                self, "{}/{}".format(self._entry_point, collection_or_name), collection_or_name)
        else:
            collection = collection_or_name
        entity = Entity(
            collection,
            {"href": "{}/{}".format(collection._href, entity_id)},
            attributes=attributes)
        return entity

    def api_version(self, version):
        return type(self)(
            self._versions[version],
            self._auth,
            logger=self.logger,
            verify_ssl=self._verify_ssl,
            ca_bundle_path=self._ca_bundle_path)

    @property
    def versions(self):
        return sorted(self._versions.keys(), reverse=True, key=LooseVersion)

    @property
    def latest_version(self):
        return self.versions[0]

    @property
    def on_latest_version(self):
        return self.version == self.latest_version

class CollectionsIndex(object):
    def __init__(self, api, data):
        self._api = api
        self._data = data
        self._collections = []
        self._load_data()

    def _load_data(self):
        for collection in self._data:
            c = Collection(
                self._api, collection["href"], collection["name"], collection["description"])
            setattr(self, collection["name"], c)
            self._collections.append(c)

    @property
    def all(self):
        return self._collections

    @property
    def all_names(self):
        return [c.name for c in self.all]

    def __contains__(self, collection):
        if isinstance(collection, six.string_types):
            return collection in self.all_names
        else:
            return collection in self.all

class SearchResult(object):
    def __init__(self, collection, data):
        self.collection = collection
        self.count = data.pop("count", 0)
        self.subcount = data.pop("subcount", 0)
        self.name = data.pop("name")
        self.resources = []
        for resource in data["resources"]:
            self.resources.append(Entity(collection, resource))

    def __iter__(self):
        for resource in self.resources:
            resource.reload()
            yield resource

    def __getitem__(self, position):
        entity = self.resources[position]
        entity.reload()
        return entity

    def __len__(self):
        return self.subcount

    def __repr__(self):
        return "<SearchResult for {!r}>".format(self.collection)

class Collection(object):
    def __init__(self, api, href, name, description=None):
        self._api = api
        self._href = href
        self._data = None
        self._subcollections = None
        self._subcollections_loaded = False
        self.action = ActionContainer(self)
        self.name = name
        self.description = description

    @property
    def api(self):
        return self._api

    @property
    def subcollections(self):
        # it's enought to try to load the subcollections list once
        if not self._subcollections_loaded:
            self._subcollections_loaded = True
            try:
                opts = self._api.options(self._href)
                self._subcollections = opts['subcollections']
            except Exception:
                # OPTIONS are supported only in new versions of API and subcollections are
                # returned only in even newer versions. Many things can go wrong here,
                # hence this broad except.
                self._subcollections = None
                self._api.logger.warning(
                    "[RESTAPI] failed to get subcollections list for %s", self._href)
        return self._subcollections

    def reload(self, expand=False, attributes=None):
        if expand is True:
            kwargs = {"expand": "resources"}
        elif expand:
            kwargs = {"expand": expand}
        else:
            kwargs = {}
        if attributes is not None:
            if isinstance(attributes, six.string_types):
                attributes = [attributes]
            kwargs.update(attributes=",".join(attributes))
        self._data = self._api.get(self._href, **kwargs)
        self._resources = self._data["resources"]
        self._count = self._data.get("count", 0)
        self._subcount = self._data.get("subcount", 0)
        self._actions = self._data.pop("actions", [])
        if self._data["name"] != self.name:
            raise ValueError("Name mishap!")

    def reload_if_needed(self):
        if self._data is None:
            self.reload()

    def query_string(self, **params):
        """Specify query string to use with the collection.

        Returns: :py:class:`SearchResult`
        """
        return SearchResult(self, self._api.get(self._href, **params))

    def raw_filter(self, filters):
        """Sends all filters to the API.

        No fancy, just a wrapper. Any advanced functionality shall be implemented as another method.

        Args:
            filters: List of filters (strings)

        Returns: :py:class:`SearchResult`
        """
        return SearchResult(self, self._api.get(self._href, **{"filter[]": filters}))

    def filter(self, q):
        """Access the ``filter[]`` functionality of ManageIQ.

        Args:
            q: An instance of :py:class:`filters.Q`

        Returns: :py:class:`SearchResult`
        """
        return self.raw_filter(q.as_filters)

    def find_by(self, **params):
        """Searches in ManageIQ using the ``filter[]`` get parameter.

        This method only supports logical AND so all key/value pairs are considered as equality
        comparision and all are logically anded.
        """
        return self.filter(Q.from_dict(params))

    def get(self, **params):
        try:
            return self.find_by(**params)[0]
        except IndexError:
            raise ValueError("No such '{}' matching query {!r}!".format(self.name, params))

    def options(self, **params):
        return self._api.options(self._href, **params)

    @property
    def count(self):
        self.reload_if_needed()
        return self._count

    @property
    def subcount(self):
        self.reload_if_needed()
        return self._subcount

    @property
    def all(self):
        self.reload(expand=True)
        return [Entity(self, r) for r in self._resources]

    def all_include_attributes(self, attributes):
        """Returns all entities present in the collection with ``attributes`` included."""
        self.reload(expand=True, attributes=attributes)
        entities = [Entity(self, r, attributes=attributes) for r in self._resources]
        self.reload()
        return entities

    def __repr__(self):
        return "<Collection {!r} ({!r})>".format(self.name, self.description)

    def __call__(self, entity_id, attributes=None):
        return self._api.get_entity(self, entity_id, attributes=attributes)

    def __iter__(self):
        self.reload(expand=True)
        for resource in self._resources:
            yield Entity(self, resource)

    def __getitem__(self, position):
        self.reload_if_needed()
        entity = Entity(self, self._resources[position])
        entity.reload()
        return entity

    def __len__(self):
        return self.subcount

class Entity(object):
    # TODO: Extend these fields
    TIME_FIELDS = {
        "updated_on", "created_on", "last_scan_attempt_on", "state_changed_on", "lastlogon",
        "updated_at", "created_at", "last_scan_on", "last_sync_on", "last_refresh_date",
        "retires_on"}
    COLLECTION_MAPPING = dict(
        ems_id="providers",
        storage_id="data_stores",
        zone_id="zones",
        host_id="hosts",
        current_group_id="groups",
        miq_user_role_id="roles",
        evm_owner_id="users",
        task_id="tasks",
    )
    EXTENDED_COLLECTIONS = dict(
        roles={"features"},
    )

    def __init__(self, collection, data, incomplete=False, attributes=None):
        self.collection = collection
        self.action = ActionContainer(self)
        self._data = data
        self._incomplete = incomplete
        self._attributes = attributes
        self._href = None
        self._load_data()

    def _load_data(self):
        if "id" in self._data:  # We have complete data
            self.reload(get=False)
        elif "href" in self._data:  # We have only href
            self._href = self._data["href"]
            # self._data = None
        else:  # Malformed
            raise ValueError("Malformed data: {!r}".format(self._data))

    def reload(self, expand=None, get=True, attributes=None):
        kwargs = {}
        if expand:
            if isinstance(expand, (list, tuple)):
                expand = ",".join(map(str, expand))
            kwargs.update(expand=expand)
        if attributes is None:
            attributes = self._attributes
        if attributes:
            if isinstance(attributes, six.string_types):
                attributes = [attributes]
            kwargs.update(attributes=",".join(attributes))
        if "href" in self._data:
            self._href = self._data["href"]
        if get and self._href:
            new = self.collection._api.get(self._href, **kwargs)
            if self._data is None:
                self._data = new
            else:
                self._data.update(new)
        self._actions = self._data.pop("actions", [])
        for key, value in self._data.items():
            if value is None:
                continue
            if key in self.TIME_FIELDS:
                setattr(self, key, iso8601.parse_date(value))
            elif key in self.COLLECTION_MAPPING.keys():
                setattr(
                    self,
                    re.sub(r"_id$", "", key),
                    self.collection._api.get_entity(self.COLLECTION_MAPPING[key], value)
                )
                setattr(self, key, value)
            elif (isinstance(value, dict) and self._href and
                  "count" in value and "resources" in value):
                href = self._href
                if not href.endswith("/"):
                    href += "/"
                subcol = Collection(self.collection._api, href + key, key)
                setattr(self, key, subcol)
            elif (isinstance(value, list) and self._href and
                  key in self.EXTENDED_COLLECTIONS.get(self.collection.name, set([]))):
                href = self._href
                if not href.endswith("/"):
                    href += "/"
                subcol = Collection(self.collection._api, href + key, key)
                setattr(self, key, subcol)
            else:
                setattr(self, key, value)

    @property
    def exists(self):
        try:
            self.reload()
        except APIException:
            return False
        else:
            return True

    @property
    def subcollections(self):
        return self.collection.subcollections

    def wait_for_existence(self, existence, **kwargs):
        return wait_for(
            lambda: self.exists, fail_condition=not existence, **kwargs)

    def wait_exists(self, **kwargs):
        return self.wait_for_existence(True, **kwargs)

    def wait_not_exists(self, **kwargs):
        return self.wait_for_existence(False, **kwargs)

    def reload_if_needed(self):
        if self._data is None or self._incomplete or not hasattr(self, "_actions"):
            self.reload()
            self._incomplete = False

    def __getattr__(self, attr):
        self.reload()
        if attr in self.__dict__:
            # It got loaded
            return self.__dict__[attr]
        if self.subcollections is not None:
            if attr not in self.subcollections:
                raise AttributeError("No such attribute/subcollection {}".format(attr))
        if not self._href:
            raise AttributeError("Can't get URL of attribute/subcollection {}".format(attr))
        # Try to get subcollection
        href = self._href
        if not href.endswith("/"):
            href += "/"
        subcol = Collection(self.collection._api, href + attr, attr)
        try:
            subcol.reload()
        except APIException:
            raise AttributeError("No such attribute/subcollection {}".format(attr))
        else:
            return subcol

    def __getitem__(self, item):
        # Backward compatibility
        return getattr(self, item)

    def __repr__(self):
        return "<Entity {!r}>".format(self._href if self._href else self._data["id"])

    def _ref_repr(self):
        return {"href": self._href} if self._href else {"id": self._data["id"]}

class ActionContainer(object):
    def __init__(self, obj):
        self._obj = obj

    def reload(self):
        self._obj.reload_if_needed()
        reloaded_actions = []
        for action in self._obj._actions:

            def _add_method(method=None):
                """Adds HTTP method to Action, e.g. ``.delete.POST()``."""
                method = method or action["method"]
                new_name = method.upper()
                base_action_obj = self.__dict__[action["name"]]
                if not hasattr(base_action_obj, new_name):
                    action_obj = base_action_obj if method == base_action_obj._method else Action(
                        self, action["name"], method, action["href"])
                    setattr(base_action_obj, new_name, action_obj)

            # There can be multiple actions with the same name and different HTTP methods
            # (e.g. action "delete" with HTTP methods POST or DELETE).
            # Create action ``.name()`` with default (first) method.
            # For each method, create action ``.name.METHOD()``.
            # E.g. default action ``.delete()`` with method POST and actions
            # ``.delete.POST()`` and ``.delete.DELETE()``.
            if action["name"] not in reloaded_actions:
                reloaded_actions.append(action["name"])
                action_obj = Action(self, action["name"], action["method"], action["href"])
                setattr(self, action["name"], action_obj)
            _add_method()

            # Edit actions on entities can be performed using PATCH and PUT methods as well.
            # These methods are not listed in "actions", therefore adding
            # them explicitly - see https://bugzilla.redhat.com/show_bug.cgi?id=1491336
            if action["name"] == "edit" and isinstance(self._obj, Entity):
                for edit_method in ("patch", "put"):
                    _add_method(method=edit_method)

    def execute_action(self, action_name, *args, **kwargs):
        # To circumvent bad method names, like `import`, you can use this one directly
        action = getattr(self, action_name)
        action_method = kwargs.pop('action_method', None)
        if action_method:
            action = getattr(action, action_method)
        return action(*args, **kwargs)

    @property
    def all(self):
        self.reload()
        return [a["name"] for a in self._obj._actions]

    @property
    def collection(self):
        if isinstance(self._obj, Collection):
            return self._obj
        elif isinstance(self._obj, Entity):
            return self._obj.collection
        else:
            raise ValueError("ActionContainer assigned to wrong object!")

    def __getattr__(self, attr):
        self.reload()
        if attr not in self.__dict__:
            raise AttributeError("No such action {}".format(attr))
        return self.__dict__[attr]

    def __contains__(self, action):
        return action in self.all

class Action(object):
    def __init__(self, container, name, method, href):
        self._container = container
        self._method = method
        self._href = href
        self._name = name

    @property
    def collection(self):
        return self._container.collection

    @property
    def api(self):
        return self.collection.api

    def __call__(self, *args, **kwargs):
        # TODO: for backwards compatibility - can be removed when no longer used
        # possibility to override HTTP method that will be used with the action
        # (e.g. force_method='delete')
        force_method = kwargs.pop('force_method', None)
        if force_method:
            return getattr(self, force_method.upper())(*args, **kwargs)

        if self._method == "post":
            resources = []
            # We got resources to post
            for res in args:
                if isinstance(res, Entity):
                    resources.append(res._ref_repr())
                else:
                    resources.append(res)
            query_dict = {"action": self._name}
            if resources:
                query_dict["resources"] = []
                for resource in resources:
                    new_res = dict(resource)
                    if kwargs:
                        new_res.update(kwargs)
                    query_dict["resources"].append(new_res)
            else:
                if kwargs:
                    query_dict["resource"] = kwargs

            result = self.api.post(self._href, **query_dict)
        elif self._method == "patch":
            result = self.api.patch(self._href, *args)
        elif self._method in ("delete", "put"):
            result = getattr(self.api, self._method)(self._href, **kwargs)
        else:
            raise NotImplementedError
        if result is None:
            return None
        # Make sure that HTTP response from action is not overriden during result processing
        action_response = self.api.response
        try:
            if "results" in result:
                outcome = [self._process_result(r) for r in result["results"]]
            else:
                outcome = self._process_result(result)
        finally:
            self.api.response = action_response
        return outcome

    def _get_entity_from_href(self, result):
        """Returns entity in correct collection.

        If the "href" value in result doesn't match the current collection,
        try to find the collection that the "href" refers to.
        """
        href_result = result['href']

        if self.collection._href.startswith(href_result):
            return Entity(self.collection, result, incomplete=True)

        href_match = re.match(r"(https?://.+/api[^?]*)/([a-z_-]+)", href_result)
        if not href_match:
            raise ValueError("Malformed href: {}".format(href_result))
        collection_name = href_match.group(2)
        entry_point = href_match.group(1)
        new_collection = Collection(
            self.collection.api,
            "{}/{}".format(entry_point, collection_name),
            collection_name
        )
        return Entity(new_collection, result, incomplete=True)

    def _process_result(self, result):
        if result is None:
            return None
        elif "href" in result:
            return self._get_entity_from_href(result)
        elif "id" in result:
            d = copy(result)
            d["href"] = "{}/{}".format(self.collection._href, result["id"])
            return Entity(self.collection, d, incomplete=True)
        elif "task_href" in result:
            collection = self.api.collections.tasks
            # reuse task_href and task_id, no other data is relevant
            d = {"href": result.get("task_href"), "id": result.get("task_id")}
            return Entity(collection, d, incomplete=True)
        elif "message" in result:
            return result
        else:
            raise NotImplementedError

    def __repr__(self):
        return "<Action {} {}#{}>".format(self._method, self._container._obj._href, self._name)


# -*- coding: utf-8 -*-
import six

QUOTES = {u'"', u"'"}

def give_another_quote(q):
    """When you pass a quote character, returns you an another one if possible"""
    for qc in QUOTES:
        if qc != q:
            return qc
    else:
        raise ValueError(u'Could not find a different quote for {}'.format(q))

def escape_filter(o):
    """Tries to escape the values that are passed to filter as correctly as possible.

    No standard way is followed, but at least it is simple.
    """
    if o is None:
        return u'NULL'
    if isinstance(o, int):
        return str(o)
    if not isinstance(o, six.string_types):
        raise ValueError('Filters take only None, int or a string type')
    if not o:
        # Empty string
        return u"''"
    # Now enforce unicode
    o = unicode_process(o)
    if u'"' not in o:
        # Simple case, just put the quote that does not exist in the string
        return u'"' + o + u'"'
    elif u"'" not in o:
        # Simple case, just put the quote that does not exist in the string
        return u"'" + o + u"'"
    else:
        # Both are there, so start guessing
        # Empty strings are sorted out, so the string must contain something.
        # String with length == 1 are sorted out because if they have a quote, they would be quoted
        # with the another quote in preceeding branch. Therefore the string is at least 2 chars long
        # here which allows us to NOT check the length here.
        first_char = o[0]
        last_char = o[-1]
        if first_char in QUOTES and last_char in QUOTES:
            # The first and last chars definitely are quotes
            if first_char == last_char:
                # Simple, just put another ones around them
                quote = give_another_quote(first_char)
                return quote + o + quote
            else:
                # I don't like this but the nature of the escape is like that ...
                # Since now it uses both of the quotes, just pick the simple ones and surround it
                return u"'" + o + u"'"
        elif first_char not in QUOTES and last_char not in QUOTES:
            # First and last chars are not quotes, so a simple solution
            return u"'" + o + u"'"
        else:
            # One of the first or last chars is not a quote
            if first_char in QUOTES:
                quote = give_another_quote(first_char)
            else:
                # last_char
                quote = give_another_quote(last_char)
            return quote + o + quote

def unicode_process(s):
    if not isinstance(s, six.string_types):
        if six.PY2:
            if hasattr(s, '__unicode__'):
                s = unicode(s)
            else:
                s = str(s)
        else:
            s = str(s)
    if (six.PY3 and isinstance(s, bytes)) or (six.PY2 and isinstance(s, str)):
        s = s.decode('utf-8')
    # Here we have Unicode!
    return s


#!/usr/bin/env python

# Plot parallax errors as a function of V according to the interpolation
# formulae on the Gaia science performance web-pages.

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from pygaia.errors.astrometric import parallaxErrorSkyAvg
from pygaia.errors.astrometric import parallaxMinError, parallaxMaxError
from pygaia.photometry.utils import vminiFromSpt
from pygaia.photometry.transformations import gminvFromVmini

from os import environ as env
from matplotlib import rc
try:
  import argparse
except ImportError:
  raise ImportError("""The argparse module does not seem to be available.\n Either upgrade to python 2.7
  or adapt the script to using the optparse module instead.""")

# Configure matplotlib
rc('text', usetex=True)
rc('font', family='serif', size=16)
rc('xtick.major', size='12')
rc('xtick.minor', size='6')
rc('ytick.major', size='12')
rc('ytick.minor', size='6')
rc('lines', linewidth=2)
rc('axes', linewidth=2)

def makePlot(args):
  """
  Make the plot with parallax performance predictions.

  :argument args: command line arguments
  """
  gmag=np.linspace(5.7,20.0,101)

  vminiB1V=vminiFromSpt('B1V')
  vminiG2V=vminiFromSpt('G2V')
  vminiM6V=vminiFromSpt('M6V')

  vmagB1V=gmag-gminvFromVmini(vminiB1V)
  vmagG2V=gmag-gminvFromVmini(vminiG2V)
  vmagM6V=gmag-gminvFromVmini(vminiM6V)

  sigparB1V=parallaxErrorSkyAvg(gmag,vminiB1V)
  sigparB1Vmin=parallaxMinError(gmag,vminiB1V)
  sigparB1Vmax=parallaxMaxError(gmag,vminiB1V)

  sigparG2V=parallaxErrorSkyAvg(gmag,vminiG2V)
  sigparG2Vmin=parallaxMinError(gmag,vminiG2V)
  sigparG2Vmax=parallaxMaxError(gmag,vminiG2V)

  sigparM6V=parallaxErrorSkyAvg(gmag,vminiM6V)
  sigparM6Vmin=parallaxMinError(gmag,vminiM6V)
  sigparM6Vmax=parallaxMaxError(gmag,vminiM6V)

  fig=plt.figure(figsize=(10,6.5))

  if (args['gmagAbscissa']):
    plt.semilogy(gmag, sigparB1V, 'b', label='B1V')
    plt.semilogy(gmag, sigparG2V, 'g', label='G2V')
    plt.semilogy(gmag, sigparM6V, 'r', label='M6V')
    plt.xlim((5,20))
    plt.ylim((4,1000))
    plt.legend(loc=4)
    plt.xlabel('$G$ [mag]')
  else:
    ax=fig.add_subplot(111)
    plt.semilogy(vmagB1V, sigparB1V, 'b', label='B1V')
    #plt.semilogy(vmagG2V, sigparG2V, 'g', label='G2V')
    plt.semilogy(vmagM6V, sigparM6V, 'r', label='M6V')
    plt.fill_between(vmagB1V, sigparB1Vmin, sigparB1Vmax, color='b', alpha=0.3)
    plt.fill_between(vmagM6V, sigparM6Vmin, sigparM6Vmax, color='r', alpha=0.3)
    plt.xlim((5,22.5))
    plt.ylim((4,1000))
    plt.text(17.2,190,'B1V',color='b')
    plt.text(18,20,'M6V',color='r')
    plt.xlabel('$V$ [mag]')
    plt.text(7,17,'calibration noise floor', size=12, bbox=dict(boxstyle="round,pad=0.3",
                       ec=(0.0, 0.0, 0.0),
                       fc=(1.0, 1.0, 1.0),
                       ))
    plt.text(14.75,80,'photon noise', rotation=45, size=12, bbox=dict(boxstyle="round,pad=0.3",
                       ec=(0.0, 0.0, 0.0),
                       fc=(1.0, 1.0, 1.0),
                       ))
    ax.annotate('non-uniformity\nover the sky', xy=(21.5, 320),  xycoords='data',
                  xytext=(21.5,80), textcoords='data', ha='center', size='12',
                  bbox=dict(boxstyle="round,pad=0.3",ec=(0,0,0),fc=(1,1,1)),
                  arrowprops=dict(facecolor='black', shrink=0.15, width=1,
                    headwidth=6),
                  horizontalalignment='right', verticalalignment='top',
                  )
    ax.annotate('', xy=(21.5, 500),  xycoords='data',
                  xytext=(21.5,950), textcoords='data', ha='center', size='12',
                  arrowprops=dict(facecolor='black', shrink=0.15, width=1,
                    headwidth=6),
                  horizontalalignment='right', verticalalignment='bottom',
                  )

  plt.xticks(np.arange(6,24,2))
  ax = plt.gca().yaxis
  ax.set_major_formatter(matplotlib.ticker.ScalarFormatter())
  plt.ticklabel_format(axis='y',style='plain')
  plt.grid(which='both')
  plt.ylabel('End-of-mission parallax standard error [$\mu$as]')

  if (args['pdfOutput']):
    plt.savefig('ParallaxErrors.pdf')
  elif (args['pngOutput']):
    plt.savefig('ParallaxErrors.png')
  else:
    plt.show()

def parseCommandLineArguments():
  """
  Set up command line parsing.
  """
  parser = argparse.ArgumentParser(description="Plot predicted Gaia sky averaged parallax errors as a function of V")
  parser.add_argument("-p", action="store_true", dest="pdfOutput", help="Make PDF plot")
  parser.add_argument("-b", action="store_true", dest="pngOutput", help="Make PNG plot")
  parser.add_argument("-g", action="store_true", dest="gmagAbscissa", help="Plot performance vs G instead of V")
  args=vars(parser.parse_args())
  return args

if __name__ in ('__main__'):
  args=parseCommandLineArguments()
  makePlot(args)


"""
For a given G-band bright limit for Gaia detection/confirmation, plot the bright limit in the V-band as a
function of (V-I).
"""

from pygaia.photometry.transformations import gminvFromVmini
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import rc
try:
  import argparse
except ImportError:
  raise ImportError("""The argparse module does not seem to be available.\n Either upgrade to python 2.7
  or adapt the script to using the optparse module instead.""")

# Configure matplotlib
rc('text', usetex=True)
rc('font', family='serif', size=16)
rc('xtick.major', size='12')
rc('xtick.minor', size='6')
rc('ytick.major', size='12')
rc('ytick.minor', size='6')
rc('lines', linewidth=2)
rc('axes', linewidth=2)

def plotBrightLimitInV(gBright, pdf=False, png=False):
  """
  Plot the bright limit of Gaia in V as a function of (V-I).

  Parameters
  ----------

  gBright - The bright limit of Gaia in G
  """
  vmini=np.linspace(0.0,6.0,1001)
  gminv=gminvFromVmini(vmini)
  vBright=gBright-gminv

  fig=plt.figure(figsize=(10,6.5))
  plt.plot(vmini,vBright,'b-')
  plt.xlabel('$(V-I)$')
  plt.ylabel('Bright limit of Gaia in $V$')
  plt.xlim(0,6)
  plt.ylim(5,11)
  plt.grid(which='both')
  plt.title("Bright limit in $G$: {0}".format(gBright))

  if (pdf):
    plt.savefig('VBandBrightLimit.pdf')
  elif (png):
    plt.savefig('VBandBrightLimit.png')
  else:
    plt.show()

def parseCommandLineArguments():
  """
  Set up command line parsing.
  """
  parser = argparse.ArgumentParser(description="""Plot the bright limit in the V-band, as a function of
  (V-I) for a given bright limit in G""")
  parser.add_argument("gmag", help="G-band bright limit", type=float)
  parser.add_argument("-p", action="store_true", dest="pdfOutput", help="Make PDF plot")
  parser.add_argument("-b", action="store_true", dest="pngOutput", help="Make PNG plot")

  args=vars(parser.parse_args())
  return args

if __name__ in ('__main__'):
  args=parseCommandLineArguments()
  plotBrightLimitInV(args['gmag'], pdf=args['pdfOutput'], png=args['pngOutput'])


__all__ = [ 'sphericalToCartesian', 'cartesianToSpherical', 'normalTriad', 'elementaryRotationMatrix',
'phaseSpaceToAstrometry', 'astrometryToPhaseSpace']

from numpy import cos, sin, sqrt, arctan2, zeros_like, isscalar
from numpy import array, transpose, dot, any, where, pi

# Astronomical Unit in meter, IAU constant and defining length
_auInMeter = 149597870700.0

# AU expressed in mas*pc or muas*kpc
_auMasParsec = 1000.0

# Number of seconds in Julian year
_julianYearSeconds = 365.25 * 86400.0

# AU expressed in km*yr/s
_auKmYearPerSec = _auInMeter/(_julianYearSeconds*1000.0)

def sphericalToCartesian(r, phi, theta):
  """
  Convert spherical to Cartesian coordinates. The input can be scalars or 1-dimensional numpy arrays.
  Note that the angle coordinates follow the astronomical convention of using elevation (declination,
  latitude) rather than its complement (pi/2-elevation), where the latter is commonly used in the
  mathematical treatment of spherical coordinates.

  Parameters
  ----------

  r     - length of input Cartesian vector.
  phi   - longitude-like angle (e.g., right ascension, ecliptic longitude) in radians
  theta - latitide-like angle (e.g., declination, ecliptic latitude) in radians

  Returns
  -------

  The Cartesian vector components x, y, z
  """
  ctheta=cos(theta)
  x=r*cos(phi)*ctheta
  y=r*sin(phi)*ctheta
  z=r*sin(theta)
  return x, y, z

def cartesianToSpherical(x, y, z):
  """
  Convert Cartesian to spherical coordinates. The input can be scalars or 1-dimensional numpy arrays.
  Note that the angle coordinates follow the astronomical convention of using elevation (declination,
  latitude) rather than its complement (pi/2-elevation), which is commonly used in the mathematical
  treatment of spherical coordinates.

  Parameters
  ----------

  x - Cartesian vector component along the X-axis
  y - Cartesian vector component along the Y-axis
  z - Cartesian vector component along the Z-axis

  Returns
  -------

  The spherical coordinates r=sqrt(x*x+y*y+z*z), longitude phi, latitude theta.

  NOTE THAT THE LONGITUDE ANGLE IS BETWEEN 0 AND +2PI. FOR r=0 AN EXCEPTION IS RAISED.
  """
  rCylSq=x*x+y*y
  r=sqrt(rCylSq+z*z)
  if any(r==0.0):
    raise Exception("Error: one or more of the points is at distance zero.")
  phi = arctan2(y,x)
  phi = where(phi<0.0, phi+2*pi, phi)
  return r, phi, arctan2(z,sqrt(rCylSq))

def normalTriad(phi, theta):
  """
  Calculate the so-called normal triad [p, q, r] which is associated with a spherical coordinate system .
  The three vectors are:

  p - The unit tangent vector in the direction of increasing longitudinal angle phi.
  q - The unit tangent vector in the direction of increasing latitudinal angle theta.
  r - The unit vector toward the point (phi, theta).

  Parameters
  ----------

  phi   - longitude-like angle (e.g., right ascension, ecliptic longitude) in radians
  theta - latitide-like angle (e.g., declination, ecliptic latitude) in radians

  Returns
  -------

  The normal triad as the vectors p, q, r
  """
  sphi=sin(phi)
  stheta=sin(theta)
  cphi=cos(phi)
  ctheta=cos(theta)
  p=array([-sphi, cphi, zeros_like(phi)])
  q=array([-stheta*cphi, -stheta*sphi, ctheta])
  r=array([ctheta*cphi, ctheta*sphi, stheta])
  return p, q, r

def elementaryRotationMatrix(axis, rotationAngle):
  """
  Construct an elementary rotation matrix describing a rotation around the x, y, or z-axis.

  Parameters
  ----------

  axis          - Axis around which to rotate ("x", "y", or "z")
  rotationAngle - the rotation angle in radians

  Returns
  -------

  The rotation matrix

  Example usage
  -------------

  rotmat = elementaryRotationMatrix("y", pi/6.0)
  """
  if (axis=="x" or axis=="X"):
    return array([[1.0, 0.0, 0.0], [0.0, cos(rotationAngle), sin(rotationAngle)], [0.0,
      -sin(rotationAngle), cos(rotationAngle)]])
  elif (axis=="y" or axis=="Y"):
    return array([[cos(rotationAngle), 0.0, -sin(rotationAngle)], [0.0, 1.0, 0.0], [sin(rotationAngle),
      0.0, cos(rotationAngle)]])
  elif (axis=="z" or axis=="Z"):
    return array([[cos(rotationAngle), sin(rotationAngle), 0.0], [-sin(rotationAngle),
      cos(rotationAngle), 0.0], [0.0, 0.0, 1.0]])
  else:
    raise Exception("Unknown rotation axis "+axis+"!")

def phaseSpaceToAstrometry(x, y, z, vx, vy, vz):
  """
  From the given phase space coordinates calculate the astrometric observables, including the radial
  velocity, which here is seen as the sixth astrometric parameter. The phase space coordinates are
  assumed to represent barycentric (i.e. centred on the Sun) positions and velocities.

  This function has no mechanism to deal with units. The velocity units are always assumed to be km/s,
  and the code is set up such that for positions in pc, the return units for the astrometry are radians,
  milliarcsec, milliarcsec/year and km/s. For positions in kpc the return units are: radians,
  microarcsec, microarcsec/year, and km/s.

  NOTE that the doppler factor k=1/(1-vrad/c) is NOT used in the calculations. This is not a problem for
  sources moving at typical velocities of Galactic stars.

  Parameters
  ----------

  x - The x component of the barycentric position vector (in pc or kpc).
  y - The y component of the barycentric position vector (in pc or kpc).
  z - The z component of the barycentric position vector (in pc or kpc).
  vx - The x component of the barycentric velocity vector (in km/s).
  vy - The y component of the barycentric velocity vector (in km/s).
  vz - The z component of the barycentric velocity vector (in km/s).

  Returns
  -------

  phi       - The longitude-like angle of the position of the source (radians).
  theta     - The latitude-like angle of the position of the source (radians).
  parallax  - The parallax of the source (in mas or muas, see above)
  muphistar - The proper motion in the longitude-like angle, multiplied by cos(theta) (mas/yr or muas/yr,
  see above)
  mutheta   - The proper motion in the latitude-like angle (mas/yr or muas/yr, see above)
  vrad      - The radial velocity (km/s)
  """
  u, phi, theta = cartesianToSpherical(x, y, z)
  parallax = _auMasParsec/u
  p, q, r = normalTriad(phi, theta)
  velocitiesArray=array([vx,vy,vz])
  if  isscalar(u):
    muphistar=dot(p,velocitiesArray)*parallax/_auKmYearPerSec
    mutheta=dot(q,velocitiesArray)*parallax/_auKmYearPerSec
    vrad=dot(r,velocitiesArray)
  else:
    muphistar=zeros_like(parallax)
    mutheta=zeros_like(parallax)
    vrad=zeros_like(parallax)
    for i in range(parallax.size):
      muphistar[i]=dot(p[:,i],velocitiesArray[:,i])*parallax[i]/_auKmYearPerSec
      mutheta[i]=dot(q[:,i],velocitiesArray[:,i])*parallax[i]/_auKmYearPerSec
      vrad[i]=dot(r[:,i],velocitiesArray[:,i])

  return phi, theta, parallax, muphistar, mutheta, vrad

def astrometryToPhaseSpace(phi, theta, parallax, muphistar, mutheta, vrad):
  """
  From the input astrometric parameters calculate the phase space coordinates. The output phase space
  coordinates represent barycentric (i.e. centred on the Sun) positions and velocities.

  This function has no mechanism to deal with units. The code is set up such that for input astrometry
  with parallaxes and proper motions in mas and mas/yr, and radial velocities in km/s, the phase space
  coordinates are in pc and km/s. For input astrometry with parallaxes and proper motions in muas and
  muas/yr, and radial velocities in km/s, the phase space coordinates are in kpc and km/s. Only positive
  parallaxes are accepted, an exception is thrown if this condition is not met.

  NOTE that the doppler factor k=1/(1-vrad/c) is NOT used in the calculations. This is not a problem for
  sources moving at typical velocities of Galactic stars.

  THIS FUNCTION SHOULD NOT BE USED WHEN THE PARALLAXES HAVE RELATIVE ERRORS LARGER THAN ABOUT 20 PER CENT
  (see http://arxiv.org/abs/1507.02105 for example). For astrometric data with relatively large parallax
  errors you should consider doing your analysis in the data space and use forward modelling of some
  kind.

  Parameters
  ----------

  phi       - The longitude-like angle of the position of the source (radians).
  theta     - The latitude-like angle of the position of the source (radians).
  parallax  - The parallax of the source (in mas or muas, see above)
  muphistar - The proper motion in the longitude-like angle, multiplied by cos(theta) (mas/yr or muas/yr,
  see above)
  mutheta   - The proper motion in the latitude-like angle (mas/yr or muas/yr, see above)
  vrad      - The radial velocity (km/s)

  Returns
  -------

  x - The x component of the barycentric position vector (in pc or kpc).
  y - The y component of the barycentric position vector (in pc or kpc).
  z - The z component of the barycentric position vector (in pc or kpc).
  vx - The x component of the barycentric velocity vector (in km/s).
  vy - The y component of the barycentric velocity vector (in km/s).
  vz - The z component of the barycentric velocity vector (in km/s).
  """
  if any(parallax<=0.0):
    raise Exception("One or more of the input parallaxes is non-positive")
  x, y, z = sphericalToCartesian(_auMasParsec/parallax, phi, theta)
  p, q, r = normalTriad(phi, theta)
  transverseMotionArray = array([muphistar*_auKmYearPerSec/parallax, mutheta*_auKmYearPerSec/parallax,
    vrad])
  if isscalar(parallax):
    velocityArray=dot(transpose(array([p, q, r])),transverseMotionArray)
    vx = velocityArray[0]
    vy = velocityArray[1]
    vz = velocityArray[2]
  else:
    vx = zeros_like(parallax)
    vy = zeros_like(parallax)
    vz = zeros_like(parallax)
    for i in range(parallax.size):
      velocityArray = dot(transpose(array([p[:,i], q[:,i], r[:,i]])), transverseMotionArray[:,i])
      vx[i] = velocityArray[0]
      vy[i] = velocityArray[1]
      vz[i] = velocityArray[2]
  return x, y, z, vx, vy, vz


#!/usr/bin/env python

# Plot proper motion errors as a function of V according to the interpolation formulae on the Gaia
# science performance web-pages.

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from pygaia.errors.astrometric import properMotionErrorSkyAvg, properMotionMinError, properMotionMaxError
from pygaia.photometry.utils import vminiFromSpt
from pygaia.photometry.transformations import gminvFromVmini

from os import environ as env
from matplotlib import rc
try:
  import argparse
except ImportError:
  raise ImportError("""The argparse module does not seem to be available.\n Either upgrade to python 2.7
  or adapt the script to using the optparse module instead.""")

# Configure matplotlib
rc('text', usetex=True)
rc('font', family='serif', size=16)
rc('xtick.major', size='12')
rc('xtick.minor', size='6')
rc('ytick.major', size='12')
rc('ytick.minor', size='6')
rc('lines', linewidth=2)
rc('axes', linewidth=2)

def makePlot(args):
  """
  Make the plot with proper motion performance predictions. The predictions are for the TOTAL proper
  motion under the assumption of equal components mu_alpha* and mu_delta.

  :argument args: command line arguments
  """
  gmag=np.linspace(5.7,20.0,101)

  vminiB1V=vminiFromSpt('B1V')
  vminiG2V=vminiFromSpt('G2V')
  vminiM6V=vminiFromSpt('M6V')

  vmagB1V=gmag-gminvFromVmini(vminiB1V)
  vmagG2V=gmag-gminvFromVmini(vminiG2V)
  vmagM6V=gmag-gminvFromVmini(vminiM6V)

  sigmualphaB1V, sigmudeltaB1V = properMotionErrorSkyAvg(gmag,vminiB1V)
  sigmuB1V = np.sqrt(0.5*sigmualphaB1V**2+0.5*sigmudeltaB1V**2)
  sigmualphaB1V, sigmudeltaB1V = properMotionMinError(gmag,vminiB1V)
  sigmuB1Vmin = np.sqrt(0.5*sigmualphaB1V**2+0.5*sigmudeltaB1V**2)
  sigmualphaB1V, sigmudeltaB1V = properMotionMaxError(gmag,vminiB1V)
  sigmuB1Vmax = np.sqrt(0.5*sigmualphaB1V**2+0.5*sigmudeltaB1V**2)

  sigmualphaG2V, sigmudeltaG2V = properMotionErrorSkyAvg(gmag,vminiG2V)
  sigmuG2V = np.sqrt(0.5*sigmualphaG2V**2+0.5*sigmudeltaG2V**2)
  sigmualphaG2V, sigmudeltaG2V = properMotionMinError(gmag,vminiG2V)
  sigmuG2Vmin = np.sqrt(0.5*sigmualphaG2V**2+0.5*sigmudeltaG2V**2)
  sigmualphaG2V, sigmudeltaG2V = properMotionMaxError(gmag,vminiG2V)
  sigmuG2Vmax = np.sqrt(0.5*sigmualphaG2V**2+0.5*sigmudeltaG2V**2)

  sigmualphaM6V, sigmudeltaM6V = properMotionErrorSkyAvg(gmag,vminiM6V)
  sigmuM6V = np.sqrt(0.5*sigmualphaM6V**2+0.5*sigmudeltaM6V**2)
  sigmualphaM6V, sigmudeltaM6V = properMotionMinError(gmag,vminiM6V)
  sigmuM6Vmin = np.sqrt(0.5*sigmualphaM6V**2+0.5*sigmudeltaM6V**2)
  sigmualphaM6V, sigmudeltaM6V = properMotionMaxError(gmag,vminiM6V)
  sigmuM6Vmax = np.sqrt(0.5*sigmualphaM6V**2+0.5*sigmudeltaM6V**2)

  fig=plt.figure(figsize=(10,6.5))

  if (args['gmagAbscissa']):
    plt.semilogy(gmag, sigmuB1V, 'b', label='B1V')
    plt.semilogy(gmag, sigmuG2V, 'g', label='G2V')
    plt.semilogy(gmag, sigmuM6V, 'r', label='M6V')
    plt.xlim((5,20))
    plt.ylim((1,500))
    plt.legend(loc=4)
  else:
    ax=fig.add_subplot(111)
    plt.semilogy(vmagB1V, sigmuB1V, 'b', label='B1V')
    #plt.semilogy(vmagG2V, sigmuG2V, 'g', label='G2V')
    plt.semilogy(vmagM6V, sigmuM6V, 'r', label='M6V')
    plt.fill_between(vmagB1V, sigmuB1Vmin, sigmuB1Vmax, color='b', alpha=0.3)
    plt.fill_between(vmagM6V, sigmuM6Vmin, sigmuM6Vmax, color='r', alpha=0.3)
    plt.xlim((5,22.5))
    plt.ylim((1,500))
    plt.text(17.5,100,'B1V',color='b')
    plt.text(18,10,'M6V',color='r')
    plt.text(7,11,'calibration noise floor', size=12, bbox=dict(boxstyle="round,pad=0.3",
                       ec=(0.0, 0.0, 0.0),
                       fc=(1.0, 1.0, 1.0),
                       ))
    plt.text(14.75,50,'photon noise', rotation=45, size=12, bbox=dict(boxstyle="round,pad=0.3",
                       ec=(0.0, 0.0, 0.0),
                       fc=(1.0, 1.0, 1.0),
                       ))
    ax.annotate('non-uniformity\nover the sky', xy=(21.5, 80),  xycoords='data',
                  xytext=(21.5,30), textcoords='data', ha='center', size='12',
                  bbox=dict(boxstyle="round,pad=0.3",ec=(0,0,0),fc=(1,1,1)),
                  arrowprops=dict(facecolor='black', shrink=0.15, width=1,
                    headwidth=6),
                  horizontalalignment='right', verticalalignment='top',
                  )
    ax.annotate('', xy=(21.5, 170),  xycoords='data',
                  xytext=(21.5,380), textcoords='data', ha='center', size='12',
                  arrowprops=dict(facecolor='black', shrink=0.15, width=1,
                    headwidth=6),
                  horizontalalignment='right', verticalalignment='bottom',
                  )

  plt.xticks(np.arange(6,24,2))
  ax = plt.gca().yaxis
  ax.set_major_formatter(matplotlib.ticker.ScalarFormatter())
  plt.ticklabel_format(axis='y',style='plain')
  plt.grid(which='both')
  plt.xlabel('$V$ [mag]')
  plt.ylabel('End-of-mission $\\sigma_\\mu$ [$\mu$as/yr]')

  basename = 'ProperMotionErrors'
  if (args['pdfOutput']):
    plt.savefig(basename+'.pdf')
  elif (args['pngOutput']):
    plt.savefig(basename+'.png')
  else:
    plt.show()

def parseCommandLineArguments():
  """
  Set up command line parsing.
  """
  parser = argparse.ArgumentParser(description="Plot predicted Gaia sky averaged proper motion errors as a function of V")
  parser.add_argument("-p", action="store_true", dest="pdfOutput", help="Make PDF plot")
  parser.add_argument("-b", action="store_true", dest="pngOutput", help="Make PNG plot")
  parser.add_argument("-g", action="store_true", dest="gmagAbscissa", help="Plot performance vs G instead of V")
  args=vars(parser.parse_args())
  return args

if __name__ in ('__main__'):
  args=parseCommandLineArguments()
  makePlot(args)


__all__ = ['enum', 'degreesToRadians', 'radiansToDegrees']

import numpy as np

from pygaia.astrometry.constants import auKmYearPerSec

def enum(typename, field_names):
    """
    Create a new enumeration type.

    Code is copyright (c) Gabriel Genellina, 2010, MIT License.

    Parameters
    ----------

    typename - Name of the enumerated type
    field_names - Names of the fields of the enumerated type
    """

    if isinstance(field_names, str):
        field_names = field_names.replace(',', ' ').split()
    d = dict((reversed(nv) for nv in enumerate(field_names)), __slots__ = ())
    return type(typename, (object,), d)()

def degreesToRadians(angle):
    """
    Convert from degrees to radians.

    Parameters
    ----------

    angle - angle in degrees

    Returns
    -------

    Angle in radians.
    """
    return angle/180.0*np.pi

def radiansToDegrees(angle):
    """
    Convert from radians to degrees.

    Parameters
    ----------

    angle - angle in radians.

    Returns
    -------

    Angle in degrees.
    """
    return angle/np.pi*180.0

def construct_covariance_matrix(cvec, parallax, radial_velocity, radial_velocity_error):
    """
    Take the astrometric parameter standard uncertainties and the uncertainty correlations as quoted in
    the Gaia catalogue and construct the covariance matrix.

    Parameters
    ----------

    cvec : array_like
        Array of shape (15,) (1 source) or (n,15) (n sources) for the astrometric parameter standard
        uncertainties and their correlations, as listed in the Gaia catalogue [ra_error, dec_error,
        parallax_error, pmra_error, pmdec_error, ra_dec_corr, ra_parallax_corr, ra_pmra_corr,
        ra_pmdec_corr, dec_parallax_corr, dec_pmra_corr, dec_pmdec_corr, parallax_pmra_corr,
        parallax_pmdec_corr, pmra_pmdec_corr]. Units are (mas^2, mas^2/yr, mas^2/yr^2).

    parallax : array_like (n elements)
        Source parallax (mas).

    radial_velocity : array_like (n elements)
        Source radial velocity (km/s, does not have to be from Gaia RVS!). If the radial velocity is not
        known it can be set to zero.

    radial_velocity_error : array_like (n elements)
        Source radial velocity  uncertainty (km/s). If the radial velocity is not know this can be set to
        the radial velocity dispersion for the population the source was drawn from.

    Returns
    -------

    Covariance matrix as a 6x6 array.
    """

    if np.ndim(cvec)==1:
        cmat = np.zeros((1,6,6))
        nsources = 1
        cv = np.atleast_2d(cvec)
    else:
        nsources = cvec.shape[0]
        cmat = np.zeros((nsources,6,6))
        cv = cvec
    for k in range(nsources):
        cmat[k,0:5,0:5] = cv[k,0:5]**2

    iu = np.triu_indices(5,k=1)
    for k in range(10):
        i = iu[0][k]
        j = iu[1][k]
        cmat[:,i,j] = cv[:,i]*cv[:,j]*cv[:,k+5]
        cmat[:,j,i] = cmat[:,i,j]

    for k in range(nsources):
        cmat[k,0:5,5] = cmat[k,0:5,2]*np.atleast_1d(radial_velocity)[k]/auKmYearPerSec
    cmat[:,5,0:5] = cmat[:,0:5,5]
    cmat[:,5,5] = cmat[:,2,2]*(radial_velocity**2 + radial_velocity_error**2)/auKmYearPerSec**2 + \
            (parallax*radial_velocity_error/auKmYearPerSec)**2

    return np.squeeze(cmat)


#!/usr/bin/env python

# Plot the lines of constant distance for stars at G=20 in a colour magnitude diagram (Mv vs V-I).

import numpy as np
import matplotlib.pyplot as plt
from pygaia.photometry.transformations import gminvFromVmini, vminGrvsFromVmini

from os import environ as env
from matplotlib import rc
try:
  import argparse
except ImportError:
  raise ImportError("""The argparse module does not seem to be available.\n Either upgrade to python 2.7
  or adapt the script to using the optparse module instead.""")

# Configure matplotlib
rc('text', usetex=True)
rc('font', family='serif', size=16)
rc('xtick.major', size='12')
rc('xtick.minor', size='6')
rc('ytick.major', size='12')
rc('ytick.minor', size='6')
rc('lines', linewidth=2)
rc('axes', linewidth=2)

def makePlot(gmag, pdf=False, png=False, rvs=False):
  """
  Make a plot of a Mv vs (V-I) colour magnitude diagram containing lines of constant distance for stars
  at G=20. This will give an idea of the reach of Gaia.

  Parameters
  ----------

  args - command line arguments
  """
  vmini = np.linspace(-0.5,4.0,100)
  if (rvs):
    gminv = -vminGrvsFromVmini(vmini)
  else:
    gminv = gminvFromVmini(vmini)
  mvlimit100pc = gmag-5.0*np.log10(100.0)+5.0-gminv
  mvlimit1kpc = gmag-5.0*np.log10(1000.0)+5.0-gminv
  mvlimit10kpc = gmag-5.0*np.log10(10000.0)+5.0-gminv

  fig=plt.figure(figsize=(8,8))
  plt.plot(vmini,mvlimit100pc,'b')
  plt.text(vmini[50]-0.4,mvlimit100pc[50],"$d=100$ pc", horizontalalignment='right', va='top')
  plt.plot(vmini,mvlimit1kpc,'r')
  plt.text(vmini[50]-0.4,mvlimit1kpc[50],"$d=1000$ pc", horizontalalignment='right', va='top')
  plt.plot(vmini,mvlimit10kpc,'g')
  plt.text(vmini[50]-0.4,mvlimit10kpc[50],"$d=10000$ pc", horizontalalignment='right', va='top')
  ax=plt.gca()
  ax.set_ylim(ax.get_ylim()[::-1])
  plt.xlabel("$(V-I)$")
  plt.ylabel("$M_V$")
  if (rvs):
    plt.title("Distance limits for $G_\\mathrm{RVS}"+"={0}$".format(gmag))
  else:
    plt.title("Distance limits for $G={0}$".format(gmag))

  if (args['pdfOutput']):
    plt.savefig('GaiaSurveyLimits.pdf')
  elif (args['pngOutput']):
    plt.savefig('GaiaSurveyLimits.png')
  else:
    plt.show()

def parseCommandLineArguments():
  """
  Set up command line parsing.
  """
  parser = argparse.ArgumentParser(description="""Plot lines of constant distance for stars a given G
  magnitude in the Mv vs (V-I) diagram""")
  parser.add_argument("gmag", help="G-band magnitude", type=float)
  parser.add_argument("-p", action="store_true", dest="pdfOutput", help="Make PDF plot")
  parser.add_argument("-b", action="store_true", dest="pngOutput", help="Make PNG plot")
  parser.add_argument("-r", action="store_true", dest="useGrvs", help="Magnitude limit is in Grvs")
  args=vars(parser.parse_args())
  return args

if __name__ in ('__main__'):
  args=parseCommandLineArguments()
  makePlot(args['gmag'], pdf=args['pdfOutput'], png=args['pngOutput'], rvs=args['useGrvs'])


__all__ = ['vradErrorSkyAvg']

from numpy import exp

_vradErrorACoeff={'B0V':0.90, 'B5V':0.90, 'A0V':1.0, 'A5V':1.15, 'F0V':1.15, 'G0V':1.15, 'G5V':1.15,
'K0V':1.15, 'K1IIIMP':1.15, 'K4V':1.15, 'K1III':1.15}
_vradErrorBCoeff={'B0V':50.00, 'B5V':26.00, 'A0V':5.50, 'A5V':4.00, 'F0V':1.50, 'G0V':0.70, 'G5V':0.60,
'K0V':0.50, 'K1IIIMP':0.39, 'K4V':0.29, 'K1III':0.21}
_vradCalibrationFloor = 0.5
_vradMagnitudeZeroPoint = 12.7

def vradErrorSkyAvg(vmag, spt):
  """
  Calculate radial velocity error from V and the spectral type. The value of the error is an average over
  the sky.

  Parameters
  ----------

  vmag - Value of V-band magnitude.
  spt  - String representing the spectral type of the star.

  Returns
  -------

  The radial velocity error in km/s.
  """
  return _vradCalibrationFloor + _vradErrorBCoeff[spt]*exp(_vradErrorACoeff[spt]*(vmag-_vradMagnitudeZeroPoint))


__all__ = ['plotCoordinateTransformationOnSky']

from pygaia.astrometry.coordinates import Transformations, CoordinateTransformation
from pygaia.utils import degreesToRadians, radiansToDegrees
from numpy import arange, argmax, roll, any, linspace, pi, zeros_like

from mpl_toolkits.basemap import Basemap
import matplotlib.pyplot as plt

def _orderGridlinePoints(x, y):
  """
  This code takes care of ordering the points (x,y), calculated for a sky map parallel or merdian, such
  that the drawing code can start at one end of the curve and end at the other (so no artifacts due to
  connecting the disjoint ends occur).

  Parameters
  ----------

  x - Set of x coordinates
  y - Set of y coordinates

  Returns
  -------

  x, y: Order set of (x,y) points
  """
  xroll=roll(x,1)
  yroll=roll(y,1)
  distance=(xroll-x)**2+(yroll-y)**2
  indexmax=argmax(distance)
  return roll(x,-indexmax), roll(y,-indexmax)

def plotCoordinateTransformationOnSky(transformation, outfile=None, myProjection='hammer',
    noTitle=False, noLabels=False, returnPlotObject=False):
    """
    Produce a sky-plot in a given coordinate system with the meridians and paralles for another
    coordinate system overlayed. The coordinate systems are specified through the
    pygaia.coordinates.Transformations enum. For example for Transformations.GAL2ECL the sky plot will be
    in Ecliptic coordinates with the Galactic coordinate grid overlayed.

    Keywords
    --------

    transformation - The coordinate transformation for which to make the plot (e.g.,
                     Transformations.GAL2ECL)
    outfile        - Save plot to this output file (default is to plot on screen). Make sure an extension
                     (.pdf, .png, etc) is included.
    myProjection   - Use this map projection (default is 'hammer', see basemap documentation)
    noTitle        - If true do not include the plot title.
    noLabels       - If true do not include plot labels.
    returnPlotObject - If true return the matplotlib object used for plotting. Further plot elements can
                       then be added.
    """
    ct = CoordinateTransformation(transformation)

    parallels=arange(-80.0,90.0,10.0)
    meridians=arange(0.0,375.0,15.0)
    meridianMax=degreesToRadians(85.0)
    parallelsMax=degreesToRadians(179.0)

    fig=plt.figure(figsize=(12,6))
    basemapInstance=Basemap(projection=myProjection,lon_0=0, celestial=True)
    basemapInstance.drawmapboundary()

    for thetaDeg in parallels:
      phi=linspace(-pi,pi,1001)
      theta=zeros_like(phi)+degreesToRadians(thetaDeg)
      phirot, thetarot = ct.transformSkyCoordinates(phi, theta)
      x ,y = basemapInstance(radiansToDegrees(phirot), radiansToDegrees(thetarot))

      indices=(phirot>=0.0)
      xplot=x[indices]
      yplot=y[indices]
      if any(indices):
        xplot, yplot = _orderGridlinePoints(xplot, yplot)
      plt.plot(xplot,yplot,'b-')

      indices=(phirot<0.0)
      xplot=x[indices]
      yplot=y[indices]
      if any(indices):
        xplot, yplot = _orderGridlinePoints(xplot, yplot)
      plt.plot(xplot,yplot,'b-')

    for phiDeg in meridians:
      theta=linspace(-meridianMax,meridianMax,1001)
      phi=zeros_like(theta)+degreesToRadians(phiDeg)
      phirot, thetarot = ct.transformSkyCoordinates(phi, theta)
      x ,y = basemapInstance(radiansToDegrees(phirot), radiansToDegrees(thetarot))

      indices=(phirot>=0.0)
      xplot=x[indices]
      yplot=y[indices]
      if any(indices):
        xplot, yplot = _orderGridlinePoints(xplot, yplot)
      plt.plot(xplot,yplot,'b-')

      indices=(phirot<0.0)
      xplot=x[indices]
      yplot=y[indices]
      if any(indices):
        xplot, yplot = _orderGridlinePoints(xplot, yplot)
      plt.plot(xplot,yplot,'b-')

    if (not noTitle):
      plt.title("Sky projection in " + ct.transformationStrings[1] + " coordinates with the corresponding " + ct.transformationStrings[0] + " grid overlayed")

    if (not noLabels):
      for theta in arange(-60,90,30):
        phirot, thetarot=ct.transformSkyCoordinates(0.0,degreesToRadians(theta))
        x, y = basemapInstance(radiansToDegrees(phirot), radiansToDegrees(thetarot))
        plt.text(x,y,"${0}$".format(theta),fontsize=16,va='bottom',ha='center',color='r')
      for phi in arange(-180,0,30):
        phirot, thetarot=ct.transformSkyCoordinates(degreesToRadians(phi),0.0)
        x, y = basemapInstance(radiansToDegrees(phirot), radiansToDegrees(thetarot))
        plt.text(x,y,"${0}$".format(phi),fontsize=16,va='bottom',ha='center',color='r')
      for phi in arange(30,180,30):
        phirot, thetarot=ct.transformSkyCoordinates(degreesToRadians(phi),0.0)
        x, y = basemapInstance(radiansToDegrees(phirot), radiansToDegrees(thetarot))
        plt.text(x,y,"${0}$".format(phi),fontsize=16,va='bottom',ha='center',color='r')

    if (outfile != None):
      plt.savefig(outfile)
    elif (returnPlotObject):
      return plt.gca(), basemapInstance
    else:
      plt.show()


#!/usr/bin/env python

# Predict parallax errors as a function of G and (V-I) according to the
# interpolation formula on the Gaia science performance web-pages.

import numpy as np
from pygaia.photometry.transformations import gminvFromVmini
from pygaia.errors.astrometric import parallaxErrorSkyAvg

from os import environ as env
try:
  import argparse
except ImportError:
  raise ImportError("""The argparse module does not seem to be available.\n Either upgrade to python 2.7
  or adapt the script to using the optparse module instead.""")

def calcParallaxError(args):
  """
  Calculate the parallax error for the given input source magnitude and colour.

  :argument args: command line arguments
  """
  gmag=float(args['gmag'])
  vmini=float(args['vmini'])
  sigmaPar=parallaxErrorSkyAvg(gmag, vmini)
  gminv=gminvFromVmini(vmini)
  print("G = {0}".format(gmag))
  print("V = {0}".format(gmag-gminv))
  print("(V-I) = {0}".format(vmini))
  print("(G-V) = {0}".format(gminv))
  print("standard error = {0} muas".format(sigmaPar))

def parseCommandLineArguments():
  """
  Set up command line parsing.
  """
  parser = argparse.ArgumentParser(description="Calculate parallax error for given G and (V-I)")
  parser.add_argument("gmag", help="G-band magnitude of source", type=float)
  parser.add_argument("vmini", help="(V-I) colour of source", type=float)

  args=vars(parser.parse_args())
  return args

if __name__ in ('__main__'):
  args=parseCommandLineArguments()
  calcParallaxError(args)


__all__ = ['gMagnitudeError', 'bpMagnitudeError', 'rpMagnitudeError']

from numpy import power, sqrt
from pygaia.errors.utils import calcZ, calcZBpRp

#
# Margin to include on predicted standard errors (i.e. multiply prediction by this value).
_scienceMargin = 1.2

#
# Mean number of CCDs crossed by a source in the AF field (G-band photometry)
_meanNumCcds = (7.0*9.0-1.0)/7.0

#
# End-of-mission CCD transit calibration floor on the photometric errors
_eomCalibrationFloorG = 3.0e-3/sqrt(_meanNumCcds)
_eomCalibrationFloorBP = 5.0e-3
_eomCalibrationFloorRP = 5.0e-3

def gMagnitudeError(G):
  """
  Calculate the single-field-of-view-transit photometric standard error in the G band as a function
  of G. A 20% margin is included.

  Parameters
  ----------

  G     - Value(s) of G-band magnitude.

  Returns
  -------

  The G band photometric standard error in units of magnitude.
  """
  z=calcZ(G)
  return 1.0e-3*sqrt(0.04895*z*z + 1.8633*z + 0.0001985) * _scienceMargin

def gMagnitudeErrorEoM(G, nobs=70):
  """
  Calculate the end of mission photometric standard error in the G band as a function
  of G. A 20% margin is included.

  Parameters
  ----------

  G     - Value(s) of G-band magnitude.

  Keywords
  --------

  nobs  - Number of observations collected (default 70).

  Returns
  -------

  The G band photometric standard error in units of magnitude.
  """
  return sqrt( (power(gMagnitudeError(G)/_scienceMargin,2) +
      _eomCalibrationFloorG*_eomCalibrationFloorG)/nobs ) * _scienceMargin

def bpMagnitudeError(G, vmini):
  """
  Calculate the single-field-of-view-transit photometric standard error in the BP band as a function
  of G and (V-I). Note: this refers to the integrated flux from the BP spectrophotometer. A margin of 20%
  is included.

  Parameters
  ----------

  G     - Value(s) of G-band magnitude.
  vmini - Value(s) of (V-I) colour.

  Returns
  -------

  The BP band photometric standard error in units of magnitude.
  """
  z=calcZBpRp(G)
  a	=	-0.000562*power(vmini,3) + 0.044390*vmini*vmini + 0.355123*vmini + 1.043270
  b	=	-0.000400*power(vmini,3) + 0.018878*vmini*vmini + 0.195768*vmini + 1.465592
  c	=	+0.000262*power(vmini,3) + 0.060769*vmini*vmini - 0.205807*vmini - 1.866968
  return 1.0e-3*sqrt(power(10.0,a)*z*z+power(10.0,b)*z+power(10.0,c))

def bpMagnitudeErrorEoM(G, vmini, nobs=70):
  """
  Calculate the end-of-mission photometric standard error in the BP band as a function of G and (V-I).
  Note: this refers to the integrated flux from the BP spectrophotometer. A margin of 20% is included.

  Parameters
  ----------

  G     - Value(s) of G-band magnitude.
  vmini - Value(s) of (V-I) colour.

  Keywords
  --------

  nobs  - Number of observations collected (default 70).

  Returns
  -------

  The BP band photometric standard error in units of magnitude.
  """
  return sqrt( (power(bpMagnitudeError(G, vmini)/_scienceMargin,2) +
      _eomCalibrationFloorBP*_eomCalibrationFloorBP)/nobs ) * _scienceMargin

def rpMagnitudeError(G, vmini):
  """
  Calculate the single-field-of-view-transit photometric standard error in the RP band as a function
  of G and (V-I). Note: this refers to the integrated flux from the RP spectrophotometer. A margin of 20%
  is included.

  Parameters
  ----------

  G     - Value(s) of G-band magnitude.
  vmini - Value(s) of (V-I) colour.

  Returns
  -------

  The RP band photometric standard error in units of magnitude.
  """
  z=calcZBpRp(G)
  a	=	-0.007597*power(vmini,3) + 0.114126*vmini*vmini - 0.636628*vmini + 1.615927
  b	=	-0.003803*power(vmini,3) + 0.057112*vmini*vmini - 0.318499*vmini + 1.783906
  c	=	-0.001923*power(vmini,3) + 0.027352*vmini*vmini - 0.091569*vmini - 3.042268
  return 1.0e-3*sqrt(power(10.0,a)*z*z+power(10.0,b)*z+power(10.0,c))

def rpMagnitudeErrorEoM(G, vmini, nobs=70):
  """
  Calculate the end-of-mission photometric standard error in the RP band as a function of G and (V-I).
  Note: this refers to the integrated flux from the RP spectrophotometer. A margin of 20% is included.

  Parameters
  ----------

  G     - Value(s) of G-band magnitude.
  vmini - Value(s) of (V-I) colour.

  Keywords
  --------

  nobs  - Number of observations collected (default 70).

  Returns
  -------

  The RP band photometric standard error in units of magnitude.
  """
  return sqrt( (power(rpMagnitudeError(G, vmini)/_scienceMargin,2) +
      _eomCalibrationFloorRP*_eomCalibrationFloorRP)/nobs ) * _scienceMargin


#!/usr/bin/env python

# Plot photometric errors as a function of V according to the interpolation formulae on the Gaia
# science performance web-pages.

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from pygaia.errors.photometric import gMagnitudeError, bpMagnitudeError, rpMagnitudeError
from pygaia.errors.photometric import gMagnitudeErrorEoM, bpMagnitudeErrorEoM, rpMagnitudeErrorEoM
from pygaia.photometry.transformations import gminvFromVmini

from os import environ as env
from matplotlib import rc
try:
  import argparse
except ImportError:
  raise ImportError("""The argparse module does not seem to be available.\n Either upgrade to python 2.7
  or adapt the script to using the optparse module instead.""")

# Configure matplotlib
rc('text', usetex=True)
rc('font', family='serif', size=16)
rc('xtick.major', size='12')
rc('xtick.minor', size='6')
rc('ytick.major', size='12')
rc('ytick.minor', size='6')
rc('lines', linewidth=2)
rc('axes', linewidth=2)

def makePlot(args):
  """
  Make the plot with photometry performance predictions.

  :argument args: command line arguments
  """
  gmag=np.linspace(3.0,20.0,171)

  vmini = args['vmini']

  vmag=gmag-gminvFromVmini(vmini)

  if args['eom']:
      sigmaG = gMagnitudeErrorEoM(gmag)
      sigmaGBp = bpMagnitudeErrorEoM(gmag, vmini)
      sigmaGRp = rpMagnitudeErrorEoM(gmag, vmini)
      yminmax = (1.0-4,0.1)
  else:
      sigmaG = gMagnitudeError(gmag)
      sigmaGBp = bpMagnitudeError(gmag, vmini)
      sigmaGRp = rpMagnitudeError(gmag, vmini)
      yminmax = (1.0-4,1)

  fig=plt.figure(figsize=(10,6.5))

  if (args['vmagAbscissa']):
    plt.semilogy(vmag, sigmaG, 'k', label='$\\sigma_G$')
    plt.semilogy(vmag, sigmaGBp, 'b', label='$\\sigma_{G_\\mathrm{BP}}$'+' for $(V-I)={0}$'.format(vmini))
    plt.semilogy(vmag, sigmaGRp, 'r', label='$\\sigma_{G_\\mathrm{RP}}$'+' for $(V-I)={0}$'.format(vmini))
    plt.xlim((6,20))
    #plt.ylim(yminmax)
    plt.legend(loc=0)
    plt.xlabel('$V$ [mag]')
  else:
    ax=fig.add_subplot(111)
    plt.semilogy(gmag, sigmaG, 'k', label='$\\sigma_G$')
    plt.semilogy(gmag, sigmaGBp, 'b', label='$\\sigma_{G_\\mathrm{BP}}$'+' for $(V-I)={0}$'.format(vmini))
    plt.semilogy(gmag, sigmaGRp, 'r', label='$\\sigma_{G_\\mathrm{RP}}$'+' for $(V-I)={0}$'.format(vmini))
    plt.xlim((6,20))
    #plt.ylim(yminmax)
    plt.legend(loc=0)
    plt.xlabel('$G$ [mag]')

  plt.xticks(np.arange(6,20,2))
  ax = plt.gca().yaxis
  #ax.set_major_formatter(matplotlib.ticker.ScalarFormatter())
  #plt.ticklabel_format(axis='y',style='plain')
  plt.grid(which='both')
  plt.ylabel('Photometric error [mag]')
  if args['eom']:
      plt.title('End-of-mission mean photometry: sky averaged errors for $(V-I)={0}$'.format(vmini), fontsize=14)
  else:
      plt.title('Single-FoV-transit photometry: sky averaged errors for $(V-I)={0}$'.format(vmini), fontsize=14)

  basename = 'PhotometricErrors'
  if (args['pdfOutput']):
    plt.savefig(basename+'.pdf')
  elif (args['pngOutput']):
    plt.savefig(basename+'.png')
  else:
    plt.show()

def parseCommandLineArguments():
  """
  Set up command line parsing.
  """
  parser = argparse.ArgumentParser(description="Plot predicted single-transit Gaia sky averaged photometric errors as a function of V or G")
  parser.add_argument("vmini", help="""(V-I) colour of source""", type=float)
  parser.add_argument("-e", action="store_true", dest="eom", help="Plot end-of-mission errors")
  parser.add_argument("-p", action="store_true", dest="pdfOutput", help="Make PDF plot")
  parser.add_argument("-b", action="store_true", dest="pngOutput", help="Make PNG plot")
  parser.add_argument("-v", action="store_true", dest="vmagAbscissa", help="Plot performance vs V instead of G")
  args=vars(parser.parse_args())
  return args

if __name__ in ('__main__'):
  args=parseCommandLineArguments()
  makePlot(args)


__all__ = ['calcZ', 'calcZBpRp', 'calcZAltStartGate', 'averageNumberOfTransits']

from scipy import isscalar
from numpy import sqrt, power, amax, array, floor, sin
from numpy import genfromtxt
from pkg_resources import resource_stream

_table = resource_stream('pygaia', 'data/errorFactorVariationBeta.txt')
_averageTransitNumber = genfromtxt(_table,
    skip_header=16, skip_footer=1,
    names=['sinBeta','betaMin','betaMax','nTransits',
      'alphaStar','delta','parallax','muAlphaStar','muDelta'])['nTransits']
_numStepsSinBeta = len(_averageTransitNumber)

def calcZ(G):
  """
  Calculate the value for the parameter z in the formula for parallax and G magnitude errors as a
  function of G and (V-I).

  Parameters
  ----------

  G - Value of G-band magnitude.

  Returns
  -------

  Value of z.
  """
  gatefloor=power(10.0,0.4*(12.0-15.0))
  if isscalar(G):
   result=amax((gatefloor,power(10.0,0.4*(G-15.0))))
  else :
    result=power(10.0,0.4*(G-15.0))
    indices=(result<gatefloor)
    result[indices]=gatefloor
  return result

def calcZBpRp(G):
  """
  Calculate the value for the parameter z in the formula for the BP and RP magnitude errors as a
  function of G and (V-I).

  Parameters
  ----------

  G - Value of G-band magnitude.

  Returns
  -------

  Value of z for BP/RP.
  """
  gatefloor=power(10.0,0.4*(11.0-15.0))
  if isscalar(G):
   result=amax((gatefloor,power(10.0,0.4*(G-15.0))))
  else :
    result=power(10.0,0.4*(G-15.0))
    indices=(result<gatefloor)
    result[indices]=gatefloor
  return result

def calcZAltStartGate(G):
  """
  Calculate the value of z in the formula for the parallax errors. In this case assume gating starts at
  G=13.3 (to simulate bright star worst performance)

  Parameters
  ----------

  G - Value of G-band magnitude.

  Returns
  -------

  Value of z.
  """
  gatefloor=power(10.0,0.4*(13.3-15.0))
  if isscalar(G):
   result=amax((gatefloor,power(10.0,0.4*(G-15.0))))
  else :
    result=power(10.0,0.4*(G-15.0))
    indices=(result<gatefloor)
    result[indices]=gatefloor
  return result

def averageNumberOfTransits(beta):
  """
  Returns the number of transits across the Gaia focal plane averaged over ecliptic longitude.

  Parameters
  ----------

  beta - Value(s) of the Ecliptic latitude.

  Returns
  -------

  Average number of transits for the input values of beta.
  """
  indices = array(floor(abs(sin(beta))*_numStepsSinBeta), dtype=int)
  indices[(indices==_numStepsSinBeta)] = _numStepsSinBeta-1
  return _averageTransitNumber[indices]