import pandas as pd
import numpy as np
import MetaTrader5 as mt5
from datetime import datetime, timedelta
from catboost import CatBoostClassifier, CatBoostRegressor
from sklearn.model_selection import train_test_split  # kept for reference; not used for the chronological split below
from sklearn.metrics import accuracy_score, classification_report, mean_squared_error
import math
import matplotlib
# Setting Agg backend for work without GUI
matplotlib.use('Agg')
import matplotlib.pyplot as plt

LOOKBACK_DEFAULT = 15
HORIZON = 24  # bars ahead the target looks; hardcoded throughout (angle labels, split embargo, etc.)


# MT5 initialization
def get_mt5_data(symbol='EURUSD', timeframe=mt5.TIMEFRAME_M5, days=60):
    if not mt5.initialize():
        print(f"MT5 initialization error: {mt5.last_error()}")
        return None

    start_date = datetime.now() - timedelta(days=days)
    rates = mt5.copy_rates_range(symbol, timeframe, start_date, datetime.now())

    # FIX: copy_rates_range can return None (bad symbol/timeframe, no history,
    # terminal not logged in, etc.) - guard before building a DataFrame from it.
    if rates is None or len(rates) == 0:
        print(f"No data returned for {symbol}. MT5 error: {mt5.last_error()}")
        mt5.shutdown()
        return None

    mt5.shutdown()

    df = pd.DataFrame(rates)
    df['time'] = pd.to_datetime(df['time'], unit='s')
    return df


# Calculate exact angle between two points (in degrees)
def calculate_angle(p1, p2):
    # p1 and p2 are tuples (normalized_time, price)
    x1, y1 = p1
    x2, y2 = p2

    # If X hasn't changed (vertical line), return 90 or -90 degrees
    if x2 - x1 == 0:
        return 90 if y2 > y1 else -90

    # Calculate angle in radians and convert to degrees
    angle_rad = math.atan2(y2 - y1, x2 - x1)
    angle_deg = math.degrees(angle_rad)

    return angle_deg


# Create angular price features
def create_angular_features(df):
    # Create a copy of DataFrame
    angular_df = df.copy()

    # Normalize time series for more accurate angle calculations
    # Convert time to numeric format for calculations
    angular_df['time_num'] = (angular_df['time'] - angular_df['time'].min()).dt.total_seconds()

    # Find ranges for normalization
    time_range = angular_df['time_num'].max() - angular_df['time_num'].min()
    price_range = angular_df['close'].max() - angular_df['close'].min()

    # FIX: guard against dividing by zero if every bar has the same timestamp
    # (would otherwise silently produce inf/NaN scale_factor and poison every angle).
    if time_range == 0:
        print("Warning: time_range is 0 (duplicate timestamps); using scale_factor=1.0 as a fallback.")
        scale_factor = 1.0
    else:
        scale_factor = price_range / time_range

    angular_df['time_scaled'] = angular_df['time_num'] * scale_factor

    # Calculate angles between consecutive points
    angles = []
    angles.append(np.nan)  # For the first point angle is undefined

    for i in range(1, len(angular_df)):
        current_point = (angular_df['time_scaled'].iloc[i], angular_df['close'].iloc[i])
        prev_point = (angular_df['time_scaled'].iloc[i - 1], angular_df['close'].iloc[i - 1])
        angle = calculate_angle(prev_point, current_point)
        angles.append(angle)

    angular_df['angle'] = angles

    # Add future price direction (24 bars ahead)
    future_directions = []
    for i in range(len(angular_df)):
        if i + HORIZON < len(angular_df):
            # 1 = up, 0 = down
            future_dir = 1 if angular_df['close'].iloc[i + HORIZON] > angular_df['close'].iloc[i] else 0
            future_directions.append(future_dir)
        else:
            future_directions.append(np.nan)

    angular_df['future_direction_24'] = future_directions

    # Calculate future price change magnitude (in percentage)
    future_changes = []
    for i in range(len(angular_df)):
        if i + HORIZON < len(angular_df):
            pct_change = (angular_df['close'].iloc[i + HORIZON] - angular_df['close'].iloc[i]) / angular_df['close'].iloc[i] * 100
            future_changes.append(pct_change)
        else:
            future_changes.append(np.nan)

    angular_df['future_change_pct_24'] = future_changes

    # Remove temporary columns
    angular_df = angular_df.drop(['time_num', 'time_scaled'], axis=1)

    return angular_df


