import random
from typing import List

from absl.testing import absltest
import apache_beam as beam
from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.testing.util import assert_that
from apache_beam.testing.util import equal_to
from google.protobuf import text_format

from xlab.base import time
from xlab.data.pipeline import produce_calc_fn
from xlab.data.proto import data_entry_pb2
from xlab.data.proto import data_type_pb2
from xlab.net.proto import time_util
from xlab.trading.dates import trading_days

DataEntry = data_entry_pb2.DataEntry
DataType = data_type_pb2.DataType
produce_source_calc_fn = produce_calc_fn.produce_source_calc_fn
produce_recursive_calc_fn = produce_calc_fn.produce_recursive_calc_fn


def _make_close_price(symbol: str, value: float, t: time.Time) -> DataEntry:
    return DataEntry(
        symbol=symbol,
        data_space=DataEntry.STOCK_DATA,
        data_type=data_type_pb2.DataType.CLOSE_PRICE,
        value=value,
        timestamp=time_util.from_time(t),
    )


def get_close_prices_for_ema20(t: time.Time, symbol: str) -> List[DataEntry]:
    # INTC 2019-12-03 to 2019-12-31
    price_data = [
        56.07, 56.02, 56.08, 56.81, 56.53, 56.59, 57.07, 57.55, 57.79, 57.70,
        57.23, 57.17, 57.96, 58.95, 59.23, 59.41, 59.82, 60.08, 59.62, 59.85,
    ]  # yapf: disable

    close_prices = [
        _make_close_price(symbol, value, time.FromCivil(d)) for value, d in zip(
            price_data, trading_days.get_last_n(time.ToCivil(t), 20))
    ]
    random.shuffle(close_prices)
    return close_prices


def get_recursive_inputs_for_ema20(t: time.Time,
                                   symbol: str) -> List[DataEntry]:
    return [
        DataEntry(
            symbol=symbol,
            data_space=DataEntry.STOCK_DATA,
            data_type=DataType.EMA_20D,
            value=58.31,
            timestamp=time_util.from_time(t - time.Hours(24)),
        ),
        _make_close_price(symbol, 59.85, t)
    ]


def expected_ema20d_test() -> DataEntry:
    return text_format.Parse(
        f"""
        symbol: "TEST"
        data_space: STOCK_DATA
        data_type: EMA_20D
        value: 58.255796513052346
        timestamp {{
            seconds: {time.as_seconds(2019, 12, 31)}
        }}""", DataEntry())


def expected_recursive_ema20d_test() -> DataEntry:
    return text_format.Parse(
        f"""
        symbol: "TEST"
        data_space: STOCK_DATA
        data_type: EMA_20D
        value: 58.45666666666667
        timestamp {{
            seconds: {time.as_seconds(2019, 12, 31)}
        }}""", DataEntry())


