import pytest
from datetime import datetime
from pandas import DataFrame, Series  # type: ignore

from auto_trading.strat.prop import Prop
from auto_trading.indicators.dumb import Dumb
from auto_trading.interfaces import PTFState
from auto_trading.orders import Long, Short


date = datetime.strptime("2015-03-31", "%Y-%m-%d")


@pytest.mark.parametrize(
    "max_delta, data, indicators_results, state, results",
    [
        (
            10,
            DataFrame({"close": {(date, "GOOG"): 10, (date, "GOOGL"): 10}}),
            Series({"GOOG": 1, "GOOGL": -1}),
            PTFState(50, {"GOOG": 0, "GOOGL": 0}),
            [Long("GOOG", price=10, amount=5)],
        ),
        (
            10,
            DataFrame({"close": {(date, "GOOG"): 10, (date, "GOOGL"): 10}}),
            Series({"GOOG": 1, "GOOGL": -1}),
            PTFState(5, {"GOOG": 0, "GOOGL": 0}),
            [],
        ),
        (
            1,
            DataFrame({"close": {(date, "GOOG"): 10, (date, "GOOGL"): 10}}),
            Series({"GOOG": 1, "GOOGL": -1}),
            PTFState(0, {"GOOG": 0, "GOOGL": 1}),
            [Short("GOOGL", price=10, amount=1), Long("GOOG", price=10, amount=1)],
        ),
        (
            20,
            DataFrame(
                {
                    "close": {
                        (date, "A"): 10,
                        (date, "B"): 10,
                        (date, "C"): 10,
                        (date, "D"): 10,
                    }
                }
            ),
            Series({"A": 1, "B": -1, "C": -1, "D": -1}),
            PTFState(0, {"A": 0, "B": 1, "C": 1, "D": 1}),
            [],
        ),
    ],
)
def test_prop(max_delta, data, indicators_results, state, results):
    strat = Prop(max_delta=max_delta, indicators={"ind": Dumb(indicators_results)})
    res = strat.run(data, state)
    assert len(res) == len(results)
    for result in results:
        assert result in res