import pandas as pd
import numpy as np
import plotly.graph_objects as go
from plotly.subplots import make_subplots
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import classification_report

# Generate synthetic OHLCV data with more fluctuation
np.random.seed(42)
num_points = 2000 # Increased from 500 to 2000 for robustness
dates = pd.date_range(start='2023-01-01', periods=num_points, freq='H')

# Start with a base price and apply random daily changes
base_price = 100
price_changes = np.random.normal(0, 1, num_points).cumsum() # Random walk centered around 0
# Scale and shift to get prices in a reasonable range, ensuring fluctuation
open_prices = base_price + price_changes * 5 # Scale factor for magnitude of changes
open_prices = np.maximum(open_prices, 50) # Ensure prices don't go too low

close_prices = open_prices + np.random.uniform(-2, 2, num_points)
high_prices = np.maximum(open_prices, close_prices) + np.random.uniform(0, 1, num_points)
low_prices = np.minimum(open_prices, close_prices) - np.random.uniform(0, 1, num_points)
volume = np.random.randint(1000, 5000, num_points)

df = pd.DataFrame({
    'datetime': dates,
    'open': open_prices,
    'high': high_prices,
    'low': low_prices,
    'close': close_prices,
    'volume': volume
})

print("Sample DataFrame head:")
display(df.head())
print("Sample DataFrame info:")
display(df.info())

def meta_labeling(
    df: pd.DataFrame,
    ema_fast: int = 12,
    ema_slow: int = 26,
    horizon: int = 5,
    meta_threshold: float = 0.55,
) -> pd.DataFrame:
    """
    Applies meta labeling by training a secondary model to filter a primary EMA signal.

    Parameters
    ----------
    df : pd.DataFrame
        OHLCV DataFrame with 'datetime', 'open', 'high', 'low', 'close', and 'volume' columns.
    ema_fast : int
        Period for the fast Exponential Moving Average (EMA) in the primary signal.
    ema_slow : int
        Period for the slow Exponential Moving Average (EMA) in the primary signal.
    horizon : int
        Number of bars to hold a trade for profitability calculation of the meta label.
    meta_threshold : float
        Minimum meta probability required for a primary signal to be accepted.

    Returns
    -------
    pd.DataFrame
        Original DataFrame with added columns: 'ema_fast', 'ema_slow', 'primary_signal',
        'meta_label', 'return_1', 'return_5', 'rsi', 'sma_ratio', 'vol_ratio',
        'meta_proba', and 'signal'.
    """
    df = df.copy().sort_values("datetime", ignore_index=True)

    # Primary signal (EMA cross)
    df["ema_fast"] = df["close"].ewm(span=ema_fast, adjust=False).mean()
    df["ema_slow"] = df["close"].ewm(span=ema_slow, adjust=False).mean()
    df["primary_signal"] = np.where(df["ema_fast"] > df["ema_slow"], 1, -1)

    # Meta label: was the primary signal profitable?
    future_return   = df["close"].shift(-horizon) / df["close"] - 1
    # Profitable if primary direction matches forward return sign
    df["meta_label"] = (
        (df["primary_signal"] == 1)  & (future_return > 0) |
        (df["primary_signal"] == -1) & (future_return < 0)
    ).astype(int)

    # Meta features
    df["return_1"]  = df["close"].pct_change(1)
    df["return_5"]  = df["close"].pct_change(5)

    # Improved RSI calculation to handle zero division and NaNs more robustly
    delta = df["close"].diff()
    gain  = delta.clip(lower=0)
    loss  = -delta.clip(upper=0)

    avg_gain = gain.rolling(14).mean()
    avg_loss = loss.rolling(14).mean().abs() # Ensure avg_loss is positive

    # Calculate RS, handling division by zero with a small epsilon or assigning specific values
    rs = avg_gain / avg_loss.replace(0, 1e-9) # Replace 0 with a small number to avoid div by zero
    df["rsi"] = 100 - (100 / (1 + rs))

    # Special cases: if no average loss, RSI is 100; if no average gain, RSI is 0
    df.loc[avg_loss == 0, "rsi"] = 100
    df.loc[avg_gain == 0, "rsi"] = 0
    df["rsi"] = df["rsi"].fillna(50) # Fill any remaining NaNs (e.g., initial period) with a neutral 50

    df["sma_ratio"] = df["close"].rolling(10).mean() / df["close"].rolling(20).mean()
    df["vol_ratio"] = df["volume"] / df["volume"].rolling(20).mean()

    feats = ["return_1", "return_5", "rsi", "sma_ratio", "vol_ratio"]

    print(f"DEBUG_PRE_DROPNA: df shape: {df.shape}")
    print(f"DEBUG_PRE_DROPNA: NaNs in feats and meta_label columns before dropna:")
    for col in feats + ["meta_label"]:
        print(f"  {col}: {df[col].isna().sum()} NaNs")

    df_clean = df.dropna(subset=feats + ["meta_label"]).copy()

    X = df_clean[feats].values
    y = df_clean["meta_label"].values
    split = int(len(X) * 0.8)

    print(f"DEBUG: Length of df_clean after dropna: {len(df_clean)}")
    print(f"DEBUG: Length of X (features array): {len(X)}")
    print(f"DEBUG: Calculated split point (80% for train): {split}")
    print(f"DEBUG: Number of training samples (X[:split]): {len(X[:split]) if X is not None else 0}")
    print(f"DEBUG: Number of testing samples (X[split:]): {len(X[split:]) if X is not None else 0}")

    # Robustness check for insufficient data for both training and testing sets
    num_samples_X = len(X)
    num_train_samples = split
    num_test_samples = num_samples_X - split

    if num_samples_X == 0:
        print(f"\nWarning: After dropping NaNs, the dataset for features is empty. Cannot train meta-model.")
        df["meta_proba"] = np.nan
        df["signal"] = 0
        return df

    if num_train_samples == 0:
        print(f"\nWarning: The training set is empty after splitting ({num_train_samples} samples). Cannot train meta-model.")
        df["meta_proba"] = np.nan
        df["signal"] = 0
        return df

    if num_test_samples == 0:
        print(f"\nWarning: The test set is empty after splitting ({num_test_samples} samples). Cannot evaluate meta-model.")
        df["meta_proba"] = np.nan
        df["signal"] = 0
        return df

    # Check for sufficient classes in the training set
    unique_classes_train = np.unique(y[:split])
    if len(unique_classes_train) < 2:
        print(f"\nWarning: Training set for meta-model contains only 1 class (all {unique_classes_train[0]}s). Cannot train meta-model.")
        df["meta_proba"] = np.nan
        df["signal"] = 0
        return df

    print(f"DEBUG: Before StandardScaler fit_transform. X[:split].shape: {X[:split].shape}") # New debug line
    sc = StandardScaler()
    X_tr = sc.fit_transform(X[:split])
    X_te = sc.transform(X[split:])

    meta_clf = GradientBoostingClassifier(n_estimators=100, random_state=42)
    meta_clf.fit(X_tr, y[:split])

    meta_proba = meta_clf.predict_proba(X_te)[:, 1]

    # Align meta probabilities back to DataFrame
    test_idx = df_clean.index[split:]
    df["meta_proba"] = np.nan
    df.loc[test_idx, "meta_proba"] = meta_proba

    # Final signal: primary direction, gated by meta confidence
    df["signal"] = 0
    df.loc[(df["meta_proba"] >= meta_threshold) & (df["primary_signal"] == 1),  "signal"] =  1
    df.loc[(df["meta_proba"] >= meta_threshold) & (df["primary_signal"] == -1), "signal"] = -1

    print("\nClassification Report for Meta Model (Test Set):")
    print(classification_report(y[split:], (meta_proba >= meta_threshold).astype(int)))
    return df