class ProduceCalcFnTest(absltest.TestCase):

    def test_produce_source_calc_fn_with_no_input(self):
        t = time.FromUnixSeconds(12345)
        inputs = []
        with TestPipeline() as p:
            out = (p \
                | beam.Create(inputs) \
                | produce_source_calc_fn(DataType.EMA_20D, t))

            assert_that(out, equal_to([]))

    def test_produce_source_calc_fn(self):
        t = time.FromCivil(time.CivilTime(2019, 12, 31))
        inputs = get_close_prices_for_ema20(t, 'TEST')
        with TestPipeline() as p:
            out = (p \
                | beam.Create(inputs) \
                | produce_source_calc_fn(DataType.EMA_20D, t))

            assert_that(out, equal_to([expected_ema20d_test()]))

    def test_produce_source_calc_fn_ignores_irrelevant_data(self):
        t = time.FromCivil(time.CivilTime(2019, 12, 31))
        inputs = get_close_prices_for_ema20(t, 'TEST')
        inputs.append(_make_close_price('X', 123.45, t))
        with TestPipeline() as p:
            out = (p \
                | beam.Create(inputs) \
                | produce_source_calc_fn(DataType.EMA_20D, t))

            assert_that(out, equal_to([expected_ema20d_test()]))

    def test_produce_source_calc_fn_not_enough_data(self):
        t = time.FromCivil(time.CivilTime(2019, 12, 31))
        inputs = get_close_prices_for_ema20(t, 'TEST')[:19]
        with TestPipeline() as p:
            out = (p \
                | beam.Create(inputs) \
                | produce_source_calc_fn(DataType.EMA_20D, t))

            assert_that(out, equal_to([]))

    def test_produce_source_calc_fn_multiple_symbols(self):
        t = time.FromCivil(time.CivilTime(2019, 12, 31))
        inputs = [
            *get_close_prices_for_ema20(t, 'TEST'),
            *get_close_prices_for_ema20(t, 'IBM')
        ]
        with TestPipeline() as p:
            out = (p \
                | beam.Create(inputs) \
                | produce_source_calc_fn(DataType.EMA_20D, t))

            assert_that(
                out,
                equal_to([
                    expected_ema20d_test(),
                    text_format.Parse(
                        f"""
                        symbol: "IBM"
                        data_space: STOCK_DATA
                        data_type: EMA_20D
                        value: 58.255796513052346
                        timestamp {{
                            seconds: {time.as_seconds(2019, 12, 31)}
                        }}""", DataEntry())
                ]))

    def test_produce_recursive_calc_fn(self):
        t = time.FromCivil(time.CivilTime(2019, 12, 31))
        inputs = get_recursive_inputs_for_ema20(t, 'TEST')
        with TestPipeline() as p:
            out = (p \
                | beam.Create(inputs) \
                | produce_recursive_calc_fn(DataType.EMA_20D, t))

            assert_that(out, equal_to([expected_recursive_ema20d_test()]))

    def test_produce_recursive_calc_fn_with_swapped_inputs(self):
        t = time.FromCivil(time.CivilTime(2019, 12, 31))
        inputs = reversed(get_recursive_inputs_for_ema20(t, 'TEST'))
        with TestPipeline() as p:
            out = (p \
                | beam.Create(inputs) \
                | produce_recursive_calc_fn(DataType.EMA_20D, t))

            assert_that(out, equal_to([expected_recursive_ema20d_test()]))

    def test_produce_recursive_calc_fn_with_extra_data(self):
        t = time.FromCivil(time.CivilTime(2019, 12, 31))
        inputs = [
            *get_recursive_inputs_for_ema20(t, 'TEST'),
            *get_close_prices_for_ema20(t, 'IBM')
        ]
        with TestPipeline() as p:
            out = (p \
                | beam.Create(inputs) \
                | produce_recursive_calc_fn(DataType.EMA_20D, t))

            assert_that(out, equal_to([expected_recursive_ema20d_test()]))


if __name__ == '__main__':
    absltest.main()

import random
from typing import List

from absl.testing import absltest
import apache_beam as beam
from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.testing.util import assert_that
from apache_beam.testing.util import equal_to
from google.protobuf import text_format

from xlab.base import time
from xlab.data.pipeline import produce_calc_fn
from xlab.data.proto import data_entry_pb2
from xlab.data.proto import data_type_pb2
from xlab.net.proto import time_util
from xlab.trading.dates import trading_days

DataEntry = data_entry_pb2.DataEntry
DataType = data_type_pb2.DataType
produce_source_calc_fn = produce_calc_fn.produce_source_calc_fn
produce_recursive_calc_fn = produce_calc_fn.produce_recursive_calc_fn


def _make_close_price(symbol: str, value: float, t: time.Time) -> DataEntry:
    return DataEntry(
        symbol=symbol,
        data_space=DataEntry.STOCK_DATA,
        data_type=data_type_pb2.DataType.CLOSE_PRICE,
        value=value,
        timestamp=time_util.from_time(t),
    )


def get_close_prices_for_ema20(t: time.Time, symbol: str) -> List[DataEntry]:
    # INTC 2019-12-03 to 2019-12-31
    price_data = [
        56.07, 56.02, 56.08, 56.81, 56.53, 56.59, 57.07, 57.55, 57.79, 57.70,
        57.23, 57.17, 57.96, 58.95, 59.23, 59.41, 59.82, 60.08, 59.62, 59.85,
    ]  # yapf: disable

    close_prices = [
        _make_close_price(symbol, value, time.FromCivil(d)) for value, d in zip(
            price_data, trading_days.get_last_n(time.ToCivil(t), 20))
    ]
    random.shuffle(close_prices)
    return close_prices


def get_recursive_inputs_for_ema20(t: time.Time,
                                   symbol: str) -> List[DataEntry]:
    return [
        DataEntry(
            symbol=symbol,
            data_space=DataEntry.STOCK_DATA,
            data_type=DataType.EMA_20D,
            value=58.31,
            timestamp=time_util.from_time(t - time.Hours(24)),
        ),
        _make_close_price(symbol, 59.85, t)
    ]


def expected_ema20d_test() -> DataEntry:
    return text_format.Parse(
        f"""
        symbol: "TEST"
        data_space: STOCK_DATA
        data_type: EMA_20D
        value: 58.255796513052346
        timestamp {{
            seconds: {time.as_seconds(2019, 12, 31)}
        }}""", DataEntry())


def expected_recursive_ema20d_test() -> DataEntry:
    return text_format.Parse(
        f"""
        symbol: "TEST"
        data_space: STOCK_DATA
        data_type: EMA_20D
        value: 58.45666666666667
        timestamp {{
            seconds: {time.as_seconds(2019, 12, 31)}
        }}""", DataEntry())


