// Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
//
// NVIDIA CORPORATION and its licensors retain all intellectual property
// and proprietary rights in and to this software, related documentation
// and any modifications thereto.  Any use, reproduction, disclosure or
// distribution of this software and related documentation without an express
// license agreement from NVIDIA CORPORATION is strictly prohibited.
//
#pragma once

#include "IFlatcache.h"

#include <carb/Framework.h>
#include <carb/Interface.h>
#include <carb/flatcache/AttrNameAndType.h>
#include <carb/flatcache/IPath.h>
#include <carb/flatcache/IToken.h>
#include <carb/flatcache/PrimChanges.h>
#include <carb/logging/Log.h>

#include <gsl/span>
#include <carb/flatcache/Type.h>

#include <map>
#include <stdint.h>

#include "carb/cpp17/Optional.h"
using carb::cpp17::optional;

namespace carb
{
namespace flatcache
{
// The comments in this file are intended to be read top to bottom, like a book.

// This file defines flatcache::StageWithHistory, an instance of which stores
// the current stage and a configurable number of frames of state history. It
// is thread safe, in the sense that the current state can be safely read and
// written in parallel with multiple threads reading the history. The class
// definition is towards the end of this file. We first define some basic types,
// a read/write accessor class for use by the main/game/sim thread, and two
// read-only accessor classes for use on render threads. These classes provide
// access to an interpolated state and a window of state history respectively.

// To specify paths, attribute names, and attribute types we use
// flatcache::Path, flatcache::Token and graph::Type types, rather than
// USD's SdfPath, TfToken and TfType. This allows us to access the stage and
// history without including USD headers.


// The main class is this file is StageWithHistory, which is defined towards
// the end of the file.
class StageWithHistory;

template<typename T>
class AttributeArrayResult;

/**
 * @invariant arrayBytes.size() must be a multiple of bytesPerElement
 */
class ConstArrayAsBytes
{
public:
    gsl::span<const gsl::byte> arrayBytes;
    size_t bytesPerElement;
    Type elementType;
};

// findPrims() returns a list of buckets of prims, represented by PrimBucketList.
class PrimBucketList
{
    friend class StageAtTimeInterval;
    friend class StageInProgress;

protected:
    PrimBucketListId m_primBucketListId;
    static carb::flatcache::IPrimBucketList* sIPrimBucketList();

    PrimBucketList(PrimBucketListId id) : m_primBucketListId(id)
    {
    }

public:
    // PrimBucketList is opaque, you have to use the getAttributesArray methods
    // of StageInProgress, StageAtTime or StageAtTimeInterval to read the
    // attributes of its elements.
    size_t bucketCount() const;
    size_t size() const;
    void print() const;

    PrimBucketListId getId() const
    {
        return m_primBucketListId;
    }

    ~PrimBucketList();
};

// ChangedPrimBucketList is a PrimBucketList that has changes stored for a
// particular listener. It is returned by StageInProgress::getChanges().
class ChangedPrimBucketList : public PrimBucketList
{
    ChangedPrimBucketList(PrimBucketListId id) : PrimBucketList(id) {}

    friend class StageInProgress;

public:
    BucketChanges getChanges(size_t index);
    AddedPrimIndices getAddedPrims(size_t index);
};


// The main/game/sim thread uses the following class to read and write the
// state at the current frame.
//
// StageInProgress can either be used RAII style, you construct it from a frameNumber,
// or non-RAII style, where you construct it from an existing stageInProgressId.
class StageInProgress
{
    StageInProgressId m_stageInProgress;
    bool m_createdFromId;
    UsdStageId m_usdStageId; // Only valid if m_createFromId == false

public:
    // The constructor creates a new frame and locks it for read/write
    StageInProgress(StageWithHistory& stageWithHistory, size_t simFrameNumber);

    // Create from an already locked frame
    StageInProgress(StageInProgressId stageInProgressId);

    // Returns the frame number allocated by constructor
    size_t getFrameNumber();

    // Returns the frame time allocated by constructor
    RationalTime getFrameTime();

    // Returns which mirrored array is valid: CPU, GPU, etc.
    ValidMirrors getAttributeValidBits(const Path& path, const Token& attrName) const;

    // getAttribute returns a read/write pointer to a non-array attribute
    // If it returns nullptr then the attribute doesn't exist in the stage
    template <typename T>
    T* getAttribute(const Path& path, const Token& attrName);

