from ftplib import FTP
import pandas as pd


host = 'ftpupload.net'
user = 'epiz_31014263'
passw = 'vmHOzdStGxo'


def format_df(df):
    df = pd.read_csv(df)
    df = df.drop('Unnamed: 0', axis=1, errors='ignore')
    df = df.drop_duplicates()
    df['date'] = pd.to_datetime(df['date'])
    df['evening_high_time'] = pd.to_datetime(df['evening_high_time'])
    df['evening_low_time'] = pd.to_datetime(df['evening_low_time'])
    df['morning_high_time'] = pd.to_datetime(df['morning_high_time'])
    df['morning_low_time'] = pd.to_datetime(df['morning_low_time'])
    df = df.set_index('date')
    df = df.sort_values(by=['date'])
    return df


if __name__ == '__main__':
    format_df()

def compare_dates(dt1, dt2, only_days=False, only_hours=False):

    if only_days:
        if dt1.year == dt2.year and dt1.month == dt2.month and dt1.day == dt2.day:
            return True

    if dt1 == dt2:
        return True


def sec_diff(dt1, dt2):

    if dt1 >= dt2:
        return (dt1 - dt2).seconds
    else:
        return -(dt2 - dt1).seconds

from AlgorithmImports import *
from datetime import datetime, timedelta
import pandas as pd
from io import StringIO
import traceback
import requests

from get_csv import format_df
from help_functions import *


