from QuantConnect.Data.UniverseSelection import *
import pandas as pd
num = 275
use_fine_filter = 1     # Try 0. There are no orders, why?

class BasicTemplateAlgorithm(QCAlgorithm):
    def Initialize(context):
        context.universe = []

        context.SetCash(10000)
        context.SetStartDate(2017, 1, 3)
        context.SetEndDate  (2017,12,31)

        context.UniverseSettings.Resolution = Resolution.Daily
        context.UniverseSettings.MinimumTimeInUniverse = 0
        if use_fine_filter:
            context.AddUniverse(context.universe_filter_coarse, context.universe_filter_fine)
        else:  # why does coarse alone fail with no orders?
            context.AddUniverse(context.universe_filter_coarse)
        context.Schedule.On(        # Schedule the 'trade' method to run
            context.DateRules.EveryDay(),
            context.TimeRules.At(9, 30),
            Action(context.trade))

    def universe_filter_fine(context, fine_data):
        # Only looking at price. Seems redundant since coarse does that.
        data_df = pd.DataFrame.from_records(
            [(float(s.Price),) for s in fine_data],
            index   = [s.Symbol for s in fine_data],
            columns = ['prc'],
            coerce_float=True)
        #context.Log("data: {}".format((data_df.prc.values)))
        #context.Log("fine len: {}".format(len(data_df)))
        context.universe = data_df.index.tolist()
        return context.universe


    def universe_filter_coarse(context, coarse_data):
        columns = ['Price', 'HasFundamentalData']
        data_df = pd.DataFrame.from_records(
            [[getattr(s, name) for name in columns] for s in coarse_data],
            index   = [s.Symbol for s in coarse_data],
            columns = columns,
            coerce_float=True)
        data_df = data_df.query("(Price > .50) and (Price < 4.0) and HasFundamentalData").nsmallest(num, 'Price')
        context.Log('len {}  prc min {}  max {}'.format(len(data_df), data_df.Price.min(), data_df.Price.max()))
        #context.Log(data_df.sort_values(by='Price', ascending=False).head())
        #context.Log(data_df.sort_values(by='Price', ascending=False).tail())
        context.universe = data_df.index.tolist()
        return context.universe

        return context.universe.index.tolist()

    def trade(context):
        #for u in context.UniverseManager.Keys:
        #    context.Log("universe name: {}".format(u))
        #for u in context.UniverseManager.Values:
        #    context.Log("universe count: {}".format(u.Members.Count))

        if not len(context.universe): return
        cash = float(context.Portfolio.CashBook["USD"].Amount)
        if cash < 500: return
        cnt    = 0  # count, to only log a few of them
        weight = 1.0 / len(context.universe)
        val    = weight * cash
        context.Log('weight ' + '%.4f' % weight)
        context.Log('   {} '.format(cash))
        for s in context.universe:
            if not context.Securities.ContainsKey(s): continue
            prc = float(context.Securities[s].Price)
            cb  = float(context.Portfolio[s].AveragePrice)  # cost basis i presume
            if cb:
                if cnt < 5:     # only log a sampling of them
                    cnt += 1
                    context.Log("{}  cb {}  prc {}".format(s, '%.2f' % cb, '%.2f' % prc))
                if prc > 1.25 * cb:  # close if up a lot
                    context.Log("    close {}  cb {}  prc {}".format(s, '%.2f' % cb, '%.2f' % prc))
                    context.SetHoldings(s, 0)
                    continue
            limit_prc = prc * .995
            if limit_prc <= 0: continue
            quantity  = int(val / limit_prc)
            if quantity  <  1: continue
            context.LimitOrder(s, quantity, limit_prc)