from numpy import zeros, ones
import numpy as np

from . import syntax_fail, good, fail
from Aspidites._vendor.contracts.library.array import np_int_dtypes, np_uint_dtypes, \
    np_float_dtypes

a_u8 = np.zeros((3, 4), dtype='uint8')
a_i8 = np.zeros((3, 4), dtype='int8')
a_u16 = np.zeros((3, 4), dtype='uint16')
a_i16 = np.zeros((3, 4), dtype='int16')
a_u32 = np.zeros((3, 4), dtype='uint32')
a_i32 = np.zeros((3, 4), dtype='int32')
a_f32 = np.zeros((3, 4), dtype='float32')
a_f64 = np.zeros((3, 4), dtype='float64')
a_bool = np.zeros((3, 4), dtype='bool')

# ## Strings
good('array', a_f32)
good('array', a_f64)
# synonims
good('ndarray', a_f32, exact=False)  # will be canonicalized to "array"
good('ndarray', a_f64, exact=False)  # same
fail('array', [0, 1])
# dtypes
good('array(uint8)', a_u8)
good('array(u1)', a_u8)
good('array(int8)', a_i8)
good('array(i1)', a_i8)
good('array(float32)', a_f32)
good('array(float64)', a_f64)
good('array(bool)', a_bool)
fail('array(float64)', a_f32)
fail('array(float32)', a_f64)
fail('array(uint8)', a_f32)
fail('array(float32)', a_u8)

# shapes
a0d = np.zeros((), dtype='float32')
a1d = np.zeros((2,))
a2d = np.zeros((2, 4,))
a3d = np.zeros((2, 4, 8))

good('shape[0]', a0d)
good('shape[1]', a1d)
good('shape[2]', a2d)
good('shape[3]', a3d)
fail('shape[>0]', a0d)
fail('shape[<1]', a1d)
fail('shape[>2]', a2d)
fail('shape[<3]', a3d)
good(['shape[x]', 'shape[y],x=y'], [a3d, a3d])
good(['shape[x]', 'shape[y],x=y'], [a2d, a2d])
good(['shape[x]', 'shape[y],x=y'], [a1d, a1d])
good(['shape[x]', 'shape[y],x=y'], [a0d, a0d])

good('array[2x4]', a2d)
fail('array[AxBxC]', a2d)
fail('array[2x4]', a3d)

good('array[HxW],H=2,W>3', a2d)
good('array[(=2)x(>3)]', a2d, exact=False)  # Parenthesis are unnecessary

# ellipsis to mean 0 or more dimensions 
good('array[2x4x...]', a2d)
good('array[2x4x...]', a3d)

# if we really want more, use:
good('shape[>2],array[2x4x...]', a3d)
fail('shape[>2],array[2x4x...]', a2d)

fail('shape[>2]', [2, 2, 3])

# Try some binding:
good('array[XxYx...],X=2,Y=4', a3d)

# We don't do in between yet
syntax_fail('array[2x...x3]')

# We don't do in between yet
syntax_fail('array[2x...x3]')
# Try some binding:
good('array[XxYx...],X=2,Y=4', a2d)
good('array[XxYx...],X=2,Y=4', a3d)

# Using the array syntax  array[shape desc]

v1d = np.zeros(100)
v2d = np.zeros((10, 10))
good('array[100]', v1d)
fail('array[100]', v2d)
good('array[10x10]', v2d)
fail('array[10x...]', v1d)
good('array[10x...]', v2d)

fail('shape[>0]', a0d)
fail('shape[<1]', a1d)
fail('shape[>2]', a2d)
fail('shape[<3]', a3d)
good('shape[0]', a0d)
good('shape[1]', a1d)
good('shape[2]', a2d)
good('shape[3]', a3d)
good('array[shape[3]]', a3d)

good('array[1x2]', zeros((1, 2)))

# Now: special comparisons for arrays
A = np.array([0, 1, 2])
B = np.array([10, 20, 30])
good('array(>=0)', A)
good('array(<=2)', A)
fail('array(<=20)', B)
good('array(=0)', zeros((10)))
fail('array(=0)', np.array([0, 1, 0]))
good('array(=1)', ones((10)))

good(['shape(x)', 'shape(x)'], [a2d, a2d])
fail(['shape(x)', 'shape(x)'], [a2d, a3d])

good('array[NxN](<=1)', np.ones((10, 10)))
good('array[NxN](<=1,float32)', np.ones((10, 10), dtype='float32'))
good('array[NxN](<=1,float32|float64)', np.ones((10, 10), dtype='float64'))
good('array[NxN](<=1,(float32|float64))',
     np.ones((10, 10), dtype='float64'))
good('array[NxN](<=1,>=1)', np.ones((10, 10)))

# more complicated tests

good('array[(2|3)xN]', np.ones((2, 10)))
good('array[(2|3)xN]', np.ones((3, 10)))
fail('array[(2|3)xN]', np.ones((4, 10)))
good('array[(2|3)x...]', np.ones((2, 10)))
good('array[(2|3)x...]', np.ones((3, 10)))
fail('array[(2|3)x...]', np.ones((4, 10)))

good('array[(2,*)xN]', np.ones((2, 10)))
fail('array[(2,3)xN]', np.ones((4, 10)))
good('array[(2,*)x...]', np.ones((2, 10)))
fail('array[(2,3)x...]', np.ones((4, 10)))


good('seq', np.ones(3))
good('seq[3]', np.ones(3))
fail('seq[3]', np.ones(2))
good('seq[6]', np.ones((2, 3)))
fail('seq[6]', np.ones((2, 4)))

# "finite" contract 
good('finite', 1)
good('finite', 0)
good('finite', -1)
good('finite', np.float_(1))
fail('finite', np.inf)
fail('finite', np.nan)

# some more complex stuff
good('array[N](>=-pi,<pi)', np.array([0], 'float32'))
good('array[N](>=0)', np.array([0], 'float32'))

# generalized ideas of numbers
# np_ints = ['int8', 'int16', 'int32', 'int64',
#            'uint8', 'uint16', 'uint32', 'uint64',
#            ]
# np_floats = ['float32', 'float64']

good('array(float)', np.array(1.32, 'float32'))


for dt in np_int_dtypes:
    x = np.array(1).astype(dt)
    good('number', x)
    fail('Number', x)

    good('array(int)', x)
    good('int', x)
    fail('Int', x)

    fail('array(float)', x)
    fail('float', x)
    fail('Float', x)
    
# generalized ideas of numbers
for dt in np_float_dtypes:
    x = np.array(1).astype(dt)
    good('number', x)
    fail('Number', x)
    
    fail('array(int)', x)
    fail('int', x)
    fail('Int', x)

    good('array(float)', x)
    good('float', x)
    fail('Float', x)
    
# generalized ideas of numbers
for dt in np_uint_dtypes:
    x = np.array(1).astype(dt)
    good('number', x)
    fail('Number', x)
    
    fail('array(int)', x)
    fail('int', x)
    fail('Int', x)

    good('array(uint)', x)
    fail('Int', x)

    fail('array(float)', x)
    fail('float', x)
    fail('Float', x)