def build_feature_dict(window, lookback):
    """
    Build the feature dict for one `lookback`-bar window.

    IMPORTANT: `window` must be exactly `lookback` rows with the LAST row
    being the anchor bar - i.e. the same bar that `future_direction_24` /
    `future_change_pct_24` are measured FROM. Every caller (training,
    backtesting, live prediction) must build its window this same way, or
    "current_price" and the other window-derived features stop matching
    what the label/prediction is actually anchored to. This one function
    is shared by all three call sites so that convention can't drift out
    of sync between them again.
    """
    feature_dict = {f'angle_{j}': window['angle'].iloc[j] for j in range(lookback)}

    feature_dict.update({
        'angle_mean': window['angle'].mean(),
        'angle_std': window['angle'].std(),
        'angle_min': window['angle'].min(),
        'angle_max': window['angle'].max(),
        'angle_last': window['angle'].iloc[-1],
        'angle_last_3_mean': window['angle'].iloc[-3:].mean(),
        'angle_last_5_mean': window['angle'].iloc[-5:].mean(),
        'angle_last_10_mean': window['angle'].iloc[-10:].mean(),
        'positive_angles_ratio': (window['angle'] > 0).mean(),
        'current_price': window['close'].iloc[-1],
        'price_std': window['close'].std(),
        'price_change_pct': (window['close'].iloc[-1] - window['close'].iloc[0]) / window['close'].iloc[0] * 100,
        'high_low_range': (window['high'].max() - window['low'].min()) / window['close'].iloc[-1] * 100,
        'last_tick_volume': window['tick_volume'].iloc[-1],
        'avg_tick_volume': window['tick_volume'].mean(),
        'tick_volume_ratio': window['tick_volume'].iloc[-1] / window['tick_volume'].mean() if window['tick_volume'].mean() > 0 else 1,
    })

    return feature_dict


# Feature preparation
def prepare_features(angular_df, lookback=LOOKBACK_DEFAULT):
    features = []
    targets_class = []  # For classification (direction)
    targets_reg = []    # For regression (percentage change)

    # Drop rows with NaN in angles or target variables
    filtered_df = angular_df.dropna(subset=['angle', 'future_direction_24', 'future_change_pct_24'])

    # Check if there's enough data after filtering
    if len(filtered_df) <= lookback:
        print("Not enough data after filtering NaN values")
        return None, None, None

    # FIX (off-by-one): the window must END on the same bar that the label
    # is anchored to. Originally the window was filtered_df.iloc[i-lookback:i]
    # (ending at bar i-1) while the label came from bar i - so "current_price"
    # and every other window feature described the bar BEFORE the one the
    # label actually measures 24 bars forward from. Shifting the window to
    # iloc[i-lookback+1:i+1] makes bar i (the label's anchor) the last row.
    for i in range(lookback - 1, len(filtered_df)):
        window = filtered_df.iloc[i - lookback + 1:i + 1]

        feature_dict = build_feature_dict(window, lookback)
        features.append(feature_dict)

        # Target for classification - direction after 24 bars
        targets_class.append(filtered_df.iloc[i]['future_direction_24'])

        # Target for regression - percentage change after 24 bars
        targets_reg.append(filtered_df.iloc[i]['future_change_pct_24'])

    return pd.DataFrame(features), np.array(targets_class), np.array(targets_reg)


def _predict_class(model, X):
    """
    CatBoostClassifier.predict() can come back as shape (n,) or (n,1)
    depending on catboost version / prediction settings. Always flatten
    to 1D so downstream indexing (`[0]`, sklearn metrics) behaves the
    same regardless of version.
    """
    return np.asarray(model.predict(X)).reshape(-1)


