"""

Magic Formula by Joel Greenblatt
    1. Establish a minimum market capitalization (usually greater than $50 million).
    2. Exclude utility and financial stocks.
    3. Exclude foreign companies (American Depositary Receipts).
    4. Determine company's earnings yield = EBIT / enterprise value.
    5. Determine company's return on capital = EBIT / (net fixed assets + working capital).
    6. Rank all companies above chosen market capitalization
       by highest earnings yield and highest return on capital (ranked as percentages).
    7. Invest in 20–30 highest ranked companies, accumulating 2–3 positions per month over a 12-month period.
    8. Re-balance portfolio once per year, selling losers one week before the year-mark and winners one week after the year mark.

    Continue over a long-term (5–10+ year) period.

    net fixed assets = Total Assets - Total Current Assets - Total Intangibles & Goodwill


#### Differences in the implementation here ###

Rather than rebalancing once/year, this algo accumulates positions each month
and holds them for just under or over 1 year depending on whether the return is down/up.

This means that it takes a while to get up to full capacity, but the rebalancing is spread
out over the course of the year rather than all at once.


##### To switch to the Acquirer's Multiple toggle the comments on lines 137/138  ####

"""


from datetime import timedelta
import pandas as pd
import numpy as np


def initialize(context):
    #: Setting a few variables that we're using for our picks
    context.picks = None
    context.fundamental_dict = {}
    context.fundamental_data = None

    #: Choices are mf -> Magic Formula and am -> Acquirer's Multiple
    context.ranker_type = 'mf'

    #: Time period in days
    days = 365
    quarters = 4
    context.time_periods = [(days/quarters)*i for i in range(0,quarters + 1)]

    # Total number of stocks to hold at any given time
    context.max_positions = 30

    # Number of positions to accumulate each month
    context.positions_per_month = 3

    context.minimum_market_cap = 500e6

    # Used to track the purchase dates of each security
    context.entry_dates = {}

    # In order to maximize post tax returns, losing positions should be sold before
    # the 1 year mark and winning positions should be held longer than one year
    context.hold_days = {
        'win': timedelta(days=375),
        'loss': timedelta(days=345)
    }

    # Buy Stocks at the beginning of each month
    schedule_function(func=buy_stocks,
                      time_rule=time_rules.market_open(),
                      date_rule=date_rules.week_end())

    # Look to close positions every week
    schedule_function(func=sell_stocks,
                      time_rule=time_rules.market_open(),
                      date_rule=date_rules.week_start())

    # plot record variables
    schedule_function(func=record_vars,
                      time_rule=time_rules.market_close(),
                      date_rule=date_rules.every_day())

def before_trading_start(context, data):
    excluded_sectors = [103, 207]
    sector_code = fundamentals.asset_classification.morningstar_sector_code
    fundamental_df = get_fundamentals(
        query(
            sector_code,
            fundamentals.valuation.market_cap,
            fundamentals.valuation.enterprise_value,
            fundamentals.cash_flow_statement.capital_expenditure,
            fundamentals.operation_ratios.roic,
            fundamentals.income_statement.ebit,
            fundamentals.income_statement.ebitda,
            fundamentals.balance_sheet.total_assets,
        )
        .filter(fundamentals.valuation.market_cap > context.minimum_market_cap)
        .filter(~sector_code.in_(excluded_sectors))
        .order_by(fundamentals.valuation.market_cap.desc())
        # .limit(200)
    )
    df_length = len(fundamental_df.columns.values)
    if df_length < 100:
        log.info("Length of fundies dataframe error: %s" % len(fundamental_df.columns.values))
    fundamental_df = fundamental_df.dropna(axis=1)

    context.fundamental_df = fundamental_df

    #: On the first of every month
    if get_datetime().day == 1:
        context.fundamental_dict[get_datetime()] = context.fundamental_df
        context.fundamental_data = pd.Panel(context.fundamental_dict)
        picks = run_historical_ranker(context)
        #: If it's not yet time to order, fill context.fundamental_df with an empty list
        if picks is None:
            context.picks = None
        else:
            context.picks = list(picks.index)

