"""Unit tests for context.py module."""

"""Copyright (C) 2023 Edward West

This program is free software: you can redistribute it and/or modify it under
the terms of the GNU Lesser General Public License as published by the Free
Software Foundation, either version 3 of the License, or (at your option) any
later version.

This program is distributed in the hope that it will be useful, but WITHOUT ANY
WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A
PARTICULAR PURPOSE. See the GNU Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public License along
with this program.  If not, see <https://www.gnu.org/licenses/>.
"""

import numpy as np
import pytest
import re
from .fixtures import *
from collections import deque
from decimal import Decimal
from pybroker.common import PriceType, StopType, to_datetime
from pybroker.config import StrategyConfig
from pybroker.context import (
    ExecContext,
    ExecResult,
    PosSizeContext,
    set_exec_ctx_data,
    set_pos_size_ctx_data,
)
from pybroker.portfolio import Order, Portfolio, Position, Trade


@pytest.fixture()
def portfolio():
    return Portfolio(100_000)


@pytest.fixture()
def end_index():
    return 100


@pytest.fixture()
def sym_end_index(symbols, end_index):
    return {sym: end_index for sym in symbols}


@pytest.fixture()
def session():
    return {"foo": 1, "bar": 2}


@pytest.fixture()
def foreign(symbols):
    return list(symbols)[-1]


@pytest.fixture()
def date(dates, end_index):
    return list(dates)[end_index - 1]


@pytest.fixture()
def orders(dates, symbols):
    return (
        Order(
            id=1,
            date=dates[0],
            symbol=symbols[1],
            type="buy",
            limit_price=None,
            fill_price=10,
            shares=200,
            fees=0,
        ),
        Order(
            id=2,
            date=dates[1],
            symbol=symbols[2],
            type="sell",
            limit_price=100,
            fill_price=101.1,
            shares=100,
            fees=0,
        ),
    )


@pytest.fixture()
def trades(dates, symbols):
    return Trade(
        id=1,
        type="long",
        symbol=symbols[-1],
        entry_date=dates[0],
        exit_date=dates[1],
        entry=100,
        exit=101,
        shares=100,
        pnl=Decimal(100),
        return_pct=Decimal(5),
        agg_pnl=Decimal(100),
        bars=1,
        pnl_per_bar=Decimal(100),
        stop=None,
    )


@pytest.fixture()
def ctx(
    col_scope,
    ind_scope,
    input_scope,
    pred_scope,
    pending_order_scope,
    portfolio,
    trained_models,
    sym_end_index,
    session,
    symbol,
    date,
):
    ctx = ExecContext(
        symbol=symbol,
        config=StrategyConfig(max_long_positions=5),
        portfolio=portfolio,
        col_scope=col_scope,
        ind_scope=ind_scope,
        input_scope=input_scope,
        pred_scope=pred_scope,
        pending_order_scope=pending_order_scope,
        models=trained_models,
        sym_end_index=sym_end_index,
        session=session,
    )
    set_exec_ctx_data(ctx, date)
    return ctx


@pytest.fixture()
def ctx_with_pos(
    col_scope,
    ind_scope,
    input_scope,
    pred_scope,
    pending_order_scope,
    portfolio,
    trained_models,
    sym_end_index,
    session,
    symbol,
    symbols,
    date,
):
    portfolio.long_positions = {
        sym: Position(sym, 200, "long") for sym in symbols
    }
    portfolio.short_positions = {
        sym: Position(sym, 100, "short") for sym in symbols
    }
    ctx = ExecContext(
        symbol=symbol,
        config=StrategyConfig(max_long_positions=5),
        portfolio=portfolio,
        col_scope=col_scope,
        ind_scope=ind_scope,
        input_scope=input_scope,
        pred_scope=pred_scope,
        pending_order_scope=pending_order_scope,
        models=trained_models,
        sym_end_index=sym_end_index,
        session=session,
    )
    set_exec_ctx_data(ctx, date)
    return ctx