# Train hybrid model
def train_hybrid_model(X, y_class, y_reg, test_size=0.3, lookback=LOOKBACK_DEFAULT, horizon=HORIZON):
    # FIX (data leakage): the original code used train_test_split(..., shuffle=True).
    # Each row's window overlaps (lookback-1) bars with its neighbors, and the
    # labels themselves look `horizon` bars into the future - so random rows are
    # highly correlated with rows right next to them in time. Shuffling scatters
    # those correlated rows across train AND test, which inflates accuracy/RMSE
    # in a way that will not hold up on genuinely unseen future data.
    #
    # Fix: split chronologically (train = earlier bars, test = later bars), and
    # drop a small "embargo" gap of `lookback + horizon` samples at the seam,
    # since samples immediately on either side of a naive cut still share
    # window bars or future-looking targets with each other.
    n = len(X)
    gap = lookback + horizon
    n_test = max(int(round(n * test_size)), 1)
    test_start = n - n_test
    train_end = max(test_start - gap, 0)

    X_train = X.iloc[:train_end]
    X_test = X.iloc[test_start:]
    y_class_train, y_class_test = y_class[:train_end], y_class[test_start:]
    y_reg_train, y_reg_test = y_reg[:train_end], y_reg[test_start:]

    if len(X_train) == 0 or len(X_test) == 0:
        print("Not enough data to form a chronological train/test split with the embargo gap.")
        return None, None, None, None, None

    # Parameters for classification model
    params_class = {
        'iterations': 500,
        'learning_rate': 0.03,
        'depth': 6,
        'loss_function': 'Logloss',
        'random_seed': 42,
        'verbose': False
    }

    # Parameters for regression model
    params_reg = {
        'iterations': 500,
        'learning_rate': 0.03,
        'depth': 6,
        'loss_function': 'RMSE',
        'random_seed': 42,
        'verbose': False
    }

    # Train classification model (direction prediction)
    print("Training classification model...")
    model_class = CatBoostClassifier(**params_class)
    model_class.fit(X_train, y_class_train, eval_set=(X_test, y_class_test),
                     early_stopping_rounds=50, verbose=False)

    y_class_pred = _predict_class(model_class, X_test)
    accuracy = accuracy_score(y_class_test, y_class_pred)
    print(f"Classification accuracy on test set (chronological, out-of-sample): {accuracy:.4f}")
    print(f"Accuracy in percentage: {accuracy * 100:.2f}%")
    print(classification_report(y_class_test, y_class_pred))

    # Feature importance for classifier
    importance_class = model_class.get_feature_importance(prettified=True)
    print("Top 10 important features for classification:")
    print(importance_class.head(10))

    # Train regression model (percentage change prediction)
    print("\nTraining regression model...")
    model_reg = CatBoostRegressor(**params_reg)
    model_reg.fit(X_train, y_reg_train, eval_set=(X_test, y_reg_test),
                  early_stopping_rounds=50, verbose=False)

    y_reg_pred = model_reg.predict(X_test)
    rmse = np.sqrt(mean_squared_error(y_reg_test, y_reg_pred))
    print(f"Regression RMSE on test set (chronological, out-of-sample): {rmse:.4f}")

    # Feature importance for regressor
    importance_reg = model_reg.get_feature_importance(prettified=True)
    print("Top 10 important features for regression:")
    print(importance_reg.head(10))

    # Visualize regression results
    plt.figure(figsize=(750 / 100, 750 / 200))
    plt.scatter(y_reg_test, y_reg_pred, alpha=0.5)
    plt.plot([-10, 10], [-10, 10], 'r--')
    plt.xlabel('Actual change (%)')
    plt.ylabel('Predicted change (%)')
    plt.title('Comparison of actual and predicted changes after 24 bars')
    plt.grid(True)
    plt.savefig('change_prediction.png')
    print("Chart saved in change_prediction.png")

    # Visualize classification results as confusion matrix
    from sklearn.metrics import confusion_matrix
    import seaborn as sns

    cm = confusion_matrix(y_class_test, y_class_pred)
    plt.figure(figsize=(750 / 100, 750 / 125))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')
    plt.xlabel('Predicted values')
    plt.ylabel('Actual values')
    plt.title('Confusion Matrix')
    plt.savefig('confusion_matrix.png')
    print("Confusion matrix saved in confusion_matrix.png")

    return model_class, model_reg, X_test, y_class_test, y_reg_test


