from __future__ import absolute_import
import unittest
import numpy as np

from tests.sample_data import SampleData
from pyti import stochastic


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

        self.percent_k_period_6_expected = [np.nan, np.nan, np.nan, np.nan,
        np.nan, 0.9856979405034324, 1.0, 1.0, 0.63513513513513342,
        0.27567567567568274, 1.0, 1.0, 0.68322981366460012, 0.0,
        0.15515515515516184, 0.0, 0.0, 0.0, 0.06131650135257203, 0.0, 0.0,
        0.4255711127487089, 1.0, 0.85463958582237798, 0.63201911589008342,
        0.58422939068100166, 0.67256637168141331, 0.55555555555554825, 0.0, 1.0,
        0.39352306182532032, 0.0, 0.0, 0.56253794778384958, 0.82179720704310821,
        1.0, 1.0, 0.83066712049012859, 0.23241362167536711,
        0.059955822025878437, 0.23704663212435031, 0.78950777202072531, 1.0,
        1.0, 0.94086165373294273, 1.0, 1.0, 1.0, 0.36487221315932178,
        0.23273518216421837, 0.38695960311835798, 0.0, 0.0, 0.0, 0.0,
        0.33420252064319617, 0.31533601378518206, 1.0, 0.0, 0.17607726597325543,
        0.038632986627041961, 0.15453194650816784, 0.0, 1.0,
        0.61413043478261453, 1.0, 1.0, 0.21932367149758231, 1.0, 1.0,
        0.17894736842105138, 0.0, 0.0, 0.12548638132295883, 0.2840466926070046,
        0.0, 0.0, 0.80735411670663715, 0.0, 1.0, 1.0, 1.0, 0.42937563971340847,
        0.14943705220061232, 0.0, 0.11392405063290814, 0.32856356631810901,
        0.48005698005698194, 0.24288107202678813, 0.62814070351758511, 1.0, 1.0,
        1.0, 1.0, 1.0, 0.52095130237826281, 1.0, 1.0, 1.0, 1.0,
        0.86164383561643876, 0.0, 0.52147239263801737, 0.0, 0.14857651245551226,
        0.28054740957966762, 0.3811983471074456, 0.0, 0.0, 0.0, 0.0, 0.0,
        0.052040212891779666, 0.0, 0.35317460317461002, 0.0, 0.0, 0.0,
        0.0079254079254060007, 0.0, 0.12661930631007018, 0.0, 0.0, 0.0,
        0.067722772277229157, 0.0, 0.24025100851636036]

        self.percent_k_period_8_expected = [np.nan, np.nan, np.nan, np.nan,
        np.nan, np.nan, np.nan, 1.0, 0.78084415584415301, 0.49576669802445755,
        1.0, 1.0, 0.68940316686967806, 0.0, 0.15515515515516184, 0.0, 0.0, 0.0,
        0.048909134500121687, 0.0, 0.0, 0.25598404255319046,
        0.81420233463035285, 0.79071481208548022, 0.63201911589008342,
        0.58422939068100166, 0.82317801672640178, 0.81521306252488657,
        0.0066371681415952387, 0.75649591685225837, 0.39352306182532032, 0.0,
        0.0, 0.56253794778384958, 0.82179720704310821, 1.0, 1.0,
        0.83066712049012859, 0.47447243022464258, 0.49302246426140284,
        0.41436738752174873, 0.79488797727989935, 0.93264248704663077, 1.0,
        0.94253770150806226, 1.0, 1.0, 1.0, 0.61401189689358671,
        0.45394736842105277, 0.52963567156063163, 0.22512234910277268, 0.0, 0.0,
        0.0, 0.33420252064319617, 0.23859191655801873, 0.43850499782702834, 0.0,
        0.17607726597325543, 0.038632986627041961, 0.15453194650816784, 0.0,
        0.26686004350978676, 0.16388687454677281, 1.0, 1.0, 0.21932367149758231,
        1.0, 1.0, 0.17956423741547525, 0.0, 0.0, 0.12548638132295883,
        0.2840466926070046, 0.0, 0.0, 0.61925199264255404, 0.0, 1.0, 1.0, 1.0,
        0.42937563971340847, 0.14943705220061232, 0.070112589559877536,
        0.17604912998976188, 0.32856356631810901, 0.18547055586131053,
        0.079801871216287013, 0.53418803418803562, 1.0, 1.0, 1.0, 1.0, 1.0,
        0.7004249291784771, 1.0, 1.0, 1.0, 1.0, 0.86164383561643876,
        0.55342465753424508, 0.78630136986300425, 0.0, 0.14857651245551226,
        0.25533807829181515, 0.32829181494662379, 0.0, 0.0, 0.0, 0.0, 0.0,
        0.040534315983417502, 0.0, 0.07229894394801159, 0.0, 0.0, 0.0,
        0.0071881606765310463, 0.0, 0.1097826086956511, 0.0, 0.0, 0.0,
        0.059915907498249425, 0.0, 0.19406227371469995]

        self.percent_k_period_10_expected = [np.nan, np.nan, np.nan, np.nan,
        np.nan, np.nan, np.nan, np.nan, np.nan, 0.76439560439560383, 1.0, 1.0,
        0.74727452923687354, 0.009910802775026999, 0.15515515515516184, 0.0,
        0.0, 0.0, 0.048909134500121687, 0.0, 0.0, 0.22642619094295152,
        0.55651595744680871, 0.47562056737588476, 0.51459143968871746,
        0.54053058216654259, 0.82317801672640178, 0.81521306252488657,
        0.46356033452807566, 0.86937475109517781, 0.30235988200590008, 0.0, 0.0,
        0.56253794778384958, 0.82179720704310821, 1.0, 1.0, 0.83066712049012859,
        0.47447243022464258, 0.49302246426140284, 0.59904697072838564,
        0.88938053097345127, 0.94829729057916878, 1.0, 0.94253770150806226, 1.0,
        1.0, 1.0, 0.78188608776843938, 0.70181741335587489, 0.7141440846001329,
        0.44852941176470656, 0.0, 0.0, 0.0, 0.24289324068224727,
        0.17340492735312743, 0.43850499782702834, 0.0, 0.089840788476118455,
        0.025024061597689246, 0.15453194650816784, 0.0, 0.26686004350978676,
        0.16388687454677281, 0.70195794053661897, 0.75054387237128717,
        0.21932367149758231, 1.0, 1.0, 0.2986512524084754, 0.0, 0.0,
        0.12548638132295883, 0.2840466926070046, 0.0, 0.0, 0.3709144326110913,
        0.0, 0.86767371601208776, 1.0, 1.0, 0.42937563971340847,
        0.14943705220061232, 0.070112589559877536, 0.17604912998976188,
        0.37563971340839536, 0.24257932446264166, 0.079801871216287013,
        0.2063841496973037, 0.37094111172262106, 1.0, 1.0, 1.0, 1.0,
        0.7004249291784771, 1.0, 1.0, 1.0, 1.0, 0.9124783362218376,
        0.63122171945701588, 0.78630136986300425, 0.0, 0.14857651245551226,
        0.25533807829181515, 0.32829181494662379, 0.0, 0.0, 0.0, 0.0, 0.0,
        0.040534315983417502, 0.0, 0.057382333978080118, 0.0, 0.0, 0.0,
        0.0064540622627167372, 0.0, 0.10167785234899253, 0.0, 0.0, 0.0,
        0.037053087757313918, 0.0, 0.17340666450986797]

        self.percent_d_period_6_expected = [np.nan, np.nan, np.nan, np.nan,
        np.nan, np.nan, np.nan, 0.99523264683447754, 0.87837837837837773,
        0.63693693693693865, 0.63693693693693876, 0.75855855855856091,
        0.8944099378882, 0.56107660455486663, 0.27946165627325398,
        0.051718385051720613, 0.051718385051720613, 0.0, 0.020438833784190678,
        0.020438833784190678, 0.020438833784190678, 0.14185703758290297,
        0.47519037091623634, 0.76007023285702902, 0.82888623390415372,
        0.69029603079782087, 0.62960495941749939, 0.60411710597265433,
        0.40937397574565387, 0.51851851851851605, 0.46450768727510677,
        0.46450768727510677, 0.13117435394177343, 0.18751264926128319,
        0.46144505160898591, 0.79477838494231923, 0.9405990690143694,
        0.94355570683004286, 0.68769358072183184, 0.37434552139712474,
        0.17647202527519865, 0.36217007539031804, 0.6755181347150252,
        0.9298359240069084, 0.98028721791098095, 0.98028721791098095,
        0.98028721791098095, 1.0, 0.78829073771977398, 0.53253579844118004,
        0.32818899948063268, 0.20656492842752547, 0.12898653437278598, 0.0, 0.0,
        0.11140084021439872, 0.2165128448094594, 0.54984617814279269,
        0.43844533792839407, 0.39202575532441847, 0.071570084200099124,
        0.12308073303615508, 0.064388311045069938, 0.38484398216938925,
        0.53804347826087151, 0.87137681159420488, 0.87137681159420488,
        0.73977455716586071, 0.73977455716586071, 0.73977455716586071,
        0.7263157894736838, 0.39298245614035049, 0.059649122807017126,
        0.041828793774319611, 0.13651102464332113, 0.13651102464332113,
        0.09468223086900153, 0.26911803890221236, 0.26911803890221236,
        0.60245137223554568, 0.66666666666666663, 1.0, 0.80979187990446944,
        0.52627089730467358, 0.19293756397134029, 0.08778703427784014,
        0.1474958723170057, 0.307514865669333, 0.35050053946729304,
        0.45035958520045166, 0.6236739251814577, 0.87604690117252837, 1.0, 1.0,
        1.0, 0.84031710079275435, 0.84031710079275435, 0.84031710079275435, 1.0,
        1.0, 0.95388127853881288, 0.62054794520547951, 0.46103874275148532,
        0.17382413087933912, 0.22334963503117655, 0.14304130734505996,
        0.27010742304754182, 0.22058191889570442, 0.12706611570248186, 0.0, 0.0,
        0.0, 0.017346737630593221, 0.017346737630593221, 0.13507160535546323,
        0.11772486772487001, 0.11772486772487001, 0.0, 0.0026418026418020004,
        0.0026418026418020004, 0.044848238078492059, 0.04220643543669006,
        0.04220643543669006, 0.0, 0.022574257425743052, 0.022574257425743052,
        0.10265792693119651]

        self.percent_d_period_8_expected = [np.nan, np.nan, np.nan, np.nan,
        np.nan, np.nan, np.nan, np.nan, np.nan, 0.75887028462287021,
        0.75887028462287009, 0.83192223267481913, 0.89646772228989269,
        0.56313438895655932, 0.28151944067494666, 0.051718385051720613,
        0.051718385051720613, 0.0, 0.016303044833373897, 0.016303044833373897,
        0.016303044833373897, 0.085328014184396825, 0.35672879239451444,
        0.62030039642300794, 0.74564542086863883, 0.66898777288552169,
        0.67980884109916229, 0.74087348997742997, 0.54834274913096126,
        0.52611538250624668, 0.38555204893972461, 0.38333965955919291,
        0.13117435394177343, 0.18751264926128319, 0.46144505160898591,
        0.79477838494231923, 0.9405990690143694, 0.94355570683004286,
        0.76837985023825706, 0.59938733832539137, 0.46062076066926472,
        0.56742594302101701, 0.71396595061609303, 0.9091768214421766,
        0.95839339618489772, 0.98084590050268738, 0.98084590050268738, 1.0,
        0.87133729896452883, 0.68931975510487975, 0.53253164562509037,
        0.40290179636148565, 0.25158600688780147, 0.075040783034257555, 0.0,
        0.11140084021439872, 0.19093147906707164, 0.33709981167608111,
        0.22569897146168236, 0.20486075460009459, 0.071570084200099124,
        0.12308073303615508, 0.064388311045069938, 0.14046399667265153,
        0.14358230601885319, 0.4769156393521865, 0.72129562484892418,
        0.73977455716586071, 0.73977455716586071, 0.73977455716586071,
        0.72652141247182511, 0.3931880791384918, 0.05985474580515842,
        0.041828793774319611, 0.13651102464332113, 0.13651102464332113,
        0.09468223086900153, 0.20641733088085135, 0.20641733088085135,
        0.53975066421418472, 0.66666666666666663, 1.0, 0.80979187990446944,
        0.52627089730467358, 0.21630842715796614, 0.13186625725008391,
        0.19157509528924946, 0.23002775072306048, 0.19794533113190219,
        0.26648682042187771, 0.53799663513477425, 0.84472934472934524, 1.0, 1.0,
        1.0, 0.9001416430594924, 0.9001416430594924, 0.9001416430594924, 1.0,
        1.0, 0.95388127853881288, 0.80502283105022787, 0.73378995433789607,
        0.44657534246574976, 0.31162596077283883, 0.1346381969157758,
        0.24406880189798374, 0.19454329774614632, 0.10943060498220793, 0.0, 0.0,
        0.0, 0.013511438661139167, 0.013511438661139167, 0.037611086643809695,
        0.02409964798267053, 0.02409964798267053, 0.0, 0.0023960535588436823,
        0.0023960535588436823, 0.038990256457394047, 0.036594202898550365,
        0.036594202898550365, 0.0, 0.019971969166083143, 0.019971969166083143,
        0.084659393737649788]

        self.percent_d_period_10_expected = [np.nan, np.nan, np.nan, np.nan,
        np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan,
        0.92146520146520128, 0.91575817641229118, 0.58572844400396684,
        0.30411349572235413, 0.055021985976729616, 0.051718385051720613, 0.0,
        0.016303044833373897, 0.016303044833373897, 0.016303044833373897,
        0.075475396980983836, 0.26098071612992008, 0.41952090525521496,
        0.51557598817047035, 0.5102475297437149, 0.62610001286055394,
        0.72630722047261032, 0.70065047125978808, 0.71604938271604668,
        0.54509832254305113, 0.39057821103369261, 0.10078662733530003,
        0.18751264926128319, 0.46144505160898591, 0.79477838494231923,
        0.9405990690143694, 0.94355570683004286, 0.76837985023825706,
        0.59938733832539137, 0.52218062173814372, 0.66048332198774662,
        0.81224159742700186, 0.94589260718420665, 0.96361166402907694,
        0.98084590050268738, 0.98084590050268738, 1.0, 0.92729536258947975,
        0.82790116704143812, 0.73261586190814898, 0.62149696990690473,
        0.38755783212161315, 0.14950980392156885, 0.0, 0.080964413560749085,
        0.13876605601179157, 0.284934388620801, 0.20396997506005191,
        0.17611526210104891, 0.038288283357935902, 0.089798932193991862,
        0.059852002701952366, 0.14046399667265153, 0.14358230601885319,
        0.37756828619772614, 0.53879622915155967, 0.55727516146849621,
        0.65662251462295651, 0.73977455716586071, 0.76621708413615852,
        0.43288375080282515, 0.099550417469491795, 0.041828793774319611,
        0.13651102464332113, 0.13651102464332113, 0.09468223086900153,
        0.12363814420369711, 0.12363814420369711, 0.41286271620772635,
        0.62255790533736255, 0.95589123867069592, 0.80979187990446944,
        0.52627089730467358, 0.21630842715796614, 0.13186625725008391,
        0.20726714431934493, 0.26475605595359963, 0.23267363636244134,
        0.17625511512541078, 0.21904237754540393, 0.52577508713997501,
        0.79031370390754041, 1.0, 1.0, 0.9001416430594924, 0.9001416430594924,
        0.9001416430594924, 1.0, 1.0, 0.97082611207394587, 0.84790001855961783,
        0.7766671418472858, 0.47250769644000673, 0.31162596077283883,
        0.1346381969157758, 0.24406880189798374, 0.19454329774614632,
        0.10943060498220793, 0.0, 0.0, 0.0, 0.013511438661139167,
        0.013511438661139167, 0.032638883320499211, 0.019127444659360039,
        0.019127444659360039, 0.0, 0.0021513540875722457, 0.0021513540875722457,
        0.036043971537236423, 0.033892617449664174, 0.033892617449664174, 0.0,
        0.012351029252437973, 0.012351029252437973, 0.070153250755727301]

    def test_percent_k_period_6(self):
        period = 6
        percent_k = stochastic.percent_k(self.data, period)
        np.testing.assert_array_equal(percent_k, self.percent_k_period_6_expected)

    def test_percent_k_period_8(self):
        period = 8
        percent_k = stochastic.percent_k(self.data, period)
        np.testing.assert_array_equal(percent_k, self.percent_k_period_8_expected)

    def test_percent_k_period_10(self):
        period = 10
        percent_k = stochastic.percent_k(self.data, period)
        np.testing.assert_array_equal(percent_k, self.percent_k_period_10_expected)

    def test_percent_k_invalid_period(self):
        period = 128
        with self.assertRaises(Exception):
            stochastic.percent_k(self.data, period)

    def test_percent_d_period_6(self):
        period = 6
        percent_d = stochastic.percent_d(self.data, period)
        np.testing.assert_array_equal(percent_d, self.percent_d_period_6_expected)

    def test_percent_d_period_8(self):
        period = 8
        percent_d = stochastic.percent_d(self.data, period)
        np.testing.assert_array_equal(percent_d, self.percent_d_period_8_expected)

    def test_percent_d_period_10(self):
        period = 10
        percent_d = stochastic.percent_d(self.data, period)
        np.testing.assert_array_equal(percent_d, self.percent_d_period_10_expected)

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