#region imports
from AlgorithmImports import *
#endregion
from datetime import datetime
import pandas as pd
import numpy as np
from config import *

log_table = {}




def read_log_table(algo, table, log_file_suffix):
    # 从持久化存储中恢复日志
    import pandas as pd
    import pickle
    
    logs = algo.ObjectStore.ReadBytes(f'{table}_logs_{log_file_suffix}.pk')
    logs = pickle.loads(bytes(logs))    
    return pd.DataFrame(logs)


def between_range(range_value, value):
    return range_value[0] < value and value < range_value[1]
    
def exit_price_distance_with_strikes(row):
    strikes = row['strikes']
    price = row['underlying_price_exit']
    low_strike, high_strike = strikes[0], strikes[1]
    return [round((price-low_strike)/price, 3), round((high_strike-price)/price, 3)]
    
def profit_and_loss(row):
    try:
        price_enter = sum(row['price_enter'])
        price_exit = sum(row['price_exit'])
        gain = 1 - price_exit / price_enter
        profit = gain * row['quantity'] * price_enter * 100
        return [gain, profit]
    except:
        return [0, 0]
        
def summary_order(order_df):
    pass


def summary_iron_condor_orders(order_df):
    pass

logged = {}
def log_once(algo, msg):
    if not logged.get(msg, False):
        if algo.Time.hour > 11:
            return
        if algo.Time.minute not in [15, 31, 45]:
            return
        last_msg = f'{algo.Time} {msg}'
        algo.Log(last_msg)
        logged[msg] = True

#region imports
from AlgorithmImports import *
#endregion

log_file_suffix = ''
log_to_disk_open = True

underlyings_stocks = [
    # come go
    # 'QQQ',
    'TSLA',
    # 'NVDA',
    # 'AAPL',
]


option_symbols = []

option_symbol_data_context = {}

max_hold_margin = 25000

default_before_market_liquidate = 5

stop_loss_pct = 0.10

default_risk_reward_ratio = 1

order_expire_time = 15

default_order_quantity = 10

from AlgorithmImports import *

from collections import deque
import statistics
import math

class STDVolatility(PythonIndicator):
    def __init__(self):
        self.roc_window = deque(maxlen=30)

        self.volatility_window = deque(maxlen=252)
        self.last_price = None
        self.Value = 0
        
    def Update(self, bar):
        if self.last_price is None:
            self.last_price = bar.Price
            return False
        
        roc = (bar.Price / self.last_price) - 1
        self.roc_window.append(roc)
        self.last_price = bar.Price
        
        if len(self.roc_window) == self.roc_window.maxlen:
            volatility = statistics.stdev(self.roc_window) * math.sqrt(252)
            self.volatility_window.append(volatility)
            self.Value = volatility
            return True
        else:
            return False
        
    @property
    def roc(self):
        return self.roc_window[-1] if len(self.roc_window) > 0 else 0
    
    @property
    def volatility(self):
        return self.volatility_window[-1] if len(self.volatility_window) > 0 else 0
    
    @property
    def volatility_rank(self):
        volatility = self.volatility
        if len(self.volatility_window) < 2:
            return 0
        min_value = min(self.volatility_window)
        max_value = max(self.volatility_window)
        if max_value == min_value:
            return 1.0
        else:
            return (volatility-min_value) / (max_value-min_value)
        
    @property
    def volatility_percentile(self):
        volatility = self.volatility
        if len(self.volatility_window) < 2:
            return 0
        return len([v for v in self.volatility_window if v < volatility]) / len(self.volatility_window)
        
    
    def to_json(self):
        return {
            'volatility': self.volatility,
            'volatility_rank': self.volatility_rank,
            'volatility_percentile': self.volatility_percentile,
        }

# from QuantConnect.Algorithm.Framework.Execution import ImmediateExecutionModel
import math
from AlgorithmImports import *
from QuantConnect.Orders import *

from common import *
from itertools import *
from option_utils import *
from config import *
from datetime import *
from strategy import *
from symbol_data import *
import json
from risk_manage import *
from share import *
from share.main_utils import *

default_before_market_lidiquate = 10

