import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler

class TradingModel:
    def __init__(self):
        self.scaler = StandardScaler()
        self.model = LogisticRegression(random_state=42, max_iter=1000)

    def train(self, X_train, y_train):
        # Flatten 3D sequences into 2D for scikit-learn
        n_samples, window_size, n_features = X_train.shape
        X_flat = X_train.reshape(n_samples, window_size * n_features)
        
        # Scale features
        X_scaled = self.scaler.fit_transform(X_flat)
        
        # Train model
        self.model.fit(X_scaled, y_train)

    def predict(self, X):
        # Flatten and scale
        n_samples, window_size, n_features = X.shape
        X_flat = X.reshape(n_samples, window_size * n_features)
        X_scaled = self.scaler.transform(X_flat)
        
        # Return probability of positive class
        return self.model.predict_proba(X_scaled)[:, 1]

def train_and_get_prediction(X_train, y_train, X_latest):
    """
    Utility function to quickly train on historical data and predict next movement.
    """
    tm = TradingModel()
    tm.train(X_train, y_train)
    
    pred = tm.predict(X_latest)
    return pred[0]