__all__ = ['CoordinateTransformation', 'Transformations', 'EpochPropagation']

from pygaia.astrometry.vectorastrometry import sphericalToCartesian, cartesianToSpherical, \
        elementaryRotationMatrix, normalTriad
from pygaia.utils import enum, degreesToRadians, radiansToDegrees
from pygaia.astrometry.constants import auKmYearPerSec

from numpy import ones_like, array, pi, cos, sin, zeros_like, zeros, arccos, select
from numpy import dot, transpose, cross, vstack, diag, sqrt, identity, tile, sum, arctan
from numpy import matmul, newaxis, squeeze
from numpy.linalg import norm
from scipy import isscalar

# Obliquity of the Ecliptic (arcsec)
_obliquityOfEcliptic = degreesToRadians(84381.41100/3600.0)

# Galactic pole in ICRS coordinates (see Hipparcos Explanatory Vol 1 section 1.5, and Murray, 1983,
# section 10.2)
_alphaGalPole = degreesToRadians(192.85948)
_deltaGalPole = degreesToRadians(27.12825)
# The galactic longitude of the ascending node of the galactic plane on the equator of ICRS (see
# Hipparcos Explanatory Vol 1 section 1.5, and Murray, 1983, section 10.2)
_omega = degreesToRadians(32.93192)

# Rotation matrix for the transformation from ICRS to Galactic coordinates. See equation (4.25) in
# chapter 4.5 of "Astrometry for Astrophysics", 2012, van Altena et al.
_matA = elementaryRotationMatrix("z",pi/2.0+_alphaGalPole)
_matB = elementaryRotationMatrix("x",pi/2.0-_deltaGalPole)
_matC = elementaryRotationMatrix("z",-_omega)
_rotationMatrixIcrsToGalactic=dot(_matC,dot(_matB,_matA))

# Alternative way to calculate the rotation matrix from ICRS to Galactic coordinates. First calculate
# the vectors describing the Galactic coordinate reference frame expressed within the ICRS.
# _vecN = array([0,0,1])
# _vecG3 = array([cos(_alphaGalPole)*cos(_deltaGalPole), sin(_alphaGalPole)*cos(_deltaGalPole),
#   sin(_deltaGalPole)])
# _vecG0 = cross(_vecN,_vecG3)
# _vecG0 = _vecG0/sqrt(dot(_vecG0,_vecG0))
# _vecG1 = -sin(_omega)*cross(_vecG3,_vecG0)+cos(_omega)*_vecG0
# _vecG2 = cross(_vecG3,_vecG1)
# _rotationMatrixIcrsToGalactic=array([_vecG1,_vecG2,_vecG3])

# Rotation matrix for the transformation from Galactic to ICRS coordinates.
_rotationMatrixGalacticToIcrs = transpose(_rotationMatrixIcrsToGalactic)

# Rotation matrix for the transformation from Ecliptic to ICRS coordinates.
_rotationMatrixEclipticToIcrs = elementaryRotationMatrix("x", -1*(_obliquityOfEcliptic))

# Rotation matrix for the transformation from ICRS to Ecliptic coordinates.
_rotationMatrixIcrsToEcliptic = transpose(_rotationMatrixEclipticToIcrs)

# Rotation matrix for the transformation from Galactic to Ecliptic coordinates.
_rotationMatrixGalacticToEcliptic = dot(_rotationMatrixIcrsToEcliptic,_rotationMatrixGalacticToIcrs)

# Rotation matrix for the transformation from Ecliptic to Galactic coordinates.
_rotationMatrixEclipticToGalactic = transpose(_rotationMatrixGalacticToEcliptic)

Transformations = enum('Transformations', ['GAL2ICRS', 'ICRS2GAL', 'ECL2ICRS', 'ICRS2ECL', 'GAL2ECL', 'ECL2GAL'])

_rotationMatrixMap = {Transformations.GAL2ICRS:_rotationMatrixGalacticToIcrs,
    Transformations.ICRS2GAL:_rotationMatrixIcrsToGalactic,
    Transformations.ECL2ICRS:_rotationMatrixEclipticToIcrs,
    Transformations.ICRS2ECL:_rotationMatrixIcrsToEcliptic,
    Transformations.GAL2ECL:_rotationMatrixGalacticToEcliptic,
    Transformations.ECL2GAL:_rotationMatrixEclipticToGalactic}

_transformationStringMap = {Transformations.GAL2ICRS:("galactic", "ICRS"),
    Transformations.ICRS2GAL:("ICRS","galactic"),
    Transformations.ECL2ICRS:("ecliptic","ICRS"),
    Transformations.ICRS2ECL:("ICRS","ecliptic"),
    Transformations.GAL2ECL:("galactic","ecliptic"),
    Transformations.ECL2GAL:("ecliptic","galactic")}

def angularDistance(phi1, theta1, phi2, theta2):
    """
    Calculate the angular distance between pairs of sky coordinates.

    Parameters
    ----------

    phi1 : float
        Longitude of first coordinate (radians).
    theta1 : float
        Latitude of first coordinate (radians).
    phi2 : float
        Longitude of second coordinate (radians).
    theta2 : float
        Latitude of second coordinate (radians).

    Returns
    -------

    Angular distance in radians.
    """
    # Formula below is more numerically stable than arccos( sin(theta1)*sin(theta2) +
    # cos(phi2-phi1)*cos(theta1)*cos(theta2) )
    # See: https://en.wikipedia.org/wiki/Great-circle_distance
    return arctan( sqrt((cos(theta2)*sin(phi2-phi1))**2 +
        (cos(theta1)*sin(theta2)-sin(theta1)*cos(theta2)*cos(phi2-phi1))**2) / (sin(theta1)*sin(theta2) +
            cos(phi2-phi1)*cos(theta1)*cos(theta2)) )

class CoordinateTransformation:
    """
    Provides methods for carrying out transformations between different coordinate (reference) systems.
    Currently the following are supported:

    ICRS or Equatorial (right ascension, declination)
    Galactic (longitude, latitude)
    Ecliptic (longitude, latitude)

    The transformations can be applied to sky coordinates or Cartesian coordinates.
    """

    def __init__(self, desiredTransformation):
        """
        Class constructor/initializer. Sets the type of transformation to be used.

        Parameters
        ----------

        desiredTransformation - The kind of coordinate transformation that should be provided. For example
        Transformations.GAL2ECL
        """
        self.rotationMatrix=_rotationMatrixMap[desiredTransformation]
        self.transformationStrings=_transformationStringMap[desiredTransformation]

    def transformCartesianCoordinates(self, x, y, z):
        """
        Rotates Cartesian coordinates from one reference system to another using the rotation matrix with
        which the class was initialized. The inputs  can be scalars or 1-dimensional numpy arrays.

        Parameters
        ----------

        x - Value of X-coordinate in original reference system
        y - Value of Y-coordinate in original reference system
        z - Value of Z-coordinate in original reference system

        Returns
        -------

        xrot - Value of X-coordinate after rotation
        yrot - Value of Y-coordinate after rotation
        zrot - Value of Z-coordinate after rotation
        """
        xrot, yrot, zrot = dot(self.rotationMatrix,[x,y,z])
        return xrot, yrot, zrot

    def transformSkyCoordinates(self, phi, theta):
        """
        Converts sky coordinates from one reference system to another, making use of the rotation matrix with
        which the class was initialized. Inputs can be scalars or 1-dimensional numpy arrays.

        Parameters
        ----------

        phi   - Value of the azimuthal angle (right ascension, longitude) in radians.
        theta - Value of the elevation angle (declination, latitude) in radians.

        Returns
        -------

        phirot   - Value of the transformed azimuthal angle in radians.
        thetarot - Value of the transformed elevation angle in radians.
        """
        r=ones_like(phi)
        x, y, z = sphericalToCartesian(r, phi, theta)
        xrot, yrot, zrot = self.transformCartesianCoordinates(x, y, z)
        r, phirot, thetarot = cartesianToSpherical(xrot, yrot, zrot)
        return phirot, thetarot

    def transformProperMotions(self, phi, theta, muphistar, mutheta):
        """
        Converts proper motions from one reference system to another, using the prescriptions in section
        1.5.3 of the Hipparcos Explanatory Volume 1 (equations 1.5.18, 1.5.19).

        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).
        muphistar - Value of the proper motion in the longitude-like angle, multiplied by cos(latitude).
        mutheta   - Value of the proper motion in the latitude-like angle.

        Returns
        -------

        muphistarrot - Value of the transformed proper motion in the longitude-like angle (including the
        cos(latitude) factor).
        muthetarot   - Value of the transformed proper motion in the latitude-like angle.
        """
        c, s = self._getJacobian(phi,theta)
        return c*muphistar+s*mutheta, c*mutheta-s*muphistar

    def transformSkyCoordinateErrors(self, phi, theta, sigPhiStar, sigTheta, rhoPhiTheta=0):
        """
        Converts the sky coordinate errors from one reference system to another, including the covariance
        term. Equations (1.5.4) and (1.5.20) from section 1.5 in the Hipparcos Explanatory Volume 1 are used.

        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).
        sigPhiStar  - Standard error in the longitude-like angle of the position of the source (radians or
                      sexagesimal units, including cos(latitude) term)
        sigTheta    - Standard error in the latitude-like angle of the position of the source (radians or
                      sexagesimal units)

        Keywords (optional)
        -------------------

        rhoPhiTheta - Correlation coefficient of the position errors. Set to zero if this keyword is not
                      provided.

        Retuns
        ------

        sigPhiRotStar  - The transformed standard error in the longitude-like angle (including
                         cos(latitude) factor)
        sigThetaRot    - The transformed standard error in the latitude-like angle.
        rhoPhiThetaRot - The transformed correlation coefficient.
        """
        if isscalar(rhoPhiTheta) and not isscalar(sigTheta):
            rhoPhiTheta=zeros_like(sigTheta)+rhoPhiTheta
        c, s = self._getJacobian(phi,theta)
        cSqr = c*c
        sSqr = s*s
        covar = sigPhiStar*sigTheta*rhoPhiTheta
        varPhiStar = sigPhiStar*sigPhiStar
        varTheta = sigTheta*sigTheta
        varPhiStarRot = cSqr*varPhiStar+sSqr*varTheta+2.0*covar*c*s
        varThetaRot = sSqr*varPhiStar+cSqr*varTheta-2.0*covar*c*s
        covarRot = (cSqr-sSqr)*covar+c*s*(varTheta-varPhiStar)
        return sqrt(varPhiStarRot), sqrt(varThetaRot), covarRot/sqrt(varPhiStarRot*varThetaRot)

    def transformProperMotionErrors(self, phi, theta, sigMuPhiStar, sigMuTheta, rhoMuPhiMuTheta=0):
        """
        Converts the proper motion errors from one reference system to another, including the covariance
        term. Equations (1.5.4) and (1.5.20) from section 1.5 in the Hipparcos Explanatory Volume 1 are used.

        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).
        sigMuPhiStar    - Standard error in the proper motion in the longitude-like direction (including
        cos(latitude) factor).
        sigMuTheta      - Standard error in the proper motion in the latitude-like direction.

        Keywords (optional)
        -------------------

        rhoMuPhiMuTheta - Correlation coefficient of the proper motion errors. Set to zero if this
                          keyword is not provided.

        Retuns
        ------

        sigMuPhiRotStar    - The transformed standard error in the proper motion in the longitude direction
        (including cos(latitude) factor).
        sigMuThetaRot      - The transformed standard error in the proper motion in the longitude direction.
        rhoMuPhiMuThetaRot - The transformed correlation coefficient.
        """
        return self.transformSkyCoordinateErrors(phi, theta, sigMuPhiStar, sigMuTheta,
                rhoPhiTheta=rhoMuPhiMuTheta)

    def transformCovarianceMatrix(self, phi, theta, covmat):
        """
        Transform the astrometric covariance matrix to its representation in the new coordinate system.

        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).
        covmat    - Covariance matrix (5x5) of the astrometric parameters.

        Returns
        -------

        covmat_rot - Covariance matrix in its representation in the new coordinate system.
        """

        c, s = self._getJacobian(phi,theta)
        jacobian = identity(5)
        jacobian[0][0]=c
        jacobian[1][1]=c
        jacobian[3][3]=c
        jacobian[4][4]=c
        jacobian[0][1]=s
        jacobian[1][0]=-s
        jacobian[3][4]=s
        jacobian[4][3]=-s

        return dot( dot(jacobian, covmat), jacobian.T )

    def _getJacobian(self, phi, theta):
        """
        Calculates the Jacobian for the transformation of the position errors and proper motion errors
        between coordinate systems. This Jacobian is also the rotation matrix for the transformation of
        proper motions. See section 1.5.3 of the Hipparcos Explanatory Volume 1 (equation 1.5.20). This
        matrix has the following form:

            |  c  s |
        J = |       |
            | -s  c |

        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).

        Returns
        -------

        c, s - The Jacobian matrix elements c and s corresponding to (phi, theta) and the currently
               desired coordinate system transformation.
        """

        p, q, r = normalTriad(phi, theta)

        # zRot = z-axis of new coordinate system expressed in terms of old system
        zRot = self.rotationMatrix[2,:]

        if (p.ndim == 2):
            zRotAll = tile(zRot, p.shape[1]).reshape(p.shape[1],3)
            pRot = cross(zRotAll, r.T)
            normPRot = norm(pRot,axis=1)
            for i in range(pRot.shape[0]):
                pRot[i] = pRot[i]/normPRot[i]
            c = zeros(pRot.shape[0])
            s = zeros(pRot.shape[0])
            for i in range(pRot.shape[0]):
                c[i] = dot(pRot[i], p.T[i])
                s[i] = dot(pRot[i], q.T[i])
            return c, s
        else:
            pRot = cross(zRot, r.T)
            pRot = pRot/norm(pRot)
            return dot(pRot,p), dot(pRot,q)

