import datetime
import matplotlib.pyplot as plt
import pandas as pd

from System import *
from QuantConnect import *
from QuantConnect.Algorithm import *


class BasicTemplateOptionsAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetCash(100000)
        self.SetStartDate(2018, 4, 1)
        self.SetEndDate(2018, 6, 30)
        
        self.optionData = {}
        
        # def init_option(season_tickets, self):
        #     for ind, row in season_tickets.iterrows():
        #         option = self.AddOption(row[0], Resolution.Minute)
        #         symbol = option.Symbol
        #         if symbol not in self.optionData:
        #             self.optionData[symbol] = OptionData(symbol)
        #         self.optionData[symbol].add_eps(row[6],row[5],row[4])
                                
        # #download EPS csv file from google drive and creating dataframe
        # url = "https://www.dropbox.com/s/4vtx2nn6jzty7lc/CND%20EPS%20data.csv?dl=1" 
        # file = self.Download(url).split("\n")
        # extract_data = [x.split(',') for x in file]
        # EPS_data = pd.DataFrame(extract_data)
        # EPS_data.columns = EPS_data.iloc[0]
        # EPS_data = EPS_data.drop(EPS_data.index[0])
        # EPS_data = EPS_data.drop(EPS_data.index[len(EPS_data)-1])
        # init_option(EPS_data, self)
        
        # Example Option
        option = self.AddOption("AAPL", Resolution.Minute)
        symbol = option.Symbol
        if symbol not in self.optionData:
            self.optionData[symbol] = OptionData(symbol)
        self.optionData[symbol].add_eps(2018, 5, 1)
        
        self.Schedule.On(self.DateRules.EveryDay(),
                 self.TimeRules.At(10,00),
                 self.BuyCall)
        
        self.Schedule.On(self.DateRules.EveryDay(),
                 self.TimeRules.At(9,45),
                 self.ClosePositions)
        
        

        
    def BuyCall(self):
        if self.Portfolio.Invested: return
    
        symbolsToTrade = []
        
        for symbol, data in self.optionData.items():
            earnings = data.dates
            if self.Time.date() in earnings:
                symbolsToTrade.append(symbol)
                
        
        for kvp in self.CurrentSlice.OptionChains:
            if kvp.Key in symbolsToTrade: 
                chain = kvp.Value

                # we sort the contracts to find at the money (ATM) contract with farthest expiration
                contracts = sorted(sorted(sorted(chain, \
                    key = lambda x: abs(chain.Underlying.Price - x.Strike)), \
                    key = lambda x: x.Expiry, reverse=True), \
                    key = lambda x: x.Right)

                # if found, trade it
                if len(contracts) == 0: continue
                symbol = contracts[0].Symbol
                
                # Divide total available buying power amongst symbols traded
                cashAllocated = self.Portfolio.MarginRemaining / len(symbolsToTrade)
                # Price per share for contract
                price = self.Securities[symbol].Close
                # Size of position given dollar amount for position
                size = int(cashAllocated/(price*100))
                
                self.MarketOrder(symbol, size)
                
    def ClosePositions(self):
        if self.Portfolio.Invested:
            self.Liquidate()

class OptionData:
    def __init__(self, symbol):
        self.symbol = symbol
        self.dates = []

    def add_eps(self, year, month, day):
        date = datetime.date(year, month, day)
        self.dates.append(date)