#!/usr/bin/env python

"""
some simple tests with struct arrays
"""

import numpy as np

dt = np.dtype ( [('unit',int),
                 ('price',float),
                 ('amount',float),
                 ])

data = np.array([ (1, 2.2, 0.0),
                 (3, 4.5, 0.0)
                 ],
               dtype= dt,
               )

print "the array:"
print data

print "pulling out a field"
print data['price']

print "pulling out a row"
print  "yields a tuple, not a rec array"
print data[0]

print "pulling out a slice:"
print data[0:1]

print "appending:"
a2 = np.array( [(3,5.5, 3)], dtype=dt)
data = np.concatenate( (data, a2) )
print data

print "some math"
data['amount'] = data['unit'] * data['price']
print repr(data)

print "adding a column"
dt2 = np.dtype ( [('unit',int),
                  ('price',float),
                  ('amount',float),
                  ('discount_price', float)
                 ])
# create a new array
data2 = np.zeros(len(data), dtype=dt2) 

# fill the array:
for field_name in dt.fields.keys():
    data2[field_name] = data[field_name]

# now some calculations:
data2['discount_price'] = data2['price'] * 0.9

print "the new array with a new column"
print repr(data2)

import accumulator

data3 = accumulator.accumulator(dtype = dt2)
data3.append((1, 2.2, 0.0, 0.0))
data3.append((3, 4.5, 0.0, 0.0))
data3.append((2, 1.2, 0.0, 0.0))
data3.append((5, 4.2, 0.0, 0.0))
print repr(data3)

# convert to regular array for calculations:
data3 = np.array(data3)
# now some calculations:
data3['discount_price'] = data3['price'] * 0.9
print repr(data3)