class EpochPropagation:
    """
    Provides methods for transforming the astrometry and radial velocity of a given source to a different
    epoch. The formulae for rigorous epoch transformation can be found on the Gaia documentation pages:
    http://gea.esac.esa.int/archive/documentation/GDR2/Data_processing/chap_cu3ast/sec_cu3ast_intro/ssec_cu3ast_intro_tansforms.html#SSS3.P1
    """

    def __init__(self):
        self.mastorad = pi/(180*3600*1000)

    def propagate_astrometry(self, phi, theta, parallax, muphistar, mutheta, vrad, t0, t1):
        """
        Propagate the position of a source from the reference epoch t0 to the new epoch t1.

        Parameters
        ----------

        phi : float
            Longitude at reference epoch (radians).
        theta : float
            Latitude at reference epoch (radians).
        parallax : float
            Parallax at the reference epoch (mas).
        muphistar : float
            Proper motion in longitude (including cos(latitude) term) at reference epoch (mas/yr).
        mutheta : float
            Proper motion in latitude at reference epoch (mas/yr).
        vrad : float
            Radial velocity at reference epoch (km/s).
        t0 : float
            Reference epoch (Julian years).
        t1 : float
            New epoch (Julian years).

        Returns
        -------

        Astrometric parameters, including the "radial proper motion" (NOT the radial velocity), at the new epoch.
        phi1, theta1, parallax1, muphistar1, mutheta1, mur1 = epoch_prop_pos(..., t0, t1)
        """

        t = t1-t0
        p0, q0, r0 = normalTriad(phi, theta)

        # Convert input data to units of radians and Julian year. Use ICRS coordinate names internally to
        # avoid errors in translating the formulae to code.
        pmra0 = muphistar*self.mastorad
        pmdec0 = mutheta*self.mastorad
        plx0 = parallax*self.mastorad
        pmr0 = vrad*parallax/auKmYearPerSec*self.mastorad
        pmtot0sqr = (muphistar**2 + mutheta**2) * self.mastorad**2

        # Proper motion vector
        pmvec0 = pmra0*p0+pmdec0*q0

        f = (1 + 2*pmr0*t + (pmtot0sqr+pmr0**2)*t**2)**(-0.5)
        u = (r0*(1+pmr0*t) + pmvec0*t)*f

        _, phi1, theta1 = cartesianToSpherical(u[0], u[1], u[2])
        parallax1 = parallax*f
        pmr1 = (pmr0+(pmtot0sqr + pmr0**2)*t)*f**2
        pmvec1 = (pmvec0*(1+pmr0*t) - r0*pmr0**2*t)*f**3
        p1, q1, r1 = normalTriad(phi1, theta1)
        muphistar1 = sum(p1*pmvec1/self.mastorad, axis=0)
        mutheta1 = sum(q1*pmvec1/self.mastorad, axis =0)
        murad1 = pmr1/self.mastorad

        return phi1, theta1, parallax1, muphistar1, mutheta1, murad1

    def propagate_pos(self, phi, theta, parallax, muphistar, mutheta, vrad, t0, t1):
        """
        Propagate the position of a source from the reference epoch t0 to the new epoch t1.

        Parameters
        ----------

        phi : float
            Longitude at reference epoch (radians).
        theta : float
            Latitude at reference epoch (radians).
        parallax : float
            Parallax at the reference epoch (mas).
        muphistar : float
            Proper motion in longitude (including cos(latitude) term) at reference epoch (mas/yr).
        mutheta : float
            Proper motion in latitude at reference epoch (mas/yr).
        vrad : float
            Radial velocity at reference epoch (km/s).
        t0 : float
            Reference epoch (Julian years).
        t1 : float
            New epoch (Julian years).

        Returns
        -------

        Coordinates phi and theta at new epoch (in radians)
        """
        phi1, theta1, parallax1, muphistar1, mutheta1, vrad1 = self.epoch_prop_astrometry(phi, theta, parallax, muphistar, mutheta, vrad, t0, t1)
        return phi1, theta1

    def propagate_astrometry_and_covariance_matrix(self, a0, c0, t0, t1):
        """
        Propagate the covariance matrix of the astrometric parameters and radial proper motion of a
        source from epoch t0 to epoch t1.

        Code based on the Hipparcos Fortran implementation by Lennart Lindegren.

        Parameters
        ----------

        a0 : array_like
            6-element vector: (phi, theta, parallax, muphistar, mutheta, vrad) in units of (radians,
            radians, mas, mas/yr, mas/yr, km/s). Shape of a should be (6,) or (6,N), with N the number of
            sources for which the astrometric parameters are provided.

        c0 : array_like
            Covariance matrix stored in a 6x6 element array. This can be constructed from the columns
            listed in the Gaia catalogue. The units are [mas^2, mas^2/yr, mas^2/yr^2] for the various
            elements. Note that the elements in the 6th row and column should be:
            c[6,i]=c[i,6]=c[i,3]*vrad/auKmYearPerSec for i=1,..,5 and
            c[6,6]=c[3,3]*(vrad^2+vrad_error^2)/auKmYearPerSec^2+(parallax*vrad_error/auKmYearPerSec)^2

            Shape of c0 should be (6,6) or (N,6,6).

        t0 : float
            Reference epoch (Julian years).

        t1 : float
            New epoch (Julian years).

        Returns
        -------

        Astrometric parameters, including the "radial proper motion" (NOT the radial velocity), and
        covariance matrix at the new epoch as a 2D matrix with the new variances on the diagional and the
        covariance in the off-diagonal elements.
        """

        zero, one, two, three = 0, 1, 2, 3
        tau = t1-t0

        # Calculate the normal triad [p0 q0 r0] at t0
        p0, q0, r0 = normalTriad(a0[0], a0[1])

        # Convert to internal units (radians, Julian year)
        par0 = a0[2]*self.mastorad
        pma0 = a0[3]*self.mastorad
        pmd0 = a0[4]*self.mastorad
        pmr0 = a0[5]*a0[2]/auKmYearPerSec*self.mastorad

        # Proper motion vector
        pmvec0 = pma0*p0+pmd0*q0

        # Auxiliary quantities
        tau2 = tau*tau
        pm02 = pma0**2 + pmd0**2
        w = one + pmr0*tau
        f2 = one/(one + two*pmr0*tau + (pm02+pmr0**2)*tau2)
        f = sqrt(f2)
        f3 = f2*f
        f4 = f2*f2

        # Position vector and parallax at t1
        u = (r0*w + pmvec0*tau)*f
        _, ra, dec = cartesianToSpherical(u[0], u[1], u[2])
        par = par0*f

        # Proper motion vector and radial proper motion at t1
        pmvec = (pmvec0*(one+pmr0*tau) - r0*pmr0**2*tau)*f3
        pmr = (pmr0+(pm02 + pmr0**2)*tau)*f2

        # Normal triad at t1
        p, q, r = normalTriad(ra, dec)

        # Convert parameters at t1 to external units (mas, Julian year)
        pma = sum(p*pmvec, axis=0)
        pmd = sum(q*pmvec, axis =0)

        a = zeros_like(a0)
        a[0] = ra
        a[1] = dec
        a[2] = par/self.mastorad
        a[3] = pma/self.mastorad
        a[4] = pmd/self.mastorad
        a[5] = pmr/self.mastorad

        # Auxiliary quantities for the partial derivatives

        pmz = pmvec0*f - three*pmvec*w
        pp0 = sum(p*p0, axis=0)
        pq0 = sum(p*q0, axis=0)
        pr0 = sum(p*r0, axis=0)
        qp0 = sum(q*p0, axis=0)
        qq0 = sum(q*q0, axis=0)
        qr0 = sum(q*r0, axis=0)
        ppmz = sum(p*pmz, axis=0)
        qpmz = sum(q*pmz, axis=0)

        J = zeros_like(c0)
        if (c0.ndim==2):
            J = J[newaxis,:,:]

        # Partial derivatives
        J[:,0,0] = pp0*w*f - pr0*pma0*tau*f
        J[:,0,1] = pq0*w*f - pr0*pmd0*tau*f
        J[:,0,2] = zero
        J[:,0,3] = pp0*tau*f
        J[:,0,4] = pq0*tau*f
        J[:,0,5] = -pma*tau2

        J[:,1,0] = qp0*w*f - qr0*pma0*tau*f
        J[:,1,1] = qq0*w*f - qr0*pmd0*tau*f
        J[:,1,2] = zero
        J[:,1,3] = qp0*tau*f
        J[:,1,4] = qq0*tau*f
        J[:,1,5] = -pmd*tau2

        J[:,2,0] = zero
        J[:,2,1] = zero
        J[:,2,2] = f
        J[:,2,3] = -par*pma0*tau2*f2
        J[:,2,4] = -par*pmd0*tau2*f2
        J[:,2,5] = -par*w*tau*f2

        J[:,3,0] = -pp0*pm02*tau*f3 - pr0*pma0*w*f3
        J[:,3,1] = -pq0*pm02*tau*f3 - pr0*pmd0*w*f3
        J[:,3,2] = zero
        J[:,3,3] = pp0*w*f3 - two*pr0*pma0*tau*f3 - three*pma*pma0*tau2*f2
        J[:,3,4] = pq0*w*f3 - two*pr0*pmd0*tau*f3 - three*pma*pmd0*tau2*f2
        J[:,3,5] = ppmz*tau*f2

        J[:,4,0] = -qp0*pm02*tau*f3 - qr0*pma0*w*f3
        J[:,4,1] = -qq0*pm02*tau*f3 - qr0*pmd0*w*f3
        J[:,4,2] = zero
        J[:,4,3] = qp0*w*f3 - two*qr0*pma0*tau*f3 - three*pmd*pma0*tau2*f2
        J[:,4,4] = qq0*w*f3 - two*qr0*pmd0*tau*f3 - three*pmd*pmd0*tau2*f2
        J[:,4,5] = qpmz*tau*f2

        J[:,5,0] = zero
        J[:,5,1] = zero
        J[:,5,2] = zero
        J[:,5,3] = two*pma0*w*tau*f4
        J[:,5,4] = two*pmd0*w*tau*f4
        J[:,5,5] = (w**2 - pm02*tau2)*f4

        JT = zeros_like(J)
        for i in range(J.shape[0]):
            JT[i] = J[i].T

        if (c0.ndim==2):
            c = matmul(J,matmul(c0[newaxis,:,:],JT))
        else:
            c = matmul(J,matmul(c0,JT))

        return a, squeeze(c)