    // getAttribute returns a read-only pointer to a non-array attribute
    // If it returns nullptr then the attribute doesn't exist in the stage
    template <typename T>
    const T* getAttributeRd(const Path& path, const Token& attrName);

    // getAttribute returns a write-only pointer to a non-array attribute
    // If it returns nullptr then the attribute doesn't exist in the stage
    template <typename T>
    T* getAttributeWr(const Path& path, const Token& attrName);

    // getAttribute returns a read/write pointer to a non-array attribute
    // If it returns nullptr then the attribute doesn't exist in the stage
    template <typename T>
    T* getAttributeGpu(const Path& path, const Token& attrName);

    // getAttribute returns a read-only pointer to a non-array attribute
    // If it returns nullptr then the attribute doesn't exist in the stage
    template <typename T>
    const T* getAttributeRdGpu(const Path& path, const Token& attrName);

    // getAttribute returns a write-only pointer to a non-array attribute
    // If it returns nullptr then the attribute doesn't exist in the stage
    template <typename T>
    T* getAttributeWrGpu(const Path& path, const Token& attrName);

    // getOrCreateAttributeWr returns a write-only pointer to a non-array
    // attribute. If the attribute doesn't exist, then it will create it.
    // The return type is a reference rather than a pointer because the
    // attribute is guaranteed to exist on exit
    template <typename T>
    T& getOrCreateAttributeWr(const Path& path, const Token& attrName, Type type);

    // getAttribute returns a read/write span of an array attribute
    // The span allows the array size to be read, but not written
    // To set the array size, use setArrayAttributeSize
    template <typename T>
    gsl::span<T> getArrayAttribute(const Path& path, const Token& attrName);

    // getAttributeRd returns a read-only span of an array attribute
    // The array size is also read only
    template <typename T>
    gsl::span<const T> getArrayAttributeRd(const Path& path, const Token& attrName);

    // getAttributeRd returns a write-only span of an array attribute
    // The array size is read only, to resize use setArrayAttributeSize
    template <typename T>
    gsl::span<T> getArrayAttributeWr(const Path& path, const Token& attrName);

    // Get the size of an array attribute. When writing CPU code, it isn't
    // normally necessary to use this method, as getArrayAttribute returns a
    // span containing the data pointer and the size.
    // However, when writing mixed CPU/GPU code it is wasteful to copy the
    // array data from GPU to CPU when just the size is required, so use this
    // method in that case.
    size_t getArrayAttributeSize(const Path& path, const Token& attrName);

    // Set the size of an array attribute
    void setArrayAttributeSize(const Path& path, const Token& attrName, size_t elemCount);

    template <typename T>
    gsl::span<T> setArrayAttributeSizeAndGet(const PrimBucketList& primBucketList,
                                             size_t primBucketListIndex,
                                             size_t indexInBucket,
                                             const Token& attrName,
                                             size_t newElemCount);

    // createPrim, destroyPrim, createAttribute and destroyAttribute do what
    // you'd expect
    void createPrim(const Path& path);
    void destroyPrim(const Path& path);
    void createAttribute(const Path& path, const Token& attrName, Type type);
    template<int n>
    void createAttributes(const Path& path, std::array<AttrNameAndType, n> attributes);

    // Deprecated: type argument is not used.
    void destroyAttribute(const Path& path, const Token& attrName, Type type);
    void destroyAttribute(const Path& path, const Token& attrName);
    template <int n>
    void destroyAttributes(const Path& path, const std::array<Token, n>& attributes);
    void destroyAttributes(const Path& path, const std::vector<Token>& attributes);

    // findPrims() finds prims that have all the attributes in "all", and any
    // of the attributes in "any", and none of the attributes in "none".
    // The attributes of the resulting prims can be accessed as piecewise
    // contiguous arrays, using getAttributeArray() below, which is typically
    // faster than calling getAttribute for each prim.
    PrimBucketList findPrims(const carb::flatcache::set<AttrNameAndType>& all,
                             const carb::flatcache::set<AttrNameAndType>& any = {},
                             const carb::flatcache::set<AttrNameAndType>& none = {});

    /**
     * Tell a listener to log changes for an attribute.
     * Attaches listener to stage if not already attached
     *
     * @param[in] attrName The attribute's name
     * @param[in] listenerId The listener
     */
    void attributeEnableChangeTracking(const Token& attrName, ListenerId listenerId);

    /**
     * Tell a listener to stop logging changes for an attribute.
     * Attaches listener to stage if not already attached
     *
     * @param[in] attrName The attribute's name
     * @param[in] listenerId The listener
     */
    void attributeDisableChangeTracking(const Token& attrName, ListenerId listenerId);