@pytest.fixture()
def ctx_with_orders(
    col_scope,
    ind_scope,
    input_scope,
    pred_scope,
    pending_order_scope,
    portfolio,
    trained_models,
    sym_end_index,
    session,
    symbol,
    date,
    orders,
    trades,
):
    portfolio.orders = deque(orders)
    portfolio.trades = deque(trades)
    portfolio.win_rate = 1
    portfolio.lose_rate = 0
    ctx = ExecContext(
        symbol=symbol,
        config=StrategyConfig(max_long_positions=5),
        portfolio=portfolio,
        col_scope=col_scope,
        ind_scope=ind_scope,
        input_scope=input_scope,
        pred_scope=pred_scope,
        pending_order_scope=pending_order_scope,
        models=trained_models,
        sym_end_index=sym_end_index,
        session=session,
    )
    set_exec_ctx_data(ctx, date)
    return ctx


def test_config(ctx):
    assert ctx.config.max_long_positions == 5


def test_dt(ctx, date):
    assert ctx.dt == to_datetime(date)


def test_win_rate(ctx_with_orders):
    assert ctx_with_orders.win_rate == 1


def test_loss_rate(ctx_with_orders):
    assert ctx_with_orders.loss_rate == 0


@pytest.mark.parametrize(
    "field", ["date", "open", "high", "low", "close", "volume", "adj_close"]
)
def test_fields(scope, ctx, field, end_index):
    scope.register_custom_cols("adj_close")
    assert len(getattr(ctx, field)) == end_index


def test_field_not_found_then_error(ctx):
    with pytest.raises(
        AttributeError, match=re.escape("Attribute 'foo' not found.")
    ):
        ctx.foo


def test_empty_field(ctx):
    assert ctx.vwap is None


@pytest.mark.parametrize(
    "field, port_field",
    [
        ("total_equity", "equity"),
        ("cash", "cash"),
        ("total_margin", "margin"),
        ("total_market_value", "market_value"),
    ],
)
def test_portfolio_field(ctx, portfolio, field, port_field):
    assert getattr(ctx, field) == getattr(portfolio, port_field)


def test_sell_all_shares(ctx_with_pos):
    ctx_with_pos.sell_all_shares()
    assert ctx_with_pos.sell_shares == 200


def test_sell_all_shares_when_no_position(ctx):
    ctx.sell_all_shares()
    assert ctx.sell_shares is None


def test_cover_all_shares(ctx_with_pos):
    ctx_with_pos.cover_all_shares()
    assert ctx_with_pos.buy_shares == 100


def test_cover_all_shares_when_no_position(ctx):
    ctx.cover_all_shares()
    assert ctx.buy_shares is None


def test_model(ctx, trained_models, symbol):
    assert (
        ctx.model(MODEL_NAME) == trained_models[(MODEL_NAME, symbol)].instance
    )


def test_model_when_not_found_then_error(ctx, symbol):
    with pytest.raises(
        ValueError,
        match=re.escape(f"Model 'undefined_model' not found for {symbol}."),
    ):
        ctx.model("undefined_model")


def test_indicator(ctx, end_index):
    assert len(ctx.indicator("hhv")) == end_index


def test_indicator_when_not_found_then_error(ctx, symbol):
    with pytest.raises(
        ValueError,
        match=re.escape(
            f"Indicator 'undefined_indicator' not found for {symbol}."
        ),
    ):
        ctx.indicator("undefined_indicator")


def test_input(ctx, end_index):
    assert len(ctx.input(MODEL_NAME)["hhv"]) == end_index


def test_input_when_not_found_then_error(ctx):
    with pytest.raises(
        ValueError, match=re.escape("Model 'undefined_model' not found.")
    ):
        ctx.input("undefined_model")


def test_preds(ctx, end_index):
    assert len(ctx.preds(MODEL_NAME)) == end_index


def test_preds_when_not_found_then_error(ctx):
    with pytest.raises(
        ValueError, match=re.escape("Model 'undefined_model' not found.")
    ):
        ctx.preds("undefined_model")