class CasualFluorescentOrangeKitten(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2023, 5, 2)
        self.SetEndDate(2023, 5, 3)
        # self.SetEndDate(2023, 2, 2)

        # self.SetStartDate(2023, 2, 9)
        # self.SetEndDate(2023, 2, 10)

        # self.SetEndDate(2023, 1, 22)
        # self.SetEndDate(2023, 2, 1)
        
        self.SetCash(3 * 10000)

        self.Portfolio.MarginCallModel = MarginCallModel.Null
        # In Initialize
        self.Portfolio.SetPositions(SecurityPositionGroupModel.Null)

        # self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage)

        # self.SetExecution(ImmediateExecutionModel())
        
        self.option_symbols = []
        
        self.init_universal()

        # self.AddRiskManagement(NormalTrailingStopRiskManagementModel(0.05, 0.05))
        self.AddRiskManagement(StopLossRiskManagementModel())
        self.AddRiskManagement(StopProfitRiskManagementModel())
      
        self.Settings.DataSubscriptionLimit = 200
        
        prop = InteractiveBrokersOrderProperties()
        prop.OutsideRegularTradingHours = True
        self.DefaultOrderProperties = prop
        self.DefaultOrderProperties.TimeInForce = TimeInForce.Day

        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.Every(timedelta(seconds=10)), lambda:cancel_order_if_expire(self, order_expire_time))

        self.Schedule.On(self.DateRules.EveryDay(),self.TimeRules.At(15, 55, 0), lambda: liquidate_combo_mkt_order(self))

        
    def init_universal(self):
        self.UniverseSettings.Resolution = Resolution.Minute
        symbols = []
        for stock in underlyings_stocks:
            equity = self.AddEquity(ticker=stock, resolution=Resolution.Minute)
            equity.SetBuyingPowerModel(BuyingPowerModel.Null)

            option = self.AddOption(stock)
            option.SetBuyingPowerModel(BuyingPowerModel.Null)
            
            option.SetFilter(lambda universe: 
                universe.IncludeWeeklys()\
                    .Strikes(-5, 5)\
                    .Expiration(0, 30)
            )

            self.option_symbols.append(option.Symbol)
            
            symbol_data_context[equity.Symbol] = SymbolData(self, equity.Symbol)
            
            symbols.extend([equity.Symbol, option.Symbol])
        
        self.AddUniverseSelection(ManualUniverseSelectionModel(symbols))
        

    def OnSecuritiesChanged(self, changes):
        for security in changes.AddedSecurities:
            if security.Symbol.ID.SecurityType == SecurityType.Option:
                # self.Securities[security.Symbol.ID].MarginModel = PatternDayTradingMarginModel()
                pass
        for security in changes.RemovedSecurities:
            if security.Symbol.ID.SecurityType == SecurityType.Option and security.Symbol in option_symbol_data_context:
                del option_symbol_data_context[security.Symbol]

    def OnOrderEvent(self, orderEvent):
        if orderEvent.Status == OrderStatus.Filled\
            and orderEvent.Symbol.ID.SecurityType == SecurityType.Option\
            and (orderEvent.IsAssignment or orderEvent.IsInTheMoney):
            order = self.Transactions.GetOrderById(orderEvent.OrderId)
            self.Log(f'ohly shit, {orderEvent.Symbol.Value}. {order.Tag}')
    

    def option_invested(self):
        for symbol, holding in self.Portfolio.items():
            if symbol.ID.SecurityType == SecurityType.Option and holding.Invested:
                return True
        return False


    def OnData(self, slice):
        
        # peng_orders
        peng_orders = self.Transactions.GetOpenOrders()
        if len(peng_orders) > 0:
            return
        
        if self.IsWarmingUp:
            return
                
        if self.option_invested():
            return
        
        now = self.Time
        
        # open_mkrt = now.replace(hour=9, minute=30)
        # just_open_mkrt = now.replace(hour=10, minute=0)
        # open_mkt_soon = open_mkrt < now < just_open_mkrt
        # if not open_mkt_soon:
        #     return
        day = self.Time.strftime("%y%m%d")
        
        for canonical_symbol, chain in slice.OptionChains.items():
            for symbol, contract in chain.Contracts.items():
                exchange = self.Securities[symbol].Exchange
                if exchange.IsClosingSoon(default_before_market_liquidate+1):
                    return
            
            cs = [contract for symbol, contract in chain.Contracts.items()]
            if len(cs) == 0:
                continue
                        
            candidates = try_setup_calendar(self, cs)
            candidates = [c for c in candidates if c.validate()]
            candidates = sorted(candidates, key=lambda c: c.score())
            for c in candidates:
                log_to_disk(self, 'candidates', c.underlying_symbol, c.to_json())
            if len(candidates) > 0:
                import json

                candidate = candidates[0]

                quantity = candidate.order_quantity()

                msg =f'place order {json.dumps(candidate.to_json())} at {quantity}'
                if self.LiveMode:
                    self.Notify.Email('violet_day@live.com', "Notify", msg)
                
                self.Log(msg)
                
                legs, price = candidate.as_combo_limit_orders()
                self.ComboLimitOrder(legs, quantity, price, tag='Enter')
                
    

    def OnEndOfAlgorithm(self):
        dump_to_storage(self, log_table)

