from joblib import Parallel, delayed
import numpy as np

from AlgorithmImports import *

#################################################################
# VERY SIMPLE TRADING LOGIC TO TEST OUT PARALLEL THREADING:
# Calculates overall mean of the chunk means
# Compares this to the current closing price
# Generates buy signal if mean > current price
# Takes full position (100%) when signal is positive (if not invested)
# Holds position until next signal
# NO SELL LOGIC IN THIS BASIC EXAMPLE 
#(it will only take one trade when it is below the 100 day close mean)
#################################################################


class ParallelProcessingAlgorithm(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2023, 1, 1)
        self.SetEndDate(2025, 1, 1)
        self.SetCash(100000)
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol
        
    def parallel_calculation(self, data_chunk):
        return np.mean(data_chunk)
    
    def OnData(self, data):
        if not self.Portfolio.Invested:
            spyhist = self.History(self.spy, 100, Resolution.Daily) #100 days of SPY closing prices
            
            # split the data into 4 equal chunks for parallel processing:
            # aka 100 days of SPY Closing prices will look like:
            #
            #  [100] --> [25][25][25][25] -->      ... 
# variables: spyhist      data_chunks          SEE LINE 46
            #
            data_chunks = np.array_split(spyhist['close'].values, 4)

            
            #################### Using Threading Backend: ####################
            # Breakdown of joblib functions
            # Parallel: Manages thread pool and task distribution
            # delayed: Postpones function execution until thread is ready
            # n_jobs=2: Two worker threads since small example
            # backend="threading": Supports threading or multiprocessing
            ###################################################################

            results = Parallel(n_jobs=2, backend="threading")( # create 2 parallel threads (n_jobs=2)
                delayed(self.parallel_calculation)(chunk) # calculates mean price for each chunk
                for chunk in data_chunks # each thread processes different data chunks simultaneously
            )
            # -->: -->  thread1: [25][25]  : -->  [   mean values   ]
            #      -->  thread2: [25][25]         [ forall 100 days ]
# variables:              (n_jobs=2)                   results
            
            if np.mean(results) > spyhist['close'].iloc[-1]:
                self.SetHoldings(self.spy, 1.0)