@pytest.mark.parametrize("pos_type", ["long", "short"])
def test_position(ctx_with_pos, pos_type, portfolio, symbol):
    assert (
        getattr(ctx_with_pos, f"{pos_type}_pos")()
        == getattr(portfolio, f"{pos_type}_positions")[symbol]
    )


@pytest.mark.parametrize("pos_fn", ["long_pos", "short_pos"])
def test_position_when_empty(ctx, pos_fn):
    assert getattr(ctx, pos_fn)() is None


def test_position_when_invalid_pos_type_then_error(ctx, symbol):
    with pytest.raises(
        ValueError, match=re.escape("Unknown pos_type: 'invalid'.")
    ):
        ctx.pos(symbol, "invalid")


@pytest.mark.parametrize("pos_type", ["long", "short"])
def test_position_with_foreign_when_empty(ctx, pos_type, foreign):
    assert getattr(ctx, f"{pos_type}_pos")(foreign) is None


@pytest.mark.parametrize("pos_type", ["long", "short", None])
def test_positions(ctx_with_pos, pos_type, portfolio):
    positions = ctx_with_pos.positions(None, pos_type)
    if pos_type is None:
        expected_positions = set(portfolio.long_positions.keys()) | set(
            portfolio.short_positions.keys()
        )
    else:
        expected_positions = set(
            getattr(portfolio, f"{pos_type}_positions").keys()
        )
    assert set(map(lambda p: p.symbol, positions)) == expected_positions


@pytest.mark.parametrize("pos_type", ["long", "short", None])
def test_positions_when_empty(ctx, pos_type):
    assert not len(list(ctx.positions(None, pos_type)))


@pytest.mark.parametrize("pos_type", ["long", "short", None])
def test_positions_with_symbol(ctx_with_pos, pos_type, foreign):
    positions = list(ctx_with_pos.positions(foreign, pos_type))
    if pos_type is None:
        assert len(positions) == 2
        assert positions[0].symbol == foreign
        assert positions[1].symbol == foreign
    else:
        assert len(positions) == 1
        assert positions[0].symbol == foreign


@pytest.mark.parametrize("pos_type", ["long", "short", None])
def test_positions_with_symbol_when_empty(ctx, pos_type, foreign):
    assert not len(list(ctx.positions(foreign, pos_type)))


@pytest.mark.parametrize(
    "col", ["date", "open", "high", "low", "close", "volume", "adj_close"]
)
def test_foreign(ctx, col, foreign, data_source_df, end_index):
    df = data_source_df[data_source_df["symbol"] == foreign]
    assert (ctx.foreign(foreign, col) == df[col].values[:end_index]).all()
    df = data_source_df[data_source_df["symbol"] == foreign]
    assert (ctx.foreign(foreign, col) == df[col].values[:end_index]).all()


def test_foreign_with_custom_col(
    scope, ctx, foreign, data_source_df, end_index
):
    scope.register_custom_cols("adj_close")
    df = data_source_df[data_source_df["symbol"] == foreign]
    assert (
        ctx.foreign(foreign, "adj_close") == df["adj_close"].values[:end_index]
    ).all()


def test_foreign_when_empty(ctx, foreign):
    assert ctx.foreign(foreign, "foo") is None


def test_foreign_when_symbol_not_found_then_error(ctx):
    with pytest.raises(ValueError, match=re.escape("Symbol 'FOO' not found.")):
        ctx.foreign("FOO", "close")


def test_foreign_with_empty_col(
    scope, ctx, foreign, data_source_df, end_index
):
    scope.register_custom_cols("adj_close")

    df = data_source_df[data_source_df["symbol"] == foreign]

    def verify_bar_data(bar_data):
        for col in (
            "date",
            "open",
            "high",
            "low",
            "close",
            "volume",
            "adj_close",
        ):
            assert (getattr(bar_data, col) == df[col].values[:end_index]).all()

    verify_bar_data(ctx.foreign(foreign))
    verify_bar_data(ctx.foreign(foreign))


