import numpy as np
from numba import njit

import vectorbt as vbt
from vectorbt.portfolio import nb
from vectorbt.portfolio.enums import Direction, SizeType

# debug
import os
os.environ['NUMBA_DISABLE_JIT'] = '1'



@njit
def pre_group_func_nb(c):
    order_value_out = np.empty(c.group_len, dtype=np.float_)
    return (order_value_out,)

@njit
def pre_segment_func_nb(c, order_value_out, size, price, size_type, direction):
    for col in range(c.from_col, c.to_col):
        c.last_val_price[col] = nb.get_col_elem_nb(c, col, price)
    nb.sort_call_seq_nb(c, size, size_type, direction, order_value_out)
    return ()

@njit
def order_func_nb(c, size, price, size_type, direction, fees, fixed_fees, slippage):
    return nb.order_nb(
        size=nb.get_elem_nb(c, size),
        price=nb.get_elem_nb(c, price),
        size_type=nb.get_elem_nb(c, size_type),
        direction=nb.get_elem_nb(c, direction),
        fees=nb.get_elem_nb(c, fees),
        fixed_fees=nb.get_elem_nb(c, fixed_fees),
        slippage=nb.get_elem_nb(c, slippage)
    )

np.random.seed(42)
close = np.random.uniform(1, 10, size=(5, 3))
size_template = vbt.RepEval('np.asarray(1 / group_lens[0])')

pf = vbt.Portfolio.from_order_func(
    close,
    order_func_nb,
    size_template,  # order_args as *args
    vbt.Rep('price'),
    vbt.Rep('size_type'),
    vbt.Rep('direction'),
    vbt.Rep('fees'),
    vbt.Rep('fixed_fees'),
    vbt.Rep('slippage'),
    segment_mask=2,  # rebalance every second tick
    pre_group_func_nb=pre_group_func_nb,
    pre_segment_func_nb=pre_segment_func_nb,
    pre_segment_args=(
        size_template,
        vbt.Rep('price'),
        vbt.Rep('size_type'),
        vbt.Rep('direction')
    ),
    broadcast_named_args=dict(  # broadcast against each other
        price=close,
        size_type=SizeType.TargetPercent,
        direction=Direction.LongOnly,
        fees=0.001,
        fixed_fees=1.,
        slippage=0.001
    ),
    template_mapping=dict(np=np),  # required by size_template
    cash_sharing=True, group_by=True,  # one group with cash sharing
)

pf.asset_value(group_by=False).vbt.plot()