# region imports
from AlgorithmImports import *
from collections import deque
# endregion

class VolumeProfileComparisonAlgorithm(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2023, 10, 1)
        self.SetEndDate(2023, 12, 31)
        self.SetCash(100_000)

        res = Resolution.Minute
        self.symbol = self.AddEquity("AAPL", res).Symbol
        self.previous_day = None
        self.bar_count = 0

        if False:
            # Create an Market Profile indicator for the symbol with Volume Profile (VOL) mode
            # https://github.com/QuantConnect/Lean/blob/master/Indicators/MarketProfile.cs#L39
            # VolumeProfile(string name, int period, decimal valueAreaVolumePercentage = 0.70m, decimal priceRangeRoundOff = 0.05m)
            # period: The period of the VP
            # valueAreaVolumePercentage: The percentage of volume contained in the value area
            # priceRangeRoundOff: How many digits you want to round and the precision. i.e 0.01 round to two digits exactly.
            # resolution: The resolution
            # selector: Selects a value from the BaseData to send into the indicator, if null defaults to casting the input value to a TradeBar
            self.vp = self.VP(
                symbol=self.symbol,
                period=380, # for plotting the full day of open market hours at 15:50
                valueAreaVolumePercentage=0.85,
                priceRangeRoundOff=0.01,
                resolution=res)

        else:
            # Initialize MyVolumeProfile indicator
            self.vp = MyVolumeProfile (
                name="MP",
                period=380, # for plotting the full day of open market hours at 15:50
                valueAreaVolumePercentage=0.85,
                priceRangeRoundOff=0.01)
            # Register the indicator for automatic updates
            self.RegisterIndicator(self.symbol, self.vp, res)

    def OnData(self, slice):
        '''
        Profile High: The highest price level within the volume profile.
        Profile Low: The lowest price level within the volume profile.
        Point of Control Price (POCPrice): The price level with the highest trading volume.
        Value Area High: The upper price level of the value area.
        Value Area Low: The lower price level of the value area.
        Point of Control Volume (POCVolume): The volume at the Point of Control price level.
        Value Area Volume: The total volume within the value area, which is the range where a specified percentage (e.g., 70%) of the total volume is traded.
        Num Buckets: The number of buckets resulting from priceRangeRoundOff.
        '''
        if self.Time.day != self.previous_day:
            self.previous_day = self.Time.day
            self.daily_bar_count = 0

        self.bar_count += 1

        if not slice.ContainsKey(self.symbol) or (self.Time.hour == 0 and self.Time.minute == 0):
            return

        close = slice[self.symbol].Close

        # Check for plotting time
        if (self.Time.hour < 16 and self.Time.minute % 5 == 0) or (self.Time.hour == 15 and self.Time.minute == 59):
            if self.vp.IsReady:
                #self.Debug(f'{self.Time} {self.bar_count}')
                self.Plot("VP1", "close", close)
                self.Plot("VP1", "vp", self.vp.Current.Value)
                self.Plot("VP1", "profilehigh", self.vp.ProfileHigh)
                self.Plot("VP1", "profilelow", self.vp.ProfileLow)
                #self.Plot("VP1", "pocprice", self.vp.POCPrice)
                self.Plot("VP1", "valueareahigh", self.vp.ValueAreaHigh)
                self.Plot("VP1", "valuearealow", self.vp.ValueAreaLow)
                self.Plot("VP2", "pocvolume", self.vp.POCVolume)
                self.Plot("VP2", "valueareavolume", self.vp.ValueAreaVolume)
                if hasattr(self.vp, 'Time'):
                    self.Plot("VP0", "hour", self.vp.Time.hour)
                if hasattr(self.vp, 'NumBuckets'):
                    self.Plot("VP3", "number of buckets", self.vp.NumBuckets)