def test_long_positions(ctx_with_pos, symbol):
    positions = tuple(ctx_with_pos.long_positions(symbol))
    assert len(positions) == 1
    assert positions[0] == Position(symbol, 200, "long")


def test_long_positions_when_empty(ctx, symbol):
    assert not len(tuple(ctx.long_positions(symbol)))


def test_short_positions(ctx_with_pos, symbol):
    positions = tuple(ctx_with_pos.short_positions(symbol))
    assert len(positions) == 1
    assert positions[0] == Position(symbol, 100, "short")


def test_short_positions_when_empty(ctx, symbol):
    assert not len(tuple(ctx.short_positions(symbol)))


def test_calc_target_shares(ctx):
    assert ctx.calc_target_shares(0.5, 33.50) == 50_000 // 33.5


def test_calc_target_shares_with_cash(ctx):
    assert ctx.calc_target_shares(1 / 3, 20, 10_000) == 166


def test_to_result(ctx, symbol, date):
    ctx.buy_fill_price = PriceType.AVERAGE
    ctx.buy_shares = 20
    ctx.buy_limit_price = 99.99
    ctx.sell_fill_price: PriceType = PriceType.HIGH
    ctx.sell_shares = 20
    ctx.sell_limit_price = 110.11
    ctx.hold_bars = 2
    ctx.score = 7
    result = ctx.to_result()
    assert result.symbol == symbol
    assert result.date == date
    assert result.buy_fill_price == PriceType.AVERAGE
    assert result.buy_shares == 20
    assert result.buy_limit_price == Decimal("99.99")
    assert result.sell_fill_price == PriceType.HIGH
    assert result.sell_shares == 20
    assert result.sell_limit_price == Decimal("110.11")
    assert result.hold_bars == 2
    assert result.score == 7
    assert len(result.long_stops) == 1
    assert result.short_stops is None
    stop = next(iter(result.long_stops))
    assert stop.symbol == symbol
    assert stop.stop_type == StopType.BAR
    assert stop.pos_type == "long"
    assert stop.percent is None
    assert stop.points is None
    assert stop.bars == 2
    assert stop.fill_price == PriceType.HIGH
    assert stop.limit_price is None


@pytest.mark.parametrize("pos_type", ["long", "short"])
@pytest.mark.parametrize(
    "stop_attr, expected_stop_type",
    [
        ("stop_loss", StopType.LOSS),
        ("stop_loss_pct", StopType.LOSS),
        ("stop_profit", StopType.PROFIT),
        ("stop_profit_pct", StopType.PROFIT),
        ("stop_trailing", StopType.TRAILING),
        ("stop_trailing_pct", StopType.TRAILING),
    ],
)
def test_to_result_when_stop(
    ctx, symbol, date, pos_type, stop_attr, expected_stop_type
):
    stop_limit = 200
    stop_amount = 20
    percent = None
    points = None
    if stop_attr.endswith("_pct"):
        percent = stop_amount
    else:
        points = stop_amount
    buy_shares = None
    sell_shares = None
    if pos_type == "long":
        buy_shares = 100
    else:
        sell_shares = 100
    ctx.buy_shares = buy_shares
    ctx.sell_shares = sell_shares
    setattr(ctx, stop_attr, stop_amount)
    setattr(ctx, f"{stop_attr.replace('_pct', '')}_limit", stop_limit)
    result = ctx.to_result()
    assert result.symbol == symbol
    assert result.date == date
    assert result.buy_fill_price == PriceType.MIDDLE
    assert result.buy_limit_price is None
    assert result.sell_fill_price == PriceType.MIDDLE
    assert result.sell_limit_price is None
    assert result.hold_bars is None
    assert result.score is None
    if pos_type == "long":
        assert result.buy_shares == 100
        assert result.sell_shares is None
        assert len(result.long_stops) == 1
        assert result.short_stops is None
        stop = next(iter(result.long_stops))
    else:
        assert result.sell_shares == 100
        assert result.buy_shares is None
        assert len(result.short_stops) == 1
        assert result.long_stops is None
        stop = next(iter(result.short_stops))
    assert stop.symbol == symbol
    assert stop.stop_type == expected_stop_type
    assert stop.pos_type == pos_type
    assert stop.percent == percent
    assert stop.points == points
    assert stop.bars is None
    assert stop.fill_price is None
    assert stop.limit_price == stop_limit