    /**
     * Tell a listener to log prim creates
     * Attaches listener to stage if not already attached
     *
     * @param[in] attrName The attribute's name
     * @param[in] listenerId The listener
     */
    void enablePrimCreateTracking(ListenerId listenerId);

    /**
     * Pause change tracking.
     *
     * @param[in] listenerId The listener to pause
     */
    void pauseChangeTracking(ListenerId listenerId);

    /**
     * Resume change tracking.
     *
     * @param[in] listenerId The listener to resume
     */
    void resumeChangeTracking(ListenerId listenerId);

    /**
     * Is change tracking paused?
     *
     * @param[in] listenerId The listener
     * @return Whether the listener is paused
     */
    bool isChangeTrackingPaused(ListenerId listenerId);

    /**
     * Get changes
     *
     * @param[in] listenerId The listener
     * @return The changes that occured since the last time the listener was popped
     */
    ChangedPrimBucketList getChanges(ListenerId listenerId);

    /**
     * Clear the list of changes 
     *
     * @param[in] listenerId The listener
     */
    void popChanges(ListenerId listenerId);

    /**
     * Get the number of listeners
     *
     * @return The number of listeners listening to this stage
     */
    size_t getListenerCount();

    /**
     * Is the listener attached to this stage
     *
     * @return Whether the listener is attached to this stage
     */
    bool isListenerAttached(ListenerId listenerId);

    /**
     * Detach the listener from the stage. Future changes will not be logged for this listener.
     *
     * @param[in] listenerId The listener
     * @return Whether the listener is attached to this stage
     */
    void detachListener(ListenerId listenerId);

    // getAttributeArray(primBucketList, index, attrName) returns a read/write
    // contiguous array of the values of attribute "attrName" for each prim of
    // bucket "index" of "primBucketList".
    // "index" must be in the range [0..primBucketList.getBucketCount())
    template <typename T>
    gsl::span<T> getAttributeArray(const PrimBucketList& primBucketList, size_t primBucketListIndex, const Token& attrName);

    template <typename T>
    gsl::span<const T> getAttributeArrayRd(const PrimBucketList& primBucketList,
                                           size_t primBucketListIndex,
                                           const Token& attrName) const;

    template <typename T>
    gsl::span<T> getAttributeArrayWr(const PrimBucketList& primBucketList,
                                     size_t primBucketListIndex,
                                     const Token& attrName);

    template <typename T>
    gsl::span<T> getAttributeArrayGpu(const PrimBucketList& primBucketList, size_t primBucketListIndex, const Token& attrName);

    template <typename T>
    gsl::span<const T> getAttributeArrayRdGpu(const PrimBucketList& primBucketList,
                                              size_t primBucketListIndex,
                                              const Token& attrName) const;

    template <typename T>
    gsl::span<T> getAttributeArrayWrGpu(const PrimBucketList& primBucketList,
                                        size_t primBucketListIndex,
                                        const Token& attrName);

    template <typename T>
    gsl::span<T> getOrCreateAttributeArrayWr(const PrimBucketList& primBucketList,
                                             size_t primBucketListIndex,
                                             const Token& attrName,
                                             Type type);

    // getAttributeArray(primBucketList, index, attrName) returns a vector of
    // array-valued attributes "attrName" for the prims of bucket "index" of
    // "primBucketList". "index" must be in the range [0..primBucketList.getBucketCount())
    // It gives read/write access to the values of each prim's array
    template <typename T>
    std::vector<gsl::span<T>> getArrayAttributeArray(const PrimBucketList& primBucketList,
                                                     size_t primBucketListIndex,
                                                     const Token& attrName) const;

    // getAttributeArray(primBucketList, index, attrName) returns a vector of
    // array-valued attributes "attrName" for the prims of bucket "index" of
    // "primBucketList". "index" must be in the range [0..primBucketList.getBucketCount())
    // It gives read-only access to the values of each prim's array
    template <typename T>
    std::vector<gsl::span<const T>> getArrayAttributeArrayRd(const PrimBucketList& primBucketList,
                                                             size_t primBucketListIndex,
                                                             const Token& attrName) const;

    // getAttributeArray(primBucketList, index, attrName) returns a vector of
    // array-valued attributes "attrName" for the prims of bucket "index" of
    // "primBucketList". "index" must be in the range [0..primBucketList.getBucketCount())
    // It gives write-only access to the values of each prim's array
    template <typename T>
    std::vector<gsl::span<T>> getArrayAttributeArrayWr(const PrimBucketList& primBucketList,
                                                       size_t primBucketListIndex,
                                                       const Token& attrName) const;

