from pyutil.mongo.engine.strategy import Strategy, strategies, configuration
from pyutil.mongo.engine.symbol import Symbol, Group
from pyutil.performance.drawdown import drawdown
from pyutil.performance.month import monthlytable
from pyutil.performance.return_series import from_nav
from pyutil.portfolio.portfolio import  similar
import pandas.testing as pt
from test.config import *


@pytest.fixture()
def group():
    Group.objects.delete()
    return Group(name="US Equity").save()


@pytest.fixture()
def symbols(group, portfolio):
    Symbol.objects.delete()
    # add the symbols to database
    for symbol in portfolio.assets:
        Symbol(name=symbol, group=group).save()


def test_strategy(symbols, portfolio):
    Strategy.objects.delete()

    s = Strategy(name="mdt", type="mdt", active=True, source="AAA")

    assert s.source == "AAA"
    assert s.type == "mdt"
    assert s.active

    assert s.portfolio is None
    assert s.last_valid_index is None

    # empty dictionary as portfolio hasn't been set
    assert Strategy.portfolios(strategies=[s]) == {}

    s.save()

    frame = Strategy.reference_frame()
    assert frame.index.name == "strategy"

    s.portfolio = portfolio
    pt.assert_frame_equal(s.portfolio.weights, portfolio.weights)
    pt.assert_frame_equal(s.portfolio.prices, portfolio.prices)

    s.save()

    similar(Strategy.portfolios(strategies=[s])["mdt"], portfolio)

    navs = Strategy.navs()
    assert not navs["mdt"].empty

    frame = Strategy.sectors(strategies=[s])
    assert frame.index.name == "Portfolio"
    assert set(frame.keys()) == {"US Equity", "Total"}
    assert frame.loc["mdt"]["US Equity"] == pytest.approx(0.308974, abs=1e-5)


def test_source(portfolio):
    with open(resource("source.py"), "r") as f:
        s = Strategy(name="Peter", source=f.read(), active=True, type="wild")
        # construct the configuration based on the strategy (and it's source code)
        c = configuration(strategy=s)
        # verify the names of the configuration
        assert c.names == portfolio.assets
        # also possible to ask the strategy directly
        assert s.assets == portfolio.assets


def test_last_valid(portfolio):
    s = Strategy(name="Maffay", source="AAA", active=True, type="wild2")
    s.portfolio = portfolio
    assert s.last_valid_index == portfolio.prices.last_valid_index()
    assert similar(s.portfolio, portfolio)


def test_strategies():
    folder = resource(name="strat")
    for name, source in strategies(folder=folder):
        assert name in {"P1", "P2"}


def test_active():
    Strategy.objects.delete()
    Strategy(name="A", active=False).save()
    Strategy(name="B", active=True).save()

    assert len(Strategy.active_strategies()) == 1
    assert len(Strategy.objects) == 2


def test_drawdown(portfolio):
    Strategy.objects.delete()
    s = Strategy(name="Maffay", source="")
    s.portfolio = portfolio
    pt.assert_series_equal(drawdown(portfolio.nav), s.drawdown)


def test_volatility(portfolio):
    Strategy.objects.delete()
    s = Strategy(name="Maffay", source="")
    s.portfolio = portfolio
    pt.assert_series_equal(from_nav(portfolio.nav).ewm_volatility().dropna(), s.ewm_volatility())


def test_monthlytable(portfolio):
    Strategy.objects.delete()
    s = Strategy(name="Maffay", source="")
    s.portfolio = portfolio
    pt.assert_frame_equal(monthlytable(portfolio.nav.pct_change()), s.monthlytable)