from AlgorithmImports import *

from stop_loss_rm import StopLossRiskManagementModel
from stop_profit_rm import StopProfitRiskManagementModel

from trailing_stop_rm import NormalTrailingStopRiskManagementModel

from AlgorithmImports import *
from symbol_data import *
from config import *
from common import *
from option_utils import *
from share import *

class StopLossRiskManagementModel(RiskManagementModel):
    def __init__(self):
        pass

    def ManageRisk(self, algorithm, targets):
        targets = list()
        
        if not algorithm.Portfolio.Invested:
            return targets
        for underlying_symbol, holdings in group_by_underlying(algorithm).items():
            pnl = sum([h.UnrealizedProfit for h in holdings])
            cost = sum([h.HoldingsCost for h in holdings])
            pct = abs(pnl/cost)
            
            if pnl < 0 and pct > stop_loss_pct:
                # FIXME
                for symbol, security in algorithm.Securities.items():
                    if security.Invested:
                        targets.append(PortfolioTarget(symbol, 0))
                log_to_disk(algorithm, 'liquidate', security.Symbol.Value, {
                    'tag': 'stop_loss'
                })
                algorithm.Log(f'stop loss {underlying_symbol.Value}, pnl is {pnl}')            
        return targets

from AlgorithmImports import *
from symbol_data import *
from config import *
from common import *
from option_utils import *

from share import *

class StopProfitRiskManagementModel(RiskManagementModel):
    def __init__(self):
        pass

    def ManageRisk(self, algorithm, targets):
        targets = list()

        if not algorithm.Portfolio.Invested or algorithm.Portfolio.TotalUnrealizedProfit < 0:
            return targets

        for underlying_symbol, holdings in group_by_underlying(algorithm).items():
            pnl = sum([h.UnrealizedProfit for h in holdings])
            cost = sum([h.HoldingsCost for h in holdings])
            pct = abs(pnl/cost)
            
            if pnl> 0 and pct > stop_loss_pct * default_risk_reward_ratio:
                for symbol, security in algorithm.Securities.items():
                    if security.Invested:
                        targets.append(PortfolioTarget(symbol, 0))
                log_to_disk(algorithm, 'liquidate', security.Symbol.Value, {
                    'tag': 'stop_profit'
                })
                algorithm.Log(f'stop profit, pnl is {pnl}')            
        return targets

from AlgorithmImports import *
from common import *
from symbol_data import *

from share import *

class NormalTrailingStopRiskManagementModel(RiskManagementModel):
    
    def __init__(self, min_profit=0.01, maximum_drawdown = 0.2):      
        self.min_profit = min_profit
        self.maximum_drawdown = maximum_drawdown
        
        self.trailing_highs = dict()
        self.reach_threhshold = dict()

    def ManageRisk(self, algorithm, targets):
        targets = list()

        for symbol, security in algorithm.Securities.items():
            
            # Remove if not invested
            if not security.Invested:
                self.trailing_highs.pop(symbol, None)
                self.reach_threhshold.pop(symbol, None)
                continue

            pnl = security.Holdings.UnrealizedProfit
            max_unrealized_profit = security.Holdings.HoldingsCost * self.min_profit
            # 亏损则忽略
            if pnl < 0 and symbol not in self.reach_threhshold:
                continue

            # Add newly invested securities
            if symbol not in self.trailing_highs:
                self.trailing_highs[symbol] = pnl
                continue

            # 收益增大的情况下，覆盖缓存中的symbol最大收益
            if self.trailing_highs[symbol] < pnl:
                self.trailing_highs[symbol] = pnl
            

            if pnl >= max_unrealized_profit and symbol not in self.reach_threhshold:
                # 当前symbol已经触发追踪止损, 首次命中
                self.reach_threhshold[symbol] = True
                algorithm.Log(f'{symbol.Value}@{security.Price} has reach traliling stop')
            
            if not self.reach_threhshold.get(symbol, False):
                continue
            
            if symbol in self.reach_threhshold:
                max_profit = self.trailing_highs[symbol]
                abs_drawdown_percent = (max_profit - pnl) / max_profit

                if abs_drawdown_percent > self.maximum_drawdown:
                    log_to_disk(algorithm, 'liquidate', security.Symbol.Value, {'tag': 'trailing_stop'})
                    targets.append(PortfolioTarget(symbol, 0))
                    
                    self.trailing_highs.pop(symbol, None)
                    self.reach_threhshold.pop(symbol, None)

                    continue
            

        return targets

