# region imports
from AlgorithmImports import *
from universe import TopologicalGraphUniverseSelectionModel
from portfolio import EqualClustersWeightingPortfolioConstructionModel
# endregion
np.random.seed(0)

class TopologicalPortfolio(QCAlgorithm):
    def initialize(self):
        self.set_end_date(datetime.now())
        self.set_start_date(self.end_date - timedelta(5*365))
        self.set_cash(1000000)

        # Lookback window to construct and analyze the topological structure.
        history_lookback = self.get_parameter("history_lookback", 250)
        # Set the period to reconstruct the topological complex.
        recalibrate_periods = [
            self.date_rules.every_day(),
            self.date_rules.week_start(),
            self.date_rules.month_start(),
            self.date_rules.year_start()
        ]
        recalibrate_period = recalibrate_periods[self.get_parameter("recalibrate_period", 0)]
        # Construct a portfolio with SPY constituents.
        universe_model = TopologicalGraphUniverseSelectionModel(
            "SPY",
            history_lookback,
            recalibrate_period,
            lambda u: [x.symbol for x in sorted(
                [x for x in u if x.weight], 
                key=lambda x: x.weight, 
                reverse=True
            )[:200]]
        )
        self.add_universe_selection(universe_model)
        self.add_alpha(ConstantAlphaModel(InsightType.PRICE, InsightDirection.UP, timedelta(1)))
        self.set_portfolio_construction(EqualClustersWeightingPortfolioConstructionModel(timedelta(1), universe_model))

        # We would like to compare with SPY for the correlation and risk-adjusted return.
        self.set_benchmark("SPY")

# region imports
from AlgorithmImports import *
from Portfolio.EqualWeightingPortfolioConstructionModel import EqualWeightingPortfolioConstructionModel
# endregion

class EqualClustersWeightingPortfolioConstructionModel(EqualWeightingPortfolioConstructionModel):
    def __init__(self, rebalance, universe_model):
        self.universe_model = universe_model
        super().__init__(rebalance)

    def determine_target_percent(self, active_insights):
        results = {}
        if self.universe_model.clustered_symbols is not None:
            # Equal weight distribution between and within each cluster.
            series = self.weight_distribution(self.universe_model.clustered_symbols)
            # Connect insight with the assigned weight.
            return {insight: series[insight.symbol] if insight.symbol in series.index else 0 for insight in active_insights}
        return {}

    def weight_distribution(self, clustered_symbols):
        # Assign weight between and within giant and small clusters. Note that we do not invest in outliers.
        weights = {}
        def assign_weights(nested_list, level=1):
            num_elements = len(nested_list)
            if num_elements == 0:
                return
            weight_per_element = 1 / num_elements
            for item in nested_list:
                if isinstance(item, list):
                    assign_weights(item, level + 1)
                else:
                    weights[item] = weights.get(item, 0) + weight_per_element / (2 ** (level - 1))
        # Calculate the overall weights.
        assign_weights(clustered_symbols)
        return pd.Series(weights) / sum(weights.values())

# region imports
from AlgorithmImports import *
from Selection.ETFConstituentsUniverseSelectionModel import ETFConstituentsUniverseSelectionModel
import kmapper as km
from sklearn.cluster import DBSCAN
from sklearn.decomposition import PCA
from umap import UMAP
# endregion
np.random.seed(0)

class TopologicalGraphUniverseSelectionModel(ETFConstituentsUniverseSelectionModel):
    def __init__(self, etf_symbol, lookback_window = 250, recalibration_period = None, universe_filter_func = None):
        self.lookback_window = lookback_window
        self.recalibration_period = recalibration_period
        self.clustered_symbols = None
        super().__init__(etf_symbol, None, universe_filter_func)

    def create_universes(self, algorithm: QCAlgorithm) -> list[Universe]:
        universe_list = super().create_universes(algorithm)
        # Set schedule event to reconstruct the topological structure.
        # Initial warm up of the universe (the maximum days until the next open is 4 days).
        algorithm.schedule.on(
            algorithm.date_rules.on([algorithm.time + timedelta(5)]),
            algorithm.time_rules.at(23, 59),
            lambda: self.get_graph_symbols(algorithm)
        )
        algorithm.schedule.on(
            self.recalibration_period if self.recalibration_period is not None else algorithm.date_rules.every_day(),
            algorithm.time_rules.at(0, 1),
            lambda: self.get_graph_symbols(algorithm)
        )
        return universe_list

    def get_graph_symbols(self, algorithm):
        # Construct simplicial complex.
        graph, symbol_list = self.construct_simplicial_complex(algorithm, self.lookback_window)
        if len(symbol_list) > 0:
            self.clustered_symbols = self.clustering_symbols(graph, symbol_list)

    def construct_simplicial_complex(self, algorithm: QCAlgorithm, lookback_window: int) -> Dict[str, object]:
        if not self.universe.selected:
            return {}, []
        # Obtain historical data to construct a graph of stock relationship.
        prices = algorithm.history(self.universe.selected, lookback_window, Resolution.DAILY).unstack(0).close
        # Calculate daily log return. Then, transpose the data since we're relating stocks.
        log_returns = np.log(prices / prices.shift(1)).dropna().T
        if log_returns.empty:
            return {}, []

        # Initialize the mapper algorithm.
        mapper = km.KeplerMapper()
        # Project the data into a 2d subspace via 2 transformation, PCA and UMAP.
        # PCA: since it can retain the most variance while denoising, as well as fast.
        # UMAP: handles non-linear relationships well and preserves both local and global structures.
        # MDS and Isomap are not included due to their potential sensitivity to noise and outliers in financial data.
        projected_data = mapper.fit_transform(log_returns, projection=[PCA(n_components=0.8, random_state=1), UMAP(n_components=1, random_state=1, n_jobs=-1)])
        # Cluster the data with DBSCAN since it is better in handling noise.
        # We are interested in the correlation distance to cluster and form a portfolio.
        graph = mapper.map(projected_data, log_returns, clusterer=DBSCAN(metric='correlation', n_jobs=-1))
        return graph, prices.columns

    def clustering_symbols(self, graph, symbol_list):
        # Each connected structure as a giant cluster.
        linked_clusters = []
        for x, y in graph['links'].items():
            isin = False
            for i in range(len(linked_clusters)):
                if x in linked_clusters[i] or y in linked_clusters[i]:
                    linked_clusters[i] = list(set(linked_clusters[i] + [x] + y))
                    isin = True
            if isin:
                continue
            linked_clusters.append([x] + y)
        linked_clusters += [[x] for x in graph['nodes'] if x not in [z for y in linked_clusters for z in y]]
        # Convert the node into symbol.
        return [[list([symbol_list[graph['nodes'][x]]][0]) for x in linked_cluster] for linked_cluster in linked_clusters]