import { DataInstance } from 'arc-ml';
import { BaseAgent } from './base-agent';

/**
 * Represents the training configuration options.
 */
interface TrainingConfig {
  // Simple Model Configs
  updatableCells?: string[]; // Which cells to update for this training session.
  multiplier?: number; // Intensity of the update (larger = bigger change).

  // Neural Network Configs
  epochs?: number; // Epochs (Number of iterations through the dataset).
  batchSize?: number; // Batch Size (Size of each mini-batch used in training).
  learningRate?: number; // Learning rate (Step size for each update).
  focus?: number[]; // The indices of the features to focus on.
  lambdas?: { [key: string]: number } // Memory retention hyperparameter for each action head.
  cleaning?: { 
    balance: { 
      oversampling: boolean, // Balance training data by duplicating instances
      multiStream: boolean // Balance gradient updates directly
    }, 
    removeSparsity: boolean // Whether or not to remove idle frames.
  }
}

export declare class ImitationLearningAgent extends BaseAgent {
  /**
   * Sends the collected data to the trainer platform.
   * @returns A promise that resolves to a boolean indicating success.
   */
  sendDataToPlatform(): Promise<boolean>;

  /**
   * Trains the model with the provided data and configuration.
   * @param trainingData - The training data.
   * @param config - The training configuration.
   * @returns A promise that resolves to a boolean indicating success.
   */
  train(trainingData: DataInstance[], config?: TrainingConfig): Promise<boolean>;

  /**
   * Discards the trained model and optionally resets the training data.
   * @param discardData - Whether to discard training data.
   */
  discardTraining(discardData?: boolean): void;
}