import { AxiosInstance, AxiosRequestConfig, AxiosResponse } from 'axios';
import { TabularModel, NeuralNetworkMultihead, DataCollector, RawDataInstance, DataInstance } from 'nrn-ml';

/**  Represents a one-dimensional numeric array. */
type Vector = number[];
/**  Represents a two-dimensional numeric matrix. */
type Matrix = number[][];
/** Represents the position of an entity in an environment. */
interface Position {
    x: number;
    y: number;
}
/** Represents an entity in an environment which has position and size. */
interface Entity extends Position {
    width: number;
    height: number;
}
/** Represents the probability output from a model (action head -> probability matrix). */
type Probabilities = Record<string, Matrix>;

/** Represents the config for the cosine similarity type. */
type CosineSimilarityConfig = {
    type: "cosineSimilarity";
    keys: {
        vector1: string;
        vector2: string;
    };
};
/** Represents the config for the raycast feature type. */
type RaycastConfig = {
    type: "raycast";
    keys: {
        origin: string;
        colliders: string;
        maxDistance: string;
    };
    setup?: {
        numRays?: number;
    };
};
/** Represents the config for the angle feature type. */
type AngleConfig = {
    type: "angle";
    keys: {
        entity1: string;
        entity2: string;
    };
};
/** Represents the config for the relativePosition feature type. */
type RelativePositionConfig = {
    type: "relativePosition";
    keys: {
        entity1: string;
        entity2: string;
        maxDistance: string;
    };
};
/** Represents the config for the relativePositionToCluster feature type. */
type RelativePositionToClusterConfig = {
    type: "relativePositionToCluster";
    keys: {
        origin: string;
        clusterEntities: string;
        maxDistance: string;
    };
};
/** Represents the config for the onehot feature type. */
type OneHotConfig = {
    type: "onehot";
    keys: {
        value: string;
    };
    setup: {
        options: string[];
    };
};
/** Represents the config for the binary feature type. */
type BinaryConfig = {
    type: "binary";
    keys: {
        value: number | string;
    };
    setup: {
        operator: "=" | ">" | "<" | "!=";
        comparison: number | string;
    };
};
/** Represents the config for the rescale feature type. */
type RescaleConfig = {
    type: "rescale";
    keys: {
        value: string;
        scaleFactor: string;
    };
};
/** Represents the config for the normalize feature type. */
type NormalizeConfig = {
    type: "normalize";
    keys: {
        value: string;
    };
    setup: {
        mean: number;
        stdev: number;
    };
};
/** Represents a configuration for a specific feature in the state space. */
type FeatureConfigType = CosineSimilarityConfig | RaycastConfig | AngleConfig | RelativePositionConfig | RelativePositionToClusterConfig | OneHotConfig | BinaryConfig | RescaleConfig | NormalizeConfig;
/** Represents the world object used to extract features. */
type GameWorldType = Record<string, any>;
/** Feature functions offered through this module. */
type ValidFeatures = "cosineSimilarity" | "raycast" | "angle" | "relativePosition" | "relativePositionToCluster" | "onehot" | "binary" | "rescale" | "normalize";
/** Required data keys and setup for each feature type. */
type RequiredDataType = Record<ValidFeatures, {
    keys: Record<string, boolean>;
    setup?: Record<string, boolean>;
}>;
/** Cosine similarity input parameters. */
type CosineSimilarityParams = {
    vector1: Vector;
    vector2: Vector;
};
/** Raycast input parameters. */
type RaycastParams = {
    origin: any;
    colliders: any[];
    maxDistance: number;
    numRays?: number;
};
/** Relative position input parameters. */
type AngleParams = {
    entity1: {
        x: number;
        y: number;
    };
    entity2: {
        x: number;
        y: number;
    };
};
/** Relative position input parameters. */
type RelativePositionParams = {
    entity1: {
        x: number;
        y: number;
    };
    entity2: {
        x: number;
        y: number;
    };
    maxDistance: number;
};
/** Relative position to cluster input parameters. */
type RelativePositionToClusterParams = {
    origin: {
        x: number;
        y: number;
    };
    clusterEntities: Entity[];
    maxDistance: number;
};
/** Onehot encoding input parameters. */
type OneHotParams = {
    value: string;
    options: string[];
};
/** Binary input parameters. */
type BinaryParams = {
    value: string | number;
    operator: "=" | ">" | "<" | "!=";
    comparison: string | number;
};
/** Rescaling input parameters. */
type RescaledParams = {
    value: number;
    scaleFactor: number;
};
/** Normalizing input parameters. */
type NormalizedParams = {
    value: number;
    mean: number;
    stdev: number;
};

/**
 * The FeatureEngineering class provides methods for configuring and processing feature extraction
 * for a state space.
 */
