import pandas as pd
import numpy as np
import plotly.graph_objects as go
from plotly.subplots import make_subplots
from typing import Tuple, List, Dict, Union
from datetime import datetime

class TrendChannel:
    def __init__(self, length: int = 100, multi: float = 3, extend: int = 50, 
                 color_up: str = '#16d897', color_dn: str = '#da853f'):
        self.length = length
        self.multi = multi
        self.extend = extend
        self.color_up = color_up
        self.color_dn = color_dn
        
    def calculate_atr(self, high: pd.Series, low: pd.Series, close: pd.Series, 
                     period: int = 14) -> pd.Series:
        high = high.squeeze()
        low = low.squeeze()
        close = close.squeeze()
        
        tr1 = high - low
        tr2 = abs(high - close.shift(1))
        tr3 = abs(low - close.shift(1))
        tr = pd.DataFrame({'tr1': tr1, 'tr2': tr2, 'tr3': tr3}).max(axis=1)
        atr = tr.rolling(window=period).mean()
        
        return atr

    def calculate_trend(self, close: pd.Series, atr: pd.Series) -> pd.Series:
        close = close.squeeze()
        atr = atr.squeeze()
        
        sma = close.rolling(window=self.length).mean()
        upper = sma + atr
        lower = sma - atr
        trend = pd.Series(False, index=close.index)

        for i in range(1, len(close)):
            prev_trend = trend.iloc[i-1]
            if i >= self.length:
                signal_up = (close.iloc[i] > upper.iloc[i]) & (close.iloc[i-1] <= upper.iloc[i-1])
                signal_dn = (close.iloc[i] < lower.iloc[i]) & (close.iloc[i-1] >= lower.iloc[i-1])
                
                if signal_up:
                    trend.iloc[i] = True
                elif signal_dn:
                    trend.iloc[i] = False
                else:
                    trend.iloc[i] = prev_trend
            else:
                trend.iloc[i] = prev_trend
        return trend
    
    def calculate_future_price(self, x1: pd.Timestamp, x2: pd.Timestamp, 
                         y1: float, y2: float, future_index: pd.Timestamp, 
                         df: pd.DataFrame) -> float:
        """
        Calculate future price projection.
        """
        if future_index > df.index[-1]:
            future_index = df.index[-1]

        # Convert timestamps to integer indices and cast explicitly
        x1_num = int(df.index.get_loc(x1))
        x2_num = int(df.index.get_loc(x2))
        future_index_num = int(df.index.get_loc(future_index))

        # Debug to confirm indices are integer scalars
        print(f"Debug - x1_num: {x1_num}, x2_num: {x2_num}, future_index_num: {future_index_num}")
        
        y1 = float(y1.iloc[0]) if isinstance(y1, pd.Series) else y1
        y2 = float(y2.iloc[0]) if isinstance(y2, pd.Series) else y2
        
        # Calculate slope and future price
        slope = (y2 - y1) / (x2_num - x1_num)
        future_price = y1 + slope * (future_index_num - x1_num)

        return float(future_price)

    def generate_channels(self, df: pd.DataFrame) -> Tuple[List[Dict], List[Tuple]]:
        df = df.copy()
        
        # Calculate indicators
        df['atr'] = self.calculate_atr(df['High'], df['Low'], df['Close'])
        df['trend'] = self.calculate_trend(df['Close'], df['atr'])
        df['hl2'] = (df['High'] + df['Low']) / 2
        
        # Calculate channel boundaries
        df['high_src'] = df['hl2'] + df['atr'] * self.multi
        df['low_src'] = df['hl2'] - df['atr'] * self.multi
        
        channels = []
        future_prices = []
        
        # Find trend changes
        trend_changes = (df['trend'] != df['trend'].shift(1)).fillna(False).astype(bool)
        trend_change_indices = df.index[trend_changes]
        
        for i in range(len(trend_change_indices)):
            start_idx = trend_change_indices[i]
            end_idx = (trend_change_indices[i + 1] 
                      if i < len(trend_change_indices) - 1 
                      else df.index[-1])
            
            # Debugging prints
            print(f"Debug - start_idx: {start_idx}, end_idx: {end_idx}")

            # Check if start_idx and end_idx exist in the DataFrame
            if start_idx not in df.index or end_idx not in df.index:
                print(f"Warning: start_idx or end_idx does not exist in DataFrame index.")
                continue

            # Use .loc for safer indexing
            y_mid_start = df.loc[start_idx, 'hl2']
            y_mid_end = df.loc[end_idx, 'hl2']
            y_high_start = df.loc[start_idx, 'high_src']
            y_high_end = df.loc[end_idx, 'high_src']
            y_low_start = df.loc[start_idx, 'low_src']
            y_low_end = df.loc[end_idx, 'low_src']
            
            # Ensure trend_value is a single scalar
            trend_value = df.loc[start_idx, 'trend']
            trend = trend_value.item() if isinstance(trend_value, pd.Series) else bool(trend_value)
            
            channel = {
                'x_start': start_idx,
                'x_end': end_idx,
                'y_mid_start': y_mid_start,
                'y_mid_end': y_mid_end,
                'y_high_start': y_high_start,
                'y_high_end': y_high_end,
                'y_low_start': y_low_start,
                'y_low_end': y_low_end,
                'trend': trend
            }
            channels.append(channel)
            
            # Calculate future_index based on the last available index in the DataFrame
            if end_idx in df.index:
                future_index = end_idx + pd.Timedelta(days=self.extend)
            else:
                future_index = df.index[-1] + pd.Timedelta(days=self.extend)
            print(f"Debug - future_index: {future_index}")

            # Check if future_index exists in the DataFrame
            if future_index not in df.index:
                print(f"Warning: future_index {future_index} does not exist in DataFrame index.")
                continue

            # Use .loc for safer indexing
            y_mid_start = df.loc[start_idx, 'hl2']
            y_mid_end = df.loc[end_idx, 'hl2']
            y_high_start = df.loc[start_idx, 'high_src']
            y_high_end = df.loc[end_idx, 'high_src']
            y_low_start = df.loc[start_idx, 'low_src']
            y_low_end = df.loc[end_idx, 'low_src']
            
            # Ensure trend_value is a single scalar
            trend_value = df.loc[start_idx, 'trend']
            trend = trend_value.item() if isinstance(trend_value, pd.Series) else bool(trend_value)
            
            channel = {
                'x_start': start_idx,
                'x_end': end_idx,
                'y_mid_start': y_mid_start,
                'y_mid_end': y_mid_end,
                'y_high_start': y_high_start,
                'y_high_end': y_high_end,
                'y_low_start': y_low_start,
                'y_low_end': y_low_end,
                'trend': trend
            }
            channels.append(channel)
            
            future_index = end_idx + pd.Timedelta(days=self.extend)
            print(f"Debug - future_index: {future_index}")

            # Check if future_index exists in the DataFrame
            if future_index not in df.index:
                print(f"Warning: future_index {future_index} does not exist in DataFrame index.")
                continue

            future_price = self.calculate_future_price(
                start_idx, end_idx,
                y_mid_start, y_mid_end,
                future_index, df
            )
            future_prices.append((future_index, future_price, trend))
        
        return channels, future_prices

    def plot(self, df: pd.DataFrame) -> go.Figure:
        channels, future_prices = self.generate_channels(df)
        
        fig = go.Figure()
        
        fig.add_trace(go.Candlestick(
            x=df.index,
            open=df['Open'],
            high=df['High'],
            low=df['Low'],
            close=df['Close'],
            name='Price'
        ))
        
        for channel in channels:
            color = self.color_up if channel['trend'] else self.color_dn
            
            # Directly using the channel's start and end indices
            fig.add_trace(go.Scatter(
                x=[channel['x_start'], channel['x_end']],
                y=[channel['y_mid_start'], channel['y_mid_end']],
                mode='lines',
                line=dict(color=color),
                name='Mid Line'
            ))
            
            for y_start, y_end in [(channel['y_high_start'], channel['y_high_end']), (channel['y_low_start'], channel['y_low_end'])]:
                fig.add_trace(go.Scatter(
                    x=[channel['x_start'], channel['x_end']],
                    y=[y_start, y_end],
                    mode='lines',
                    line=dict(color=color, dash='dash'),
                    showlegend=False
                ))
        
        for future_idx, price, trend in future_prices:
            color = self.color_up if trend else self.color_dn
            fig.add_trace(go.Scatter(
                x=[future_idx],
                y=[float(price)],
                mode='markers+text',
                marker=dict(color=color, symbol='diamond', size=10),
                text=f'Future Price: {price:.2f}',
                textposition='top center',
                showlegend=False
            ))
        
        fig.update_layout(
            title='Future Trend Channel',
            yaxis_title='Price',
            xaxis_title='Date',
            template='plotly_dark',
            showlegend=True,
            xaxis_rangeslider_visible=False
        )
        
        return fig

