import numpy as np
import pandas as pd
import MetaTrader5 as mt5
from datetime import datetime
from sklearn.preprocessing import MinMaxScaler
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import warnings
warnings.filterwarnings('ignore')

class TrendForceIndicator:
    def __init__(self, symbol: str, timeframe: int = mt5.TIMEFRAME_H1, lookback: int = 10000):
        """
        Инициализация индикатора
        
        Args:
            symbol (str): Торговый символ (например, "EURUSD")
            timeframe (int): Таймфрейм (по умолчанию H1)
            lookback (int): Количество баров для анализа
        """
        self.symbol = symbol
        self.timeframe = timeframe
        self.lookback = lookback
        self.scaler = MinMaxScaler(feature_range=(3, 9))

    def analyze_market(self) -> pd.DataFrame:
        """Основной метод анализа рынка"""
        # Получаем данные из MT5
        rates = mt5.copy_rates_from_pos(self.symbol, self.timeframe, 0, self.lookback)
        if rates is None:
            raise ValueError(f"Failed to get data for {self.symbol}")
            
        df = pd.DataFrame(rates)
        df['time'] = pd.to_datetime(df['time'], unit='s')
        df.set_index('time', inplace=True)
        
        # Рассчитываем компоненты
        df = self._calculate_components(df)
        return df.tail(500)  # Возвращаем последние 500 баров

    def _calculate_components(self, df: pd.DataFrame) -> pd.DataFrame:
        """Расчет всех компонент индикатора"""
        # Базовые компоненты
        df['volatility'] = df['close'].pct_change().rolling(20).std()
        df['momentum'] = df['close'].pct_change(5)
        
        # Объемный профиль
        df['volume_ma'] = df['tick_volume'].rolling(20).mean()
        df['volume_trend'] = df['tick_volume'] / df['volume_ma']
        
        # Расчет силы тренда
        df['trend_force'] = df['volatility'] * df['volume_trend'] * abs(df['momentum'])
        df['trend_force_norm'] = self.scaler.fit_transform(
            df['trend_force'].values.reshape(-1, 1)
        ).flatten()
        
        # Определение направления
        df['trend_direction'] = np.sign(df['momentum'])
        
        # Расчет торговых сессий
        df['session_coef'] = 1.0
        hour = df.index.hour
        
        # Коэффициенты для разных сессий
        asian_mask = (hour >= 0) & (hour < 8)
        european_mask = (hour >= 8) & (hour < 16)
        american_mask = (hour >= 16) & (hour < 24)
        
        df.loc[asian_mask, 'session_coef'] = 0.7    # Азиатская сессия
        df.loc[european_mask, 'session_coef'] = 1.0  # Европейская сессия
        df.loc[american_mask, 'session_coef'] = 0.9  # Американская сессия
        
        # Итоговый показатель силы тренда
        df['trend_force_adjusted'] = df['trend_force_norm'] * df['session_coef']
        
        # Компоненты для определения разворотов
        df['direction_strength'] = df['trend_direction'] * df['trend_force_adjusted']
        
        # Паттерны разворота
        df['reversal_pattern'] = np.where(
            (df['trend_force_adjusted'] > 6) & 
            (df['direction_strength'].diff().rolling(3).std() > 1.5),
            1, 0
        )
        
        # Подтверждение объемом
        df['volume_confirmation'] = np.where(
            (df['reversal_pattern'] == 1) &
            (df['volume_trend'] > df['volume_trend'].rolling(20).mean() * 1.5),
            'Strong',
            'Weak'
        )
        
        return df

    def create_visualization(self, df: pd.DataFrame = None):
        """Создание интерактивной визуализации"""
        if df is None:
            df = self.analyze_market()
        
        df = df.reset_index()
        
        # Создаем график с тремя подокнами
        fig = make_subplots(rows=3, cols=1, 
                           shared_xaxes=True,
                           subplot_titles=('Price', 'Trend Force', 'Trend Direction'),
                           row_heights=[0.5, 0.25, 0.25],
                           vertical_spacing=0.05)

        # График цены
        fig.add_trace(
            go.Candlestick(
                x=df['time'],
                open=df['open'],
                high=df['high'],
                low=df['low'],
                close=df['close'],
                name='OHLC',
                hoverlabel=dict(
                    bgcolor='white',
                    font=dict(size=12)
                )
            ),
            row=1, col=1
        )

        # График силы тренда
        fig.add_trace(
            go.Scatter(
                x=df['time'],
                y=df['trend_force_adjusted'],
                mode='lines',
                line=dict(color='blue', width=2),
                name='Trend Force',
                hovertemplate="<br>".join([
                    "Time: %{x}",
                    "Force: %{y:.2f}",
                    "<extra></extra>"
                ])
            ),
            row=2, col=1
        )

        # Опорные уровни
        fig.add_hline(y=3, line_dash="dash", line_color="yellow", row=2, col=1)
        fig.add_hline(y=6, line_dash="dash", line_color="green", row=2, col=1)

        # График направления тренда
        fig.add_trace(
            go.Bar(
                x=df['time'],
                y=df['trend_direction'] * df['trend_force_adjusted'],
                name='Trend Direction',
                marker_color=np.where(df['trend_direction'] > 0, 'green', 'red'),
                hovertemplate="<br>".join([
                    "Time: %{x}",
                    "Direction: %{y:.2f}",
                    "Reversal: %{text}",
                    "<extra></extra>"
                ]),
                text=df['volume_confirmation']
            ),
            row=3, col=1
        )

        # Настройка внешнего вида
        fig.update_layout(
            height=1000,
            title_text=f"Trend Force Indicator - {self.symbol}",
            showlegend=True,
            xaxis3_title="Time",
            yaxis_title="Price",
            yaxis2_title="Force",
            yaxis3_title="Direction"
        )

        return fig

def main():
    # Инициализация MT5
    if not mt5.initialize():
        print("Failed to initialize MetaTrader 5")
        return
        
    try:
        # Создаем и тестируем индикатор
        indicator = TrendForceIndicator('AUDUSD')
        df = indicator.analyze_market()
        fig = indicator.create_visualization(df)
        fig.show()
        
    finally:
        mt5.shutdown()

if __name__ == "__main__":
    main()