declare class FeatureEngineering {
    static numFeatures?: number;
    static stateConfig: FeatureConfigType[];
    /** Mapping of feature types to their respective processing functions. */
    static conversionFunctions: {
        cosineSimilarity: typeof FeatureEngineering.getCosineSimilarity;
        raycast: typeof FeatureEngineering.getRaycasts;
        angle: typeof FeatureEngineering.getAngle;
        relativePosition: typeof FeatureEngineering.getRelativePosition;
        relativePositionToCluster: typeof FeatureEngineering.getRelativePositionToCluster;
        onehot: typeof FeatureEngineering.getOneHotEncoding;
        binary: typeof FeatureEngineering.getBinary;
        rescale: typeof FeatureEngineering.getRescaledValue;
        normalize: typeof FeatureEngineering.getNormalizedValue;
    };
    /** Number of features returned from feature engineering functions (-1 for dynamic size). */
    static featureSizes: Record<ValidFeatures, number>;
    /** Required data keys and setup for each feature type. */
    static requiredData: RequiredDataType;
    /**
     * Sets the state configuration for feature extraction.
     * @param config - The array of feature configurations.
     */
    static setStateConfig(config?: FeatureConfigType[]): void;
    /**
     * Validate that the state configuration is correct.
     * @param config - The array of feature configurations.
     * @returns The number of features in the state config
     */
    static _validateStateConfig(config: FeatureConfigType[]): number;
    /**
     * Validate that the key exists in the game world.
     * @param value - Value extracted from the world.
     * @param key The key used to extract a value.
     */
    static _validateKeyInWorld(value: any, key: string): void;
    /**
     * Extracts a value from the game world.
     * @param world - The world object containing the data for feature extraction.
     * @param key - Key to extract feature object.
     * @returns Object that will be used in feature engineering
     */
    static _parseWorldWithKey(world: GameWorldType, key: string): any;
    /**
     * Extracts the state features from the world object based on the current state configuration.
     * @param world - The world object containing the data for feature extraction.
     * @returns An array of feature values.
     */
    static getState(world: GameWorldType): Vector;
    static _dotProduct(A: Vector, B: Vector): number;
    static _normL2(vector: Vector): number;
    /**
     * Processes the cosine similarity of two vectors.
     * @param params - Parameters required for cosine similarity.
     * @returns The cosine similarity.
     */
    static getCosineSimilarity({ vector1, vector2 }: CosineSimilarityParams): Vector;
    /**
     * Processes raycast features based on the provided parameters.
     * @param params - Parameters required for raycasting.
     * @returns An array of raycast results.
     */
    static getRaycasts({ origin, colliders, maxDistance, numRays }: RaycastParams): Vector;
    /**
     * Processes the angle between an origin and another entity.
     * @param params - Parameters for angle calculation.
     * @returns An array containing the sine and cosine of the angle between the entities.
     */
    static getAngle({ entity1, entity2 }: AngleParams): Vector;
    /**
     * Processes the relative position between an origin and another entity.
     * @param params - Parameters for relative position calculation.
     * @returns An array containing the distance and directional components.
     */
    static getRelativePosition({ entity1, entity2, maxDistance }: RelativePositionParams): Vector;
    /**
     * Processes relative position features between an origin and a cluster of objects.
     * @param params - Parameters for relative position calculation.
     * @returns An array containing the distance and directional components.
     */
    static getRelativePositionToCluster({ origin, clusterEntities, maxDistance }: RelativePositionToClusterParams): Vector;
    /**
     * Processes one-hot encoding for a given value and set of options.
     * @param params - Parameters for one-hot encoding.
     * @returns An array representing the one-hot encoded value.
     */
    static getOneHotEncoding({ value, options }: OneHotParams): Vector;
    /**
     * Processes binary features based on a comparison operation.
     * @param params - Parameters for the binary operation.
     * @returns An array with the result of the comparison (1 or 0).
     */
    static getBinary({ value, operator, comparison }: BinaryParams): Vector;
    /**
     * Processes rescaled values.
     * @param params - Parameters for rescaling.
     * @returns An array with the rescaled value.
     */
    static getRescaledValue({ value, scaleFactor }: RescaledParams): Vector;
    /**
     * Processes normalized values.
     * @param params - Parameters for normalization.
     * @returns An array with the normalized value.
     */
    static getNormalizedValue({ value, mean, stdev }: NormalizedParams): Vector;
}

/**
 * Represents the progress of a chunked request.
 */
type ChunkProgress = {
    chunk: number;
    totalChunks: number;
    progress: number;
    partialScore?: number;
};
/**
 * Options for the sendChunkedRequest method.
 */
type ChunkedRequestOptions = {
    onProgress?: (progress: ChunkProgress) => void;
    retryAttempts?: number;
    retryDelay?: number;
};
/**
 * The APIClient class provides methods to interact with the backend API.
 */
interface APIClientType {
    /** API key for authentication. */
    apiKey: string;
    /** Session data for the current user. */
    session: Record<string, any>;
    /** Indicates whether the current session is valid. */
    isValidSession: boolean;
    /** Game ID associated with the client. */
    gameId: string;
    /** Game metadata retrieved from the backend. */
    game: Record<string, any>;
    /** Indicates whether the current game is valid. */
    isValidGame: boolean;
    /** Threshold for chunking large requests. */
    chunkingThreshold: number;
    /** Backend base URL for API requests. */
    backend: string;
    /** Axios client instance. */
    client: AxiosInstance;
    /**
     * Sets the API key for authentication and validates the session.
     * @param apiKey - The API key to set.
     * @returns A promise that resolves to whether the session is valid.
     */
    setApiKey(apiKey: string): Promise<boolean>;
    /**
     * Sets the game ID and validates the game.
     * @param gameId - The game ID to set.
     * @returns A promise that resolves to whether the game is valid.
     */
    setGameId(gameId: string): Promise<boolean>;
    /**
     * Overrides the backend URL and validates the session.
     * @param newBackend - The new backend URL to set.
     * @returns A promise that resolves to whether the session is valid.
     */
    overrideBackendUrl(newBackend: string): Promise<boolean>;
    /**
     * Sets whether or not to use cookie auth and validates the session.
     * @param useCookieAuth - Whether or not to use cookie auth.
     * @returns A promise that resolves to whether the session is valid.
     */
    setUseCookieAuth(useCookieAuth: boolean): Promise<boolean>;
    /**
     * Sets the threshold for chunking large requests.
     * @param threshold - The chunking threshold to set.
     */
    setChunkingThreshold(threshold: number): void;
    /**
     * Validates the current game by fetching its metadata.
     * @returns A promise that resolves to whether the game is valid.
     */
    validateGame(): Promise<boolean>;
    /**
     * Validates the current session by checking the API key.
     * @returns A promise that resolves to whether the session is valid.
     */
    validateSession(): Promise<boolean>;
    /**
     * Sends a GET request to the backend.
     * @param path - The API endpoint path.
     * @param config - Optional Axios request configuration.
     * @returns A promise that resolves to the Axios response.
     */
    get(path: string, config?: AxiosRequestConfig): Promise<AxiosResponse>;
    /**
     * Sends a POST request to the backend.
     * @param path - The API endpoint path.
     * @param data - The request payload.
     * @param config - Optional Axios request configuration.
     * @returns A promise that resolves to the Axios response.
     */
    post(path: string, data?: Record<string, any>, config?: AxiosRequestConfig): Promise<AxiosResponse>;
    /**
     * Sends a PUT request to the backend.
     * @param path - The API endpoint path.
     * @param data - The request payload.
     * @param config - Optional Axios request configuration.
     * @returns A promise that resolves to the Axios response.
     */
    put(path: string, data?: Record<string, any>, config?: AxiosRequestConfig): Promise<AxiosResponse>;
    /**
     * Sends a DELETE request to the backend.
     * @param path - The API endpoint path.
     * @param config - Optional Axios request configuration.
     * @returns A promise that resolves to the Axios response.
     */
    delete(path: string, config?: AxiosRequestConfig): Promise<AxiosResponse>;
    /**
     * Sends a PATCH request to the backend.
     * @param path - The API endpoint path.
     * @param data - The request payload.
     * @param config - Optional Axios request configuration.
     * @returns A promise that resolves to the Axios response.
     */
    patch(path: string, data?: Record<string, any>, config?: AxiosRequestConfig): Promise<AxiosResponse>;
    /**
     * Sends a chunked request to the backend.
     * @param endpoint - The API endpoint path.
     * @param body - The core request data to send across with each chunk.
     * @param dataToChunk - The data to send, split into chunks.
     * @param options - Options for tracking progress and handling retries.
     * @returns A promise that resolves when all chunks are sent successfully.
     */
    chunkedUpload(endpoint: string, body: Record<string, any>, dataToChunk?: any[], options?: ChunkedRequestOptions): Promise<any>;
}