if __name__ == "__main__":
    import yfinance as yf
    
    # Download sample data
    symbol = '^IBEX'
    df = yf.download(symbol, start='2023-01-01', end='2024-01-01')
    
    # Create and plot trend channel
    tc = TrendChannel(length=100, multi=3, extend=50)
    fig = tc.plot(df)
    fig.show()

Fix the warnings:
Debug - start_idx: 2023-01-02 00:00:00+00:00, end_idx: 2023-06-02 00:00:00+00:00
Debug - future_index: 2023-07-22 00:00:00+00:00
Warning: future_index 2023-07-22 00:00:00+00:00 does not exist in DataFrame index.
Debug - start_idx: 2023-06-02 00:00:00+00:00, end_idx: 2023-10-03 00:00:00+00:00
Debug - future_index: 2023-11-22 00:00:00+00:00
Debug - future_index: 2023-11-22 00:00:00+00:00
Debug - x1_num: 106, x2_num: 193, future_index_num: 229
Debug - start_idx: 2023-10-03 00:00:00+00:00, end_idx: 2023-11-14 00:00:00+00:00
Debug - future_index: 2024-01-03 00:00:00+00:00
Warning: future_index 2024-01-03 00:00:00+00:00 does not exist in DataFrame index.
Debug - start_idx: 2023-11-14 00:00:00+00:00, end_idx: 2023-12-29 00:00:00+00:00
Debug - future_index: 2024-02-17 00:00:00+00:00
Warning: future_index 2024-02-17 00:00:00+00:00 does not exist in DataFrame index.