# region imports
from AlgorithmImports import *
from alpha_1637 import Strategy001AlphaModel
from porfolio_1737 import Stragety001PortfolioConstructionModel
# endregion

class MeanReversionAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2009, 1, 10)
        self.set_end_date(2024, 5, 12)
        self.set_cash(30000)  # Set Strategy Cash
        
        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)

        self.add_universe_selection(ETFConstituentsUniverseSelectionModel("QQQ"))
        self.universe_settings.resolution = Resolution.DAILY

        self.add_alpha(Strategy001AlphaModel())

        self.set_portfolio_construction(Stragety001PortfolioConstructionModel())

        self.set_execution(ImmediateExecutionModel())
    

#region imports
from AlgorithmImports import *
#endregion

class Strategy001AlphaModel(AlphaModel):
    def __init__(self):
        self.entry_rolling_window_dict = {}
        self.exit_rolling_window_dict = {}

    def update(self, algorithm, data):
        entry_insights_list = self.entry_insights_update(algorithm, data)
        exit_insights_list = self.exit_insights_update(algorithm, data)
        insight_group = Insight.group(entry_insights_list + exit_insights_list)

        return insight_group
    
    def exit_insights_update(self, algorithm, data):
        exit_insights_list = []
        for symbol, trade_bar in data.items():
            self.exit_rolling_window_dict[symbol].add(trade_bar)
            
            if (self.exit_rolling_window_dict[symbol].count == 2) and algorithm.portfolio[symbol].invested:
                if (self.exit_rolling_window_dict[symbol][0] is None) or (self.exit_rolling_window_dict[symbol][1] is None):
                    algorithm.log(f"system error:{symbol.value}:self.exit_rolling_window_dict[symbol][0] or [1] is None: trade_bar is None? {trade_bar is None}")
                    return []
                if self.exit_rolling_window_dict[symbol][0].close > self.exit_rolling_window_dict[symbol][1].close:
                    exit_insights_list.append(
                        Insight(
                            symbol=symbol, 
                            period=timedelta(days=1), 
                            type=InsightType.PRICE, 
                            direction=InsightDirection.FLAT,
                            magnitude= self.exit_rolling_window_dict[symbol][0].volume,
                            confidence=self.exit_rolling_window_dict[symbol][0].volume))
                    # algorithm.log(f"exit insight:{symbol.value}")
        return exit_insights_list


    def entry_insights_update(self, algorithm, data):
        entry_insights_list = []
        if data.time.weekday() in [4, 5, 1]:
            for symbol, trade_bar in data.items():
                self.entry_rolling_window_dict[symbol].add(trade_bar)

                if (data.time.weekday() == 1) and (self.entry_rolling_window_dict[symbol].count == 3):
                    if (self.entry_rolling_window_dict[symbol][0] is None) or (self.entry_rolling_window_dict[symbol][1] is None) or (self.entry_rolling_window_dict[symbol][2] is None):
                        algorithm.log(f"system error:{symbol.value}:self.entry_rolling_window_dict[symbol][0] or [1] or [2] is None: trade_bar is None? {trade_bar is None}")
                        return []
                    flag = True
                    flag = flag and (self.entry_rolling_window_dict[symbol][0].close < self.entry_rolling_window_dict[symbol][1].close)
                    flag = flag and (self.entry_rolling_window_dict[symbol][1].close < self.entry_rolling_window_dict[symbol][2].close)
                    if flag:
                        entry_insights_list.append(
                            Insight(
                                symbol=symbol, 
                                period=timedelta(days=1), 
                                type=InsightType.PRICE, 
                                direction=InsightDirection.UP, 
                                magnitude=self.entry_rolling_window_dict[symbol][0].volume,
                                confidence=self.entry_rolling_window_dict[symbol][0].volume))
                        # algorithm.log(f"entry insight:{symbol.value}")
                        
        return entry_insights_list

    def on_securities_changed(self, algorithm, changes):
        for security in changes.added_securities:
            self.entry_rolling_window_dict.setdefault(security.symbol, RollingWindow[TradeBar](3))
            self.exit_rolling_window_dict.setdefault(security.symbol, RollingWindow[TradeBar](2))

#region imports
from AlgorithmImports import *
#endregion

 # Portfolio construction scaffolding class; basic method args.
class Stragety001PortfolioConstructionModel(PortfolioConstructionModel):
    # Create list of PortfolioTarget objects from Insights
    def create_targets(self, algorithm: QCAlgorithm, insights: List[Insight]) -> List[PortfolioTarget]:
        return super().create_targets(algorithm, insights)

    # Determines the target percent for each insight
    def determine_target_percent(self, activeInsights: List[Insight]) -> Dict[Insight, float]:
        targets = {}
        securities_count = 0
        for security_holding in self.algorithm.portfolio.values():
            if security_holding.invested:
                securities_count += 1
        self.algorithm.log(f"securities_count:{securities_count}")
        
        sorted_acive_insights = sorted(activeInsights, key=lambda x: x.confidence, reverse=True)

        # exit
        for insight in sorted_acive_insights:
            if insight.direction == InsightDirection.FLAT:
                targets[insight] = 0.0
                securities_count -= 1
            if securities_count > 5 or securities_count < 0:
                self.algorithm.log(f"securities_count:{securities_count}:exit")
                continue
            assert (securities_count >= 0) and (securities_count <= 5)

        # # entry
        for insight in sorted_acive_insights:
            if (insight.direction == InsightDirection.UP) and (securities_count < 5):
                targets[insight] = 0.2
                securities_count += 1
            if securities_count > 5 or securities_count < 0:
                self.algorithm.log(f"securities_count:{securities_count}:entry")
                continue
            assert (securities_count >= 0) and (securities_count <= 5)

        log_string = ""
        for t in targets:
            log_string += t.to_string() + "; "
        self.algorithm.log(log_string)

        return targets