#region imports
from AlgorithmImports import *
from symbol_data import *
from functools import *
from itertools import *
from config import *
from symbol_data import *
from common import *
from option_utils import *
import json

cache = {}

def gen_all_expiry(day, all_options):
    day = day.strftime("%y%m%d")
    if day in cache:
        return cache[day]

    all_expiries = list(set([c.Expiry for c in all_options]))
    all_expiries_combo = [(a,b) for a,b in permutations(all_expiries, 2) if a<b]
    
    cache[day] = all_expiries_combo
    return all_expiries_combo

def try_setup_calendar(algo, all_options):
    option_map = {(c.Expiry, c.Right, c.Strike) : c for c in all_options}
    desc = []
    
    all_expiries_combo = gen_all_expiry(algo.Time, all_options)
    
    for expiry_a, expiry_b in all_expiries_combo:
        strikes = set([c.Strike for c in all_options if c.Expiry==expiry_a])
        for strike in strikes:
            if (expiry_a, OptionRight.Put, strike) in option_map and (expiry_b, OptionRight.Put, strike) in option_map:
                option_short = option_map[(expiry_a, OptionRight.Put, strike)]
                option_long = option_map[(expiry_b, OptionRight.Put, strike)]
                desc.append(Desc(algo, option_short, option_long))
            
            if (expiry_a, OptionRight.Call, strike) in option_map and (expiry_b, OptionRight.Call, strike) in option_map:
                option_short = option_map[(expiry_a, OptionRight.Call, strike)]
                option_long = option_map[(expiry_b, OptionRight.Call, strike)]
                desc.append(Desc(algo, option_short, option_long))

    return desc


class Desc:
    def __init__(self, algo, option_short, option_long):
        self.algo = algo
        
        self.symbol_data = symbol_data_context[option_short.UnderlyingSymbol]

        self.underlying_symbol =  option_short.UnderlyingSymbol
        self.option_short, self.option_long = option_short, option_long
        self.diff_day = (self.option_long.Expiry-self.option_short.Expiry).days
        
        self.strike = option_short.Strike
        
        self.expiry = (self.option_short.Expiry - self.algo.Time).days
        self.expiry_a = self.option_short.Expiry.strftime("%y%m%d")
        self.expiry_b = self.option_long.Expiry.strftime("%y%m%d")
        self.right = 'call' if self.option_short.Right == OptionRight.Call else 'put'
        
        self.ID =f'{self.option_short.UnderlyingSymbol.Value}/{self.expiry_a}/{self.expiry_b}/{self.right}/{self.strike}'
            
        self.premium = [
            - self.option_short.BidPrice + self.option_long.AskPrice,
            - self.option_short.AskPrice + self.option_long.BidPrice,
            - self.option_short.LastPrice + self.option_long.LastPrice
        ]

        self.max_loss = self.option_long.AskPrice - self.option_short.BidPrice
        self.size = [
            self.option_short.BidSize,
            self.option_long.AskSize
        ]

    def to_json(self):
        return {
            '_id': self.ID,
            'premium': self.premium,
            'max_loss': self.max_loss,
            # 'strike_atr': self.strike_atr,

            'spread': self.spread,

            'underlying_price': self.option_short.UnderlyingLastPrice,

            # 'size': self.size,
            'delta_pct_': abs(self.norm_delta * self.symbol_data.atr.Current.Value /self.max_loss),
            
            'expiry': self.expiry,
            'internal_expiry': (self.option_long.Expiry-self.option_short.Expiry).days,

            'greek': {
                'delta': self.norm_delta,
                'gamma': self.norm_gamma,
                'theta_per_day': self.norm_theta_per_day,
            },

            'sd': self.symbol_data.to_json()
        }

    @property    
    def spread(self):
        bid_and_ask_price =  - self.option_short.BidPrice + self.option_long.AskPrice
        unrealized_price =  - self.option_short.AskPrice + self.option_long.BidPrice
        close_price = self.option_long.LastPrice - self.option_short.LastPrice

        return 1-unrealized_price/bid_and_ask_price if bid_and_ask_price!=0 else 0

    @property
    def norm_delta(self):
        return - self.option_short.Greeks.Delta + self.option_long.Greeks.Delta

    @property
    def norm_gamma(self):
        return - self.option_long.Greeks.Gamma + self.option_short.Greeks.Gamma

    @property
    def norm_theta_per_day(self):
        return self.option_long.Greeks.ThetaPerDay - self.option_short.Greeks.ThetaPerDay

    def as_orders_limit_price(self):
        return -self.option_short.LastPrice + self.option_long.LastPrice

    def as_combo_limit_orders(self):
        return [
            Leg.Create(self.option_short.Symbol, -1),
            Leg.Create(self.option_long.Symbol, +1),
        ], - self.option_short.BidPrice + self.option_long.AskPrice

    @property
    def strike_atr(self):
        if self.option_short.Right == OptionRight.Call:
            return (self.option_short.Strike - self.option_short.UnderlyingLastPrice) / self.symbol_data.atr.Current.Value
        else:
            return (self.option_short.UnderlyingLastPrice - self.option_short.Strike) / self.symbol_data.atr.Current.Value

    def validate(self):
        base_validate_items = {
            # 'premium': all([p>0.01 for p in self.premium]),
            # 'size': default_order_quantity < min(self.size),
            'spread': self.spread < 0.03,
            # 'strike_atr': abs(self.strike_atr) < 0.5
        }

        if not all(base_validate_items.values()):
            return False

        greek_validate_items = {
            '': True,
            
            'delta_pct': abs(self.norm_delta * self.symbol_data.atr.Current.Value /self.max_loss) < 0.010,
            # 'theta_per_day_pct': abs(self.norm_theta_per_day/self.max_loss) > 0.300,

            # 'gamma': self.norm_gamma > 0.0,
            # 'theta': self.norm_theta_per_day > 0,
        }

        if not all(greek_validate_items.values()):
            return False
        return True

    def score(self):
        return abs(self.norm_theta_per_day)
        
    def order_quantity(self):
        legs, price = self.as_combo_limit_orders()
        return math.floor(
            self.algo.Portfolio.Cash * 0.3 / 100 / price
        )
        
        
        # import math
        # max_cash = self.algo.Portfolio.Cash * 3

        # stop_loss = 0.1
        
        # quantity = min([
        #     math.floor(max_cash / self.underlying_price / 100),
        #     math.floor(abs(self.algo.Portfolio.Cash * stop_loss / self.payoff.max_loss)),
        # ])
            
        # return min([
        #     quantity,
        #     min(self.size)
        # ])