    // getPathArray(primBucketList, index) returns a read-only contiguous array
    // of the paths of the prims of bucket "index" of "primBucketList".
    // "index" must be in the range [0..primBucketList.getBucketCount())
    gsl::span<const Path> getPathArray(const PrimBucketList& primBucketList, size_t primBucketListIndex) const;

    flatcache::set<AttrNameAndType> getAttributeNamesAndTypes(const PrimBucketList& primBucketList,
        size_t primBucketListIndex) const;

    // TODO: replace with an iterator for iterating over bucket names
    void printBucketNames() const;

    // Record that the attribute at path.attrName has been modified. Right now this is
    // done explicitly to give a high degree of control over which attributes get
    // passed to the notice.
    void logAttributeWriteForNotice(const Path& path, const Token& attrName);

    // Construct and send a TfNotice with a vector of objects paths
    // that have changed, much like the ObjectsChanged notice from USD
    void broadcastTfNoticeForAttributesChanged() const;

    // Connection API
    /**
     * @brief Create a connection on the target prim
     *
     * @param path the target prim on which to create a connection
     * @param connectionName specifies the connections attribute name on the prim
     * @param connection specifies the target prim and attribute of the connection
     */
    void createConnection(const Path& path, const Token& connectionName, const Connection& connection);

    /**
     * @brief Create an arbitrary number of connections on the target prim
     *
     * @param path the target prim on which to create a connection
     * @param connectionNames a span of attribute names. Must match the size of the connections vector
     * @param connections a span of connections. Must match the size of the connectionNames vector
     */
    void createConnections(const Path& path, const gsl::span<Token>& connectionNames, const gsl::span<Connection>& connections );

    /**
     * @brief removes a connection from a prim
     *
     * @param path the target prim from which to remove a connection
     * @param connectionName the name of the connection to remove
     */
    void destroyConnection(const Path& path, const Token& connectionName);

    /**
     * @brief removes an arbitary number of connections from a prim
     *
     * @param path the target prim from which to remove the connections
     * @param connectionNames the names of the connections to be removed
     */
    void destroyConnections(const Path& path, const gsl::span<Token>& connectionNames);

    /**
     * @brief Get a R/W pointer to a connection on the target prim
     *
     * @param path the target prim
     * @param connectionName the target connection name
     * @return a R/W pointer to the connection
     */
    Connection* getConnection(const Path& path, const Token& connectionName);

    /**
     * @brief Get a read only pointer to a connection on the target prim
     *
     * @param path the target prim
     * @param connectionName the target connection name
     * @return a read only pointer to the connection
     */
    const Connection* getConnectionRd(const Path& path, const Token& connectionName);

    /**
     * @brief Get a write only pointer to a connection on the target prim
     *
     * @param path the target prim
     * @param connectionName the target connection name
     * @return  a write only pointer to the connection
     */
    Connection* getConnectionWr(const Path& path, const Token& connectionName);

    /**
     * @brief Copy all attributes from the source prim to the destination prim
     * Will create attributes if they do not exist on the destination prim
     * If an attribute exists on both prims they must have compatible types to copy.
     *
     * @param[in] srcPath the source prim 
     * @param[in] dstPath the destination prim
     */
    void copyAttributes(const Path& srcPath, const Path& dstPath);

    /**
     * @brief Copy the specified attributes from the source prim to the destination prim
     * Will create attributes if they do not exist on the destination prim
     * If an attribute exists on both prims they must have compatible types to copy.
     *
     * @param[in] srcPath the source prim 
     * @param[in] srcAttrs a span of attributes to be copied.
     * @param[in] dstPath the destination prim
     */
    void copyAttributes(const Path& srcPath, const gsl::span<Token>& srcAttrs, const Path& dstPath);
     
    /**
     * @brief Copy the specified attributes from the source prim to the the specified 
     * attributes on the destination prim
     * Will create attributes if they do not exist on the destination prim
     * If an attribute exists on both prims they must have compatible types to copy.
     * Note: The srcAttrs and dstAttrs must be the same size as the function assumes
     * that the copy is 1 to 1 in terms of name alignment
     *
     * @param[in] srcPath the source prim 
     * @param[in] srcAttrs a span of attributes to be copied.
     * @param[in] dstPath the destination prim
     * @param[in] dstAttrs a span of attributes to be copied.
     */
    void copyAttributes(const Path& srcPath, const gsl::span<Token>& srcAttrs, const Path& dstPath, const gsl::span<Token>& dstAttrs);


