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

class CreativeYellowGreenArmadillo(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2026, 1, 1)
        self.set_cash(100000)
        self.settings.seed_initial_prices = True
        
        # Universe settings
        self._universe_size = 20
        self._lookback_days = 15  # Need 10 days for correlation
        self._rebalance_days = 5
        
        # Add universe
        self._universe = self.add_universe(self._select_coarse)
        
        # Schedule rebalancing
        self.schedule.on(self.date_rules.every_day(), 
                        self.time_rules.after_market_open("SPY", 30), 
                        self._rebalance)
        
        # Track last rebalance time
        self._last_rebalance = self.time
    
    def _select_coarse(self, coarse):
        """Universe selection using WorldQuant Brain alpha formula"""
        # Filter for liquid stocks
        filtered = [x for x in coarse if x.has_fundamental_data 
                   and x.price > 5 
                   and x.dollar_volume > 10000000]
        
        if len(filtered) < self._universe_size:
            return [x.symbol for x in filtered]
        
        # Get historical data for all symbols
        symbols = [x.symbol for x in filtered]
        history = self.history(symbols, self._lookback_days, Resolution.DAILY)
        
        if history.empty:
            return symbols[:self._universe_size]
        
        # Calculate alpha scores
        scores = {}
        
        for symbol in symbols:
            try:
                if symbol not in history.index.get_level_values(0):
                    continue
                    
                df = history.loc[symbol]
                
                if len(df) < 10:  # Need at least 10 days
                    continue
                
                close = df['close'].values
                open_price = df['open'].values
                volume = df['volume'].values
                
                # Calculate alpha: (-rank(ts_covariance(rank(close),rank(volume),5)))+rank(-ts_corr(open,volume,10))
                alpha = self._calculate_alpha(close, open_price, volume)
                
                if not np.isnan(alpha):
                    scores[symbol] = alpha
                    
            except Exception as e:
                continue
        
        # Select top stocks by alpha score
        sorted_symbols = sorted(scores.items(), key=lambda x: x[1], reverse=True)
        selected = [symbol for symbol, score in sorted_symbols[:self._universe_size]]
        
        return selected
    
    def _calculate_alpha(self, close, open_price, volume):
        """Calculate WorldQuant Brain alpha formula"""
        # Part 1: -rank(ts_covariance(rank(close), rank(volume), 5))
        rank_close = self._rank(close)
        rank_volume = self._rank(volume)
        
        # Calculate 5-day rolling covariance
        cov_values = []
        for i in range(4, len(rank_close)):
            window_close = rank_close[i-4:i+1]
            window_volume = rank_volume[i-4:i+1]
            cov = np.cov(window_close, window_volume)[0, 1]
            cov_values.append(cov)
        
        if len(cov_values) == 0:
            return np.nan
        
        # Rank the covariance values (cross-sectional would be across stocks, here we approximate with time series)
        rank_cov = self._rank(np.array(cov_values))
        part1 = -rank_cov[-1]  # Take the most recent value and negate
        
        # Part 2: rank(-ts_corr(open, volume, 10))
        if len(open_price) < 10:
            return np.nan
            
        # Calculate 10-day rolling correlation
        corr_values = []
        for i in range(9, len(open_price)):
            window_open = open_price[i-9:i+1]
            window_volume = volume[i-9:i+1]
            corr = np.corrcoef(window_open, window_volume)[0, 1]
            corr_values.append(corr)
        
        if len(corr_values) == 0:
            return np.nan
        
        # Rank the negative correlation values
        rank_neg_corr = self._rank(-np.array(corr_values))
        part2 = rank_neg_corr[-1]  # Take the most recent value
        
        # Combine
        alpha = part1 + part2
        return alpha
    
    def _rank(self, values):
        """Rank values from 0 to 1"""
        if len(values) <= 1:
            return np.array([0.5] * len(values))
        
        # Use pandas for ranking
        series = pd.Series(values)
        ranks = series.rank(method='average', pct=True)
        return ranks.values
    
    def _rebalance(self):
        """Rebalance portfolio"""
        # Only rebalance every N days
        if (self.time - self._last_rebalance).days < self._rebalance_days:
            return
        
        self._last_rebalance = self.time
        
        # Get active universe members
        active_securities = [x for x in self._universe.selected]
        
        if len(active_securities) == 0:
            return
        
        # Equal weight portfolio
        weight = 1.0 / len(active_securities)
        targets = [PortfolioTarget(symbol, weight) for symbol in active_securities]
        
        self.set_holdings(targets, liquidate_existing_holdings=True)