# region imports
from AlgorithmImports import *
from datetime import timedelta
# endregion

class IVTSAlpha(AlphaModel):
    
    ''' Provides an implementation of ivts_alpha_model that returns insight for each security'''

    def __init__(self, type, direction, period, magnitude = None, confidence = None, weight = 0.2):
        '''Initializes a new instance of the ConstantAlphaModel class
        Args:
            type: The type of insight
            direction: The direction of the insight
            period: The period over which the insight with come to fruition
            magnitude: The predicted change in magnitude as a +- percentage
            confidence: The confidence in the insight
            weight: The portfolio weight of the insights'''
        self.type = type
        self.direction = direction
        self.period = period
        self.magnitude = magnitude
        self.confidence = confidence
        self.weight = weight
        self.securities = []
        self.insights_time_by_symbol = {}

    def update(self, algorithm, data):
        ''' Creates a constant insight for each security as specified via the constructor
        Args:
            algorithm: The algorithm instance
            data: The new data available
        Returns:
            The new insights generated'''
        insights = []

        for security in self.securities:
            # security price could be zero until we get the first data point. e.g. this could happen
            # when adding both forex and equities, we will first get a forex data point
            if security.price != 0 and self.should_emit_insight(algorithm.utc_time, security.symbol):
                insights.append(Insight(security.symbol, self.period, self.type, self.direction, self.magnitude, self.confidence, weight = self.weight))

        return insights


    def on_securities_changed(self, algorithm, changes):
        ''' Event fired each time the we add/remove securities from the data feed
        Args:
            algorithm: The algorithm instance that experienced the change in securities
            changes: The security additions and removals from the algorithm'''
        for added in changes.added_securities:
            self.securities.append(added)

        # this will allow the insight to be re-sent when the security re-joins the universe
        for removed in changes.removed_securities:
            if removed in self.securities:
                self.securities.remove(removed)
            if removed.symbol in self.insights_time_by_symbol:
                self.insights_time_by_symbol.pop(removed.symbol)


    def should_emit_insight(self, utc_time, symbol):
        if symbol.is_canonical():
            # canonical futures & options are none tradable
            return False

        generated_time_utc = self.insights_time_by_symbol.get(symbol)

        if generated_time_utc is not None:
            # we previously emitted a insight for this symbol, check it's period to see
            # if we should emit another insight
            if utc_time - generated_time_utc < self.period:
                return False

        # we either haven't emitted a insight for this symbol or the previous
        # insight's period has expired, so emit a new insight now for this symbol
        self.insights_time_by_symbol[symbol] = utc_time
        return True

# region imports
from AlgorithmImports import *
# endregion

# Your New Python File

#region imports
from AlgorithmImports import *
from universe import *
from alpha import *
from portfolio import *
#endregion

class ConservativeApgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2024, 1, 1)
        self.set_end_date(2024, 1, 7)
        self.set_cash(100_000)
        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)

        # SPY 500 companies
        spy = self.add_equity("SPY",
            resolution = Resolution.MINUTE,
            data_normalization_mode = DataNormalizationMode.RAW).symbol
        self.set_benchmark(spy)

        # Subscribe price data unadjusted for splits and dividends ("raw") into the algorithm. Required for options and useful for more accurately modeling historical periods.
        self.universe_settings.data_normalization_mode = DataNormalizationMode.RAW
        self.universe_settings.asynchronous = True
        self.add_universe_selection(
            OptionChainedUniverseSelectionModel(
                # Add 10 highest dollar trading volume securities to the universe.
                self.add_universe(self.universe.dollar_volume.top(10)),
                # Set contract filter to return only front month CALL options that have the strike price within 2 strike level.
                lambda option_filter_universe: option_filter_universe.strikes(-2, +2).front_month()
            )
        )
        self.set_alpha(IVTSAlpha(InsightType.PRICE, InsightDirection.UP, timedelta(1)))
        self.Settings.RebalancePortfolioOnInsightChanges = False
        self.Settings.RebalancePortfolioOnSecurityChanges = False
        self.set_portfolio_construction(StraddlePortfolio())
        self.set_risk_management(NullRiskManagementModel())
        self.set_execution(NullExecutionModel())

# region imports
from AlgorithmImports import *
from datetime import timedelta
# endregion

# Portfolio construction scaffolding class; basic method args.
# Note: this seems equivalent to the EqualWeightingPortfolioConstructionModel
class StraddlePortfolio(PortfolioConstructionModel):

    def __init__(self, rebalance = Resolution.DAILY, portfolio_bias = PortfolioBias.LONG_SHORT):
        '''Initialize a new instance of EqualWeightingPortfolioConstructionModel
        Args:
            rebalance: Rebalancing parameter. If it is a timedelta, date rules or Resolution, it will be converted into a function.
                              If None will be ignored.
                              The function returns the next expected rebalance time for a given algorithm UTC DateTime.
                              The function returns null if unknown, in which case the function will be called again in the
                              next loop. Returning current time will trigger rebalance.
            portfolio_bias: Specifies the bias of the portfolio (Short, Long/Short, Long)'''
        super().__init__()
        self.portfolio_bias = portfolio_bias

        # If the argument is an instance of Resolution or Timedelta
        # Redefine rebalancing_func
        rebalancing_func = rebalance
        if isinstance(rebalance, int):
            rebalance = Extensions.to_time_span(rebalance)
        if isinstance(rebalance, timedelta):
            rebalancing_func = lambda dt: dt + rebalance
        if rebalancing_func:
            self.set_rebalancing_func(rebalancing_func)

    # Determines the target percent for each insight
    def determine_target_percent(self, active_insights: List[Insight]) -> Dict[Insight, float]:

        targets = {}

        # give equal weighting to each security
        count = sum(x.direction != InsightDirection.FLAT and self.respect_portfolio_bias(x) for x in active_insights)
        percent = 0 if count == 0 else 1.0 / count
        for insight in active_insights:
            targets[insight] = (insight.direction if self.respect_portfolio_bias(insight) else InsightDirection.FLAT) * percent
        return targets

    def respect_portfolio_bias(self, insight):
        '''Method that will determine if a given insight respects the portfolio bias
        Args:
            insight: The insight to create a target for
        '''
        return self.portfolio_bias == PortfolioBias.LONG_SHORT or insight.direction == self.portfolio_bias

# region imports
from AlgorithmImports import *
# endregion

class ETFConstituentsUniverse(ETFConstituentsUniverseSelectionModel):
    def __init__(self, benchmark, universe_settings: UniverseSettings = None) -> None:
        super().__init__(benchmark, universe_settings, self.etf_constituents_filter)

    def etf_constituents_filter(self, constituents):
        return [c.symbol for c in constituents]