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


#
# Get last 30 bars, categorize current volume and body as Ultra, high, avg, low.
# Get support and ressistance levels from past 30 bars. 
# Test it on HO future first.
# Add ability to test for group of symbols, with different resolutions.
# Train on last one month optimize parameters(n).
# in V2, explore exit strategies.
#
class VolumeSpreadAnalysisStrategy(QCAlgorithm):

    def Initialize(self):
        self.set_start_date(2023, 1, 1)
        self.set_end_date(2023,2, 1)
        self.set_cash(100000)
        
        
        self.symbols = [self.add_equity(s, Resolution.MINUTE).Symbol for s in ["TSLA"]]
        self.history_window = 30
        
        self.data_history = {symbol: [] for symbol in self.symbols}
        
        self.exit_data = {}
        for symbol in self.symbols:
            self.schedule.on(self.date_rules.every_day(symbol), self.time_rules.before_market_close(symbol, 10), self._before_market_close)

    def _before_market_close(self) -> None:
        self.liquidate(tag="EOD close")

    def OnData(self, data: Slice):
        if not self.IsWithinTradingHours():
            return
        
        for symbol in self.symbols:
            if not data.bars.contains_key(symbol):
                continue

            bar = data.bars[symbol]
            history = self.data_history[symbol]

            # Calculate bar body
            body = abs(bar.close - bar.open)

            # Append new data point to history
            history.append({
                'volume': bar.volume,
                'body': body
            })

            # Keep only last 30 bars
            if len(history) > self.history_window:
                history.pop(0)

            # Ensure enough history
            if len(history) < self.history_window:
                continue
            
            # Calculate support and resistance.

            # Get historical volume and body arrays
            volume_list = [h['volume'] for h in history[:-1]]  # exclude current
            body_list = [h['body'] for h in history[:-1]]

            # Calculate thresholds
            volume_thresholds = self.GetBuckets(volume_list)
            body_thresholds = self.GetBuckets(body_list)

            # Classify current volume and body
            volume_bucket = self.Classify(bar.volume, volume_thresholds)
            body_bucket = self.Classify(body, body_thresholds)

            # Trading logic: ULTRA volume and ULTRA body
            if volume_bucket == "ULTRA" and body_bucket == "ULTRA":
                # Do not buy again if we already has holdings.
                if self.portfolio[symbol].invested:
                    continue
                
                self.debug(f"{self.time} - SELL {symbol} - Volume: {bar.volume} ({volume_bucket}), Body: {body:.2f} ({body_bucket})")
                buy_price = bar.close

                # Buy 1 stock of that symbol. 
                self.market_order(symbol, 1, tag=f"Buy order price={buy_price}");
                self.exit_data[symbol] = {
                    "buy_price" : buy_price,
                    "tp_price" : buy_price + 1,
                    "stop_loss" : buy_price * 0.9
                }
        
        for invested_symbol in self.exit_data:
            tp_price = self.exit_data[invested_symbol]['tp_price']
            stop_loss = self.exit_data[invested_symbol]['stop_loss']
            bar = data.bars[invested_symbol]
            if bar.close > tp_price:
                self.liquidate(invested_symbol, tag="take profit")
            elif bar.close < stop_loss:
                self.liquidate(invested_symbol, tag="stop loss")


    def IsWithinTradingHours(self):
        return self.time.hour >= 10 and self.time.hour < 15

    def GetBuckets(self, values):
        q25 = np.percentile(values, 25)
        q50 = np.percentile(values, 50)
        q75 = np.percentile(values, 99)
        return {
            "LOW": (float('-inf'), q25),
            "NORMAL": (q25, q50),
            "HIGH": (q50, q75),
            "ULTRA": (q75, float('inf'))
        }

    def Classify(self, value, thresholds):
        for bucket, (low, high) in thresholds.items():
            if low <= value < high:
                return bucket
        return "ULTRA"  # fallback in edge case