"""
Given a time-series of length 260 generated by:

p(t) = 100 + sum[sin(t * c_i)] for c_i in
c = [46, 64, 98, 99, 101, 104, 106, 107, 109, 110, 111, 111, 111, 115, 115]

Calculate the returns based on the following strategy:
Hold the stock when it is within the closed interval [98, 101] and
sell at the end of the time series if still holding.

So, if the time series was:
[96, 97, 98, 102, 100, 100, 100, 96]

The strategy would (N=do nothing, B=buy, S=sell):
[N, N, B, S, B, N, N, S]

What is the cumulative percentage return using this strategy:
1) simply on a percentage return basis
e.g., day 1: 10%, day 2: 10%, total return would be 21% after two days.

2) assuming $1000.00 of starting capital and only purchasing whole "shares"
"""

import math
import pytest

class DataPoints(object):
    def __init__(self, c, length, low, high, capital, own, shares):
        self.c = c
        self.length = length
        self.low = low
        self.high = high
        self.init_capital = capital
        self.current_capital = capital
        self.own = own
        self.shares = shares

    def create_data_point(self, t):
        return 100 + sum(map(lambda c_i: math.sin(c_i * t), self.c))

    def create_time_series(self):
        return [self.create_data_point(t) for t in xrange(self.length)]

    def define_strategy(self, index, value):
        # Bug here where it says self.own is True at the initial state.
        if self.own:
            if index + 1 == self.length or value >= self.high:
                self.own = False
                self.current_capital += self.shares * value
                self.shares = 0
                return 'S'
            else:
                return 'N'
        else:
            if self.low <= value <= self.high:
                self.own = True
                self.shares = math.floor(self.current_capital/value)
                self.current_capital -= self.shares * value
                return 'B'
            else:
                return 'N'

    def map_strategy(self):
        time_series = self.create_time_series()
        decisions = [
            self.define_strategy(index, value)
            for index, value in enumerate(time_series)
        ]
        return zip(time_series, decisions)

    def cumulative_percentage_returns(self):
        self.map_strategy()
        returns = (self.current_capital - self.init_capital) / self.init_capital * 100
        return 'RETURNS: ' + str(returns) + '%'


class TestDataPoints(object):
    def setup(self):
        self.c = [46, 64, 98, 99, 101, 104, 106, 107, 109, 110, 111, 111, 111, 115, 115]
        self.length = 260
        self.low = 98
        self.high = 101
        self.capital = 1000
        self.shares = 0

        self.d = DataPoints(c=self.c, length=self.length, low=self.low,
                            high=self.high, capital=self.capital,
                            own=False, shares=self.shares)
        self.d_own = DataPoints(c=self.c, length=self.length, low=self.low,
                            high=self.high, capital=self.capital,
                            own=True, shares=0)

    def teardown(self):
        pass

    def test_create_data_point(self):
        assert self.d.create_data_point(0) == 100.0
        assert self.d.create_data_point(100) == 97.54988657392637

    def test_define_strategy_start(self):
        assert self.d.define_strategy(0, 100) == 'B'
        assert self.d.define_strategy(0, 96) == 'N'
        # assert self.d.define_strategy(0, 104) == 'N'

    def test_define_strategy_middle(self):
        assert self.d_own.define_strategy(5, 100) == 'N'
        assert self.d_own.define_strategy(5, 95) == 'N'
        assert self.d_own.define_strategy(5, 104) == 'S'

    def test_define_strategy_end(self):
        assert self.d_own.define_strategy(261, 100) == 'N'
        assert self.d_own.define_strategy(260, 3) == 'N'

    def test_cumulative_percentage_returns(self):
        assert self.d.cumulative_percentage_returns() == \
            "RETURNS: 1441.62257536%"