"""
Plot the radial velocity horizons for stars of various spectral types. The radial velocity horizon
defines out to what distance the star can be seen for a given radial velocity accuracy.

Anthony Brown 2014
"""

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from pygaia.errors.spectroscopic import vradErrorSkyAvg
from pygaia.photometry.utils import vminiFromSpt, gabsFromSpt, vabsFromSpt
from pygaia.photometry.transformations import vminGrvsFromVmini

from os import environ as env
from matplotlib import rc
from matplotlib.colors import hsv_to_rgb

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=18)
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 horizons. The plot shows V-band magnitude vs distance for a number of
  spectral types and over the range 5.7<G<20. In addition a set of crudely drawn contours show the points
  where 0.1, 1, and 10 per cent relative parallax accracy are reached.

  Parameters
  ----------

  args - Command line arguments.
  """
  distances = 10.0**np.linspace(1,6,10001)

  spts=['B0V', 'A0V', 'F0V', 'G0V', 'K0V', 'K4V', 'K1III']
  twokmsRV = []
  twokmsV = []
  vabsTwokms = []
  fivekmsRV = []
  fivekmsV = []
  vabsFivekms = []
  tenkmsRV = []
  tenkmsV = []
  vabsTenkms = []

  fig=plt.figure(figsize=(11,7.8))
  deltaHue = 240.0/(len(spts)-1)
  hues = (240.0-np.arange(len(spts))*deltaHue)/360.0
  hsv=np.zeros((1,1,3))
  hsv[0,0,1]=1.0
  hsv[0,0,2]=0.9
  for hue,spt in zip(hues, spts):
    hsv[0,0,0]=hue
    vmags = vabsFromSpt(spt)+5.0*np.log10(distances)-5.0
    vmini=vminiFromSpt(spt)
    grvsmags = vmags - vminGrvsFromVmini(vmini)
    rvError = vradErrorSkyAvg(vmags, spt)
    observed = (grvsmags>=5.7) & (grvsmags<=16.1)
    rvError = rvError[observed]
    # Identify the points where the relative parallax accuracy is 0.1, 1, or 10 per cent.
    if (rvError.min()<=2.0):
      index = len(rvError[rvError<=2.0])-1
      twokmsRV.append(distances[observed][index])
      twokmsV.append(vmags[observed][index])
      vabsTwokms.append(vabsFromSpt(spt))
    if (rvError.min()<=5.0):
      index = len(rvError[rvError<=5.0])-1
      fivekmsRV.append(distances[observed][index])
      fivekmsV.append(vmags[observed][index])
      vabsFivekms.append(vabsFromSpt(spt))
    if (rvError.min()<=10.0):
      index = len(rvError[rvError<=10.0])-1
      tenkmsRV.append(distances[observed][index])
      tenkmsV.append(vmags[observed][index])
      vabsTenkms.append(vabsFromSpt(spt))
    plt.semilogx(distances[observed], vmags[observed], '-', label=spt, color=hsv_to_rgb(hsv)[0,0,:])
    plt.text(distances[observed][-1], vmags[observed][-1], spt, horizontalalignment='center',
        verticalalignment='bottom', fontsize=14)

  # Draw the "contours" of constant radial velocity accuracy.
  twokmsRV = np.array(twokmsRV)
  twokmsV = np.array(twokmsV)
  indices = np.argsort(vabsTwokms)
  plt.semilogx(twokmsRV[indices],twokmsV[indices],'k--')
  plt.text(twokmsRV[indices][-1]*0.8,twokmsV[indices][-1],"$2$ km s$^{-1}$", ha='right', size=16,
      bbox=dict(boxstyle="round, pad=0.3", ec=(0.0, 0.0, 0.0), fc=(1.0, 1.0, 1.0),))

  fivekmsRV = np.array(fivekmsRV)
  fivekmsV = np.array(fivekmsV)
  indices = np.argsort(vabsFivekms)
  plt.semilogx(fivekmsRV[indices],fivekmsV[indices],'k--')
  plt.text(fivekmsRV[indices][-1]*0.8,fivekmsV[indices][-1],"$5$ km s$^{-1}$", ha='right', size=16,
      bbox=dict(boxstyle="round, pad=0.3", ec=(0.0, 0.0, 0.0), fc=(1.0, 1.0, 1.0),))

  tenkmsRV = np.array(tenkmsRV)
  tenkmsV = np.array(tenkmsV)
  indices = np.argsort(vabsTenkms)
  plt.semilogx(tenkmsRV[indices],tenkmsV[indices],'k--')
  plt.text(tenkmsRV[indices][-1]*0.8,tenkmsV[indices][-1]+0.5,"$10$ km s$^{-1}$", ha='right', size=16,
      bbox=dict(boxstyle="round, pad=0.3", ec=(0.0, 0.0, 0.0), fc=(1.0, 1.0, 1.0),))

  plt.title('Radial velocity accuracy horizons ($A_V=0$)')

  plt.xlabel('Distance [pc]')
  plt.ylabel('V')
  plt.grid()
  #leg=plt.legend(loc=4, fontsize=14, labelspacing=0.5)
  plt.ylim(5,20)

  basename='RadialVelocityHorizons'
  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 radial velocity horizons for various spectral types.")
  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()
  makePlot(args)


__all__ = ['parallaxErrorSkyAvg', 'parallaxErrorSkyAvgAltStartGate', 'positionErrorSkyAvg',
           'properMotionErrorSkyAvg', 'parallaxError', 'positionError', 'properMotionError']

from numpy import sqrt, sin, array, floor, power
from scipy import isscalar
from pygaia.errors.utils import calcZ, calcZAltStartGate
from numpy import genfromtxt
from pkg_resources import resource_stream

_table = resource_stream('pygaia', 'data/errorFactorVariationBetaIncludingNumberOfTransits.txt')
_astrometricErrorFactors = genfromtxt(_table,
    skip_header=4, skip_footer=1,
    names=['sinBeta','alphaStar','delta','parallax','muAlphaStar','muDelta'])
_numStepsSinBeta = len(_astrometricErrorFactors['sinBeta'])

# Scaling factors for sky averaged position and proper motion errors. The errors are scaled with respect
# to the parallax error values. Note that the error are quoted in true arc terms (using phi*) for the
# longitude like component.
_scalingForPositions = {'Total':0.743, 'AlphaStar':0.787, 'Delta':0.699}
_scalingForProperMotions = {'Total':0.526, 'AlphaStar':0.556, 'Delta':0.496}

# The upper band of the parallax errors at the bright end
_parallaxErrorMaxBright = 14.0

# The nominal mission lifetime
_nominalLifeTime = 5.0

def errorScalingFactor(observable, beta):
  """
  Look up the numerical factors to apply to the sky averaged parallax error in order to obtain error
  values for a given astrometric parameter, taking the Ecliptic latitude and the number of transits into
  account.

  Parameters
  ----------

  observable - Name of astrometric observable (one of: alphaStar, delta, parallax, muAlphaStar, muDelta)
  beta       - Values(s) of the Ecliptic latitude.

  Returns
  -------

  Numerical factors to apply to the errors of the given observable.
  """
  if isscalar(beta):
    index=int(floor(abs(sin(beta))*_numStepsSinBeta))
    if index == _numStepsSinBeta:
      return _astrometricErrorFactors[observable][_numStepsSinBeta-1]
    else:
      return _astrometricErrorFactors[observable][index]
  else:
    indices = array(floor(abs(sin(beta))*_numStepsSinBeta), dtype=int)
    indices[(indices==_numStepsSinBeta)] = _numStepsSinBeta-1
    return _astrometricErrorFactors[observable][indices]

def errorScalingMissionLength(extension, p):
    """
    Calculate the factor by which to scale errors for a given Gaia mission extension.

    Parameters
    ----------

    extension - The mission extension in years (a negative extension can be used to make crude
    performance predictions part-way through the mission lifetime).
    p - The power by which the errors scale with time (error ~ t^p, p=-0.5 for parallax and celestial
    position, and -1.5 for proper motion).

    Returns
    -------

    The factor by which to scale the errors
    """
    return power((_nominalLifeTime+extension)/_nominalLifeTime, p)

def parallaxErrorSkyAvg(G, vmini, extension=0.0):
  """
  Calculate the sky averaged parallax error from G and (V-I).

  Parameters
  ----------

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

  Keywords
  --------

  extension - Add this amount of years to the mission lifetime and scale the errors accordingly.

  Returns
  -------

  The parallax error in micro-arcseconds.
  """
  factor = errorScalingMissionLength(extension, -0.5)
  z=calcZ(G)
  return sqrt(-1.631 + 680.766*z + 32.732*z*z)*(0.986 + (1.0 - 0.986)*vmini)*factor

def parallaxMinError(G, vmini, extension=0.0):
  """
  Calculate the minimum parallax error from G and (V-I). This correspond to the sky regions with the
  smallest astrometric errors.  At the bright end the parallax error is at least 14 muas due to the
  gating scheme.

  Parameters
  ----------

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

  Keywords
  --------

  extension - Add this amount of years to the mission lifetime and scale the errors accordingly.

  Returns
  -------

  The minimum parallax error in micro-arcseconds.
  """
  return _astrometricErrorFactors["parallax"].min()*parallaxErrorSkyAvg(G, vmini, extension=extension)

def parallaxMaxError(G, vmini, extension=0.0):
  """
  Calculate the maximum parallax error from G and (V-I). This correspond to the sky regions with the
  largest astrometric errors.  At the bright end the parallax error is at least 14 muas due to the
  gating scheme.

  Parameters
  ----------

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

  Keywords
  --------

  extension - Add this amount of years to the mission lifetime and scale the errors accordingly.

  Returns
  -------

  The maximum parallax error in micro-arcseconds.
  """
  errors = _astrometricErrorFactors["parallax"].max()*parallaxErrorSkyAvg(G, vmini, extension=extension)
  indices = (errors<_parallaxErrorMaxBright)
  errors[indices]=_parallaxErrorMaxBright
  return errors

def parallaxError(G, vmini, beta, extension=0.0):
  """
  Calculate the parallax error from G and (V-I) and the Ecliptic latitude beta of the source. The
  parallax error is calculated by applying numerical factors to the sky average error. These factors
  depend on beta and the average number of transits across a source for that value of beta.

  The code implements table 6 from the Gaia science performance pages:
  http://www.rssd.esa.int/SA/GAIA/docs/SciencePerformance/table6.htm

  Parameters
  ----------

  G     - Value(s) of G-band magnitude.
  vmini - Value(s) of (V-I) colour.
  beta  - Value(s) of the Ecliptic latitude.

  Keywords
  --------

  extension - Add this amount of years to the mission lifetime and scale the errors accordingly.

  Returns
  -------

  The parallax error in micro-arcseconds.
  """
  return parallaxErrorSkyAvg(G, vmini, extension=extension)*errorScalingFactor('parallax',beta)

def parallaxErrorSkyAvgAltStartGate(G, vmini, extension=0.0):
  """
  Calculate the sky averaged parallax error from G and (V-I). In this case assume gating starts at G=13.3
  (to simulate bright star worst performance)

  Parameters
  ----------

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

  Keywords
  --------

  extension - Add this amount of years to the mission lifetime and scale the errors accordingly.

  Returns
  -------

  The parallax error in micro-arcseconds.
  """
  factor = errorScalingMissionLength(extension, -0.5)
  z=calcZAltStartGate(G)
  return sqrt(-1.631 + 680.766*z + 32.732*z*z)*(0.986 + (1.0 - 0.986)*vmini)*factor

def positionErrorSkyAvg(G, vmini, extension=0.0):
  """
  Calculate the sky averaged position errors from G and (V-I).

  NOTE! THE ERRORS ARE FOR SKY POSITIONS IN THE ICRS (I.E., RIGHT ASCENSION, DECLINATION). MAKE SURE YOUR
  SIMULATED ASTROMETRY IS ALSO ON THE ICRS.

  Parameters
  ----------

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

  Keywords
  --------

  extension - Add this amount of years to the mission lifetime and scale the errors accordingly.

  Returns
  -------

  The error in alpha* and the error in delta, in that order, in micro-arcsecond.
  """
  parallaxError = parallaxErrorSkyAvg(G, vmini, extension=extension)
  return _scalingForPositions['AlphaStar']*parallaxError, \
         _scalingForPositions['Delta']*parallaxError

def positionMinError(G, vmini, extension=0.0):
  """
  Calculate the minimum position errors from G and (V-I). These correspond to the sky regions with the
  smallest astrometric errors.

  NOTE! THE ERRORS ARE FOR SKY POSITIONS IN THE ICRS (I.E., RIGHT ASCENSION, DECLINATION). MAKE SURE YOUR
  SIMULATED ASTROMETRY IS ALSO ON THE ICRS.

  Parameters
  ----------

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

  Keywords
  --------

  extension - Add this amount of years to the mission lifetime and scale the errors accordingly.

  Returns
  -------

  The minimum error in alpha* and the error in delta, in that order, in micro-arcsecond.
  """
  parallaxError = parallaxErrorSkyAvg(G, vmini, extension=extension)
  return _astrometricErrorFactors['alphaStar'].min()*parallaxError, \
         _astrometricErrorFactors['delta'].min()*parallaxError

def positionMaxError(G, vmini, extension=0.0):
  """
  Calculate the maximum position errors from G and (V-I). These correspond to the sky regions with the
  largest astrometric errors.

  NOTE! THE ERRORS ARE FOR SKY POSITIONS IN THE ICRS (I.E., RIGHT ASCENSION, DECLINATION). MAKE SURE YOUR
  SIMULATED ASTROMETRY IS ALSO ON THE ICRS.

  Parameters
  ----------

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

  Keywords
  --------

  extension - Add this amount of years to the mission lifetime and scale the errors accordingly.

  Returns
  -------

  The maximum error in alpha* and the error in delta, in that order, in micro-arcsecond.
  """
  parallaxError = parallaxErrorSkyAvg(G, vmini, extension)
  return _astrometricErrorFactors['alphaStar'].max()*parallaxError, \
         _astrometricErrorFactors['delta'].max()*parallaxError

def positionError(G, vmini, beta, extension=0.0):
  """
  Calculate the position errors from G and (V-I) and the Ecliptic latitude beta of the source.

  NOTE! THE ERRORS ARE FOR SKY POSITIONS IN THE ICRS (I.E., RIGHT ASCENSION, DECLINATION). MAKE SURE YOUR
  SIMULATED ASTROMETRY IS ALSO ON THE ICRS.

  Parameters
  ----------

  G     - Value(s) of G-band magnitude.
  vmini - Value(s) of (V-I) colour.
  beta  - Value(s) of the Ecliptic latitude.

  Keywords
  --------

  extension - Add this amount of years to the mission lifetime and scale the errors accordingly.

  Returns
  -------

  The error in alpha* and the error in delta, in that order, in micro-arcsecond.
  """
  parallaxError = parallaxErrorSkyAvg(G, vmini, extension=extension)
  return errorScalingFactor('alphaStar',beta)*parallaxError, \
         errorScalingFactor('delta',beta)*parallaxError

def properMotionErrorSkyAvg(G, vmini, extension=0.0):
  """
  Calculate the sky averaged proper motion errors from G and (V-I).

  NOTE! THE ERRORS ARE FOR PROPER MOTIONS IN THE ICRS (I.E., RIGHT ASCENSION, DECLINATION). MAKE SURE
  YOUR SIMULATED ASTROMETRY IS ALSO ON THE ICRS.

  Parameters
  ----------

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

  Keywords
  --------

  extension - Add this amount of years to the mission lifetime and scale the errors accordingly.

  Returns
  -------

  The error in mu_alpha* and the error in mu_delta, in that order, in micro-arcsecond/year.
  """
  factor = errorScalingMissionLength(extension, -1.5)
  parallaxError = parallaxErrorSkyAvg(G, vmini)*factor
  return _scalingForProperMotions['AlphaStar']*parallaxError, \
         _scalingForProperMotions['Delta']*parallaxError

def properMotionMinError(G, vmini, extension=0.0):
  """
  Calculate the minimum proper motion errors from G and (V-I). These correspond to the sky regions with
  the smallest astrometric errors.

  NOTE! THE ERRORS ARE FOR PROPER MOTIONS IN THE ICRS (I.E., RIGHT ASCENSION, DECLINATION). MAKE SURE
  YOUR SIMULATED ASTROMETRY IS ALSO ON THE ICRS.

  Parameters
  ----------

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

  Keywords
  --------

  extension - Add this amount of years to the mission lifetime and scale the errors accordingly.

  Returns
  -------

  The minimum error in mu_alpha* and the error in mu_delta, in that order, in micro-arcsecond/year.
  """
  factor = errorScalingMissionLength(extension, -1.5)
  parallaxError = parallaxErrorSkyAvg(G, vmini)*factor
  return _astrometricErrorFactors['muAlphaStar'].min()*parallaxError, \
         _astrometricErrorFactors['muDelta'].min()*parallaxError

def properMotionMaxError(G, vmini, extension=0.0):
  """
  Calculate the maximum proper motion errors from G and (V-I). These correspond to the sky regions with
  the largest astrometric errors.

  NOTE! THE ERRORS ARE FOR PROPER MOTIONS IN THE ICRS (I.E., RIGHT ASCENSION, DECLINATION). MAKE SURE
  YOUR SIMULATED ASTROMETRY IS ALSO ON THE ICRS.

  Parameters
  ----------

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

  Keywords
  --------

  extension - Add this amount of years to the mission lifetime and scale the errors accordingly.

  Returns
  -------

  The maximum error in mu_alpha* and the error in mu_delta, in that order, in micro-arcsecond/year.
  """
  factor = errorScalingMissionLength(extension, -1.5)
  parallaxError = parallaxErrorSkyAvg(G, vmini)*factor
  indices = (parallaxError<_parallaxErrorMaxBright)
  parallaxError[indices] = _parallaxErrorMaxBright
  return _astrometricErrorFactors['muAlphaStar'].max()*parallaxError, \
         _astrometricErrorFactors['muDelta'].max()*parallaxError

def properMotionError(G, vmini, beta, extension=0.0):
  """
  Calculate the proper motion errors from G and (V-I) and the Ecliptic latitude beta of the source.

  NOTE! THE ERRORS ARE FOR PROPER MOTIONS IN THE ICRS (I.E., RIGHT ASCENSION, DECLINATION). MAKE SURE
  YOUR SIMULATED ASTROMETRY IS ALSO ON THE ICRS.

  Parameters
  ----------

  G     - Value(s) of G-band magnitude.
  vmini - Value(s) of (V-I) colour.
  beta  - Value(s) of the Ecliptic latitude.

  Keywords
  --------

  extension - Add this amount of years to the mission lifetime and scale the errors accordingly.

  Returns
  -------

  The error in mu_alpha* and the error in mu_delta, in that order, in micro-arcsecond/year.
  """
  factor = errorScalingMissionLength(extension, -1.5)
  parallaxError = parallaxErrorSkyAvg(G, vmini)*factor
  return errorScalingFactor('muAlphaStar',beta)*parallaxError, \
         errorScalingFactor('muDelta',beta)*parallaxError

def totalProperMotionErrorSkyAvg(G, vmini, extension=0.0):
  """
  Calculate the sky averaged total proper motion error from G and (V-I). This refers to the error on the
  length of the proper motion vector.

  NOTE! THE ERRORS ARE FOR PROPER MOTIONS IN THE ICRS (I.E., RIGHT ASCENSION, DECLINATION). MAKE SURE
  YOUR SIMULATED ASTROMETRY IS ALSO ON THE ICRS.

  Parameters
  ----------

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

  Keywords
  --------

  extension - Add this amount of years to the mission lifetime and scale the errors accordingly.

  Returns
  -------

  The error on the total proper motion in micro-arcsecond/year.
  """
  factor = errorScalingMissionLength(extension, -1.5)
  parallaxError = parallaxErrorSkyAvg(G, vmini)*factor
  return _scalingForProperMotions['Total']*parallaxError


"""
Plot the parallax horizons for stars of various spectral types. The parallax horizon defines out to what
distance the star can be seen for a given relative parallax accuracy.