# Predict future movements
def predict_future_movement(model_class, model_reg, angular_df, lookback=LOOKBACK_DEFAULT, feature_names=None):
    # Check if there's enough data
    if len(angular_df) < lookback:
        return {"error": "Not enough data for prediction"}

    # Get last lookback bars - the last row here IS the anchor bar (the most
    # recent close), consistent with how windows are now built during training.
    window = angular_df.tail(lookback)

    feature_dict = build_feature_dict(window, lookback)
    X_pred = pd.DataFrame([feature_dict])

    # If feature names are passed, check if they exist
    if feature_names:
        for feature in feature_names:
            if feature not in X_pred.columns:
                X_pred[feature] = 0
        X_pred = X_pred[feature_names]

    # Direction prediction (classification)
    direction_prob = model_class.predict_proba(X_pred)[0]
    direction_prediction = _predict_class(model_class, X_pred)[0]

    # Percent change prediction (regression)
    change_prediction = model_reg.predict(X_pred)[0]

    # Prepare result
    result = {
        'direction_prediction': 'UP' if direction_prediction == 1 else 'DOWN',
        'direction_probability': direction_prob[int(direction_prediction)],
        'prob_up': direction_prob[1],
        'prob_down': direction_prob[0],
        'change_prediction': change_prediction,
        'current_price': window['close'].iloc[-1],
        'predicted_price': window['close'].iloc[-1] * (1 + change_prediction / 100),
        'prediction_horizon': '24 bars',
    }

    # Form trading signal
    result['signal'] = 'STRONG_BUY' if direction_prob[1] > 0.75 and change_prediction > 0.5 else \
                       'BUY' if direction_prob[1] > 0.6 else \
                       'STRONG_SELL' if direction_prob[0] > 0.75 and change_prediction < -0.5 else \
                       'SELL' if direction_prob[0] > 0.6 else \
                       'NEUTRAL'

    return result


# Visualize angles
def visualize_angles(angular_df, save_path='angle_visualization.png'):
    plt.figure(figsize=(750 / 100, 750 / 75))

    # Price chart
    plt.subplot(2, 1, 1)
    plt.plot(angular_df['time'], angular_df['close'])
    plt.title('Closing Price')
    plt.grid(True)

    # Angles chart
    plt.subplot(2, 1, 2)
    plt.plot(angular_df['time'], angular_df['angle'])
    plt.axhline(y=0, color='r', linestyle='--')
    plt.title('Price Movement Angles')
    plt.grid(True)

    plt.tight_layout()
    plt.savefig(save_path)

    return save_path


