from __future__ import absolute_import
import unittest
import numpy as np

from tests.sample_data import SampleData
from pyti import moving_average_envelope


class TestMovingAverageEnvelope(unittest.TestCase):
    def setUp(self):
        """Create data to use for testing."""
        self.data = SampleData().get_sample_close_data()

        self.center_band_period_6_expected = [np.nan, np.nan, np.nan, np.nan,
        np.nan, 804.55166666666673, 807.84333333333336, 809.89666666666665,
        811.21833333333325, 811.20333333333338, 812.51166666666666,
        813.88000000000011, 814.40333333333331, 813.18666666666661,
        812.6783333333334, 810.23333333333346, 806.20333333333338,
        799.25166666666667, 793.06499999999994, 785.82499999999993,
        778.30499999999995, 775.09000000000003, 774.75166666666667,
        776.35333333333347, 776.68333333333339, 779.10666666666668,
        782.55166666666673, 784.03833333333341, 781.79333333333341,
        781.85500000000002, 781.81833333333327, 781.17833333333328,
        775.88166666666666, 773.70666666666659, 774.42666666666662,
        777.66499999999996, 782.99833333333333, 787.4766666666668,
        792.12333333333345, 793.86333333333334, 795.21833333333336,
        795.20000000000016, 794.85333333333335, 797.77499999999998,
        803.81666666666672, 810.46833333333336, 817.15666666666664,
        822.19999999999993, 824.55999999999983, 824.90499999999986,
        826.52833333333331, 826.42666666666662, 822.80833333333339,
        817.61833333333345, 814.28833333333341, 812.64499999999998,
        809.72499999999991, 808.505, 807.48333333333323, 807.23000000000002,
        806.75500000000011, 805.25833333333321, 803.72666666666657,
        802.04166666666663, 802.36333333333334, 803.52666666666664,
        805.11000000000001, 805.08666666666659, 807.51666666666677,
        809.49833333333333, 809.89666666666665, 808.18333333333328,
        805.62666666666667, 804.84666666666669, 802.55833333333339,
        798.31000000000006, 795.5916666666667, 795.43166666666673,
        794.28000000000009, 795.0916666666667, 796.21833333333336,
        799.1450000000001, 800.50333333333344, 799.26666666666677, 799.495,
        797.67500000000007, 795.64666666666665, 793.17999999999995,
        792.25166666666667, 792.61833333333345, 793.74166666666667,
        794.58000000000004, 795.21833333333325, 796.80666666666673,
        799.15999999999997, 800.42500000000007, 801.98666666666668,
        803.67000000000007, 805.09499999999991, 806.05166666666662,
        806.39499999999987, 807.06833333333327, 807.23000000000002,
        805.59666666666669, 804.04999999999984, 802.65500000000009,
        801.56499999999994, 799.25, 792.40166666666664, 786.52166666666665,
        779.64333333333332, 772.54333333333341, 765.60000000000002,
        759.44500000000005, 757.98500000000001, 756.55833333333328,
        755.81666666666661, 752.16833333333341, 748.25500000000011,
        744.10000000000002, 740.005, 735.63666666666666, 729.73333333333323,
        725.005, 720.53333333333342, 716.43500000000006, 712.72500000000002]

        self.upper_band_period_6_expected = [np.nan, np.nan, np.nan, np.nan,
        np.nan, 885.00683333333347, 888.62766666666676, 890.88633333333337,
        892.34016666666662, 892.32366666666678, 893.76283333333345,
        895.26800000000014, 895.84366666666676, 894.5053333333334,
        893.94616666666684, 891.25666666666689, 886.82366666666678,
        879.17683333333343, 872.37149999999997, 864.40750000000003,
        856.13549999999998, 852.59900000000016, 852.22683333333339,
        853.98866666666686, 854.3516666666668, 857.01733333333345,
        860.80683333333343, 862.44216666666682, 859.97266666666678,
        860.04050000000007, 860.0001666666667, 859.29616666666664,
        853.46983333333344, 851.07733333333329, 851.86933333333332,
        855.43150000000003, 861.2981666666667, 866.22433333333356,
        871.33566666666684, 873.24966666666671, 874.74016666666682,
        874.72000000000025, 874.33866666666677, 877.55250000000001,
        884.19833333333349, 891.5151666666668, 898.87233333333336,
        904.41999999999996, 907.01599999999985, 907.39549999999997,
        909.18116666666674, 909.06933333333336, 905.08916666666676,
        899.38016666666692, 895.7171666666668, 893.90950000000009,
        890.69749999999999, 889.35550000000012, 888.23166666666668,
        887.95300000000009, 887.43050000000017, 885.78416666666658,
        884.09933333333333, 882.24583333333339, 882.59966666666674,
        883.87933333333342, 885.62100000000009, 885.59533333333331,
        888.26833333333354, 890.44816666666679, 890.88633333333337,
        889.00166666666667, 886.18933333333337, 885.33133333333342,
        882.81416666666678, 878.14100000000019, 875.15083333333348,
        874.97483333333344, 873.7080000000002, 874.60083333333341,
        875.84016666666673, 879.05950000000018, 880.55366666666691,
        879.1933333333335, 879.44450000000006, 877.44250000000011,
        875.21133333333341, 872.49800000000005, 871.47683333333339,
        871.88016666666692, 873.1158333333334, 874.03800000000012,
        874.7401666666666, 876.48733333333348, 879.07600000000002,
        880.4675000000002, 882.18533333333346, 884.03700000000015,
        885.60450000000003, 886.65683333333334, 887.03449999999998,
        887.77516666666668, 887.95300000000009, 886.15633333333346,
        884.45499999999993, 882.92050000000017, 881.72149999999999,
        879.17500000000007, 871.64183333333335, 865.17383333333339,
        857.60766666666677, 849.79766666666683, 842.16000000000008,
        835.38950000000011, 833.78350000000012, 832.21416666666664,
        831.39833333333331, 827.38516666666681, 823.08050000000014,
        818.5100000000001, 814.0055000000001, 809.20033333333333,
        802.70666666666659, 797.5055000000001, 792.58666666666682,
        788.07850000000008, 783.99750000000006]

        self.lower_band_period_6_expected = [np.nan, np.nan, np.nan, np.nan,
        np.nan, 724.09650000000011, 727.05900000000008, 728.90700000000004,
        730.09649999999999, 730.08300000000008, 731.26049999999998,
        732.49200000000008, 732.96299999999997, 731.86799999999994,
        731.41050000000007, 729.21000000000015, 725.58300000000008,
        719.32650000000001, 713.75849999999991, 707.24249999999995,
        700.47449999999992, 697.58100000000002, 697.27650000000006,
        698.71800000000019, 699.0150000000001, 701.19600000000003,
        704.29650000000004, 705.63450000000012, 703.61400000000003,
        703.66950000000008, 703.63649999999996, 703.06049999999993,
        698.29349999999999, 696.3359999999999, 696.98399999999992,
        699.89850000000001, 704.69849999999997, 708.72900000000016,
        712.91100000000017, 714.47699999999998, 715.69650000000001,
        715.68000000000018, 715.36800000000005, 717.99749999999995,
        723.43500000000006, 729.42150000000004, 735.44100000000003,
        739.9799999999999, 742.10399999999981, 742.41449999999986,
        743.87549999999999, 743.78399999999999, 740.52750000000003,
        735.8565000000001, 732.85950000000014, 731.38049999999998,
        728.75249999999994, 727.65449999999998, 726.7349999999999,
        726.50700000000006, 726.07950000000017, 724.73249999999996,
        723.35399999999993, 721.83749999999998, 722.12700000000007,
        723.17399999999998, 724.59900000000005, 724.57799999999997,
        726.7650000000001, 728.54849999999999, 728.90700000000004,
        727.36500000000001, 725.06399999999996, 724.36200000000008,
        722.30250000000012, 718.47900000000004, 716.03250000000003,
        715.88850000000002, 714.85200000000009, 715.5825000000001,
        716.59649999999999, 719.23050000000012, 720.45300000000009,
        719.34000000000015, 719.54550000000006, 717.90750000000003,
        716.08199999999999, 713.86199999999997, 713.02650000000006,
        713.3565000000001, 714.36750000000006, 715.12200000000007,
        715.6964999999999, 717.12600000000009, 719.24400000000003,
        720.38250000000005, 721.78800000000001, 723.30300000000011,
        724.58549999999991, 725.44650000000001, 725.75549999999987,
        726.36149999999998, 726.50700000000006, 725.03700000000003,
        723.64499999999987, 722.38950000000011, 721.4085, 719.32500000000005,
        713.16150000000005, 707.86950000000002, 701.67899999999997,
        695.2890000000001, 689.04000000000008, 683.5005000000001,
        682.18650000000002, 680.90249999999992, 680.23500000000001,
        676.95150000000012, 673.42950000000008, 669.69000000000005,
        666.00450000000001, 662.07299999999998, 656.75999999999988,
        652.50450000000001, 648.48000000000013, 644.79150000000004,
        641.45249999999999]

    def test_center_band_period_6(self):
        period = 6
        cb = moving_average_envelope.center_band(self.data, period)
        np.testing.assert_array_equal(cb, self.center_band_period_6_expected)

    def test_center_band_invalid_period(self):
        period = 128
        with self.assertRaises(Exception) as cm:
            moving_average_envelope.center_band(self.data, period)
        expected = "Error: data_len < period"
        self.assertEqual(str(cm.exception), expected)

    def test_upper_band_period_6(self):
        period = 6
        env_percentage = 0.1
        ub = moving_average_envelope.upper_band(self.data, period, env_percentage)
        np.testing.assert_array_equal(ub, self.upper_band_period_6_expected)

    def test_upper_band_invalid_period(self):
        period = 128
        env_percentage = 0.1
        with self.assertRaises(Exception) as cm:
            moving_average_envelope.upper_band(self.data, period, env_percentage)
        expected = "Error: data_len < period"
        self.assertEqual(str(cm.exception), expected)

    def test_lower_band_period_6(self):
        period = 6
        env_percentage = 0.1
        lb = moving_average_envelope.lower_band(self.data, period, env_percentage)
        np.testing.assert_array_equal(lb, self.lower_band_period_6_expected)

    def test_lower_band_invalid_period(self):
        period = 128
        env_percentage = 0.1
        with self.assertRaises(Exception) as cm:
            moving_average_envelope.lower_band(self.data, period, env_percentage)
        expected = "Error: data_len < period"
        self.assertEqual(str(cm.exception), expected)