@pytest.mark.parametrize(
    "stop_attr", ["stop_loss", "stop_profit", "stop_trailing"]
)
def test_to_result_when_stop_pct_and_points_then_error(ctx, stop_attr):
    ctx.buy_shares = 100
    with pytest.raises(
        ValueError,
        match=re.escape(
            f"Only one of {stop_attr} or {stop_attr}_pct can be set."
        ),
    ):
        setattr(ctx, stop_attr, 20)
        setattr(ctx, f"{stop_attr}_pct", 20)
        ctx.to_result()


def test_orders(ctx_with_orders, orders):
    assert tuple(ctx_with_orders.orders()) == orders


def test_orders_when_empty(ctx):
    assert not len(list(ctx.orders()))


def test_trades(ctx_with_orders, trades):
    assert tuple(ctx_with_orders.trades()) == trades


def test_trades_when_empty(ctx):
    assert not len(list(ctx.trades()))


@pytest.mark.parametrize(
    "cover_attr, buy_attr",
    [
        ("cover_fill_price", "buy_fill_price"),
        ("cover_shares", "buy_shares"),
        ("cover_limit_price", "buy_limit_price"),
    ],
)
def test_cover(ctx, cover_attr, buy_attr):
    setattr(ctx, cover_attr, 99)
    assert getattr(ctx, cover_attr) == 99
    assert getattr(ctx, cover_attr) == getattr(ctx, buy_attr)
    assert ctx._cover is True


def test_set_exec_ctx_data(ctx, sym_end_index):
    date = np.datetime64("2020-01-01")
    ctx._foreign = {"SPY": np.random.rand(100)}
    ctx._cover = True
    ctx.buy_fill_price = PriceType.AVERAGE
    ctx.buy_shares = 100
    ctx.buy_limit_price = 99
    ctx.sell_fill_price = PriceType.CLOSE
    ctx.sell_shares = 200
    ctx.sell_limit_price = 80
    ctx.hold_bars = 5
    ctx.score = 45.5
    ctx.stop_loss = 10
    ctx.stop_loss_pct = 20
    ctx.stop_loss_limit = 99
    ctx.stop_profit = 20
    ctx.stop_profit_pct = 30
    ctx.stop_profit_limit = 99.99
    ctx.stop_trailing = 100
    ctx.stop_trailing_pct = 15
    ctx.stop_trailing_limit = 80.8
    set_exec_ctx_data(ctx, date)
    assert ctx.dt == to_datetime(date)
    assert ctx.bars == sym_end_index[ctx.symbol]
    assert not ctx._foreign
    assert ctx._cover is False
    assert ctx.buy_fill_price == PriceType.MIDDLE
    assert ctx.buy_shares is None
    assert ctx.buy_limit_price is None
    assert ctx.sell_fill_price == PriceType.MIDDLE
    assert ctx.sell_shares is None
    assert ctx.sell_limit_price is None
    assert ctx.hold_bars is None
    assert ctx.score is None
    assert ctx.stop_loss is None
    assert ctx.stop_loss_pct is None
    assert ctx.stop_loss_limit is None
    assert ctx.stop_profit is None
    assert ctx.stop_profit_pct is None
    assert ctx.stop_profit_limit is None
    assert ctx.stop_trailing is None
    assert ctx.stop_trailing_pct is None
    assert ctx.stop_trailing_limit is None