declare class APIClient implements APIClientType {
    static instance: any;
    apiKey: string;
    session: any;
    isValidSession: boolean;
    useCookieAuth: boolean;
    gameId: string;
    game: any;
    isValidGame: boolean;
    chunkingThreshold: number;
    backend: string;
    client: AxiosInstance;
    constructor();
    static getInstance(): APIClientType;
    _createClient(backend: string): void;
    _setupInterceptors(): void;
    _handleError(error: any, throwError?: boolean): void;
    setApiKey(apiKey: string): Promise<boolean>;
    setGameId(gameId: string): Promise<boolean>;
    overrideBackendUrl(newBackend: string): Promise<boolean>;
    setUseCookieAuth(useCookieAuth: boolean): Promise<boolean>;
    setChunkingThreshold(threshold: number): void;
    validateGame(): Promise<boolean>;
    validateSession(): Promise<boolean>;
    get(path: string, config?: AxiosRequestConfig): Promise<AxiosResponse>;
    post(path: string, data?: Record<string, any>, config?: AxiosRequestConfig): Promise<AxiosResponse>;
    put(path: string, data?: Record<string, any>, config?: AxiosRequestConfig): Promise<AxiosResponse>;
    delete(path: string, config?: AxiosRequestConfig): Promise<AxiosResponse>;
    patch(path: string, data?: Record<string, any>, config?: AxiosRequestConfig): Promise<AxiosResponse>;
    chunkedUpload(endpoint: string, body: Record<string, any>, dataToChunk: any[], options?: ChunkedRequestOptions): Promise<any>;
}

/**
 * Represents a valid state representation.
 * neural-network: Matrix
 * simple: number
 */
type ValidState = Matrix | number;
/** Represents a raw state representation before formatting (Vector converts to Matrix) */
type RawState = ValidState | Vector;
/** Represents a class that is validates state inputs to the model/agent. */
interface StateValidatorType {
    /** The model types for each of the models in the agent wrapper. */
    modelTypes: ModelType[];
    /** The expected dimensionalities of the models' input. */
    expectedInputDims: Vector;
    /** Add the input dimensionality for validation. */
    addInputDimValidation(config: any): void;
    /** Check if the state is valid and reformat if necessary. */
    validateState(inputs: RawState, idx?: number): ValidState;
    /** Get an empty state (used when skipping agents for inference). */
    getEmptyState(idx?: number): Vector | number;
}

/** Represents an agent 'input' to the environment (i.e. the action). */
type InputType = Record<string, boolean>;
/** Represents the cooldown counter for actions that an agent takes. */
type InputCooldownType = Record<string, number>;
/** The actions which are currently locked for a total of 'cooldown' frames. */
type LockedActionType = {
    name: string;
    head: string;
    cooldown: number;
};
/** The core agent wrapper which wrappers are built on top of. */
interface AgentWrapperCoreType {
    numAgents: number;
    model: Model;
    frameDelay: number;
    forcedHold: number;
    pressActions: string[];
    holdActions: string[];
    forceHoldActions: string[];
    continuousActions: Record<string, boolean>;
    actionToHead: Record<string, string>;
    inputs: InputType[];
    inputCooldown: InputCooldownType[];
    previousHoldActions: InputType[];
    lockedAction: LockedActionType[];
    /** Reset all of the active button presses. */
    reset(): void;
    /**
     * Applies cooldown to prevent very quick consecutive actions
     * @param inputs - Input for the current agent.
     * @param action - Name of the action.
     * @param agentIdx - The agent index.
     */
    applyPressCooldown(inputs: InputType, action: string, agentIdx?: number): void;
    /**
     * Applies cooldown and locks certain actions in place
     * @param inputs - Input for the current agent.
     * @param action - Name of the action.
     * @param agentIdx - The agent index.
     */
    applyHoldCooldown(inputs: InputType, action: string, agentIdx?: number): void;
    /**
     * Checks whether or not an action is locked
     * @param action - Name of the action.
     * @param agentIdx - The agent index.
     * @returns boolean for whether or not an action is locked
     */
    checkLockedAction(action: string, agentIdx?: number): boolean;
    /**
     * Copy over previous actions for the locked head
     * @param inputs - Input for the current agent.
     * @param agentIdx - The agent index.
     */
    copyPrevHoldActions(inputs: InputType, agentIdx?: number): void;
    /**
     * Apply frame delay to all press and hold actions
     * @param inputs - Input for the current agent.
     * @param agentIdx - The agent index.
     */
    applyFrameDelay(inputs: InputType, agentIdx?: number): void;
    /**
     * Keep track of the hold actions from the previous timestep
     * @param inputs - Input for the current agent.
     * @param agentIdx - The agent index.
     */
    trackPreviousHoldInputs(inputs: InputType, agentIdx?: number): void;
    /**
     * Check whether a given agent is currently locked in its action
     * @param agentIdx - The agent index.
     * @returns boolean of whether it is locked or not
     */
    isActionLocked(agentIdx?: number): boolean;
}