from AlgorithmImports import *
from indicator import *

class DailyIndicator(PythonIndicator):
    def __init__(self):
        self.Value = 0
        self.today = None
        self.today_highest_high = None
        self.today_lowest_low = None
        self.today_volume = 0
        # self.yesterday_close = 0
    
    def Update(self, input):
        day = input.EndTime.strftime('%y-%M-%d')
        if day != self.today:
            self.today_highest_high = input.High
            self.today_lowest_low = input.Low
            self.today = day
            self.today_volume = input.Volume
            self.Value = self.today_highest_high - self.today_lowest_low
        else:
            self.today_highest_high = max(input.High, self.today_highest_high)
            self.today_lowest_low = max(input.Low, self.today_lowest_low)
            self.today_volume = self.today_volume + input.Volume
            
            self.Value = self.today_highest_high - self.today_lowest_low
        return True
    
    def to_json(self):
        return {
            'atr': self.atr.Current.Value,
        }
        
    def is_ready(self):
        return all([
            self.atr.IsReady()
        ])

    
class SymbolData:
    def __init__(self, algorithm, symbol):
        self.algorithm = algorithm
        self.symbol = symbol
        
        self.atr = AverageTrueRange(1)
        self.di = DailyIndicator()

        self.daily_consolidator = TradeBarConsolidator(timedelta(days=1))

        algorithm.SubscriptionManager.AddConsolidator(self.symbol, self.daily_consolidator)
        
        self.daily_consolidator.DataConsolidated += self.consolidation_handler

        self.warmup()


    def consolidation_handler(self, sender, input):
        if input is  None:
            return
        # point = IndicatorDataPoint(index, row.close)
        # self.algorithm.Debug(input.EndTime)
        self.atr.Update(input)
        self.di.Update(input)

    def to_json(self):
        return {
            'tr': self.atr.Current.Value
        }
    
    def warmup(self): 
        history = self.algorithm.History(
            self.symbol, 
            self.algorithm.Time - timedelta(days=5), 
            self.algorithm.Time, 
            Resolution.Daily, 
            dataNormalizationMode=DataNormalizationMode.Raw, 
            ).reset_index(level=0)

        for index, row in history.iterrows():
            if self.algorithm.LiveMode:
                self.algorithm.Log(row.to_dict())
            bar = TradeBar(index, self.symbol, row.open, row.high, row.low, row.close, row.volume)            
            self.consolidation_handler(None, bar)
    
symbol_data_context = {}            

feature_snap = {}