class FatOrangeSnake(QCAlgorithm):
    """
    Setup:
        0 - buy call at low, sell at high
        1 - buy put at high, sell at low
    OTM:
        How many perc contract out of money
    max_DTE:
        Max days before expiration to buy call/put
    """

    def Initialize(self):
        self.SetStartDate(2012, 1, 1)  # Set Start Date
        #self.SetEndDate(2014, 1, 1)
        self.start_cash = 1000000
        self.setup = 0
        self.OTM = 0.01  # How many perc contract out of money
        self.max_logs = 6  # not used now
        self.max_DTE = 1  # Max days before expiration to buy call/put
        self.perc_buy = 0.01  # how many percent from portfolio to buy with
        self.exponential_gain = 1
        self.logging = 0  # if 1, logs data, be careful, you only have 3MB per day

        self.SetCash(self.start_cash)  # Set Strategy Cash
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)
        self.equity = self.AddEquity('SPY', Resolution.Minute)
        self.equity.SetDataNormalizationMode(DataNormalizationMode.Raw)
        self.symbol = self.equity.Symbol
        self.contract = ''
        self.contracts_subscribed = {}  # option contract
        self.underlying_price = 0
        self.SetBenchmark("SPY")
        self.contracts_added = set()
        self.strikes = ''
        self.all_options = 0
        # if self.logging == 0:
        #     self.Log = lambda x: None

        df = self.Download('https://www.dropbox.com/s/rttn6rq5wmqzt2y/extremes_new.csv?dl=1')
        self.df = format_df(StringIO(df))  # downloaded data of max/min daily prices

    def is_action_necessary(self):
        """
        Find out if we need to prepare or buy or liquidate option.
        If option within 30 minutes of buying, prepare the option
        If time of buying/selling, send info
        """

        if self.Time.weekday() == 4:
            df_row = self.df[(self.df.index.year == self.Time.year) &
                             (self.df.index.month == self.Time.month) &
                             (self.df.index.day == self.Time.day)]

            if not df_row.empty:
                morning_low = 0  # 0 - not, 1 - within 30 mins, 2 - exact
                morning_high = 0
                evening_low = 0
                evening_high = 0

                sec1 = sec_diff(df_row.iloc[0]['morning_low_time'], self.Time)
                sec2 = sec_diff(df_row.iloc[0]['morning_high_time'], self.Time)
                sec3 = sec_diff(df_row.iloc[0]['evening_low_time'], self.Time)
                sec4 = sec_diff(df_row.iloc[0]['evening_high_time'], self.Time)

                num = 60

                if 0 <= sec1 <= num:
                    morning_low = 2 if sec1 == 0 else 1
                if 0 <= sec2 <= num:
                    morning_high = 2 if sec2 == 0 else 1
                if 0 <= sec3 <= num:
                    evening_low = 2 if sec3 == 0 else 1
                if 0 <= sec4 <= num:
                    evening_high = 2 if sec4 == 0 else 1

                return morning_low, morning_high, evening_low, evening_high

    def prepare_option(self, data, is_call):
        """ Finds and stubscribes to option which will later be bought """

        if not self.contracts_subscribed:
            self.contracts_subscribed = self.options_filter(data, is_call)

    def sort_contracts(self, contracts):
        # sort contracts by expiration date (2nd priority)
        contracts = sorted(contracts, key=lambda x: (x.ID.Date - self.Time).days)
        # sort contract by closest strike price to underlying (1st priority)
        return sorted(contracts, key=lambda x: x.ID.StrikePrice)

    def buy_option(self, data, is_call):
        """
        Go through all subscribed options (subscribed during option filtering)
        Select first option out of money by percentage given in input
        Buy the option and log
        Remove all options from subscribing that are not the bought option (we still need to close it)
        """

        if not self.contracts_subscribed:
            if self.logging: self.Log('Contract not prepared')
        elif not self.Portfolio.Invested:
            contracts = [self.contracts_subscribed[k] for k in self.contracts_subscribed]
            contracts = self.filter_contracts(contracts, data, is_call)
            self.contract = self.sort_contracts(contracts)[0]

            if data.ContainsKey(self.contract):
                sec = self.Securities[self.contract.Value]

                if self.logging: self.Log(f'SPY: {self.underlying_price} Strike: {self.contract.ID.StrikePrice} '
                                          f'ASK: {sec.AskPrice} SIZE: {sec.AskSize} '
                                          f'BID: {sec.BidPrice} SIZE: {sec.BidSize} ')

                if sec.AskPrice > 0:
                    cash = self.Portfolio.Cash if self.exponential_gain else self.start_cash
                    quantity = round((cash * self.perc_buy) / (sec.AskPrice * 100))
                    self.all_options += quantity
                    self.Buy(self.contract, quantity)
                    # self.Buy(self.symbol, 1)
            else:
                if self.logging: self.Log('Contract not found')

        else:
            if self.logging: self.Log('Portfolio already invested')

        # after buying option, we no longer need to subscribe to all options data, only track the bought option
        if self.contracts_subscribed:
            for k in self.contracts_subscribed:
                if self.contracts_subscribed[k] in self.contracts_added and self.contract.ID != \
                        self.contracts_subscribed[k].ID:
                    self.RemoveOptionContract(self.contracts_subscribed[k])
            self.contracts_subscribed = {}

    def liquidate_option(self):
        """
        Sell option, log what happened
        Remove option from subscribed data, we're done with it
        """
        if self.contract:
            sec = self.Securities[self.contract.Value]

            if self.logging: self.Log(f'SPY: {self.underlying_price}'
                                      f'ASK: {sec.AskPrice} SIZE: {sec.AskSize} '
                                      f'BID: {sec.AskPrice} SIZE: {sec.AskSize}')
            self.Liquidate(self.symbol)
            self.Liquidate(self.contract)
            if self.logging: self.Log('Contract closed')
            self.RemoveOptionContract(self.contract)
            self.contract = ''

    def filter_contracts(self, contracts, data, is_call):
        contracts = [i for i in contracts if self.max_DTE >= (i.ID.Date - data.Time).days]

        if is_call:
            # filter only calls/puts
            contracts = [i for i in contracts if i.ID.OptionRight == OptionRight.Call]
            # filter out all contracts that are not out of the money by percentage given in input
            contracts = [i for i in contracts if
                         i.ID.StrikePrice >= round(self.underlying_price * (1+self.OTM))]
            # strikes = set([i.ID.StrikePrice for i in contracts])
        else:
            contracts = [i for i in contracts if i.ID.OptionRight == OptionRight.Put]
            # strikes = set([i.ID.StrikePrice for i in contracts])
            contracts = [i for i in contracts if
                         self.underlying_price >= round(i.ID.StrikePrice * (1+self.OTM))]

        return contracts

    def options_filter(self, data, is_call):
        """
        Go through options available, and filter the right one to buy
        Subscribe to all options with strike prices, to later pick correctly during buying
        """

        contracts = self.OptionChainProvider.GetOptionContractList(self.symbol, data.Time)
        # filter contracts that have longer expiration day than given
        contracts = self.filter_contracts(contracts, data, is_call)

        if contracts:
            contracts = self.sort_contracts(contracts)[:4]

            contracts = {i.ID.StrikePrice: i for i in contracts}
            # self.Log(f'strikes: {[i for i in contracts]}')

            for k in contracts:
                if contracts[k] not in self.contracts_added:
                    self.contracts_added.add(contracts[k])
                    self.AddOptionContract(contracts[k], Resolution.Minute)

            return contracts
        else:
            return ''

    def OnData(self, data):
        '''OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.
            Arguments:
                data: Slice object keyed by symbol containing the stock data
        '''

        is_necessary = self.is_action_necessary()
        self.underlying_price = self.Securities[self.symbol].Price

        if is_necessary:

            morning_low, morning_high, evening_low, evening_high = is_necessary

            if self.setup == 0:
                if morning_low == 1:
                    self.prepare_option(data, is_call=1)
                if morning_low == 2:
                    self.buy_option(data, is_call=1)
                    # self.Buy('SPY', 1)

                if evening_high == 2:
                    self.liquidate_option()
            if self.setup == 1:
                if morning_high == 1:
                    self.prepare_option(data, is_call=0)
                if morning_high == 2:
                    self.buy_option(data, is_call=0)
                    # self.Buy('SPY', -1)

                if evening_low == 2:
                    self.liquidate_option()
                    # self.Liquidate('SPY')

        # if self.Portfolio['SPY'].Invested:
        #     self.Liquidate('SPY')

    def OnOrderEvent(self, orderEvent):
        if self.logging: self.Log(str(orderEvent))

    def OnEndOfAlgorithm(self):
        if self.logging: self.Log(f'all options: {self.all_options} fees: {self.Portfolio.TotalFees / self.all_options}$ per option')