/**  Represents the models types offered from NRN ML. */
type ModelType = "neural-network" | "simple";
/**  Represents the type of actions for a model head. */
type ActionType = "discrete" | "continuous";
/**  Represents the activation function for a model head. */
type ActionActivationType = "linear" | "softmax" | "tanh" | "sigmoid";
/** Represents the valid policy options. */
type Policy = "argmaxPolicy" | "probabilisticSampling";
/** Represents the configuration a model's action head */
type ActionHeadMetadata = {
    policyMapping: Policy;
    order: string[];
    actionType?: ActionType;
    activationName?: ActionActivationType;
};
/** Represents the raw configuration (when creating a new model) base for all model types */
interface RawModelConfigBase {
    modelType: ModelType;
    inputDim: number;
    actionOrder: string[] | string[][];
    actionNames?: string | string[];
    actionTypes?: ActionType | ActionType[];
    actionActivations?: ActionActivationType | ActionActivationType[];
    actionPolicies?: Policy | Policy[];
    multiheadBool?: boolean;
}
/** Represents the formatted configuration (when loading in a model) base for all model types */
interface FormattedModelConfigBase {
    modelType?: ModelType;
    inputDim?: number;
    actionHeads: string[];
    actionMetadata: Record<string, ActionHeadMetadata>;
    multiheadBool?: boolean;
}
/** The method used to initialize a new cell. */
type TabularInitializationMethod = "empty" | "random";
/** Represents the formatted configuration for tabular models */
interface FormattedModelConfigTabular extends FormattedModelConfigBase {
    initializationMethod: TabularInitializationMethod;
    numDiscreteStates: number;
}
/** Represents the raw configuration for neural-network models */
interface RawModelConfigNeuralNetwork extends RawModelConfigBase {
    neurons?: number[];
    activationFunctionName?: "elu" | "relu";
}
/** Represents the configuration a model's output */
type OutputConfig = {
    activation: ActionActivationType;
    outputType: "mean" | "quantileRegression";
    quantiles: number;
};
/** Represents the formatted configuration for neural-network models */
interface FormattedModelConfigNeuralNetwork extends FormattedModelConfigBase {
    nFeatures: number;
    neurons: number[];
    activationFunctionName?: "elu" | "relu";
    movingAverageType?: "Simple" | "Exponential";
    decimalPrecision?: number;
    outputConfig?: OutputConfig;
}
/** Represents the model data for the agent */
interface ModelDataType {
    config?: RawModelConfigBase | RawModelConfigNeuralNetwork | FormattedModelConfigTabular | FormattedModelConfigNeuralNetwork;
    parameters?: Record<string, Matrix>;
    frequencies?: any[];
    trainable?: boolean;
}
/** Represents a model's action metadata configuration for all heads. */
type ActionMetadata = Record<string, ActionHeadMetadata>;
/** Represents a valid model from the nrn-ml package. */
type Model = TabularModel | NeuralNetworkMultihead;

/** Represents the configuration options for an agent. */
interface AgentConfig {
    useAgentWrapper?: boolean;
    frameDelay?: number;
    forcedHold?: number;
    numSamples?: {
        [key: string]: number;
    };
    forceHoldActions?: string[];
    holdActions?: string[];
    policySimilarityThreshold?: number;
}

/** Selection inputs for probabilistic sampling */
type ProbabilisticSelectionFunctionInputs = {
    probabilities: Probabilities;
    agentIdx?: number;
    row?: number;
};
/** Probabilistic sampling wrapper that builds upon the core agent wrapper. */
interface ProbabilisticAgentWrapperType extends AgentWrapperCoreType {
    numSamples: Record<string, number>;
    framesRemaining: Record<string, number[]>;
    previousPolicy: Probabilities;
    currentAction: Record<string, string[][]>;
    actionSubkeys: Record<string, string[]>;
    policySimilarityThreshold: number;
    /**
     * Monte carlo sampling from the softmax probability output
     * @param probabilities - The probability output for all model heads.
     * @param actionKey - The action head to sample from.
     * @param row - The row to sample from in the probability matrix
     * @returns array of sampled inputs
     */
    monteCarloSampling(probabilities: Probabilities, actionKey: string, row?: number): InputType[];
    /**
     * Convert the selected input to a string
     * @param input - The sampled input (action).
     * @param actionKey - The action head being evaluated.
     * @returns The conversion of the input as a string
     */
    convertInputToString(input: InputType, actionKey: string): string | undefined;
    /**
     * Assign the actions that were sampled to the agent's current input
     * @param actionKey - The action head being evaluated.
     * @param agentIdx - The agent index
     * @returns The assigned inputs
     */
    assignSampledInput(actionKey: string, agentIdx?: number): InputType;
    /**
     * Sample the actions from the probability output of the model
     * @param probabilities - The probability output for all model heads.
     * @param actionKey - The action head to sample from.
     * @param agentIdx - The agent index
     * @param row - The row to sample from in the probability matrix
     * @returns The input (action) for the given agent
     */
    sampleAction(probabilities: Probabilities, actionKey: string, agentIdx?: number, row?: number): InputType;
    /**
     * Force all actions to be false for an action head
     * @param actionGroup - The action head's name.
     * @returns The input (action) of false for all actions in a given action head
     */
    forceNoAction(actionGroup: string): InputType;
    /**
     * Probabilistic sampling function
     * @param data - All of the relevant inputs to select an action for the current frame.
     * @returns The sampled inputs
     */
    selectionFunction(data: ProbabilisticSelectionFunctionInputs): InputType;
}

