import pandas as pd
import numpy as np
import os
from sklearn.model_selection import train_test_split

def load_sequences_from_csv(csv_path, max_len=30, test_size=0.2, random_state=42):
    """
    Загружает последовательности из одного или нескольких CSV-файлов.
    Возвращает X_train, X_test, y_train, y_test, le (LabelEncoder).
    """
    if os.path.isdir(csv_path):
        dfs = []
        for f in os.listdir(csv_path):
            if f.endswith('.csv'):
                dfs.append(pd.read_csv(os.path.join(csv_path, f)))
        df = pd.concat(dfs, ignore_index=True)
    else:
        df = pd.read_csv(csv_path)

    sequences = {}
    for seq_id, group in df.groupby('sequence_id'):
        group = group.sort_values('frame')
        label = group['label'].iloc[0]
        features = group[[f'f{i}' for i in range(99)]].values
        # Обрезаем или падинг
        if len(features) > max_len:
            features = features[:max_len]
        elif len(features) < max_len:
            pad = np.zeros((max_len - len(features), 99))
            features = np.vstack([features, pad])
        sequences[seq_id] = (label, features)

    labels = []
    X = []
    for label, feats in sequences.values():
        labels.append(label)
        X.append(feats)

    X = np.array(X)
    from sklearn.preprocessing import LabelEncoder
    le = LabelEncoder()
    y = le.fit_transform(labels)

    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=test_size, stratify=y, random_state=random_state
    )
    return X_train, X_test, y_train, y_test, le