import numpy as np
import tensorflow as tf
from tensorflow.keras import layers, models, callbacks

def build_lstm_model(input_shape):
    """
    Builds a robust LSTM classification network with regularizers 
    to handle noisy financial time-series data.
    """
    model = models.Sequential([
        # 1. First LSTM Layer: Returns sequences to feed into the next layer
        layers.Input(shape=input_shape),
        layers.LSTM(units=64, return_sequences=True, 
                    kernel_regularizer=tf.keras.regularizers.l2(0.001)),
        layers.Dropout(0.3),  # High dropout to reduce overfitting
        
        # 2. Second LSTM Layer: Returns only the last step's hidden state
        layers.LSTM(units=32, return_sequences=False,
                    kernel_regularizer=tf.keras.regularizers.l2(0.001)),
        layers.Dropout(0.3),
        
        # 3. Dense Hidden Layer for abstract feature synthesis
        layers.Dense(units=16, activation='relu'),
        layers.BatchNormalization(),  # Stabilizes training weights
        
        # 4. Output Layer: Single neuron outputting a probability (0.0 to 1.0)
        layers.Dense(units=1, activation='sigmoid')
    ])
    
    # Compile with Binary Crossentropy for the binary UP/DOWN classification task
    model.compile(
        optimizer=tf.keras.optimizers.Adam(learning_rate=0.0005),
        loss='binary_crossentropy',
        metrics=['accuracy', tf.keras.metrics.AUC(name='auc')]
    )
    
    return model

def train_trading_model(X, y):
    """Handles data splitting, model creation, and training with early stopping."""
    
    # 1. Chronological Train/Test Split (Crucial for Time-Series! Do NOT random split)
    # Using the last 20% of chronological data for evaluation to avoid data leakage
    split_idx = int(len(X) * 0.8)
    X_train, X_val = X[:split_idx], X[split_idx:]
    y_train, y_val = y[:split_idx], y[split_idx:]
    
    # 2. Extract dimensions for input shape: (Time Steps, Features)
    input_shape = (X.shape[1], X.shape[2])
    model = build_lstm_model(input_shape)
    model.summary()
    
    # 3. Set up Training Callbacks
    # Early Stopping prevents the network from memorizing noise when val loss stalls
    custom_callbacks = [
        callbacks.EarlyStopping(
            monitor='val_loss', 
            patience=5, 
            restore_best_weights=True
        ),
        callbacks.ReduceLROnPlateau(
            monitor='val_loss', 
            factor=0.5, 
            patience=3
        )
    ]
    
    # 4. Fit the Model
    print("\nStarting model training...")
    history = model.fit(
        X_train, y_train,
        validation_data=(X_val, y_val),
        epochs=30,
        batch_size=64,
        callbacks=custom_callbacks,
        verbose=1
    )
    
    return model, history

# --- PIPELINE INTEGRATION EXAMPLE ---
# trained_model, training_history = train_trading_model(X, y)
# trained_model.save("portfolio_lstm_bot.h5")