Anthony Brown 2013
"""

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

from os import environ as env
from matplotlib import rc
from matplotlib.colors import hsv_to_rgb

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=18)
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 horizons. The plot shows V-band magnitude vs distance for a number of
  spectral types and over the range 5.7<G<20. In addition a set of crudely drawn contours show the points
  where 0.1, 1, and 10 per cent relative parallax accracy are reached.

  Parameters
  ----------

  args - Command line arguments.
  """
  distances = 10.0**np.linspace(1,6,10001)
  av = args['extinction']
  ai = 0.479*av #Cardelli et al R=3.1

  spts = ['B0I', 'B1V', 'G2V', 'K4V', 'M0V', 'M6V', 'K1III', 'M0III']
  pointOnePercD = []
  pointOnePercV = []
  onePercD = []
  onePercV = []
  tenPercD = []
  tenPercV = []
  vabsPointOnePerc = []
  vabsOnePerc = []
  vabsTenPerc = []

  fig=plt.figure(figsize=(11,7.8))
  deltaHue = 240.0/(len(spts)-1)
  hues = (240.0-np.arange(len(spts))*deltaHue)/360.0
  hsv=np.zeros((1,1,3))
  hsv[0,0,1]=1.0
  hsv[0,0,2]=0.9
  for hue,spt in zip(hues, spts):
    hsv[0,0,0]=hue
    vmags = vabsFromSpt(spt)+5.0*np.log10(distances)-5.0+av
    vmini=vminiFromSpt(spt)+av-ai
    #gmags = gabsFromSpt(spt)+5.0*np.log10(distances)-5.0
    gmags = vmags + gminvFromVmini(vmini)
    relParErr = parallaxErrorSkyAvg(gmags,vmini)*distances/1.0e6
    observed = (gmags>=5.7) & (gmags<=20.0)
    relParErrObs = relParErr[observed]
    # Identify the points where the relative parallax accuracy is 0.1, 1, or 10 per cent.
    if (relParErrObs.min()<0.001):
      index = len(relParErrObs[relParErrObs<=0.001])-1
      pointOnePercD.append(distances[observed][index])
      pointOnePercV.append(vmags[observed][index])
      vabsPointOnePerc.append(vabsFromSpt(spt))
    if (relParErrObs.min()<0.01):
      index = len(relParErrObs[relParErrObs<=0.01])-1
      onePercD.append(distances[observed][index])
      onePercV.append(vmags[observed][index])
      vabsOnePerc.append(vabsFromSpt(spt))
    if (relParErrObs.min()<0.1):
      index = len(relParErrObs[relParErrObs<=0.1])-1
      tenPercD.append(distances[observed][index])
      tenPercV.append(vmags[observed][index])
      vabsTenPerc.append(vabsFromSpt(spt))
    plt.semilogx(distances[observed], vmags[observed], '-', label=spt, color=hsv_to_rgb(hsv)[0,0,:])
    if (spt=='B0I'):
      plt.text(distances[observed][-1]-1.0e5, vmags[observed][-1], spt, horizontalalignment='right',
          verticalalignment='bottom', fontsize=14)
    else:
      plt.text(distances[observed][-1], vmags[observed][-1], spt, horizontalalignment='center',
          verticalalignment='bottom', fontsize=14)

  # Draw the "contours" of constant relative parallax accuracy.
  pointOnePercD = np.array(pointOnePercD)
  pointOnePercV = np.array(pointOnePercV)
  indices = np.argsort(vabsPointOnePerc)
  plt.semilogx(pointOnePercD[indices],pointOnePercV[indices],'k--')
  plt.text(pointOnePercD[indices][-1]*1.2,pointOnePercV[indices][-1]-2.5,"$0.1$\\%", ha='right', size=16,
      bbox=dict(boxstyle="round, pad=0.3", ec=(0.0, 0.0, 0.0), fc=(1.0, 1.0, 1.0),))

  onePercD = np.array(onePercD)
  onePercV = np.array(onePercV)
  indices = np.argsort(vabsOnePerc)
  plt.semilogx(onePercD[indices],onePercV[indices],'k--')
  plt.text(onePercD[indices][-1]*1.2,onePercV[indices][-1]-2.5,"$1$\\%", ha='right', size=16,
      bbox=dict(boxstyle="round, pad=0.3", ec=(0.0, 0.0, 0.0), fc=(1.0, 1.0, 1.0),))

  tenPercD = np.array(tenPercD)
  tenPercV = np.array(tenPercV)
  indices = np.argsort(vabsTenPerc)
  plt.semilogx(tenPercD[indices],tenPercV[indices],'k--')
  plt.text(tenPercD[indices][-1]*1.5,tenPercV[indices][-1]-2.5,"$10$\\%", ha='right', size=16,
      bbox=dict(boxstyle="round, pad=0.3", ec=(0.0, 0.0, 0.0), fc=(1.0, 1.0, 1.0),))

  plt.title('Parallax relative accuracy horizons ($A_V={0}$)'.format(av))

  plt.xlabel('Distance [pc]')
  plt.ylabel('V')
  plt.grid()
  #leg=plt.legend(loc=4, fontsize=14, labelspacing=0.5)
  plt.ylim(5,26)

  basename='ParallaxHorizons'
  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 parallax horizons for various spectral types.")
  parser.add_argument("-a", dest="extinction", type=float, help="Value of extinction Av", default=0.0)
  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()
  makePlot(args)