/** Single agent implementation of the probabilistic agent. */
interface SingleAgentWrapperType extends ProbabilisticAgentWrapperType {
    /**
     * Selecting an action for the agent
     * @param input - The inputs to the model for action selection.
     * @returns The selected input (action)
     */
    selectAction(input: any): InputType;
}

/** Represents the various raw action formats for data collection. */
type RawActionFormats = Vector | Matrix | Probabilities | InputType;
interface GeneralAgentCoreType extends StateValidatorType {
    /** The number of agents. Each element in the array refers to the quantity of agents for a particular model */
    numAgents: Vector;
    /** Whether the agent has been initialized. */
    initializedBool: boolean;
    /** The agent's configuration options. */
    agentConfigs: AgentConfig[];
    /** Data collector for managing training data. */
    dataCollectors: DataCollector[][];
    /** The raw model configs when first initialized. */
    initialModelConfigs: any[];
    /** The raw model configs when first initialized. */
    modelTypes: ModelType[];
    /** The share configuration for actions that can be taken by agents in this instance. */
    actionConfig: {
        heads: string[];
        metadata: ActionMetadata;
    };
    /** Initialize data collectors for each of the agents in the instance. */
    initializeDataCollectors(): void;
    /** Validate that the model type in the configuration is correct. */
    validateModelType(modelConfig: FormattedModelConfigBase, idx?: number): void;
    /** Adding validation params given the model configuration. */
    addValidation(configs: FormattedModelConfigBase[]): void;
    /**
     * Sets the frame interval for collection checking.
     * @param interval - The interval at which we collect data.
     */
    setCollectionInterval(interval: number): void;
    /**
     * Sets the reward threshold for collection checking.
     * @param threshold - The absolute value reward threshold.
     */
    setCollectionRewardThreshold(threshold: number): void;
    /**
     * Collects training data.
     * @param dataInstance - The data instance to collect.
     * @param groupIdx - The index of the agent group that corresponds to a specific model
     * @param agentIdx - The index of the agent within the group
     * @returns Whether or not data was eligible to be collected this frame.
     */
    collect(dataInstance: RawDataInstance, groupIdx?: number, agentIdx?: number): boolean;
    /**
     * Retrieves the training data collected by an agent.
     * @param groupIdx - The index of the agent group that corresponds to a specific model
     * @param agentIdx - The index of the agent within the group*
     * @returns The training data.
     */
    getTrainingData(groupIdx?: number, agentIdx?: number): DataInstance[];
    /** Empties the array that stores the training data. */
    clearTrainingData(): void;
    /**
     * Converts actions to a continuous representation.
     * @param rawAction - The action to be collected.
     * @param actionGroup - the action heads to convert
     * @returns A matrix of continuous actions
     */
    createContinuousActionArray(rawAction: RawActionFormats, actionGroup: string): Matrix;
    /**
     * Converts actions to one-hot encoding if they are formatted as a dictionary of bools.
     * @param rawAction - The action to be collected.
     * @param actionGroup - the action heads to convert
     * @returns A matrix of onehot encoded actions
     */
    convertActionToOneHot(rawAction: RawActionFormats, actionGroup: string): Matrix;
    /**
     * Converts all actions to the appropriate format for collection
     * @param rawAction - The action to be collected for all action heads
     * @returns The formatted action
     */
    convertActionsForCollect(rawAction: RawActionFormats): Vector | Matrix | Record<string, Matrix>;
    /**
     * Converts a continuous action (with x and y) to an action that is compatible with an analog stick
     * @param action - The continuous action that will be converted to analog
     * @param allowInsideCircle - Whether we allow coordinates to remain inside the circle border
     * @returns The X and Y coordinates on the analog stick
     */
    convertToAnalog(action: {
        x: number;
        y: number;
    }, allowInsideCircle?: boolean): {
        x: number;
        y: number;
    };
}

/** Represents the combined output from multiple action heads */
type CombinedActionOutput = Record<string, boolean | number>;
interface AgentCoreType extends GeneralAgentCoreType {
    /** The model used for inference in the agent instance. */
    model: Model;
    /** Temorary instance of the model after training (to compare against the baseline model). */
    trainedModel?: Model;
    /** The agent wrapper used around the underlying model. */
    agent: SingleAgentWrapperType;
    /** Temorary instance of the agent after training. */
    trainedAgent?: SingleAgentWrapperType;
    /**
     * Initializes the agent and its model.
     * @param modelData - Configuration and parameters (optional) for the model being created
     * @returns Whether or not a new random model was created.
     */
    createAgent(modelData: ModelDataType): boolean;
    /**
     * Gets the probabilities for the provided inputs.
     * @param inputs - The currrent state.
     * @param postTrainingBool - Whether to use the trained model.
     * @returns The probabilities for each action.
     */
    getProbabilities(inputs: RawState, postTrainingBool?: boolean): Record<string, Matrix>;
    /**
     * Selects an action based on the provided inputs.
     * @param inputs - The currrent state.
     * @param postTrainingBool - Whether to use the trained model.
     * @returns An object which shows the actions that were selected.
     */
    selectAction(inputs: RawState, postTrainingBool?: boolean): CombinedActionOutput;
    /** Removes any input tracking that was left over from a previous game. */
    clearInputTracker(): void;
}

