import functools
import nose
import numpy
import apogee.tools.read as apread
from apogee.tools import bitmask
_DATA= None #such that we can re-use it in different tests

# Decorator for known failure
def known_failure(test):
    @functools.wraps(test)
    def inner(*args, **kwargs):
        try:
            test(*args, **kwargs)
        except Exception:
            raise nose.SkipTest
        else:
            raise AssertionError('Test is expected to fail, but passed instead')
    return inner

def test_read():
    global _DATA
    _DATA= apread.allStar(raw=True)
    assert not _DATA is None, '_DATA was not successfully read'
    return None

@known_failure
def test_telescope():
    #Test the telescope tag against the APSTAR_ID
    onemIndx= numpy.array(['apogee.apo1m' in s for s in _DATA['APSTAR_ID']])
    assert numpy.sum(onemIndx*(True-(_DATA['TELESCOPE'] == 'apo1m'))) == 0,\
        'TELESCOPE tag does not correspond to APSTAR_ID for 1m data'
    return None

def test_targflags_apogee_target1():
    # Test that TARGFLAGS corresponds to the bits in APOGEE_TARGET
    targ1bits= range(31) #don't check 31, bc always set
    targ1bits.pop(14) #14 not populated
    for targbit in targ1bits:
        name= bitmask.apogee_target1_string(targbit)
        targindx= numpy.array([name in s for s in _DATA['TARGFLAGS']],
                              dtype='bool')
        if targbit == 0:
            targindx*= \
                numpy.array([not 'APOGEE_FAINT_EXTRA' in s for s in _DATA['TARGFLAGS']],
                            dtype='bool')
        badindx= ((_DATA['APOGEE_TARGET1'] & 2**targbit) != 0)*(True-targindx)
        assert numpy.sum(badindx) == 0, 'Some objects with bit %i set in apogee_target1 do not have the corresponding flag name in TARGFLAGS set' % targbit
    return None

def test_targflags_apogee_target2():
    # Test that TARGFLAGS corresponds to the bits in APOGEE_TARGET
    targ2bits= [1,2,3,4,9,10,11,12,13,14,15,16,17]
    for targbit in targ2bits:
        name= bitmask.apogee_target2_string(targbit)
        targindx= numpy.array([name in s for s in _DATA['TARGFLAGS']],
                              dtype='bool')
        if targbit == 0:
            targindx*= \
                numpy.array([not 'APOGEE_FAINT_EXTRA' in s for s in _DATA['TARGFLAGS']],
                            dtype='bool')
        badindx= ((_DATA['APOGEE_TARGET2'] & 2**targbit) != 0)*(True-targindx)
        assert numpy.sum(badindx) == 0, 'Some objects with bit %i set in apogee_target2 do not have the corresponding flag name in TARGFLAGS set' % targbit
    return None

@known_failure
def test_extratarg():
    #Test that extratarg tag is 
    # 0 for main survey targets, 
    # 1 for commissioning,
    # 2 for tellurics
    mainIndx= (((_DATA['APOGEE_TARGET1'] & 2**11) != 0)\
                   +((_DATA['APOGEE_TARGET1'] & 2**12) != 0)
               +((_DATA['APOGEE_TARGET1'] & 2**13) != 0))
    assert numpy.sum(mainIndx*_DATA['EXTRATARG'] != 0) == 0, '%i main survey targets have EXTRATARG > 0' % numpy.sum(mainIndx*_DATA['EXTRATARG'] > 0)
    commIndx= _DATA['COMMISS'] == 1
    assert numpy.sum(commIndx*(_DATA['EXTRATARG'] != 1)) == 0, '%i commissioning targets have EXTRATARG neq 1' % numpy.sum(commIndx*(_DATA['EXTRATARG'] != 1))
    tellIndx= (_DATA['APOGEE_TARGET2'] & 2**9) != 0
    assert numpy.sum(tellIndx*(_DATA['EXTRATARG'] != 2)) == 0, '%i telluric targets have EXTRATARG neq 2' % numpy.sum(tellIndx*(_DATA['EXTRATARG'] != 2))
    return None

@@ -2582,32 +2582,33 @@
    return None%0A%0A
+#
 @known_failure%0Ad
@@ -2659,17 +2659,16 @@
 g tag is
- 
 %0A    # 0
@@ -2719,17 +2719,24 @@
 ssioning
-,
+ (bit 1)
 %0A    # 2
@@ -2749,16 +2749,24 @@
 ellurics
+ (bit 2)
 %0A    mai
@@ -2934,24 +2934,25 @@
  != 0))%0A    
+#
 assert numpy
@@ -2957,32 +2957,33 @@
 py.sum(mainIndx*
+(
 _DATA%5B'EXTRATARG
@@ -2990,16 +2990,17 @@
 '%5D != 0)
+)
  == 0, '
@@ -3037,17 +3037,19 @@
 TRATARG 
-%3E
+neq
  0' %25 nu
@@ -3133,58 +3133,172 @@
     
-assert numpy.sum(commIndx*(_DATA%5B'EXTRATARG'%5D != 1
+commBitSet= numpy.array(%5Bbitmask.bit_set(1,e) for e in _DATA%5B'EXTRATARG'%5D%5D,%0A                            dtype='bool')%0A    assert numpy.sum(commIndx*(True-commBitSet
 )) =
@@ -3320,36 +3320,52 @@
 sioning targets 
-have
+do not have bit 1 in
  EXTRATARG neq 1
@@ -3363,13 +3363,11 @@
 ARG 
-neq 1
+set
 ' %25 
@@ -3390,33 +3390,30 @@
 dx*(
-_DATA%5B'EXTRATARG'%5D != 1))
+True-commBitSet)) == 0
 %0A   
@@ -3469,58 +3469,172 @@
     
-assert numpy.sum(tellIndx*(_DATA%5B'EXTRATARG'%5D != 2
+tellBitSet= numpy.array(%5Bbitmask.bit_set(2,e) for e in _DATA%5B'EXTRATARG'%5D%5D,%0A                            dtype='bool')%0A    assert numpy.sum(tellIndx*(True-tellBitSet
 )) =
@@ -3651,36 +3651,52 @@
 elluric targets 
-have
+do not have bit 2 in
  EXTRATARG neq 2
@@ -3694,13 +3694,11 @@
 ARG 
-neq 2
+set
 ' %25 
@@ -3721,31 +3721,23 @@
 dx*(
-_DATA%5B'EXTRATARG'%5D != 2
+True-tellBitSet
 ))%0A 

fix test of extratarg after Jon's input