# df of closing prices
import pandas as pd
import numpy as np
import talib as ta
import math
from collections import deque
from datetime import datetime, timedelta

class customRSI(object):
    def __init__(self, name):
        self.Name = name
        self.Time = datetime.min #minimum/earliest time, NOT MINUTE
        self.cyclelen = 14
        self.h1 = 30.0
        self.h2 = 70.0
        self.Value = 0
        self.IsReady = False
        self.queue = RollingWindow[float](self.cyclelen + 1) # +1 for difference cal
        self.diffdeque = RollingWindow[float](self.cyclelen)
        
        
    def __repr__(self):
        return "{0} -> IsReady: {1}. Time: {2}. Value: {3}".format(self.Name, self.IsReady, self.Time, self.Value)
        
    def Update(self, input):
        
        #fill the queue
        self.queue.Add(input.Close) # 15 day closing prices
        
        # calc rsi
        self.pricechange = self.change()

        # initiate
        count = self.queue.Count
        self.Time = input.EndTime
        self.IsReady = count == self.cyclelen + 1
        
        
        # if its ready iniate these
        if self.IsReady == True:
            self.uplist = [0.0 if i < 0 else i for i in self.pricechange]
            self.downlist = [0.0 if i > 0 else -i for i in self.pricechange]
            self.rs = np.mean(self.uplist)/np.mean(self.downlist)
            self.Value = rsi = 100.0 - 100.0 / (1 + self.rs)
            
        else:
            self.Value = 0
        
        
    def change(self): # pine change fx
        try:
            self.diff = self.queue[0]-self.queue[1]
            self.diffdeque.Add(self.diff)
        except:
            return
        return self.diffdeque

from customRSI import customRSI


class UncoupledVentralCoil(QCAlgorithm):
    stopMarketTicket = None
    stopMarketOrderFillTime = datetime.min
    highestSPYPrice = 0
    symbolsList = []
    
    def Initialize(self):
        self.SetStartDate(2015, 1, 1)  # Set Start Date
        self.SetEndDate(2018,1,1)    #Set End Date
        self.SetCash(100000)  # Set Strategy Cash
        # self.AddEquity("SPY", Resolution.Daily)
        
        self.buylist = []
        self.selllist = []
        
        # self.Securities["SPY"].SetDataNormalizationMode(DataNormalizationMode.Raw)

        # Create a chart for the indicator
        crsiChart = Chart("Awesome", ChartType.Stacked)
        crsiChart.AddSeries(Series('custom', SeriesType.Line))
        
        self.coarse_count = 10
        self.AddUniverse(self.CoarseSelectionFunction)
        self.UniverseSettings.Resolution = Resolution.Daily
        self.SetRiskManagement(TrailingStopRiskManagementModel(0.2))
        
    def OnSecuritiesChanged(self, changes):
        for security in changes.AddedSecurities:
            self.SetHoldings(security.Symbol, 0.01)
            self.stopMarketTicket = self.StopMarketOrder(security.Symbol, -1.0 * self.CalculateOrderQuantity(security.Symbol, 0.01), 0.85 * self.Securities[security.Symbol].Close)
            
        
    def CoarseSelectionFunction(self, coarse):
        
        sym = sorted([x for x in coarse if x.HasFundamentalData and x.Price>0], 
                                key = lambda x: x.DollarVolume, reverse = True)
        
        syms = sym[:self.coarse_count]
        
        for i in syms:
            if not self.Securities.ContainsKey(i.Symbol):
                self.symbolsList.append(i.Symbol)
        
        for symbol in self.symbolsList:
            if symbol not in self.buylist:
                self.AddEquity(f"{symbol}", Resolution.Daily)
                self.custom = customRSI(f"{symbol}")
                self.RegisterIndicator(symbol, self.custom, Resolution.Daily)
                if self.custom.Value < self.custom.h1:
                    self.buylist.append(symbol)
                # if self.custom.Value > self.custom.h2:
                # self.selllist.append(symbol)
            
        return self.buylist