type AgentType$1 = "base" | "reinforcement" | "imitation";
type ArchitectureConfig = {
    inputDim: number;
    stateBounds: any[];
    actionNames: string[];
    actionTypes: string[];
    actionActivations: string[];
    actionOrder: string[][];
    stateSpaceConfig?: {
        config: FeatureConfigType[];
    };
};
/** Architecture metadata. */
interface Architecture {
    id: string;
    slug: string;
    modelType: ModelType;
    name: string;
    config: ArchitectureConfig;
}
/** Agent metadata. */
type AgentData = {
    id: string;
    name?: string;
    architecture?: Architecture;
};
/** Owner of the agent. */
type AgentOwner = {
    type: string;
    id: string;
};
interface BaseAgentType extends AgentCoreType {
    /** Instance of the API client for all networking. */
    api: APIClientType;
    /** The architecture metadata associated with this agent. */
    architecture: Architecture;
    /** The public Id of the agent. */
    id: string;
    /** The name of the agent. */
    name: string;
    /** The owner of this agent. */
    owner: AgentOwner;
    /** The type of agent. */
    agentType: AgentType$1;
    /** The model data for initialization (it gets deleted after use). */
    modelData?: ModelDataType;
    /** Change the agent's name */
    setName(newAgentName: string): void;
    /**
     * Initializes the agent and its model.
     * @returns A promise that resolves when initialization is complete.
     */
    initialize(): Promise<void>;
    /**
     * Uploads training data to the backend.
     * @param contributionMapId - The contribution map that the data will be going towards
     * @returns A promise that resolves to a boolean indicating success.
     */
    uploadData(contributionMapId?: string): Promise<boolean>;
    /**
     * Saves the current model state to the backend.
     * @param newModelBool - Whether to save as a new model.
     */
    save(newModelBool?: boolean): Promise<void>;
    /** Deletes this agent's model. */
    delete(newModelBool?: boolean): Promise<void>;
}

/**
 * Represents the training configuration options.
 */
interface TrainingConfig {
    updatableCells?: string[];
    multiplier?: number;
    epochs?: number;
    batchSize?: number;
    learningRate?: number;
    focus?: number[];
    lambdas?: {
        [key: string]: number;
    };
    cleaning?: {
        balance: {
            oversampling: boolean;
            multiStream: boolean;
        };
        removeSparsity: boolean;
    };
}
interface ImitationLearningAgentType extends BaseAgentType {
    /**
     * 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>;
    /**
     * Saves the current model state to the backend.
     * @param newModelBool - Whether to save as a new model.
     */
    save(newModelBool?: boolean): Promise<void>;
    /**
     * Discards the trained model and optionally resets the training data.
     * @param discardData - Whether to discard training data.
     */
    discardTraining(discardData?: boolean): void;
}

interface ReinforcementLearningAgentType extends BaseAgentType {
}

interface DemoAgentType extends AgentCoreType {
    /** Downloads the current model parameters as a JSON file */
    downloadParameters(): Promise<void>;
}

interface TrainableDemoAgentType extends DemoAgentType {
    /** Instance of the API client for all networking. */
    api: APIClientType;
    /**
     * Trains the agent with imitation learning
     * @param config - The hyperparameter configuration for training
     * @returns boolean for whether or not training was successful
     */
    train(config?: TrainingConfig): Promise<boolean>;
}

/** Valid production agent classes */
type ProductionAgent = ImitationLearningAgentType | ReinforcementLearningAgentType;
/** Agent contructor for production agents */
type ProductionAgentConstructor<T extends ProductionAgent> = new (agentData: AgentData) => T;
/** Valid demo agent classes */
type ValidDemoAgent = DemoAgentType | TrainableDemoAgentType;
/** Valid agent types offered through the SDK */
type AgentType = "reinforcement" | "imitation";
/** Inputs to use when creating a new agent */
type AgentCreationInputs = {
    architectureId: string;
    agentType: AgentType;
    name?: string;
};
/** Inputs to use when loading in an existing agent */
type AgentLoadingInputs = {
    agentType: AgentType;
    agentId: string;
};
/** Agent information overview */
type AgentOverview = {
    architecture: {
        id: string;
    };
    id: string;
    name: string;
    stats?: any;
};
/** Agent summary info for a user */
type AgentSummaryInfo = {
    agents: AgentOverview[];
    maxAgents: number;
};
/** Architecture details */
interface ArchitectureDetails extends Architecture {
    description: string;
    config: any;
    modelType: ModelType;
}
/** Architecture summary info for a user */
type ArchitectureSummaryInfo = {
    architectures: ArchitectureDetails[];
    maxArchitectures: number;
};

/** Parameters for staggering agent inference. */
type StaggerDataType = {
    interval: number;
    cooldowns: number[];
};
/** Multi-agent implementation of the probabilistic agent. */
interface MultiAgentWrapperType extends ProbabilisticAgentWrapperType {
    staggerData: StaggerDataType;
    /**
     * Sets the interval to stagger inference
     * E.g. if there are 20 agents and the interval = 5, and 4 agents perform inference per frame
     * @param interval - The stagger interval.
     */
    setStaggerInterval(interval: number): void;
    /** Gets a mapping from agent idx to whether or not that agent is locked */
    getLockedAgents(): Record<number, boolean>;
    /**
     * Selecting actions for all the specified agents
     * @param input - The inputs to the model for action selection.
     * @param agentsToSkip - Mapping of the agents to skip (typically the ones with locked actions)
     * @returns The selected inputs (actions)
     */
    selectAction(input: any, agentsToSkip?: Record<number, boolean>): InputType;
}