# Function for backtesting signals
def backtest_signals(angular_df, model_class, model_reg, feature_names, lookback=LOOKBACK_DEFAULT):
    print("\nPerforming signal backtesting...")

    # Filter out NaN values
    valid_df = angular_df.dropna(subset=['angle', 'future_direction_24', 'future_change_pct_24'])

    # Check if there's enough data
    if len(valid_df) <= lookback + 50:  # Need at least 50 points for backtesting
        print("Not enough data for backtesting")
        return

    signals = []
    actual_changes = []
    timestamps = []

    # Generate predictions for each point in history.
    # FIX (same off-by-one as prepare_features): window now ends on bar i,
    # the same bar future_change_pct_24 is measured from, instead of i-1.
    for i in range(lookback - 1, len(valid_df) - HORIZON):
        window = valid_df.iloc[i - lookback + 1:i + 1]

        feature_dict = build_feature_dict(window, lookback)
        X_pred = pd.DataFrame([feature_dict])

        # Check feature names match
        if feature_names:
            for feature in feature_names:
                if feature not in X_pred.columns:
                    X_pred[feature] = 0
            X_pred = X_pred[feature_names]

        # Make prediction
        signal = 1 if _predict_class(model_class, X_pred)[0] == 1 else -1  # 1 = BUY, -1 = SELL
        signals.append(signal)

        # Save actual change
        actual_changes.append(valid_df.iloc[i]['future_change_pct_24'])
        timestamps.append(valid_df.iloc[i]['time'])

    # Results analysis
    signals = np.array(signals)
    actual_changes = np.array(actual_changes)

    # Profit/loss by signals (in percentage)
    pnl = signals * actual_changes

    # Statistics
    win_rate = np.sum(pnl > 0) / len(pnl)
    avg_win = np.mean(pnl[pnl > 0]) if np.any(pnl > 0) else 0
    avg_loss = np.mean(pnl[pnl < 0]) if np.any(pnl < 0) else 0
    profit_factor = abs(np.sum(pnl[pnl > 0]) / np.sum(pnl[pnl < 0])) if np.sum(pnl[pnl < 0]) != 0 else float('inf')

    print(f"Total signals: {len(signals)}")
    print(f"Win rate: {win_rate:.2%}")
    print(f"Average win: {avg_win:.2f}%")
    print(f"Average loss: {avg_loss:.2f}%")
    print(f"Profit Factor: {profit_factor:.2f}")
    print(f"Total return: {np.sum(pnl):.2f}%")

    # NOTE (still worth knowing): this backtest reuses the SAME model_class that
    # was trained on data covering this same historical period (including, for
    # the training-window portion, these exact bars). It is a sanity check that
    # the model fits its own training era, not a true out-of-sample backtest.
    # For an honest backtest, only run this loop over the chronological TEST
    # portion (or on fresh data the model never trained on).

    # Plot results
    cumulative_pnl = np.cumsum(pnl)

    plt.figure(figsize=(750 / 100, 750 / 150))
    plt.plot(cumulative_pnl)
    plt.title('Model Equity Curve')
    plt.xlabel('Number of trades')
    plt.ylabel('Cumulative return (%)')
    plt.grid(True)
    plt.savefig('backtest_results.png')
    print("Backtesting results chart saved in backtest_results.png")


# Main function
def main():
    # Get EURUSD data
    print("Loading EURUSD data from MetaTrader5...")
    df = get_mt5_data(symbol='EURUSD', days=90)  # Take more data for 24-bar horizon

    if df is None or len(df) == 0:
        print("Failed to get data")
        return

    print(f"Loaded {len(df)} bars")

    # Create angular features
    print("Creating angular features...")
    angular_df = create_angular_features(df)
    print(f"Created {len(angular_df)} records with angular features")

    # Visualize angles
    print("Visualizing angles...")
    viz_path = visualize_angles(angular_df)
    print(f"Visualization saved in {viz_path}")

    # Prepare features
    print("Preparing features...")
    X, y_class, y_reg = prepare_features(angular_df, lookback=LOOKBACK_DEFAULT)

    if X is None:
        print("Not enough data to train the model")
        return

    print(f"Prepared {len(X)} samples")

    # Train hybrid model
    print("Training hybrid model...")
    model_class, model_reg, X_test, y_class_test, y_reg_test = train_hybrid_model(X, y_class, y_reg)

    # FIX: train_hybrid_model can now return all-None if there isn't enough
    # data for a chronological split with the embargo gap - handle that
    # instead of crashing on the next line.
    if model_class is None:
        print("Training aborted: not enough data for a clean chronological split.")
        return

    # Predict future movement
    feature_names = X.columns.tolist()
    prediction = predict_future_movement(model_class, model_reg, angular_df, feature_names=feature_names)

    print("\nPREDICTION FOR MOVEMENT AFTER 24 BARS:")
    for k, v in prediction.items():
        print(f"{k}: {v}")

    # Information about last bars
    print("\nLast 5 bars with angular features:")
    print(angular_df.tail(5)[['time', 'open', 'close', 'angle', 'tick_volume']])

    # Analyze signals on historical data
    backtest_signals(angular_df, model_class, model_reg, feature_names, lookback=LOOKBACK_DEFAULT)


if __name__ == "__main__":
    main()