__all__ = ['vminiFromSpt', 'vabsFromSpt', 'gabsFromSpt']

from scipy import isscalar
from numpy import sqrt, power, amax
from pygaia.photometry.transformations import gminvFromVmini

_sptToVminiVabsDictionary={'B0V':(-0.31, -3.5), 'B1V':(-0.24,-2.7), 'B5V':(-0.08,0.0), 'A0V':(0.01,0.0),
'A5V':(0.16,1.69), 'F0V':(0.38,2.98), 'G0V':(0.67,4.24), 'G2V':(0.72,4.7), 'G5V':(0.74,4.78),
'K0V':(0.87,5.58), 'K1IIIMP':(0.99,1.53), 'K4V':(1.23,7.21), 'K1III':(1.04,2.16), 'M0V':(1.71,8.62),
'M2V':(2.02,9.48), 'M6V':(3.69,14.2), 'M0III':(1.65,-0.66), 'B0I':(-0.22,-7.1), 'B1I':(-0.16,-6.7)}

def vminiFromSpt(spt):
  """
  Obtain (V-I) for the input spectral type.

  Parameters
  ----------

  spt - String representing the spectral type of the star.

  Returns
  -------

  The value of (V-I).
  """
  if spt in _sptToVminiVabsDictionary:
    return _sptToVminiVabsDictionary[spt][0]
  else:
    message="Unknown spectral type. Allowed values are: "
    for key in _sptToVminiVabsDictionary.keys():
      message += key+" "
    raise Exception(message)

def vabsFromSpt(spt):
  """
  Obtain M_V (absolute magnitude in V-band) for the input spectral type.

  Parameters
  ----------

  spt - String representing the spectral type of the star.

  Returns
  -------

  The value of M_V.
  """
  if spt in _sptToVminiVabsDictionary:
    return _sptToVminiVabsDictionary[spt][1]
  else:
    message="Unknown spectral type. Allowed values are: "
    for key in _sptToVminiVabsDictionary.keys():
      message += key+" "
    raise Exception(message)

def gabsFromSpt(spt):
  """
  Obtain M_G (absolute magnitude in G-band) for the input spectral type.

  Parameters
  ----------

  spt - String representing the spectral type of the star.

  Returns
  -------

  The value of M_G.
  """
  if spt in _sptToVminiVabsDictionary:
    return vabsFromSpt(spt) + gminvFromVmini(vminiFromSpt(spt))
  else:
    message="Unknown spectral type. Allowed values are: "
    for key in _sptToVminiVabsDictionary.keys():
      message += key+" "
    raise Exception(message)


#!/usr/bin/env python

# Plot relative parallax errors as a function of distance for stars of a given spectral type.