interface MultiAgentCoreType extends GeneralAgentCoreType {
    /** The model used for inference in the agent instance. */
    models: Model[];
    /** The agent wrapper used around the underlying model. */
    agents: MultiAgentWrapperType[];
    /**
     * Sets the interval to stagger inference for all agent groups
     * @param interval - The stagger interval.
     */
    setStaggerIntervals(intervals: Vector): void;
    /**
     * Initializes the agent and its model.
     * @param modelData - Configuration and parameters (optional) for the model being created
     * @returns Whether or not a new random model was created.
     */
    createAgents(modelDataArray: ModelDataType[]): void;
    collectAgentGroup(groupedDataInstances: RawDataInstance[], groupIdx?: number): void;
    collectAll(allDataInstances: RawDataInstance[][]): void;
    getAllTrainingData(): DataInstance[][][];
    getProbabilitiesAgentGroup(groupedInputs: Matrix, groupIdx?: number): Record<string, Matrix>;
    getAllProbabilities(allInputs: Matrix[]): Record<string, Matrix>[];
    getAllLockedAgents(): Record<number, boolean>[];
    selectAllActions(allInputs: any[], agentsToSkip: Record<number, boolean>[]): CombinedActionOutput[];
}

declare class AgentFactory {
    api: APIClientType;
    static api: APIClient;
    /** Gets the api key */
    static getApiKey(): string;
    /** Gets the game id */
    static getGameId(): string;
    /** Gets the backend url being used by the API Client */
    static getBackendUrl(): string;
    /** Overrides the default backend url. */
    static overrideBackendUrl(newBackend: string): Promise<boolean>;
    /** Sets the static game id */
    static setGameId(gameId: string): Promise<boolean>;
    /** Sets the api key */
    static setApiKey(apiKey: string): Promise<boolean>;
    /** Enables cookie auth */
    static enableCookieAuth(): Promise<boolean>;
    /** Disables cookie auth */
    static disableCookieAuth(): Promise<boolean>;
    /** Gets the base url for agents */
    static _getAgentUrl(architectureId?: string): string;
    /** Gets the base url for architectures */
    static _getArchitectureUrl(architectureId?: string): string;
    /** Gets all the high-level agent information for a user (and optionally for an architecture) */
    static getAgentInfo(architectureId?: string): Promise<AgentSummaryInfo>;
    /** Gets all the high-level architecture information for a user (and optionally for a specific architecture) */
    static getArchitectureInfo(architectureId?: string): Promise<ArchitectureSummaryInfo | ArchitectureDetails>;
    /** Gets the class for a given agent type */
    static _getAgentClass(agentType: AgentType$1): ProductionAgentConstructor<ProductionAgent>;
    /** Instantiate and initialize a new agent */
    static _setupAgent(agentType: AgentType$1, agentData: AgentData, modelData?: ModelDataType): Promise<ProductionAgent>;
    /**
     * Creates a new agent in the NRN agents database and returns an instance of the agent.
     * @param inputs - The architecture, type, and name of the agent.
     * @returns An instance of the specified agent type.
     * @throws Will throw an error if the agent type is invalid.
     */
    static createNewAgent({ architectureId, agentType, name }: AgentCreationInputs): Promise<ProductionAgent>;
    /**
     * Loads in an existing agent and returns an instance of it.
     * @param inputs - The type and id of the agent to load.
     * @returns An instance of the specified agent type.
     * @throws Will throw an error if the agent type is invalid.
     */
    static loadAgent({ agentType, agentId }: AgentLoadingInputs): Promise<ProductionAgent>;
    /**
     * Creates a demo instance for simple experimentation
     * @param modelData - Model data, which includes a mix of hyperparameters and parameters
     * @param agentConfig - Configuration options for the agent.
     * @returns An instance of a demo agent.
     */
    static createDemoAgent(modelData: ModelDataType, agentConfig?: AgentConfig): ValidDemoAgent;
    /**
     * Creates a demo instance for simple experimentation
     * @param modelDataArray - Model data for each model in the agent cluster
     * @param numAgents - Number of agents for each model.
     * @param agentConfigs - Configuration options for each if the agents.
     * @returns A cluster of agents to run multi-agents in various environments.
     */
    static createAgentArmy(modelDataArray: ModelDataType[], numAgents: Vector, agentConfigs: AgentConfig[]): MultiAgentCoreType;
}

/** Represents a the bounds for a feature in the state space. */
type StateBounds = {
    type: "range" | "set";
    range?: {
        min: number;
        max: number;
    };
    set?: number[];
};
/** Represents the actions arrays section of the registry slot. */
interface RegistryActionArrays {
    actionNames: string[];
    actionTypes: ActionType[];
    actionActivations: ActionActivationType[];
    actionOrder: string[][];
}
type RegistryActionFormatted = {
    order: string[];
    type?: ActionType;
    activation?: ActionActivationType;
};
/** Represents a slot in the registry. */
interface RegistrySlot extends RegistryActionArrays {
    name?: string;
    gameId: string;
    architectureId: string;
    modelType: ModelType;
    inputDim: number;
    stateBounds?: StateBounds[];
}
/** Represents the inputs required to register a new model architecture. */
type RegistrationInputs = {
    modelType: ModelType;
    architectureId: string;
    name: string;
    inputDim: number;
    actions: Record<string, RegistryActionFormatted>;
    stateBounds?: StateBounds[];
};
/**
 * The Registry class provides methods to interact with the backend for managing registry slots.
 */