class MyVolumeProfile(PythonIndicator):
    """
    Represents a Volume Profile Indicator.

    The Volume Profile indicator displays trading activity over a specified period, showing the volume at different price levels. 
    This indicator is useful for identifying support and resistance levels and understanding where significant trading activity has occurred.

    Args:
        name (str): The name of the indicator.
        period (int): The lookback period for calculating the volume profile.
        valueAreaVolumePercentage (float): Percentage of total volume contained in the Value Area.
        priceRangeRoundOff (float): Precision for rounding off price levels.
    """
    def __init__(self, name, period, valueAreaVolumePercentage=0.70, priceRangeRoundOff=0.05):
        super().__init__(name)
        self.period = period
        self.valueAreaVolumePercentage = valueAreaVolumePercentage # Percentage of total volume contained in the ValueArea
        self.priceRangeRoundOff = 1 / priceRangeRoundOff # The range of roundoff to the prices. i.e two decimal places, three decimal places
        self.is_ready = False
        self.Value = 0
        self.Time = datetime.min
        self.volumePerPrice = {} # Buckets with Close values and Volume values in the given period of time
        self.oldDataPoints = deque(maxlen=period) # Rolling window filled with tuple of (close, volume)
        self.ProfileHigh = 0  # Highest price level in the volume profile
        self.ProfileLow = 0   # Lowest price level in the volume profile
        self.ValueAreaHigh = 0  # Upper boundary of the value area
        self.ValueAreaLow = 0   # Lower boundary of the value area
        self.ValueAreaVolume = 0  # Total volume within the value area
        self.POCPrice = 0  # Price level with the highest volume (Point of Control)
        self.POCVolume = 0  # Volume at the Point of Control
        self.NumBuckets = 0  # Number of distinct price levels considered

    def Update(self, input: TradeBar) -> bool:
        if input is None:
            return False

        # Get time, close, volume
        self.Time = input.Time
        rounded_close = self._round(input.Close)

        # Update rolling window
        self.oldDataPoints.append((rounded_close, input.Volume))

        # Update buckets with key=Close and value=Volume
        if rounded_close not in self.volumePerPrice:
            self.volumePerPrice[rounded_close] = input.Volume
        else:
            self.volumePerPrice[rounded_close] += input.Volume

        # Remove old data from bucket and rolling window
        if len(self.oldDataPoints) == self.period:
            self.is_ready = True
            # Remove and return the leftmost (or first) item from rolling window
            removed_close, removed_volume = self.oldDataPoints.popleft()
            # Remove old data point's volume
            if removed_close in self.volumePerPrice:
                self.volumePerPrice[removed_close] -= removed_volume
                # Remove old data point, if no volume left
                if self.volumePerPrice[removed_close] <= 0:
                    del self.volumePerPrice[removed_close]

        # Update derived values
        if self.is_ready and len(self.volumePerPrice) > 0:
            self.UpdatePOC()
            self.UpdateValueArea()

        return self.is_ready

    def _round(self, value: float) -> float:
        # Round the value to the specified precision
        return round(value * self.priceRangeRoundOff) / self.priceRangeRoundOff

    def UpdatePOC(self):
        # Determine the Point of Control (POC)
        poc_price, poc_volume = max(self.volumePerPrice.items(), key=lambda x: x[1])
        self.Value = poc_price
        self.POCPrice = poc_price
        self.POCVolume = poc_volume

        # Determine Profile High and Low
        self.ProfileHigh = max(self.volumePerPrice.keys())
        self.ProfileLow = min(self.volumePerPrice.keys())

        # Determine the number of price buckets for control purposes
        self.NumBuckets = len(self.volumePerPrice)

    def UpdateValueArea(self):
        # Calculate the value area and its volume
        total_volume = sum(volume for _, volume in self.oldDataPoints)
        border_area_volume_target = total_volume * (1 - self.valueAreaVolumePercentage) / 2

        # Create list of (Close, Volume) sorted by Volume
        sorted_volume_data = sorted(self.volumePerPrice.items(), key=lambda x: x[0], reverse=False)

        # Determine ValueAreaLow price
        current_volume = 0
        value_area_prices = []
        for price, volume in sorted_volume_data:
            if current_volume + volume > border_area_volume_target:
                break
            value_area_prices.append(price)
            current_volume += volume
        self.ValueAreaLow = max(value_area_prices) if len(value_area_prices) > 0 else self.ProfileLow
        border_area_volume_low = current_volume

        # Determine ValueAreaHigh price
        current_volume = 0
        value_area_prices = []
        for price, volume in reversed(sorted_volume_data):
            if current_volume + volume > border_area_volume_target:
                break
            value_area_prices.append(price)
            current_volume += volume
        self.ValueAreaHigh = min(value_area_prices) if len(value_area_prices) > 0 else self.ProfileHigh
        border_area_volume_high = current_volume

        # Calculate Value Area Volume
        self.ValueAreaVolume = total_volume - border_area_volume_low - border_area_volume_high

    @property
    def IsReady(self) -> bool:
        return self.is_ready