    StageInProgressId getId() const
    {
        return m_stageInProgress;
    }

    /**
     * @brief Check whether a prim exists at a given path
     * @param[in] the path
     * @return true if a prim exists at the path
     */
    bool primExists(const Path& path); 

    // If StageInProgress was created from an Id, then do nothing
    // Else unlock the current sim frame, allowing it to be read by
    // other threads
    ~StageInProgress();
};

// The following two classes, StageAtTime and StageAtTimeInterval
// are used by reader threads to read the history. StageAtTime is
// used when the state of a stage is needed at a particular point in time.
// StageAtTimeInterval is used when we need all the stage history in a given time
// window.
//
// There can be multiple threads reading the history buffer, for example
// multiple sensor renderers running at different rates. We use shared locks
// to allow multiple threads to read the same frame of history.
//
// StageAtTimeInterval takes an RAII approach to locking, constructing one locks
// a range of slots for reading, and destructing unlocks them.
class StageAtTimeInterval
{
    StageAtTimeIntervalId m_stageAtTimeInterval;
    static carb::flatcache::IStageAtTimeInterval* sIStageAtTimeInterval();

public:
    // The constructor locks frames of history
    StageAtTimeInterval(StageWithHistory& stageWithHistory,
                        RationalTime beginTime,
                        RationalTime endTime,
                        bool includeEndTime = false);

    StageAtTimeInterval(StageWithHistoryId stageWithHistoryId,
                        RationalTime beginTime,
                        RationalTime endTime,
                        bool includeEndTime = false);

    ValidMirrors getAttributeValidBits(const PathC& path, const TokenC& attrName) const;

    // Get values of locked elements
    template <typename T>
    std::vector<const T*> getAttributeRd(const Path& path, const Token& attrName) const;

    // Get GPU pointer and size of locked elements
    template <typename T>
    std::vector<const T*> getAttributeRdGpu(const Path& path, const Token& attrName) const;

    // Get the size of an array attribute. When writing CPU code, it isn't
    // normally necessary to use this method, as getArrayAttribute returns a
    // span containing the data pointer and the size.
    // However, when writing mixed CPU/GPU code it is wasteful to copy the
    // array data from GPU to CPU when just the size is required, so use this
    // method in that case.
    std::vector<size_t> getArrayAttributeSize(const Path& path, const Token& attrName) const;

    /**
     * @brief Get an array-valued attribute for reading from a single prim
     *
     * @param path The path of the prim
     * @param attrName The name of the attribute
     * 
     * @return a vector of array spans, one for each time sample within the current StageAtTimeInterval
     */
    template <typename T>
    std::vector<gsl::span<const T>> getArrayAttributeRd(const Path& path, const Token& attrName) const;

    /**
     * @brief Get an array-valued attribute as bytes for reading from a single prim.
     *        This is useful for converting to VtValue
     *
     * @param path The path of the prim
     * @param attrName The name of the attribute
     *
     * @return a vector of array spans, one for each time sample within the
     *         current StageAtTimeInterval
     */
    std::vector<ConstArrayAsBytes> getArrayAttributeRawRd(const Path& path, const Token& attrName) const;

    // Get timestamps of locked elements
    std::vector<RationalTime> getTimestamps() const;

    size_t getTimeSampleCount() const;

    // findPrims() finds prims that have all the attributes in "all", and any
    // of the attributes in "any", and none of the attributes in "none".
    // The attributes of the resulting prims can be accessed as piecewise
    // contiguous arrays, using getAttributeArray() below, which is typically
    // faster than calling getAttribute for each prim.
    PrimBucketList findPrims(const carb::flatcache::set<AttrNameAndType>& all,
                             const carb::flatcache::set<AttrNameAndType>& any = {},
                             const carb::flatcache::set<AttrNameAndType>& none = {});

    // getAttributeArray(primBucketList, index, attrName) returns for each
    // timesample, a read-only, contiguous array of the values of attribute
    // "attrName" for each prim of bucket "index" of "primBucketList".
    // "index" must be in the range [0..primBucketList.getBucketCount())
    template <typename T>
    std::vector<gsl::span<const T>> getAttributeArrayRd(const PrimBucketList& primBucketList,
                                                        size_t primBucketListIndex,
                                                        const Token& attrName) const;