interface RegistryType {
    /** Instance of the API client for all networking. */
    api: APIClientType;
    /** List of registered slots. This property is fetched dynamically and may not be initialized. */
    registrySlots?: RegistrySlot[];
    /** Maximum number of slots available in the registry. This value is set after fetching slots. */
    maxSlots?: number;
    /** Gets the api key */
    getApiKey(): string;
    /** Gets the backend url being used by the API Client */
    getBackendUrl(): string;
    /**
     * Sets the api key for the API Client
     * @param apiKey - The api key for the user.
     */
    setApiKey(apiKey: string): Promise<boolean>;
    /**
     * Overrides the default backend url.
     * @param newBackend - The new url for the backend.
     */
    overrideBackendUrl(newBackend: string): Promise<boolean>;
    /** Enables cookie auth */
    enableCookieAuth(): Promise<boolean>;
    /** Disables cookie auth */
    disableCookieAuth(): Promise<boolean>;
    /**
     * Fetches the registered slots for the game from the backend.
     * Updates `registrySlots` and sets the maximum number of slots (`maxSlots`).
     * @returns A promise that resolves when the slots have been fetched.
     */
    fetchRegistrySlots(): Promise<void>;
    /**
     * Unregisters a model from the registry.
     * @param architectureId - The ID of the architecture to unregister.
     * @returns A promise that resolves when the model has been unregistered.
     */
    unregister(architectureId: string): Promise<void>;
    /**
     * Registers a new model architecture in the registry.
     * @param RegistrationInputs - The inputs for the new model architecture.
     * @returns A promise that resolves when the model architecture has been registered.
     */
    register(registrationInputs: RegistrationInputs): Promise<void>;
    /**
     * Prints the current registry slots to the console.
     * If `registrySlots` is undefined, it first fetches the slots from the backend.
     * Displays both registered and empty slots.
     * @returns A promise that resolves when the slots have been printed.
     */
    printSlots(): Promise<void>;
}

declare class Registry implements RegistryType {
    api: APIClientType;
    registrySlots: RegistrySlot[];
    maxSlots: number;
    constructor(gameId: string);
    getApiKey(): string;
    getBackendUrl(): string;
    setApiKey(apiKey: string): Promise<boolean>;
    overrideBackendUrl(newBackend: string): Promise<boolean>;
    enableCookieAuth(): Promise<boolean>;
    disableCookieAuth(): Promise<boolean>;
    fetchRegistrySlots(): Promise<void>;
    unregister(architectureId: string): Promise<void>;
    _convertActions(actions: Record<string, RegistryActionFormatted>): RegistryActionArrays;
    register({ architectureId, modelType, name, inputDim, stateBounds, actions }: RegistrationInputs): Promise<void>;
    _validateRegistration(modelType: ModelType, inputDim: number, actions: Record<string, RegistryActionFormatted>): string;
    printSlots(): Promise<void>;
}

/** Represents a bucket in the data factory. */
interface Bucket$1 {
    id: string;
    slug: string;
    name: string;
    description?: string;
}
/** Represents the inputs required to create a new bucket. */
interface BucketCreationInput {
    slug: string;
    name: string;
    description?: string;
}
/** Represents a user contribution data for a particular contribution map. */
interface UserContribution {
    score: number;
    stdev: number;
    count: number;
    user?: string;
    contributionMap: {
        id: string;
        slug: string;
    };
}
/** Represents an entry on the contribution leaderboard. */
interface ContributionLeaderboardRow extends UserContribution {
    rank: number;
    identifier: string;
}
/** Contribution leaderboard. */
type ContributionLeaderboard = ContributionLeaderboardRow[];
/** Represents the methods available in the DataFactory class. */
interface DataFactoryType {
    /** Instance of the API client for all networking. */
    api: APIClientType;
    /** The game ID associated with this data factory. */
    gameId: string;
    /** The architecture ID associated with this data factory. */
    architectureId: string;
    /** List of registered buckets. This property is fetched dynamically and may not be initialized. */
    buckets?: Bucket$1[];
    /** Maximum number of buckets available in the data factory. This value is set after fetching buckets. */
    maxBuckets?: number;
    /** Gets the API key */
    getApiKey(): string;
    /** Gets the backend URL */
    getBackendUrl(): string;
    /** Sets the API key */
    setApiKey(apiKey: string): Promise<boolean>;
    /** Overrides the backend URL */
    overrideBackendUrl(newBackend: string): Promise<boolean>;
    /** Enables cookie authentication */
    enableCookieAuth(): void;
    /** Disables cookie authentication */
    disableCookieAuth(): void;
    /** Get contributions for the user */
    getMyContributions(campaignId: string): any;
    /** Get the leaderboard for contributions */
    getContributionLeaderboard(campaignId: string, contributionMap: string, numRows?: number): any;
    /** Fetches buckets from the backend */
    fetchBuckets(): Promise<void>;
    /** Creates a new bucket */
    createBucket(data: BucketCreationInput): Promise<Bucket$1>;
    /**
     * Uploads data to a bucket
     * @param bucketId - The ID of the bucket to upload to
     * @param data - The data to upload
     */
    uploadData(bucketId: string, data: any): Promise<boolean>;
    /**
     * Deletes a bucket
     * @param bucketId - The ID of the bucket to delete
     */
    deleteBucket(bucketId: string): Promise<void>;
    /** Prints the registered buckets */
    printSlots(): Promise<void>;
}

type Bucket = {
    id: string;
    slug: string;
    name: string;
    description: string;
};
declare class DataFactory implements DataFactoryType {
    api: APIClientType;
    gameId: string;
    architectureId: string;
    buckets: Bucket[];
    maxBuckets: number;
    constructor(gameId: string, architectureId: string);
    getApiKey(): string;
    getBackendUrl(): string;
    setApiKey(apiKey: string): Promise<boolean>;
    overrideBackendUrl(newBackend: string): Promise<boolean>;
    enableCookieAuth(): Promise<boolean>;
    disableCookieAuth(): Promise<boolean>;
    getMyContributions(campaignId: string): Promise<UserContribution[]>;
    getContributionLeaderboard(campaignId: string, contributionMap: string, numRows?: number): Promise<ContributionLeaderboard>;
    fetchBuckets(): Promise<void>;
    createBucket(data: BucketCreationInput): Promise<Bucket>;
    uploadData(bucketId: string, data: any, contributionMapId?: string, metadata?: any): Promise<boolean>;
    deleteBucket(bucketId: string): Promise<void>;
    printSlots(): Promise<void>;
}

export { AgentFactory, DataFactory, FeatureEngineering, Registry };