import numpy as np
import matplotlib.pyplot as plt
from collections import OrderedDict
from pygaia.photometry.transformations import gminvFromVmini
from pygaia.errors.astrometric import parallaxErrorSkyAvg

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(pdf=False, png=False):
  """
  Plot relative parallax errors as a function of distance for stars of a given spectral type.

  Parameters
  ----------

  args - command line arguments
  """
  logdistancekpc = np.linspace(-1,np.log10(20.0),100)
  sptVabsAndVmini=OrderedDict([('K0V',(5.58,0.87)), ('G5V',(4.78,0.74)), ('G0V',(4.24,0.67)),
    ('F5V',(3.50,0.50)), ('F0V',(2.98,0.38)), ('RC',(0.8,1.0))])
  lines={}

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

  for spt in sptVabsAndVmini.keys():
    vmag=sptVabsAndVmini[spt][0]+5.0*logdistancekpc+10.0
    indices=(vmag>14) & (vmag<16)
    gmag=vmag+gminvFromVmini(sptVabsAndVmini[spt][1])
    parerrors=parallaxErrorSkyAvg(gmag,sptVabsAndVmini[spt][1])
    relparerrors=parerrors*10**logdistancekpc/1000.0
    plt.loglog(10**logdistancekpc, relparerrors,'--k',lw=1)
    plt.loglog(10**logdistancekpc[indices], relparerrors[indices],'-',label=spt)
  plt.xlim(0.1,20.0)
  plt.ylim(0.001,0.5)
  plt.text(0.9, 0.05,'Colours indicate $14<V<16$',
     horizontalalignment='right',
     verticalalignment='bottom',
     transform = currentAxis.transAxes)
  plt.legend(loc=2)
  plt.xlabel('distance [kpc]')
  plt.ylabel('$\\sigma_\\varpi/\\varpi$')
  plt.grid(which='both')

  if (args['pdfOutput']):
    plt.savefig('RelativeParallaxErrorsVsDist.pdf')
  elif (args['pngOutput']):
    plt.savefig('RelativeParallaxErrorsVsDist.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("-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()
  makePlot(pdf=args['pdfOutput'], png=args['pngOutput'])


#!/usr/bin/env python

# Plot radial velocity 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.spectroscopic import vradErrorSkyAvg
from pygaia.photometry.transformations import vminGrvsFromVmini
from pygaia.photometry.utils import vminiFromSpt

from os import environ as env
from matplotlib import rc
from matplotlib.colors import hsv_to_rgb
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 radial velocity performance predictions.

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

  spts=['B0V', 'B5V', 'A0V', 'A5V', 'F0V', 'G0V',
        'G5V', 'K0V', 'K1IIIMP', 'K4V', 'K1III']

  fig=plt.figure(figsize=(10,6.5))
  deltaHue = 240.0/(len(spts)-1)
  hsv=np.zeros((1,1,3))
  hsv[0,0,1]=1.0
  hsv[0,0,2]=0.9
  count=0
  for spt in spts:
    hsv[0,0,0]=(240-count*deltaHue)/360.0
    vmag = vminGrvsFromVmini(vminiFromSpt(spt)) + gRvs
    vradErrors = vradErrorSkyAvg(vmag, spt)
    plt.plot(vmag, vradErrors, '-', label=spt, color=hsv_to_rgb(hsv)[0,0,:])
    count+=1
  plt.grid(which='both')
  plt.xlim(9,17.5)
  plt.ylim(0,20)
  plt.xticks(np.arange(9,18,1))
  plt.yticks(np.arange(0,20.5,5))
  plt.xlabel('$V$ [mag]')
  plt.ylabel('End-of-mission radial velocity error [km s$^{-1}$]')
  leg=plt.legend(loc=0,  handlelength=2.0, labelspacing=0.10)
  for t in leg.get_texts():
    t.set_fontsize(12)

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

def parseCommandLineArguments():
  """
  Set up command line parsing.
  """
  parser = argparse.ArgumentParser(description="Plot predicted Gaia sky averaged radial velocity 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="grvsAbscissa", help="Plot performance vs Grvs instead of V")
  args=vars(parser.parse_args())
  return args

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


from __future__ import print_function, division, unicode_literals

from abc import ABCMeta, abstractproperty, abstractmethod
from collections import namedtuple
from distutils.version import LooseVersion
from functools import wraps
from itertools import takewhile, dropwhile
import operator
import re
import sys

try:
    from bs4 import BeautifulSoup, PageElement, NavigableString
except ImportError:  # pragma: no cover
    raise ImportError("Soupy requires beautifulsoup4")

try:
    import six
    from six.moves import map
    assert LooseVersion(six.__version__) >= LooseVersion('1.9')
except(ImportError, AssertionError):   # pragma: no cover
    raise ImportError("Soupy requires six version 1.9 or later")

__version__ = '0.4.dev'

__all__ = ['Soupy', 'Q', 'Node', 'Scalar', 'Collection',
           'Null', 'NullNode', 'NullCollection',
           'either', 'NullValueError', 'QDebug']

# extract the thing inside string reprs (eg u'abc' -> abc)
QUOTED_STR = re.compile("^[ub]?['\"](.*?)['\"]$")

QDebug = namedtuple('QDebug', ('expr', 'inner_expr', 'val', 'inner_val'))
"""Namedtuple that holds information about a failed expression evaluation."""

@six.add_metaclass(ABCMeta)
class Wrapper(object):

    @abstractmethod
    def val(self):
        pass  # pragma: no cover

    @abstractmethod
    def orelse(self, value):
        pass  # pragma: no cover

    def nonnull(self):
        """
        Require that a node is not null

        Null values will raise NullValueError, whereas nonnull
        values return self.

        useful for being strict about portions of queries.

        Examples:

           node.find('a').nonnull().find('b').orelse(3)

           This will raise an error if find('a') doesn't match,
           but provides a fallback if find('b') doesn't match.

        """
        return self

    @abstractmethod
    def isnull(self):
        pass  # pragma: no cover

    @abstractmethod
    def map(self, func):
        pass  # pragma: no cover

    @abstractmethod
    def apply(self, func):
        pass  # pragma: no cover

    @classmethod
    def wrap(cls, value):
        """
        Wrap value in the appropriate wrapper class,
        based upon its type.
        """
        if isinstance(value, Wrapper):
            return value

        if hasattr(value, 'children'):
            return Node(value)

        return Scalar(value)

    def __getitem__(self, key):
        return self.map(operator.itemgetter(key))

    def dump(self, *args, **kwargs):
        """
        Extract derived values into a Scalar(tuple) or Scalar(dict)

        The keyword names passed to this function become keys in
        the resulting dictionary, while positional arguments passed to
        this function become elements in the resulting tuple.

        The positional arguments and keyword values are functions that
        are called on this Node.

        Notes:

            - The input functions are called on the Node, **not** the
              underlying BeautifulSoup element
            - If the function returns a wrapper, it will be unwrapped
            - Only either positional arguments or keyword arguments may
              be passed, not both.

        Example:

            >>> soup = Soupy("<b>hi</b>").find('b')
            >>> data = soup.dump(name=Q.name, text=Q.text).val()
            >>> data == {'text': 'hi', 'name': 'b'}
            True

            >> name, text = soup.dump(Q.name, Q.text).val()
            >> (name, text) == ('hi', 'b')
            True
        """
        if args and kwargs:
            raise ValueError('Cannot pass both arguments and keywords to dump')

        if args:
            result = tuple(_unwrap(self.apply(func)) for func in args)
        else:
            result = dict((name, _unwrap(self.apply(func)))
                          for name, func in kwargs.items())
        return Wrapper.wrap(result)

    @abstractmethod
    def require(self, func, msg='Requirement Violated'):
        pass  # pragma: no cover

class NullValueError(ValueError):

    """
    The NullValueError exception is raised when attempting
    to extract values from Null objects
    """
    pass

class QKeyError(KeyError):

    """
    A custom KeyError subclass that better formats
    exception messages raised inside expressions
    """

    def __str__(self):
        parts = self.args[0].split('\n\n\t')
        return parts[0] + '\n\n\t' + _dequote(repr(parts[1]))

QKeyError.__name__ = str('KeyError')

@six.python_2_unicode_compatible
class BaseNull(Wrapper):

    """
    This is the base class for null wrappers. Null values are returned
    when the result of a function is ill-defined.
    """

    def val(self):
        """
        Raise :class:`NullValueError`
        """
        raise NullValueError()

    def orelse(self, value):
        """
        Wraps value and returns the result
        """
        return Wrapper.wrap(value)

    def map(self, func):
        """
        Returns :class:`Null`
        """
        return self

    def apply(self, func):
        """
        Returns :class:`Null`
        """
        return self

    def nonnull(self):
        """
        Raises :class:`NullValueError`
        """
        raise NullValueError()

    def require(self, func, msg="Requirement is violated (wrapper is null)"):
        """
        Raises :class:`NullValueError`
        """
        raise NullValueError()

    def isnull(self):
        """
        Return Scalar(True) if this item is a null value
        """
        return Scalar(True)

    def __setitem__(self, key, val):
        pass

    def __bool__(self):
        return False

    __nonzero__ = __bool__

    def __str__(self):
        return "%s()" % type(self).__name__

    __repr__ = __str__

    def __hash__(self):
        return hash(type(self))

    def __eq__(self, other):
        return type(self)()

    def __ne__(self, other):
        return type(self)()

@six.python_2_unicode_compatible
class Some(Wrapper):

    def __init__(self, value):
        self._value = value

    def map(self, func):
        """
        Call a function on a wrapper's value, and wrap the result if necessary.

        Parameters:

            func : function(val) -> val

        Examples:

            >>> s = Scalar(3)
            >>> s.map(Q * 2)
            Scalar(6)
        """
        return Wrapper.wrap(_make_callable(func)(self._value))

    def apply(self, func):
        """
        Call a function on a wrapper, and wrap the result if necessary.

        Parameters:

            func: function(wrapper) -> val

        Examples:

            >>> s = Scalar(5)
            >>> s.apply(lambda val: isinstance(val, Scalar))
            Scalar(True)
        """
        return Wrapper.wrap(_make_callable(func)(self))

    def orelse(self, value):
        """
        Provide a fallback value for failed matches.

        Examples:

            >>> Scalar(5).orelse(10).val()
            5
            >>> Null().orelse(10).val()
            10
        """
        return self

    def val(self):
        """
        Return the value inside a wrapper.

        Raises :class:`NullValueError` if called on a Null object
        """
        return self._value

    def require(self, func, msg="Requirement violated"):
        """
        Assert that self.apply(func) is True.

        Parameters:

            func : func(wrapper)
            msg : str
               The error message to display on failure

        Returns:

            If self.apply(func) is True, returns self.
            Otherwise, raises NullValueError.
        """
        if self.apply(func):
            return self
        raise NullValueError(msg)

    def isnull(self):
        """
        Return Scalar(True) if this item is a null value
        """
        return Scalar(False)

    def __str__(self):
        # returns unicode
        # six builds appropriate py2/3 methods from this
        return "%s(%s)" % (type(self).__name__, _repr(self._value))

    def __repr__(self):
        return repr(self.__str__())[1:-1]  # trim off quotes

    def __setitem__(self, key, val):
        return self.map(Q.__setitem__(key, val))

    def __hash__(self):
        return hash(self._value)

    def __eq__(self, other):
        return self.map(lambda x: x == other)

    def __ne__(self, other):
        return self.map(lambda x: x != other)

class Null(BaseNull):

    """
    The class for ill-defined Scalars.
    """

    def __getattr__(self, attr):
        return Null()

    def __call__(self, *args, **kwargs):
        return Null()

    def __gt__(self, other):
        return Null()

    def __ge__(self, other):
        return Null()

    def __lt__(self, other):
        return Null()

    def __le__(self, other):
        return Null()

    def __len__(self):
        raise TypeError("Null has no len()")

    def __add__(self, other):
        return Null()

    def __sub__(self, other):
        return Null()

    def __mul__(self, other):
        return Null()

    def __div__(self, other):
        return Null()

    def __floordiv__(self, other):
        return Null()

    def __pow__(self, other):
        return Null()

    def __mod__(self, other):
        return Null()

    def __truediv__(self, other):
        return Null()

    def __hash__(self):
        return super(Null, self).__hash__()

class Scalar(Some):

    """
    A wrapper around single values.

    Scalars support boolean testing (<, ==, etc), and
    use the wrapped value in the comparison. They return
    the result as a Scalar(bool).

    Calling a Scalar calls the wrapped value, and wraps
    the result.

    Examples:

        >>> s = Scalar(3)
        >>> s > 2
        Scalar(True)
        >>> s.val()
        3
        >>> s + 5
        Scalar(8)
        >>> s + s
        Scalar(6)
        >>> bool(Scalar(3))
        True
        >>> Scalar(lambda x: x+2)(5)
        Scalar(7)

    """

    def __getattr__(self, attr):
        return self.map(operator.attrgetter(attr))

    def __call__(self, *args, **kwargs):
        return self.map(operator.methodcaller('__call__', *args, **kwargs))

    def __gt__(self, other):
        return self.map(lambda x: x > other)

    def __ge__(self, other):
        return self.map(lambda x: x >= other)

    def __lt__(self, other):
        return self.map(lambda x: x < other)

    def __le__(self, other):
        return self.map(lambda x: x <= other)

    def __bool__(self):
        return bool(self._value)

    __nonzero__ = __bool__

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

    def __add__(self, other):
        if isinstance(other, BaseNull):
            return other
        return self.map(Q + _unwrap(other))

    def __sub__(self, other):
        if isinstance(other, BaseNull):
            return other
        return self.map(Q - _unwrap(other))

    def __mul__(self, other):
        if isinstance(other, BaseNull):
            return other
        return self.map(Q * _unwrap(other))

    def __div__(self, other):
        if isinstance(other, BaseNull):
            return other
        return self.map(Q / _unwrap(other))

    def __floordiv__(self, other):
        if isinstance(other, BaseNull):
            return other
        return self.map(Q // _unwrap(other))

    def __pow__(self, other):
        if isinstance(other, BaseNull):
            return other
        return self.map(Q ** _unwrap(other))

    def __mod__(self, other):
        if isinstance(other, BaseNull):
            return other
        return self.map(Q % _unwrap(other))

    def __truediv__(self, other):
        if isinstance(other, BaseNull):
            return other
        return self.map(Q / _unwrap(other))

class Collection(Some):

    """
    Collection's store lists of other wrappers.

    They support most of the list methods (len, iter, getitem, etc).
    """

    def __init__(self, items):
        super(Collection, self).__init__(list(items))
        self._items = self._value
        self._assert_items_are_wrappers()

    def _assert_items_are_wrappers(self):
        for item in self:
            if not isinstance(item, Wrapper):
                raise TypeError("Collection can only hold other wrappers")

    def val(self):
        """
        Unwraps each item in the collection, and returns as a list
        """
        return list(self.iter_val())

    def first(self):
        """
        Return the first element of the collection, or :class:`Null`
        """
        return self[0]

    def iter_val(self):
        """
        An iterator version of :meth:`val`
        """
        return (item.val() for item in self._items)

    def each(self, *funcs):
        """
        Call `func` on each element in the collection.

        If multiple functions are provided, each item
        in the output will be a tuple of each
        func(item) in self.

        Returns a new Collection.

        Example:

            >>> col = Collection([Scalar(1), Scalar(2)])
            >>> col.each(Q * 10)
            Collection([Scalar(10), Scalar(20)])
            >>> col.each(Q * 10, Q - 1)
            Collection([Scalar((10, 0)), Scalar((20, 1))])
        """

        funcs = list(map(_make_callable, funcs))

        if len(funcs) == 1:
            return Collection(map(funcs[0], self._items))

        tupler = lambda item: Scalar(
            tuple(_unwrap(func(item)) for func in funcs))
        return Collection(map(tupler, self._items))

    def exclude(self, func=None):
        """
        Return a new Collection excluding some items

        Parameters:

            func : function(Node) -> Scalar

                A function that, when called on each item
                in the collection, returns a boolean-like
                value. If no function is provided, then
                truthy items will be removed.

        Returns:

            A new Collection consisting of the items
            where bool(func(item)) == False
        """
        func = _make_callable(func)
        inverse = lambda x: not func(x)
        return self.filter(inverse)

    def filter(self, func=None):
        """
        Return a new Collection with some items removed.

        Parameters:

            func : function(Node) -> Scalar

                A function that, when called on each item
                in the collection, returns a boolean-like
                value. If no function is provided, then
                false-y items will be removed.

        Returns:

            A new Collection consisting of the items
            where bool(func(item)) == True

        Examples:

            node.find_all('a').filter(Q['href'].startswith('http'))
        """
        func = _make_callable(func)
        return Collection(filter(func, self._items))

    def takewhile(self, func=None):
        """
        Return a new Collection with the last few items removed.

        Parameters:

            func : function(Node) -> Node

        Returns:

            A new Collection, discarding all items
            at and after the first item where bool(func(item)) == False

        Examples:

            node.find_all('tr').takewhile(Q.find_all('td').count() > 3)

        """
        func = _make_callable(func)
        return Collection(takewhile(func, self._items))

    def dropwhile(self, func=None):
        """
        Return a new Collection with the first few items removed.

        Parameters:

            func : function(Node) -> Node

        Returns:

            A new Collection, discarding all items
            before the first item where bool(func(item)) == True
        """
        func = _make_callable(func)
        return Collection(dropwhile(func, self._items))

    def __getitem__(self, key):
        if isinstance(key, int):
            try:
                return self._items[key]
            except IndexError:
                return NullNode()

        # slice
        return Collection(list(self._items).__getitem__(key))

    def dump(self, *args, **kwargs):
        """
        Build a list of dicts, by calling :meth:`Node.dump`
        on each item.

        Each keyword provides a function that extracts a value
        from a Node.

        Examples:

            >>> c = Collection([Scalar(1), Scalar(2)])
            >>> c.dump(x2=Q*2, m1=Q-1).val()
            [{'x2': 2, 'm1': 0}, {'x2': 4, 'm1': 1}]
        """
        return self.each(Q.dump(*args, **kwargs))

    def __len__(self):
        return self.map(len).val()

    def count(self):
        """
        Return the number of items in the collection, as a :class:`Scalar`
        """
        return Scalar(len(self))

    def zip(self, *others):
        """
        Zip the items of this collection with one or more
        other sequences, and wrap the result.

        Unlike Python's zip, all sequences must be the same length.

        Parameters:

            others: One or more iterables or Collections

        Returns:

            A new collection.

        Examples:

            >>> c1 = Collection([Scalar(1), Scalar(2)])
            >>> c2 = Collection([Scalar(3), Scalar(4)])
            >>> c1.zip(c2).val()
            [(1, 3), (2, 4)]
        """
        args = [_unwrap(item) for item in (self,) + others]
        ct = self.count()
        if not all(len(arg) == ct for arg in args):
            raise ValueError("Arguments are not all the same length")
        return Collection(map(Wrapper.wrap, zip(*args)))

    def dictzip(self, keys):
        """
        Turn this collection into a Scalar(dict), by zipping keys and items.

        Parameters:

            keys: list or Collection of NavigableStrings
                The keys of the dictionary

        Examples:

            >>> c = Collection([Scalar(1), Scalar(2)])
            >>> c.dictzip(['a', 'b']).val() == {'a': 1, 'b': 2}
            True
        """
        return Scalar(dict(zip(_unwrap(keys), self.val())))

    def __iter__(self):
        for item in self._items:
            yield item

    def all(self):
        """
        Scalar(True) if all items are truthy, or collection is empty.
        """
        return self.map(all)

    def any(self):
        """
        Scalar(True) if any items are truthy. False if empty.
        """
        return self.map(any)

    def none(self):
        """
        Scalar(True) if no items are truthy, or collection is empty.
        """
        return self.map(lambda items: not any(items))

    def __bool__(self):
        return bool(self._items)

    __nonzero__ = __bool__

class NullCollection(BaseNull, Collection):

    """
    Represents in invalid Collection.

    Returned by some methods on other Null objects.
    """

    def __init__(self):
        pass

    def iter_val(self):
        raise NullValueError()

    def each(self, func):
        return self

    def filter(self, func=None):
        return self

    def takewhile(self, func=None):
        return self

    def dropwhile(self, func=None):
        return self

    def first(self):
        return NullNode()  # XXX don't like this assumption

    def __getitem__(self, key):
        if isinstance(key, int):
            return NullNode()  # XXX don't like this assumption

        # slice
        return self

    def dump(self, *args, **kwargs):
        return NullCollection()

    def count(self):
        return Scalar(0)

@six.add_metaclass(ABCMeta)
class NodeLike(object):

    # should return NodeLike
    parent = abstractproperty()
    next_sibling = abstractproperty()
    previous_sibling = abstractproperty()

    # should return scalar-like
    text = abstractproperty()
    attrs = abstractproperty()
    name = abstractproperty()

    # should return CollectionLike
    children = abstractproperty()
    contents = abstractproperty()
    descendants = abstractproperty()
    parents = abstractproperty()
    next_siblings = abstractproperty()
    previous_siblings = abstractproperty()

    @abstractmethod
    def find(self, name=None, attrs={}, recursive=True,
             text=None, **kwargs):
        pass  # pragma: no cover

    @abstractmethod
    def find_all(self, name=None, attrs={}, recursive=True,
                 text=None, **kwargs):
        pass  # pragma: no cover

    @abstractmethod
    def select(self, selector):
        pass  # pragma: no cover

    @abstractmethod
    def find_next_sibling(self, *args, **kwargs):
        pass  # pragma: no cover

    @abstractmethod
    def find_previous_sibling(self, *args, **kwargs):
        pass  # pragma: no cover

    @abstractmethod
    def find_parent(self, *args, **kwargs):
        pass  # pragma: no cover

    @abstractmethod
    def find_next_siblings(self, *args, **kwargs):
        pass  # pragma: no cover

    @abstractmethod
    def find_previous_siblings(self, *args, **kwargs):
        pass  # pragma: no cover

    @abstractmethod
    def find_parents(self, *args, **kwargs):
        pass  # pragma: no cover

    @abstractmethod
    def prettify(self):
        pass  # pragma: no cover

    def __iter__(self):
        for item in self.children:
            yield item

    def __call__(self, *args, **kwargs):
        return self.find_all(*args, **kwargs)

class Node(NodeLike, Some):

    """
    The Node class is the main wrapper around
    BeautifulSoup elements like Tag. It implements many of the
    same properties and methods as BeautifulSoup for navigating
    through documents, like find, select, parents, etc.
    """
    def __new__(cls, value, *args, **kwargs):
        if isinstance(value, NavigableString):
            return object.__new__(NavigableStringNode)

        return object.__new__(cls)

    def _wrap_node(self, func):
        val = func(self._value)
        return NullNode() if val is None else Node(val)

    def _wrap_multi(self, func):
        vals = func(self._value)
        return Collection(map(Node, vals))

    def _wrap_scalar(self, func):
        val = func(self._value)
        return Scalar(val)

    @property
    def children(self):
        """
        A :class:`Collection` of the child elements.
        """
        return self._wrap_multi(operator.attrgetter('children'))

    @property
    def parents(self):
        """
        A :class:`Collection` of the parents elements.
        """
        return self._wrap_multi(operator.attrgetter('parents'))

    @property
    def contents(self):
        """
        A :class:`Collection` of the child elements.
        """
        return self._wrap_multi(operator.attrgetter('contents'))

    @property
    def descendants(self):
        """
        A :class:`Collection` of all elements nested inside this Node.
        """
        return self._wrap_multi(operator.attrgetter('descendants'))

    @property
    def next_siblings(self):
        """
        A :class:`Collection` of all siblings after this node
        """
        return self._wrap_multi(operator.attrgetter('next_siblings'))

    @property
    def previous_siblings(self):
        """
        A :class:`Collection` of all siblings before this node
        """
        return self._wrap_multi(operator.attrgetter('previous_siblings'))

    @property
    def parent(self):
        """
        The parent :class:`Node`, or :class:`NullNode`
        """
        return self._wrap_node(operator.attrgetter('parent'))

    @property
    def next_sibling(self):
        """
        The :class:`Node` sibling after this, or :class:`NullNode`
        """
        return self._wrap_node(operator.attrgetter('next_sibling'))

    @property
    def previous_sibling(self):
        """
        The :class:`Node` sibling prior to this, or :class:`NullNode`
        """
        return self._wrap_node(operator.attrgetter('previous_sibling'))

    @property
    def attrs(self):
        """
        A :class:`Scalar` of this Node's attribute dictionary

        Example:

            >>> Soupy("<a val=3></a>").find('a').attrs
            Scalar({u'val': u'3'})
        """
        return self._wrap_scalar(operator.attrgetter('attrs'))

    @property
    def text(self):
        """
        A :class:`Scalar` of this Node's text.

        Example:

            >>> node = Soupy('<p>hi there</p>').find('p')
            >>> node
            Node(<p>hi there</p>)
            >>> node.text
            Scalar(u'hi there')

        """
        return self._wrap_scalar(operator.attrgetter('text'))

    @property
    def name(self):
        """
        A :class:`Scalar` of this Node's tag name.

        Example:

            >>> node = Soupy('<p>hi there</p>').find('p')
            >>> node
            Node(<p>hi there</p>)
            >>> node.name
            Scalar(u'p')
        """
        return self._wrap_scalar(operator.attrgetter('name'))

    def find(self, *args, **kwargs):
        """
        Find a single Node among this Node's descendants.

        Returns :class:`NullNode` if nothing matches.

        This inputs to this function follow the same semantics
        as BeautifulSoup. See http://bit.ly/bs4doc for more info.

        Examples:

         - node.find('a')  # look for `a` tags
         - node.find('a', 'foo') # look for `a` tags with class=`foo`
         - node.find(func) # find tag where func(tag) is True
         - node.find(val=3)  # look for tag like <a, val=3>
        """
        op = operator.methodcaller('find', *args, **kwargs)
        return self._wrap_node(op)

    def find_next_sibling(self, *args, **kwargs):
        """
        Like :meth:`find`, but searches through :attr:`next_siblings`
        """
        op = operator.methodcaller('find_next_sibling', *args, **kwargs)
        return self._wrap_node(op)

    def find_parent(self, *args, **kwargs):
        """
        Like :meth:`find`, but searches through :attr:`parents`
        """
        op = operator.methodcaller('find_parent', *args, **kwargs)
        return self._wrap_node(op)

    def find_previous_sibling(self, *args, **kwargs):
        """
        Like :meth:`find`, but searches through :attr:`previous_siblings`
        """
        op = operator.methodcaller('find_previous_sibling', *args, **kwargs)
        return self._wrap_node(op)

    def find_all(self, *args, **kwargs):
        """
        Like :meth:`find`, but selects all matches (not just the first one).

        Returns a :class:`Collection`.

        If no elements match, this returns a Collection with no items.
        """
        op = operator.methodcaller('find_all', *args, **kwargs)
        return self._wrap_multi(op)

    def find_next_siblings(self, *args, **kwargs):
        """
        Like :meth:`find_all`, but searches through :attr:`next_siblings`
        """
        op = operator.methodcaller('find_next_siblings', *args, **kwargs)
        return self._wrap_multi(op)

    def find_parents(self, *args, **kwargs):
        """
        Like :meth:`find_all`, but searches through :attr:`parents`
        """
        op = operator.methodcaller('find_parents', *args, **kwargs)
        return self._wrap_multi(op)

    def find_previous_siblings(self, *args, **kwargs):
        """
        Like :meth:`find_all`, but searches through :attr:`previous_siblings`
        """
        op = operator.methodcaller('find_previous_siblings', *args, **kwargs)
        return self._wrap_multi(op)

    def select(self, selector):
        """
        Like :meth:`find_all`, but takes a CSS selector string as input.
        """
        op = operator.methodcaller('select', selector)
        return self._wrap_multi(op)

    def prettify(self):
        return self.map(Q.prettify()).val()

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

    def __bool__(self):
        return True

    __nonzero__ = __bool__

class NavigableStringNode(Node):

    """
    The NavigableStringNode is a special case Node that wraps
    BeautifulSoup NavigableStrings. This class implements sensible
    versions of properties and methods that are missing from
    the NavigableString object.
    """

    @property
    def attrs(self):
        """
        An empty :class:`Scalar` dict
        """
        return Scalar({})

    @property
    def text(self):
        """
        A :class:`Scalar` of the string value
        """
        return Scalar(self._value.string)

    @property
    def name(self):
        """
        An empty :class:`Scalar` dict
        """
        return Scalar('')

    @property
    def children(self):
        """
        An empty :class:`Collection`
        """
        return Collection([])

    @property
    def contents(self):
        """
        An empty :class:`Collection`
        """
        return Collection([])

    @property
    def descendants(self):
        """
        An empty :class:`Collection`
        """
        return Collection([])

    def find(self, *args, **kwargs):
        """
        Returns :class:`NullNode`
        """
        return NullNode()

    def find_all(self, *args, **kwargs):
        """
        Returns an empty :class:`Collection`
        """
        return Collection([])

    def select(self, selector):
        """
        Returns an empty :class:`Collection`
        """
        return Collection([])

    def prettify(self):
        return self.text.val()

class NullNode(NodeLike, BaseNull):

    """
    NullNode is returned when a query doesn't match any node
    in the document.
    """

    def _get_null(self):
        """
        Returns the NullNode
        """
        return NullNode()

    def _get_null_set(self):
        """
        Returns the NullCollection
        """
        return NullCollection()

    children = property(_get_null_set)
    parents = property(_get_null_set)
    contents = property(_get_null_set)
    descendants = property(_get_null_set)
    next_siblings = property(_get_null_set)
    previous_siblings = property(_get_null_set)

    parent = property(_get_null)
    next_sibling = property(_get_null)
    previous_sibling = property(_get_null)

    attrs = property(lambda self: Null())
    text = property(lambda self: Null())
    name = property(lambda self: Null())

    def find(self, *args, **kwargs):
        """
        Returns :class:`NullNode`
        """
        return NullNode()

    def find_parent(self, *args, **kwargs):
        """
        Returns :class:`NullNode`
        """
        return NullNode()

    def find_previous_sibling(self, *args, **kwargs):
        """
        Returns :class:`NullNode`
        """
        return NullNode()

    def find_next_sibling(self, *args, **kwargs):
        """
        Returns :class:`NullNode`
        """
        return NullNode()

    def find_all(self, *args, **kwargs):
        """
        Returns :class:`NullCollection`
        """
        return NullCollection()

    def find_parents(self, *args, **kwargs):
        """
        Returns :class:`NullCollection`
        """
        return NullCollection()

    def find_next_siblings(self, *args, **kwargs):
        """
        Returns :class:`NullCollection`
        """
        return NullCollection()

    def find_previous_siblings(self, *args, **kwargs):
        """
        Returns :class:`NullCollection`
        """
        return NullCollection()

    def select(self, selector):
        """
        Returns :class:`NullCollection`
        """
        return NullCollection()

    def dump(self, *args, **kwargs):
        """
        Returns :class:`Null`
        """
        return Null()

    def prettify(self):
        return "Null Node"

    def __len__(self):
        return 0

def either(*funcs):
    """
    A utility function for selecting the first non-null query.

    Parameters:

      funcs: One or more functions

    Returns:

       A function that, when called with a :class:`Node`, will
       pass the input to each `func`, and return the first non-Falsey
       result.

    Examples:

       >>> s = Soupy("<p>hi</p>")
       >>> s.apply(either(Q.find('a'), Q.find('p').text))
       Scalar('hi')
    """
    def either(val):
        for func in funcs:
            result = val.apply(func)
            if result:
                return result
        return Null()

    return either

def _helpful_failure(method):
    """
    Decorator for eval_ that prints a helpful error message
    if an exception is generated in a Q expression
    """

    @wraps(method)
    def wrapper(self, val):
        try:
            return method(self, val)
        except:
            exc_cls, inst, tb = sys.exc_info()

            if hasattr(inst, '_RERAISE'):
                _, expr, _, inner_val = Q.__debug_info__
                Q.__debug_info__ = QDebug(self, expr, val, inner_val)
                raise

            if issubclass(exc_cls, KeyError):  # Overrides formatting
                exc_cls = QKeyError

            # Show val, unless it's too long
            prettyval = repr(val)
            if len(prettyval) > 150:
                prettyval = "<%s instance>" % (type(val).__name__)

            msg = "{0}\n\n\tEncountered when evaluating {1}{2}".format(
                inst, prettyval, self)

            new_exc = exc_cls(msg)
            new_exc._RERAISE = True
            Q.__debug_info__ = QDebug(self, self, val, val)

            six.reraise(exc_cls, new_exc, tb)

    return wrapper

@six.python_2_unicode_compatible
class Expression(object):

    """
    Soupy expressions are a shorthand for building single-argument functions.

    Users should use the ``Q`` object, which is just an instance of Expression.
    """

    def __str__(self):
        return 'Q'

    def __repr__(self):
        return repr(str(self))[1:-1]  # trim quotes

    def __iter__(self):
        yield self

    def _chain(self, other):
        return Chain(tuple(iter(self)) + tuple(iter(other)))

    def __getattr__(self, key):
        return self._chain(Attr(key))

    def __getitem__(self, key):
        return self._chain(GetItem(key))

    def __call__(self, *args, **kwargs):
        return self._chain(Call(args, kwargs))

    def __gt__(self, other):
        return BinaryOp(operator.gt, '>', self, other)

    def __ge__(self, other):
        return BinaryOp(operator.ge, '>=', self, other)

    def __lt__(self, other):
        return BinaryOp(operator.lt, '<', self, other)

    def __le__(self, other):
        return BinaryOp(operator.le, '<=', self, other)

    def __eq__(self, other):
        return BinaryOp(operator.eq, '==', self, other)

    def __ne__(self, other):
        return BinaryOp(operator.ne, '!=', self, other)

    def __add__(self, other):
        return BinaryOp(operator.add, '+', self, other)

    def __sub__(self, other):
        return BinaryOp(operator.sub, '-', self, other)

    def __div__(self, other):
        return BinaryOp(operator.__div__, '/', self, other)

    def __floordiv__(self, other):
        return BinaryOp(operator.floordiv, '//', self, other)

    def __truediv__(self, other):
        return BinaryOp(operator.truediv, '/', self, other)

    def __mul__(self, other):
        return BinaryOp(operator.mul, '*', self, other)

    def __rmul__(self, other):
        return BinaryOp(operator.mul, '*', other, self)

    def __pow__(self, other):
        return BinaryOp(operator.pow, '**', self, other)

    def __mod__(self, other):
        return BinaryOp(operator.mod, '%', self, other)

    @_helpful_failure
    def eval_(self, val):
        """
        Pass the argument ``val`` to the function, and return the result.

        This special method is necessary because the ``__call__`` method
        builds a new function stead of evaluating the current one.
        """
        return val

    def debug_(self):
        """
        Returns debugging information for the previous error raised
        during expression evaluation.

        Returns a QDebug namedtuple with four fields:

          - expr is the last full expression to have raised an exception
          - inner_expr is the specific sub-expression that raised the exception
          - val is the value that expr tried to evaluate.
          - inner_val is the value that inner_expr tried to evaluate

        If no exceptions have been triggered from expression evaluation,
        then each field is None.

        Examples:

            >>> Scalar('test').map(Q.upper().foo)
            Traceback (most recent call last):
            ...
            AttributeError: 'str' object has no attribute 'foo'
            ...
            >>> dbg = Q.debug_()
            >>> dbg.expr
            Q.upper().foo
            >>> dbg.inner_expr
            .foo
            >>> dbg.val
            'test'
            >>> dbg.inner_val
            'TEST'
        """
        result = self.__debug_info__
        if isinstance(result, QDebug):
            return result
        return QDebug(None, None, None, None)

@six.python_2_unicode_compatible
class Call(Expression):

    """An expression for calling a function or method"""

    def __init__(self, args, kwargs):
        self._args = args
        self._kwargs = kwargs

    @_helpful_failure
    def eval_(self, val):
        return val.__call__(*self._args, **self._kwargs)

    def __str__(self):
        result = list(map(_uniquote, self._args))
        if self._kwargs:
            result.append('**%s' % _uniquote(self._kwargs))
        return '(%s)' % (', '.join(result))

@six.python_2_unicode_compatible
class BinaryOp(Expression):

    """A binary operation"""

    def __init__(self, op, symbol, left, right):
        self.op = op
        self.left = left
        self.right = right
        self.symbol = symbol

    @_helpful_failure
    def eval_(self, val):
        left = self.left
        right = self.right
        if isinstance(left, Expression):
            left = left.eval_(val)

        if isinstance(right, Expression):
            right = right.eval_(val)

        return self.op(left, right)

    def __str__(self):
        l, r = self.left, self.right
        if isinstance(l, BinaryOp):
            l = '(%s)' % str(l)
        if isinstance(r, BinaryOp):
            r = '(%s)' % str(r)

        return "%s %s %s" % (l, self.symbol, r)

@six.python_2_unicode_compatible
class Attr(Expression):

    """An expression for fetching an attribute (eg, obj.item)"""

    def __init__(self, attribute_name):
        self._name = attribute_name

    @_helpful_failure
    def eval_(self, val):
        return operator.attrgetter(self._name)(val)

    def __str__(self):
        return '.%s' % self._name

@six.python_2_unicode_compatible
class GetItem(Expression):

    """An expression for getting an item (eg, obj['item'])"""

    def __init__(self, key):
        self._name = key

    @_helpful_failure
    def eval_(self, val):
        return operator.itemgetter(self._name)(val)

    def __str__(self):
        return "[%s]" % _uniquote(self._name)

@six.python_2_unicode_compatible
class Chain(Expression):

    """An chain of expressions (eg a.b.c)"""

    def __init__(self, items):
        self._items = items

    def __iter__(self):
        for item in self._items:
            yield item

    @_helpful_failure
    def eval_(self, val):
        for item in self._items:
            val = item.eval_(val)
        return val

    def __str__(self):
        return ''.join(map(_uniquote, self._items))

def _make_callable(func):
    # If func is an expression, we call via eval_
    # otherwise, we call func directly
    if func is None:
        func = Q
    return getattr(func, 'eval_', func)

def _unwrap(val):
    if isinstance(val, Wrapper):
        return val.val()
    return val

def _dequote(str):
    try:
        return QUOTED_STR.findall(str)[0]
    except IndexError:
        raise AssertionError("Not a quoted string")

def _uniquote(value):
    """
    Convert to unicode, and add quotes if initially a string
    """
    if isinstance(value, six.binary_type):
        try:
            value = value.decode('utf-8')
        except UnicodeDecodeError:  # Not utf-8. Show the repr
            value = six.text_type(_dequote(repr(value)))  # trim quotes

    result = six.text_type(value)

    if isinstance(value, six.text_type):
        result = "'%s'" % result
    return result

def _repr(value):
    value = repr(value)
    if isinstance(value, six.binary_type):
        value = value.decode('utf-8')
    return value

class Soupy(Node):

    def __init__(self, val, *args, **kwargs):
        if not isinstance(val, PageElement):
            val = BeautifulSoup(val, *args, **kwargs)
        super(Soupy, self).__init__(val)

Q = Expression()