# === Вот эта связка нужна для всех исполняемых скриптов внутри библиотеки, которая будет использоваться как сабмодуль ===

import sys
from pathlib import Path

ROOT = Path(__file__).parent.parent
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

# === перед вмеми импортами ===

import numpy as np
import joblib
import json
import argparse
import tensorflow as tf
from tensorflow.keras import layers, models
from sklearn.metrics import classification_report, confusion_matrix, accuracy_score
from ml_gestures_dynamic.sequence_utils import load_sequences_from_csv

def train(data_path, model_path, max_len=30, lstm_units=64, epochs=50, batch_size=16, test_size=0.2):
    # Загрузка данных
    X_train, X_test, y_train, y_test, le = load_sequences_from_csv(data_path, max_len, test_size)
    num_classes = len(le.classes_)
    print(f"Classes: {le.classes_}")
    print(f"Train samples: {len(X_train)}, Test samples: {len(X_test)}")

    # Построение модели
    model = models.Sequential([
        layers.LSTM(lstm_units, input_shape=(max_len, 99), return_sequences=True),
        layers.Dropout(0.3),
        layers.LSTM(lstm_units),
        layers.Dropout(0.3),
        layers.Dense(num_classes, activation='softmax')
    ])
    model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
    model.fit(X_train, y_train, epochs=epochs, batch_size=batch_size, validation_data=(X_test, y_test))

    # Оценка
    y_pred = np.argmax(model.predict(X_test), axis=1)
    acc = accuracy_score(y_test, y_pred)
    report = classification_report(y_test, y_pred, target_names=le.classes_, output_dict=True)
    cm = confusion_matrix(y_test, y_pred).tolist()

    print(f"\nTest accuracy: {acc:.4f}")
    print("\nClassification Report:")
    for cls in le.classes_:
        print(f"{cls}: precision={report[cls]['precision']:.3f}, recall={report[cls]['recall']:.3f}, f1={report[cls]['f1-score']:.3f}")
    print("\nConfusion Matrix:")
    print(cm)

    # Сохранение модели и метаданных
    model.save(model_path)
    classes_path = model_path.replace('.h5', '_classes.pkl')
    joblib.dump(le.classes_, classes_path)

    report_data = {
        'model_type': 'LSTM',
        'max_len': max_len,
        'lstm_units': lstm_units,
        'epochs': epochs,
        'batch_size': batch_size,
        'test_size': test_size,
        'accuracy': acc,
        'classification_report': report,
        'confusion_matrix': cm,
        'classes': le.classes_.tolist(),
        'train_samples': len(X_train),
        'test_samples': len(X_test)
    }
    report_path = model_path.replace('.h5', '_report.json')
    with open(report_path, 'w') as f:
        json.dump(report_data, f, indent=2)

    print(f"\nModel saved to {model_path}")
    print(f"Classes saved to {classes_path}")
    print(f"Report saved to {report_path}")

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Обучение LSTM для динамических жестов')
    parser.add_argument('--data', required=True, help='Путь к CSV-файлу или папке с CSV-файлами')
    parser.add_argument('--model', required=True, help='Путь для сохранения модели (.h5)')
    parser.add_argument('--max_len', type=int, default=30, help='Длина последовательности')
    parser.add_argument('--lstm_units', type=int, default=64, help='Количество нейронов в LSTM')
    parser.add_argument('--epochs', type=int, default=50, help='Количество эпох')
    parser.add_argument('--batch_size', type=int, default=16, help='Размер батча')
    parser.add_argument('--test_size', type=float, default=0.2, help='Доля тестовой выборки')
    args = parser.parse_args()
    train(args.data, args.model, args.max_len, args.lstm_units, args.epochs, args.batch_size, args.test_size)