df_signals = meta_labeling(df, meta_threshold=0.55)
print("\n--- Final Signal Distribution ---")
print(df_signals["signal"].value_counts())

buy_signals  = df_signals[df_signals["signal"] ==  1]
sell_signals = df_signals[df_signals["signal"] == -1]

fig = make_subplots(rows=2, cols=1, shared_xaxes=True,
    subplot_titles=["Price + Meta-Labeled Signals", "Meta Probability"],
    row_heights=[0.65, 0.35])

fig.add_trace(go.Candlestick(x=df_signals["datetime"],
    open=df_signals["open"], high=df_signals["high"],
    low=df_signals["low"], close=df_signals["close"], name="Price"), row=1, col=1)

fig.add_trace(go.Scatter(x=buy_signals["datetime"],  y=buy_signals["low"]  * 0.999,
    mode="markers", marker=dict(symbol="triangle-up",   size=10, color="green"), name="Buy"),  row=1, col=1)

fig.add_trace(go.Scatter(x=sell_signals["datetime"], y=sell_signals["high"] * 1.001,
    mode="markers", marker=dict(symbol="triangle-down", size=10, color="red"),   name="Sell"), row=1, col=1)

fig.add_trace(go.Scatter(x=df_signals["datetime"], y=df_signals["meta_proba"],
    mode="lines", name="Meta P(correct)", line=dict(color="purple")), row=2, col=1)

fig.add_hline(y=0.55, line_dash="dot", line_color="orange", row=2, col=1)

fig.update_layout(title_text="Meta Labeling Results",
    xaxis_rangeslider_visible=False, height=700, xaxis2_title="Datetime",
    legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1))

fig.show()