def buy_stocks(context, data):
    # Only accumulate positions if there's room in the portfolio
    stocks_owned = position_count(context, data)
    now = get_datetime()
    stocks_bought = 0

    open_orders = get_open_orders()

    if context.picks is None:
        return

    for stock in context.picks:
        if not data.can_trade(stock) or stock in open_orders:
            continue
        if stocks_bought >= context.positions_per_month or stocks_owned >= context.max_positions:
            return
        # Skip stocks already owned
        if stock in context.portfolio.positions:
            continue
        # Some securities throw an error, not sure why???
        #try:
        order_target_percent(stock, 1.0 / context.max_positions)
        context.entry_dates[stock] = now
        stocks_bought += 1
        stocks_owned += 1
        # except Exception as e:
        #     log.debug(e)

def sell_stocks(context, data):
    now = get_datetime()
    open_orders = get_open_orders()
    for stock in context.portfolio.positions:
        if data.can_trade(stock) and stock not in open_orders:
            cost_basis = context.portfolio.positions[stock].cost_basis
            returns = data.current(stock, 'price') / cost_basis - 1
            if returns >= 0:
                entry_date = context.entry_dates[stock]
                if now >= entry_date + context.hold_days['win']:
                    order_target(stock, 0)
            elif returns < 0:
                entry_date = context.entry_dates[stock]
                if now >= entry_date + context.hold_days['loss']:
                    order_target(stock, 0)


def position_count(context, data):
    return sum(1 for stock, position in context.portfolio.positions.items()
               if position.amount > 0)

def run_historical_ranker(context):
    """
    Ranks our stocks and returns the context.max_positions amount of them if available.
    Returns None if we currently don't have enough historical data
    """
    #: Instantiate a few variables that we need
    trailing_metrics = {}
    fund_hist = context.fundamental_data
    time_periods = context.time_periods

    #: Check that we have data for the earliest date or before then
    if not check_earliest_date(time_periods, fund_hist):
        return None

    #: Create our ranks by putting them into a dict keyed by time period and the ranking
    for t in time_periods:
        temp_fund_df = find_closest_date(fund_hist, get_datetime() - timedelta(t))
        if context.ranker_type == 'am':
            ranks = acquirers_multiple(context, temp_fund_df)
        elif context.ranker_type == 'mf':
            ranks = magic_formula_ranks(context, temp_fund_df)
        trailing_metrics[t] = ranks

    #: Create a DataFrame that we can simply call the mean() on to get a historical average
    trailing_metrics_df = pd.DataFrame(trailing_metrics).transpose()
    trailing_mean = trailing_metrics_df.mean()

    #: Magic Formula takes highest ranked stocks
    trailing_mean.sort(ascending=True)
    return trailing_mean.head(context.max_positions)

def magic_formula_ranks(context, df):
    """
    Creates the magic formula ranking for a given DataFrame
    """
    #: Create our sorting mechanisms
    earnings_yield_ranking = df.ix['ebit'] / df.ix['enterprise_value']
    roic_ranking = df.ix['roic'].copy()

    #: Sort our rankings
    earnings_yield_ranking.sort(ascending=False)
    roic_ranking.sort(ascending=False)

    #: Create our ranking mechanisms
    earnings_yield_ranking = pd.Series(range(len(earnings_yield_ranking)), index=earnings_yield_ranking.index)
    roic_ranking = pd.Series(range(len(roic_ranking)), index=roic_ranking.index)

    return earnings_yield_ranking + roic_ranking

def acquirers_multiple(context, df):
    """
    Creates the acquirer's multiple ranking for a given DataFrame
    """
    multiples = df.ix['enterprise_value'] / (df.ix['ebitda'] + df.ix['capital_expenditure'])
    multiples.sort(ascending=True)
    return pd.Series(range(len(multiples)), index=multiples.index)

def find_closest_date(df_panel, date):
    """
    Finds the closest date if an exact match isn't possible, otherwise finds the exact match
    """
    date_index = df_panel.items.searchsorted(date)
    date = df_panel.items[date_index]
    return df_panel[date]

def check_earliest_date(time_periods, fund_history):
    """
    Finds the earliest date and makes sure that we have data for that time
    """
    earliest_date = get_datetime() - timedelta(max(time_periods))
    if earliest_date < min(fund_history.items):
        return False
    return True

def record_vars(context, data):
    record(leverage=context.account.leverage,
           stocks_owned=position_count(context, data))