    template <typename T>
    std::vector<gsl::span<const T>> getAttributeArrayRdGpu(const PrimBucketList& primBucketList,
                                                           size_t primBucketListIndex,
                                                           const Token& attrName) const;

    template <typename T>
    std::vector<std::vector<gsl::span<const T>>> getArrayAttributeArrayRd(const PrimBucketList& primBucketList,
                                                                          size_t primBucketListIndex,
                                                                          const Token& attrName) const;

    /**
     * @brief Read a raw byte representation for a given attribute from a given bucket. This is useful for doing things such as batched type conversions.
     *
     * @param primBucketList the list of buckets
     * @param primBucketListIndex the specific bucket to search
     * @param attrName the token describing the desired attribute
     *
     * @return a vector of byte arrays, one for each time sample within the current StageAtTimeInterval
     */
    std::vector<gsl::span<const char>> getAttributeArrayRawRd(const PrimBucketList& primBucketList,
                                                              size_t primBucketListIndex,
                                                              const Token& attrName) const;

    // getPathArray(primBucketList, index) returns for each timesample a
    // read-only contiguous array of the paths of the prims of bucket "index"
    // of "primBucketList".
    // "index" must be in the range [0..primBucketList.getBucketCount())
    // The reason a separate path array is returned per sample is that prims
    // can be added and deleted from frame to frame, and we need to check which
    // prim a sample corresponds to when interpolating.
    std::vector<gsl::span<const Path>> getPathArray(const PrimBucketList& primBucketList,
                                                    size_t primBucketListIndex) const;

    /**
     * @brief Get a Connection on a target prim
     * 
     * @param path the target prim
     * @param connectionName the connection name
     * @return a vector of read only pointers to connections
     */
    std::vector<const Connection*> getConnectionRd(const Path& path, const Token& connectionName);

    // TODO: replace with an iterator for iterating over bucket names
    void printBucketNames() const;

    /**
     * @brief write the current data for this stageInProgress to the specified UsdStage
     *        this will write all attributes at the currentTime in getFrameNumber()
     *
     * @param usdStageId Valid usdStage in the stage cache
     *
     * @return none
     */
    void exportUsd(UsdStageId usdStageId) const;

    // Get the number of attributes for a given bucket.
    std::vector<size_t> getAttributeCounts(const PrimBucketList& primBucketList, size_t primBucketListIndex) const;
    
    // Get the name and type of each attribute for a given bucket.
    std::pair< std::vector<std::vector<Token>>, std::vector<std::vector<Type>>>
        getAttributeNamesAndTypes(const PrimBucketList& primBucketList, size_t primBucketListIndex) const;

    /** @brief Write a cache file to disk at a specified location
     *
     * @note Many parameters to this function are optional
     * @note This currently only writes the first time in the interval
     *
     * @file[in The location the file is desired to be written to
     * @workingBuffer[in] [Optional] In order to avoid costly reallocations
     *                    the code will attempt to re-use the memory at the buffer
     *                    location if it is large enough. If the buffer isn't larg
     *                    enough the cost of allocation, and re-traversal may be paid
     * @workingBufferSize[in] [Optional] If workingBuffer is non null, then this desrcibes the length
     *                        of the buffer
     * @return The amount of data needed to serialize the cache, a return value of 0 indicates an error
     *
     */
    uint64_t writeCacheToDisk(const char* file, uint8_t* workingBuffer, uint64_t workingBufferSize)  const;

    /** @brief Add a ref count to any data backed by the StageAtTimeIntercal
     *
     * @note The ref count will not enforce any behavior currently, but will
     *       print a warning if backing data is deleted before all ref counts
     *       are cleared
     *
     * @return None
     *
     */
    void addRefCount();

    /** @brief Remove a ref count from an existing timeInterval
     *
     * @return True if ref count was removed successfully, failure conditions may
     *         include
     *           (1) StageAtTimeInterval doesn't exist
     *           (2) RefCount was already 0
     */
    bool removeRefCount();

    /** @brief Query ref count for a stage at time
    *
    * @note A stage at time might be represented by multiple actual data sources
    *       in that case we return the largest refcount of all the data sources
    *
    * @return number of reference counts
    */
    unsigned int getRefCount();