class ProduceCalcFnTest(absltest.TestCase):

    def test_produce_source_calc_fn_with_no_input(self):
        t = time.FromUnixSeconds(12345)
        inputs = []
        with TestPipeline() as p:
            out = (p \
                | beam.Create(inputs) \
                | produce_source_calc_fn(DataType.EMA_20D, t))

            assert_that(out, equal_to([]))

    def test_produce_source_calc_fn(self):
        t = time.FromCivil(time.CivilTime(2019, 12, 31))
        inputs = get_close_prices_for_ema20(t, 'TEST')
        with TestPipeline() as p:
            out = (p \
                | beam.Create(inputs) \
                | produce_source_calc_fn(DataType.EMA_20D, t))

            assert_that(out, equal_to([expected_ema20d_test()]))

    def test_produce_source_calc_fn_ignores_irrelevant_data(self):
        t = time.FromCivil(time.CivilTime(2019, 12, 31))
        inputs = get_close_prices_for_ema20(t, 'TEST')
        inputs.append(_make_close_price('X', 123.45, t))
        with TestPipeline() as p:
            out = (p \
                | beam.Create(inputs) \
                | produce_source_calc_fn(DataType.EMA_20D, t))

            assert_that(out, equal_to([expected_ema20d_test()]))

    def test_produce_source_calc_fn_not_enough_data(self):
        t = time.FromCivil(time.CivilTime(2019, 12, 31))
        inputs = get_close_prices_for_ema20(t, 'TEST')[:19]
        with TestPipeline() as p:
            out = (p \
                | beam.Create(inputs) \
                | produce_source_calc_fn(DataType.EMA_20D, t))

            assert_that(out, equal_to([]))

    def test_produce_source_calc_fn_multiple_symbols(self):
        t = time.FromCivil(time.CivilTime(2019, 12, 31))
        inputs = [
            *get_close_prices_for_ema20(t, 'TEST'),
            *get_close_prices_for_ema20(t, 'IBM')
        ]
        with TestPipeline() as p:
            out = (p \
                | beam.Create(inputs) \
                | produce_source_calc_fn(DataType.EMA_20D, t))

            assert_that(
                out,
                equal_to([
                    expected_ema20d_test(),
                    text_format.Parse(
                        f"""
                        symbol: "IBM"
                        data_space: STOCK_DATA
                        data_type: EMA_20D
                        value: 58.255796513052346
                        timestamp {{
                            seconds: {time.as_seconds(2019, 12, 31)}
                        }}""", DataEntry())
                ]))

    def test_produce_recursive_calc_fn(self):
        t = time.FromCivil(time.CivilTime(2019, 12, 31))
        inputs = get_recursive_inputs_for_ema20(t, 'TEST')
        with TestPipeline() as p:
            out = (p \
                | beam.Create(inputs) \
                | produce_recursive_calc_fn(DataType.EMA_20D, t))

            assert_that(out, equal_to([expected_recursive_ema20d_test()]))

    def test_produce_recursive_calc_fn_with_swapped_inputs(self):
        t = time.FromCivil(time.CivilTime(2019, 12, 31))
        inputs = reversed(get_recursive_inputs_for_ema20(t, 'TEST'))
        with TestPipeline() as p:
            out = (p \
                | beam.Create(inputs) \
                | produce_recursive_calc_fn(DataType.EMA_20D, t))

            assert_that(out, equal_to([expected_recursive_ema20d_test()]))

    def test_produce_recursive_calc_fn_with_extra_data(self):
        t = time.FromCivil(time.CivilTime(2019, 12, 31))
        inputs = [
            *get_recursive_inputs_for_ema20(t, 'TEST'),
            *get_close_prices_for_ema20(t, 'IBM')
        ]
        with TestPipeline() as p:
            out = (p \
                | beam.Create(inputs) \
                | produce_recursive_calc_fn(DataType.EMA_20D, t))

            assert_that(out, equal_to([expected_recursive_ema20d_test()]))


if __name__ == '__main__':
    absltest.main()

# INTC 2019-12-03 to 2019-12-31 # yapf: disable symbol: "TEST"         data_space: STOCK_DATA         data_type: EMA_20D         value: 58.255796513052346         timestamp {{             seconds: {time.as_seconds(2019, 12, 31)}         }} symbol: "TEST"         data_space: STOCK_DATA         data_type: EMA_20D         value: 58.45666666666667         timestamp {{             seconds: {time.as_seconds(2019, 12, 31)}         }} symbol: "IBM"                         data_space: STOCK_DATA                         data_type: EMA_20D                         value: 58.255796513052346                         timestamp {{                             seconds: {time.as_seconds(2019, 12, 31)}                         }}