# region imports
from AlgorithmImports import *

from sklearn.cluster import KMeans
from sklearn.preprocessing import MinMaxScaler
# endregion

class MuscularFluorescentPinkPanda(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2020, 1, 1)
        self.set_cash(1_000_000)
        self.universe_settings.extended_market_hours = True
        self._universe = self.add_universe(self._select_assets)
        self._spy = self.add_equity('SPY')
        self.schedule.on(self.date_rules.every_day(self._spy.symbol), self.time_rules.before_market_open(self._spy.symbol, 30), self._rebalance)

        self._liquidity_filter_size = self.get_parameter('liquidity_filter_size', 100)
        self._roc_threshold = self.get_parameter('roc_threshold', 0.1)
        self._clusters = self.get_parameter('clusters', 5)
        self._training_data_period = timedelta(self.get_parameter('training_data_days', 365))  # Calendar days
        self._hold_duration = self.get_parameter('hold_duration', 3)  # Trading days
        self._max_positions = 3

        self._scaler = MinMaxScaler(feature_range=(0, 1))
        self._kmeans = KMeans(n_clusters=self._clusters, random_state=42)

        self._trades = []
        self._columns = ['end_time', 'symbol', 'price_volatility', 'volume_volatility', 'pct_above_vwap', 'vwap_deviation', 'fraction_of_vol_below_vwap']
        self._factor_history = pd.DataFrame(columns=self._columns).set_index(['end_time', 'symbol'])
        self.set_warm_up(self._training_data_period)

    def _select_assets(self, fundamentals):
        # Select the most liquid assets.
        symbols = [f.symbol for f in sorted(fundamentals, key=lambda f: f.dollar_volume)[-self._liquidity_filter_size:]]

        # Select the subset of symbols that had >=10% growth yesterday from open to close.
        history = self.history(symbols, 1, Resolution.DAILY)
        symbols = [idx[0] for idx in history[(history.close / history.open - 1) >= self._roc_threshold].index]

        # Get the data we'll need to calculate the factors.
        vwap_by_symbol = {s: IntradayVwap('') for s in symbols}
        df_by_symbol = {s: pd.DataFrame(columns=['close', 'vwap', 'volume']) for s in symbols}
        for bars in self.history[TradeBar](symbols, self._spy.exchange.hours.get_previous_market_open(self.time, False), self.time, Resolution.MINUTE, extended_market_hours=False):
            for symbol, bar in bars.items():
                vwap = vwap_by_symbol[symbol]
                if vwap.update(bar):
                    df_by_symbol[symbol].loc[bar.end_time] = [bar.close, vwap.current.value, bar.volume]
        # Calculate factors for the selected assets.
        # The clustering process will try to minimize these.
        factors = pd.DataFrame(columns=self._columns)
        for symbol, df in df_by_symbol.items():
            if df.empty:
                continue
            factors.loc[len(factors)] = [
                self.time, 
                symbol, 
                df.close.std() / df.close.mean(),                      # price_volatility
                df.volume.std() / df.volume.mean(),                    # volume_volatility
                (df.close > df.vwap).astype(int).sum() / len(df),      # pct_above_vwap
                np.mean(np.abs(df.close - df.vwap)) / df.vwap.mean(),  # vwap_deviation
                df[df.close < df.vwap].volume.sum() / df.volume.sum()  # fraction_of_vol_below_vwap
            ]
        if factors.empty:
            return []
        factors.set_index(['end_time', 'symbol'], inplace=True)

        # Define the universe.
        universe = []
        if not self.is_warming_up:
            # Fit the scaler and k-means model to the training data.
            self._kmeans.fit(self._scaler.fit_transform(self._factor_history)) # Drop the `time` column
            # Find the cluster that's closest to the origin.
            closest_cluster_idx = np.argmin(np.linalg.norm(self._kmeans.cluster_centers_, axis=1))
            # Predict the cluster of the current universe.
            cluster_labels = self._kmeans.predict(self._scaler.transform(factors))
            # Select the subset of asset in the universe that are in the cluster closest to the origin.
            universe = [factors.index.levels[1][i] for i, label in enumerate(cluster_labels) if label == closest_cluster_idx]

        # Append the latest factor samples to the history.
        self._factor_history = pd.concat([self._factor_history, factors], axis=0).reset_index()
        # Trim off samples that have fallen out of the lookback window (1 year).
        self._factor_history = self._factor_history[self._factor_history.end_time >= self.time-self._training_data_period].set_index(['end_time', 'symbol'])
        # Return the assets selected for the universe.
        return universe

    def _rebalance(self):
        # Scan for exits.
        closed_trades = []
        for i, trade in enumerate(self._trades):
            trade.scan(self)
            if trade.closed:
                closed_trades.append(i)
        # Delete closed trades
        for i in closed_trades[::-1]:
            del self._trades[i]
        # Scan for entries.
        if self._universe.selected and len(self._trades) < self._max_positions:
            for symbol in self._universe.selected:
                self._trades.append(Trade(self, symbol, self._hold_duration, -1/self._max_positions))
        self.plot('Trades', 'Open', len(self._trades))


class Trade:

    def __init__(self, algorithm, symbol, hold_duration, weight):
        self._symbol = symbol
       
        # Determine position size
        self.closed = True
        price = algorithm.securities[symbol].price
        if not price:
            return
        self._quantity = algorithm.calculate_order_quantity(symbol, weight)
        if self._quantity == 0:
            return
        self.closed = False
       
        # Enter trade
        algorithm.market_order(symbol, self._quantity)
       
        # Set variable for exit logic
        self._hold_duration = hold_duration
       
    def scan(self, algorithm):
        if self.closed:
            return
        self._hold_duration -= 1
        if self._hold_duration <= 0:
           algorithm.market_order(self._symbol, -self._quantity)    
           self.closed = True