def test_set_pos_ctx_data(
    date,
    portfolio,
    col_scope,
    ind_scope,
    input_scope,
    pred_scope,
    pending_order_scope,
    trained_models,
    sym_end_index,
):
    buy_results = [
        ExecResult(
            symbol="SPY",
            date=date,
            buy_shares=100,
            buy_fill_price=99,
            buy_limit_price=99,
            sell_shares=None,
            sell_fill_price=None,
            sell_limit_price=None,
            hold_bars=None,
            score=1,
            long_stops=None,
            short_stops=None,
        ),
        ExecResult(
            symbol="AAPL",
            date=date,
            buy_shares=200,
            buy_fill_price=90,
            buy_limit_price=90,
            sell_shares=None,
            sell_fill_price=None,
            sell_limit_price=None,
            hold_bars=None,
            score=2,
            long_stops=None,
            short_stops=None,
        ),
    ]
    sell_results = [
        ExecResult(
            symbol="TSLA",
            date=date,
            buy_shares=None,
            buy_fill_price=None,
            buy_limit_price=None,
            sell_shares=100,
            sell_fill_price=80,
            sell_limit_price=80,
            hold_bars=None,
            score=1,
            long_stops=None,
            short_stops=None,
        ),
    ]
    sessions = {"SPY": {}, "AAPL": {}, "TSLA": {"foo": 1}}
    ctx = PosSizeContext(
        StrategyConfig(max_long_positions=1),
        portfolio,
        col_scope,
        ind_scope,
        input_scope,
        pred_scope,
        pending_order_scope,
        trained_models,
        sessions,
        sym_end_index,
    )
    set_pos_size_ctx_data(ctx, buy_results, sell_results)
    assert ctx.sessions == sessions
    buy_signals = list(ctx.signals("buy"))
    assert len(buy_signals) == 1
    assert buy_signals[0].id == 0
    assert buy_signals[0].symbol == "SPY"
    assert buy_signals[0].shares == 100
    assert buy_signals[0].score == 1
    assert buy_signals[0].type == "buy"
    assert buy_signals[0].bar_data is not None
    sell_signals = list(ctx.signals("sell"))
    assert len(sell_signals) == 1
    assert sell_signals[0].id == 2
    assert sell_signals[0].symbol == "TSLA"
    assert sell_signals[0].shares == 100
    assert sell_signals[0].score == 1
    assert sell_signals[0].type == "sell"
    assert sell_signals[0].bar_data is not None
    all_signals = list(ctx.signals())
    assert len(all_signals) == 2
    assert all_signals[0].id == buy_signals[0].id
    assert all_signals[0].symbol == buy_signals[0].symbol
    assert all_signals[0].shares == buy_signals[0].shares
    assert all_signals[0].score == buy_signals[0].score
    assert all_signals[0].type == buy_signals[0].type
    assert all_signals[0].bar_data is not None
    assert all_signals[1].id == sell_signals[0].id
    assert all_signals[1].symbol == sell_signals[0].symbol
    assert all_signals[1].shares == sell_signals[0].shares
    assert all_signals[1].score == sell_signals[0].score
    assert all_signals[1].type == sell_signals[0].type
    assert all_signals[1].bar_data is not None


def test_cancel_pending_order(ctx, pending_orders):
    assert ctx.cancel_pending_order(pending_orders[0].id)
    orders = tuple(ctx.pending_orders())
    assert orders == tuple([pending_orders[1]])


def test_cancel_all_pending_orders(ctx):
    ctx.cancel_all_pending_orders()
    assert not tuple(ctx.pending_orders())


def test_cancel_all_pending_orders_when_symbol(ctx, pending_orders):
    ctx.cancel_all_pending_orders("SPY")
    assert tuple(ctx.pending_orders()) == tuple([pending_orders[1]])


def test_pending_orders(ctx, pending_orders):
    assert tuple(ctx.pending_orders()) == pending_orders
    assert tuple(ctx.pending_orders("AAPL")) == tuple([pending_orders[1]])
