import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

def ema(data, key, alpha=0.5, ic=None):
	result = []
	if ic is None:
		r = data[key][0]
	else:
		r = ic
	for v in data[key]:
		r = alpha*r + (1-alpha)*v
		result.append(r)

	return np.array(result)

def colsum(data, keys):
	result = np.zeros(data.shape[0])
	for key in keys:
		result += data[key]

	return result

store = pd.HDFStore('test.h5')
if 1:
	index = pd.date_range('1/1/2001', periods=4, freq='D')
	df = pd.DataFrame(np.array([[1, 2, 3, 101], [4, 5, 6, 102], [7, 8, 9, 103], [10, 11, 12, 104]]),columns=['high', 'low', 'open', 'close'], index=index)
	store['df'] = df
else:
	df = store['df']

df['ema(close)'] = ema(df, 'close')
df['colsum'] = colsum(df, ['high', 'low', 'open', 'close'])
df['ema(colsum)'] = ema(df, 'colsum')

# Print table
print(df.loc['2001-01-01':'2001-01-02'])

# Plot graph
df.plot()
plt.legend()
plt.grid()
plt.show()

store.close()