    // Unlocks elements to allow them to be reused.
    ~StageAtTimeInterval();
};

// StageAtTime is used when the state of a stage is needed at
// a particular point in time, which may or may not be one of the times sampled
// in the history. If it is, then getAttributeRd returns the exact value sampled.
// If not, it linearly interpolates using the two closest samples in the history.
//
// StageAtTime takes an RAII approach to locking, constructing one
// locks one or two frames in the history (depending on whether interpolation
// is needed), and destructing unlocks them.
class StageAtTime
{
    // Invariants:
    // I0: if sampleTimes.size()==2, m_theta = (m_time - sampleTimes[0]) /
    //                                         (sampleTimes[1] - sampleTimes[0])
    //       where sampleTimes = m_historyWindow.getTimestamps()
    //
    //     In particular, m_theta increases linearly from 0 to 1 as m_time
    //     increases from sampleTimes[0] to sampleTimes[1]
    //
    // TODO: do we need to delay conversion from rational number to double?
    StageAtTimeInterval m_historyWindow;
    RationalTime m_time;
    double m_theta;

    void initInterpolation()
    {
        std::vector<RationalTime> sampleTimes = m_historyWindow.getTimestamps();

        if (sampleTimes.size() == 2)
        {
            if ((double)sampleTimes[0].denominator == 0.0 || (double)sampleTimes[1].denominator == 0.0)
            {
                CARB_LOG_WARN_ONCE("StageWithHistory initInterpolation(): cannot divide by a denominator with a value of zero.");
                m_theta = 0.0;
            }
            else
            {
                double a_t = (double)sampleTimes[0].numerator / (double)sampleTimes[0].denominator;
                double b_t = (double)sampleTimes[1].numerator / (double)sampleTimes[1].denominator;

                if (a_t == b_t)
                    m_theta = 0.0;

                else
                {
                    double c_t = (double)m_time.numerator / (double)m_time.denominator;
                    m_theta = (c_t - a_t) / (b_t - a_t);
                }
            }    
        }
        else if (sampleTimes.size() == 1)
            m_theta = 0.0;
    }

public:
    // Locks one or two history elements for read.
    StageAtTime(StageWithHistory& stageWithHistory, RationalTime time)
        : m_historyWindow(stageWithHistory, time, time, true), m_time(time)
    {
        initInterpolation();
    }

    StageAtTime(StageWithHistoryId stageWithHistoryId, RationalTime time)
        : m_historyWindow(stageWithHistoryId, time, time, true), m_time(time)
    {
        initInterpolation();
    }
  
    // Auxiliary method to communicate attributes of types which will not be interpolated
    // Supported types: bool, int, uint
    // no samples found: return nullopt 
    // samples found:    return pair{value of sample in frame n, value of sample in frame n+1}
    template <typename T>
    optional<std::pair<optional<T>, optional<T>>> getNonInterpolatableAttributeRd(const Path& path, const Token& attrName) const;

    ValidMirrors getAttributeValidBits(const PathC& path, const TokenC& attrName) const;

    // Read interpolated elements
    template <typename T>
    optional<T> getAttributeRd(const Path& path, const Token& attrName) const;

    // Read GPU elements (interpolation not supported yet!)
    template <typename T>
    const T* getAttributeRdGpu(const Path& path, const Token& attrName) const;

    // Get array attribute size, useful for GPU attributes
    size_t getArrayAttributeSize(const Path& path, const Token& attrName) const;

    // Get arrau attribute read
    template <typename T>
    gsl::span<const T> getArrayAttributeRd(const Path& path, const Token& attrName);

    // findPrims() finds prims that have all the attributes in "all", and any
    // of the attributes in "any", and none of the attributes in "none".
    // The attributes of the resulting prims can be accessed as piecewise
    // contiguous arrays, using getAttributeArray() below, which is typically
    // faster than calling getAttribute for each prim.
    PrimBucketList findPrims(const carb::flatcache::set<AttrNameAndType>& all,
                             const carb::flatcache::set<AttrNameAndType>& any = {},
                             const carb::flatcache::set<AttrNameAndType>& none = {})
    {
        return m_historyWindow.findPrims(all, any, none);
    }

    // getAttributeArray(primBucketList, index, attrName) returns a read-only
    // contiguous array of the values of attribute "attrName" for each prim of
    // bucket "index" of "primBucketList".
    // "index" must be in the range [0..primBucketList.getBucketCount())
    template <typename T>
    AttributeArrayResult<T> getAttributeArrayRd(const PrimBucketList& primBucketList,
                                                size_t primBucketListIndex,
                                                const Token& attrName) const;

    template <typename T>
    AttributeArrayResult<T> getAttributeArrayRdGpu(const PrimBucketList& primBucketList,
                                                   size_t primBucketListIndex,
                                                   const Token& attrName) const;

