import os
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
from binance.client import Client
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report
import ta

# Connect to Binance
client = Client()

# Settings
DATA_FILE = "btc_data_4h_full.csv"
symbol = "BTCUSDT"
interval = Client.KLINE_INTERVAL_4HOUR

# Load or download data
if os.path.exists(DATA_FILE):
    print("Loading existing data...")
    df = pd.read_csv(DATA_FILE, index_col=0, parse_dates=True)
    last_timestamp = df.index[-1]
    start_time = last_timestamp + timedelta(hours=4)
    start_str = start_time.strftime("%d %B %Y %H:%M:%S")
    
    print(f"Downloading new data from {start_str}...")
    new_klines = client.get_historical_klines(symbol, interval, start_str)
    if new_klines:
        new_df = pd.DataFrame(new_klines, columns=['timestamp', 'open', 'high', 'low', 'close', 'volume',
                                                    'close_time', 'quote_av', 'trades', 'tb_base_av', 'tb_quote_av', 'ignore'])
        new_df = new_df[['timestamp', 'open', 'high', 'low', 'close', 'volume']]
        new_df[['open', 'high', 'low', 'close', 'volume']] = new_df[['open', 'high', 'low', 'close', 'volume']].astype(float)
        new_df['timestamp'] = pd.to_datetime(new_df['timestamp'], unit='ms')
        new_df = new_df.set_index('timestamp')
        df = pd.concat([df, new_df])
        df = df[~df.index.duplicated(keep='first')]
        df.to_csv(DATA_FILE)
else:
    print("Downloading all data from scratch...")
    klinesT = client.get_historical_klines(symbol, interval, "01 December 2021")
    df = pd.DataFrame(klinesT, columns=['timestamp', 'open', 'high', 'low', 'close', 'volume',
                                        'close_time', 'quote_av', 'trades', 'tb_base_av', 'tb_quote_av', 'ignore'])
    df = df[['timestamp', 'open', 'high', 'low', 'close', 'volume']]
    df[['open', 'high', 'low', 'close', 'volume']] = df[['open', 'high', 'low', 'close', 'volume']].astype(float)
    df['timestamp'] = pd.to_datetime(df['timestamp'], unit='ms')
    df = df.set_index('timestamp')
    df.to_csv(DATA_FILE)

# Feature Engineering - Maximum Indicators
# (RSI, MACD, EMAs, SMAs, BB, ATR, ADX, Stochastic, Williams %R, CCI, Momentum)

# RSI (Relative Strength Index)
df['rsi'] = ta.momentum.RSIIndicator(df['close'], window=14).rsi()

# MACD (Moving Average Convergence Divergence)
df['macd'] = ta.trend.MACD(df['close']).macd()

# EMA (Exponential Moving Averages)
df['ema_10'] = df['close'].ewm(span=10, adjust=False).mean()
df['ema_20'] = df['close'].ewm(span=20, adjust=False).mean()
df['ema_50'] = df['close'].ewm(span=50, adjust=False).mean()
df['ema_100'] = df['close'].ewm(span=100, adjust=False).mean()

# SMA (Simple Moving Averages)
df['sma_10'] = df['close'].rolling(window=10).mean()
df['sma_20'] = df['close'].rolling(window=20).mean()
df['sma_50'] = df['close'].rolling(window=50).mean()
df['sma_100'] = df['close'].rolling(window=100).mean()

# Bollinger Bands
bb_indicator = ta.volatility.BollingerBands(df['close'], window=20, window_dev=2)
df['bb_bbm'] = bb_indicator.bollinger_mavg()
df['bb_bbh'] = bb_indicator.bollinger_hband()
df['bb_bbl'] = bb_indicator.bollinger_lband()
df['bb_width'] = (df['bb_bbh'] - df['bb_bbl']) / df['bb_bbm']

# Average True Range (ATR)
df['atr'] = ta.volatility.AverageTrueRange(df['high'], df['low'], df['close'], window=14).average_true_range()

# ADX - Average Directional Index (Trend strength)
df['adx'] = ta.trend.ADXIndicator(df['high'], df['low'], df['close'], window=14).adx()

# Stochastic Oscillator
stoch = ta.momentum.StochasticOscillator(df['high'], df['low'], df['close'], window=14)
df['stoch_k'] = stoch.stoch()
df['stoch_d'] = stoch.stoch_signal()

# Williams %R
df['williams_r'] = ta.momentum.WilliamsRIndicator(df['high'], df['low'], df['close'], lbp=14).williams_r()

# CCI (Commodity Channel Index)
df['cci'] = ta.trend.CCIIndicator(df['high'], df['low'], df['close'], window=20).cci()

# Momentum (manual calculation)
df['momentum'] = df['close'] - df['close'].shift(10)  # Simple momentum calculation

# Ichimoku Cloud Indicators
ichimoku = ta.trend.IchimokuIndicator(df['high'], df['low'], window1=9, window2=26, window3=52)
df['ichimoku_tenkan_sen'] = ichimoku.ichimoku_conversion_line()
df['ichimoku_kijun_sen'] = ichimoku.ichimoku_base_line()
df['ichimoku_senkou_span_a'] = ichimoku.ichimoku_a()
df['ichimoku_senkou_span_b'] = ichimoku.ichimoku_b()
df['ichimoku_chikou_span'] = df['close'].shift(-26)

# Create Uptrend/Downtrend labels based on cloud (1 = Uptrend, 0 = Downtrend, -1 = Neutral)
def ichimoku_trend_label(row):
    if row['close'] > row['ichimoku_senkou_span_a'] and row['close'] > row['ichimoku_senkou_span_b']:
        return 1  # Uptrend
    elif row['close'] < row['ichimoku_senkou_span_a'] and row['close'] < row['ichimoku_senkou_span_b']:
        return 0  # Downtrend
    else:
        return -1  # Neutral

# Apply function to create 'cloud_trend' labels
df['cloud_trend'] = df.apply(ichimoku_trend_label, axis=1)

# Drop rows with NaN values
df = df.dropna()

# Features and Target
features = df.drop(columns=['open', 'high', 'low', 'close', 'volume', 'cloud_trend']).columns
X = df[features]
y = df['cloud_trend']  # Now predicting cloud trend: up, down, or neutral

# Train/Test Split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, shuffle=False)

# Train Random Forest (with class balancing)
model = RandomForestClassifier(n_estimators=200, class_weight="balanced", random_state=42)
model.fit(X_train, y_train)

# Evaluate
y_pred = model.predict(X_test)
print(classification_report(y_test, y_pred))

# Predict latest movement
latest_features = X.iloc[-1].values.reshape(1, -1)
predicted_trend = model.predict(latest_features)
trend_label = predicted_trend[0]

# Print trend prediction
if trend_label == 1:
    print("Predicted next trend: Uptrend")
elif trend_label == 0:
    print("Predicted next trend: Downtrend")
else:
    print("Predicted next trend: Neutral (inside the cloud)")