    template <typename T>
    AttributeArrayResult<std::vector<T>> getArrayAttributeArrayRd(const PrimBucketList& primBucketList,
                                                             size_t primBucketListIndex,
                                                             const Token& attrName) const;

    /**
     * @brief Read a raw byte representation for a given attribute from a given bucket. This is useful for doing things such as batched type conversions.
     *
     * @param primBucketList the list of buckets
     * @param primBucketListIndex the specific bucket to search
     * @param attrName the token describing the desired attribute
     *
     * @return a vector of byte arrays, one for each time sample underlying the current StageAtTime. Note: Does not perform any interpolation.
     */
    std::vector<gsl::span<const char>> getAttributeArrayRawRd(const PrimBucketList& primBucketList,
                                                              size_t primBucketListIndex,
                                                              const Token& attrName) const;

    // getPathArray(primBucketList, index) returns a read-only contiguous array
    // of the paths of the prims of bucket "index" of "primBucketList".
    // "index" must be in the range [0..primBucketList.getBucketCount())
    gsl::span<const Path> getPathArray(const PrimBucketList& primBucketList, size_t primBucketListIndex) const;

    /**
     * @brief Get a read only pointer to a connection on a prim
     * 
     * @param path the target prim
     * @param connectionName the target connection
     * @return returns a vector read only pointers to connections, one per time sample
     */
    std::vector<const Connection*> getConnectionRd(const Path& path, const Token& connectionName);

    // TODO: replace with an iterator for iterating over bucket names
    void printBucketNames() const;
    
    // Get the number of attributes for a given bucket.
    size_t getAttributeCount(const PrimBucketList& primBucketList, size_t primBucketListIndex) const;
    
    // Get the name and type of each attribute for a given bucket.
    std::pair< std::vector<Token>, std::vector<Type>>
        getAttributeNamesAndTypes(const PrimBucketList& primBucketList, size_t primBucketListIndex) const;

    // Unlocks elements to allow them to be reused.
    ~StageAtTime() = default;

    /** @brief Write a cache file to disk at a specified location
     *
     * @note Many parameters to this function are optional
     *
     * @file[in The location the file is desired to be written to
     * @workingBuffer[in] [Optional] In order to avoid costly reallocations
     *                    the code will attempt to re-use the memory at the buffer
     *                    location if it is large enough. If the buffer isn't larg
     *                    enough the cost of allocation, and re-traversal may be paid
     * @workingBufferSize[in] [Optional] If workingBuffer is non null, then this desrcibes the length
     *                        of the buffer
     * @return The amount of data needed to serialize the cache, a return value of 0 indicates an error
     *
     */
    uint64_t writeCacheToDisk(const char* file, uint8_t* workingBuffer, uint64_t workingBufferSize)  const;

    /** @brief Add a ref count to any data backed by the StageAtTimeInterval
     *
     * @note The ref count will not enforce any behavior currently, but will
     *       print a warning if backing data is deleted before all ref counts
     *       are cleared
     *
     * @return None
     *
     */
    void addRefCount();

    /** @brief Remove a ref count from an existing timeInterval
     *
     * @return True if ref count was removed successfully, failure conditions may
     *         include
     *           (1) StageAtTimeInterval doesn't exist
     *           (2) RefCount was already 0
     */
    bool removeRefCount();

    /** @brief Query ref count for a stage at time
    *
    * @note A stage at time might be represented by multiple actual data sources
    *       in that case we return the largest refcount of all the data sources
    *
    * @return number of reference counts
    */
    unsigned int getRefCount();
};

// Finally, here is the main class, StageWithHistory.
class StageWithHistory
{
    StageWithHistoryId m_stageWithHistory;
    UsdStageId m_usdStageId;

    friend class StageInProgress;
    friend class StageAtTimeInterval;

public:
    StageWithHistory(UsdStageId usdStageId, size_t historyFrameCount, RationalTime simPeriod, bool withCuda=false);
    ~StageWithHistory();

    /**
     * Create a listener
     * This just creates a listener ID, you have to attach it to a stage to use it.
     * Note that there is no destroyListener method. To stop using an ID, detach it from all stages it is attached to.
     * @return The listener
     */
    ListenerId createListener();
};

const ListenerId kInvalidListenerId = { 0 };

} // namespace flatcache
} // namespace carb

// Implement above C++ methods by calling C-ABI interfaces
#include "WrapperImpl.h"
