#pragma once
/*
 * Prepacking: Group together technology-mapped netlist blocks before packing.
 * This gives hints to the packer on what groups of blocks to keep together
 * during packing.
 *
 * Primary uses: 1) "Forced" packs (eg LUT+FF pair)
 *               2) Carry-chains
 */

#include <algorithm>
#include <map>
#include <unordered_set>
#include "atom_netlist_fwd.h"
#include "pack_patterns.h"
#include "vtr_assert.h"
#include "vtr_range.h"
#include "vtr_strong_id.h"
#include "vtr_vector.h"
#include "vtr_vector_map.h"

// Forward declarations
class t_pack_molecule;
class LogicalModels;
struct t_logical_block_type;
struct t_pb_graph_node;

// A unique ID used to identify a molecule generated by the prepacker.
typedef vtr::StrongId<struct pack_molecule_id_tag, size_t> PackMoleculeId;

// A unique ID used to identify a chain of molecules generated by the prepacker.
typedef vtr::StrongId<struct molecule_chain_id_tag, size_t> MoleculeChainId;

/**
 * @brief Holds general information to be shared between molecules that
 *        represent the same chained pack pattern.
 *
 * For example, molecules that are representing a long carry chain that spans
 * multiple logic blocks.
 */
struct t_chain_info {
    /// @brief Is this a long chain that is divided on multiple clusters
    ///        (divided on multiple molecules).
    bool is_long_chain = false;
};

/**
 * @brief Describes the molecule pack pattern type.
 */
enum class e_pack_pattern_molecule_type : bool {
    MOLECULE_SINGLE_ATOM, ///<single atom forming a molecule (no pack pattern associated)
    MOLECULE_FORCED_PACK  ///<more than one atom representing a packing pattern forming a large molecule
};

/**
 * @brief Represents a grouping of atom blocks that match a pack_pattern,
 *        these groups are intended to be placed as a single unit during packing
 *
 * A chain is a special type of pack pattern since it can extend across multiple
 * logic blocks. The prepacker segments the chain into molecules that each fit
 * in a logic block by identifying the atom that forms the root of the chain,
 * and starting the first molecule from it. Long chains can lead to multiple
 * molecules; a new molecule is created as the chain is traversed every time we
 * exceed the maximum number of bits a single logic block can implement. The
 * MoleculeChainId can be used to identify the molecules that are part of the
 * same chain; it is set to Invalid if a given molecule did not come from a
 * chain.
 */
class t_pack_molecule {
  public:
    // =========================================================================
    // General molecule info
    // =========================================================================

    /// @brief Intrinsic "goodness" score for molecule independent of rest of netlist.
    float base_gain;

    /// @brief Whether this molecule represents a single atom or more atoms
    ///        representing a packing pattern.
    enum e_pack_pattern_molecule_type type;

    // =========================================================================
    // Large molecules info
    // =========================================================================

    /// @brief If not a single atom, this is the pack pattern representing this molecule.
    t_pack_patterns* pack_pattern;

    /// @brief Index of the pack_pattern->root_block in the atom_blocks_ids.
    ///        root_block_id = atom_block_ids[root]
    int root;

    /// @brief [0..num_blocks-1] IDs of atom blocks that implements this molecule,
    ///        indexed by t_pack_pattern_block->block_id.
    ///
    /// This vector may contain invalid atom block ids (when the molecule does
    /// not completely fill the pattern).
    std::vector<AtomBlockId> atom_block_ids;

    /// @brief The unique ID of the chain this molecule is a part of if is_chain.
    ///        If this molecule is not part of a chain, this would be invalid.
    ///
    /// Multiple molecules may point to the same chain.
    MoleculeChainId chain_id;

    // =========================================================================
    // Class methods
    // =========================================================================

    // A molecule is a chain if it is a forced pack and its pack pattern is a chain
    inline bool is_chain() const {
        return type == e_pack_pattern_molecule_type::MOLECULE_FORCED_PACK && pack_pattern->is_chain;
    }
};

/**
 * @brief Statistics on a molecule.
 *
 * This is used during packing to quickly look up information on a molecule to
 * compute gain information.
 *
 * TODO: The prepacker should precompute this information per molecule.
 */
struct t_molecule_stats {
    /// @brief Number of blocks across all primitives in molecule.
    int num_blocks = 0;

    /// @brief Number of pins across all primitives in molecule.
    int num_pins = 0;
    /// @brief Number of input pins across all primitives in molecule.
    int num_input_pins = 0;
    /// @brief Number of output pins across all primitives in molecule.
    int num_output_pins = 0;

    /// @brief Number of *used external* (i.e. come from outside of the
    /// molecule) pins across all primitives in molecle.
    int num_used_ext_pins = 0;
    /// @brief Number of *used external* input pins across all primitives in
    ///        molecule.
    int num_used_ext_inputs = 0;
    /// @brief Number of *used external* output pins across all primitives in
    ///        molecule.
    int num_used_ext_outputs = 0;
};

/// @brief Sets of nets that cross the boundary of a molecule, partitioned by pin type.
struct t_molecule_external_nets {
    std::unordered_set<AtomNetId> ext_clock_nets;  ///< Nets driving a clock pin from outside the molecule.
    std::unordered_set<AtomNetId> ext_input_nets;  ///< Nets driving an input pin from outside the molecule.
    std::unordered_set<AtomNetId> ext_output_nets; ///< Nets driven by the molecule with at least one sink outside.
};

/**
 * @brief Class that performs prepacking.
 *
 * This class maintains the prepacking state, allowing the use of molecules
 * (prepacked atoms) while this object exists. After prepacking, every atom will
 * be part of a molecule (with a large number being part of single-atom
 * molecules).
 *
 * Molecules currently come from pack patterns in the architecture file. For
 * example, a 3-bit carry chain in most architectures would turn into a molecule
 * containing the 3 atoms forming the carry chain.
 *
 * To use the prepacker, call the init method with a complete atom netlist.
 * Then maintain this object (do not reset or destroy it) so long as the
 * molecules are needed.
 *
 * // Initialize device and atom netlist
 * // ...
 * Prepacker prepacker;
 * prepacker.init(atom_ctx.netlist(), device_ctx.logical_block_types);
 * // ...
 * // Use the prepacked molecules.
 * // ...
 * prepacker.reset(); // Or if the prepacker object is destroyed.
 * // Prepacked molecules can no longer be used beyond this point.
 *
 */
class Prepacker {
  public:
    // Iterator for the pack molecule IDs
    typedef typename vtr::vector_map<PackMoleculeId, PackMoleculeId>::const_iterator molecule_iterator;

    // Range for the pack molecule IDs
    typedef typename vtr::Range<molecule_iterator> molecule_range;

    // This class maintains pointers to internal data structures, and as such
    // should not be copied or moved (prevents unsafe accesses).
    Prepacker(const Prepacker&) = delete;
    Prepacker& operator=(const Prepacker&) = delete;

    // Destructor of the class.
    ~Prepacker();

    /**
     * @brief Constructor. Performs prepacking.
     *
     * Initializes the prepacker by performing prepacking and allocating the
     * necessary data structures.
     *
     *  @param atom_nlist           The atom netlist to prepack.
     *  @param models
     *  @param logical_block_types  A list of the logical block types on the device.
     */
    Prepacker(const AtomNetlist& atom_nlist,
              const LogicalModels& models,
              const std::vector<t_logical_block_type>& logical_block_types);

    /**
     * @brief A range of all prepacked molecules. Every atom should exist in one
     *        of these molecules.
     */
    inline molecule_range molecules() const {
        return vtr::make_range(pack_molecule_ids_.begin(), pack_molecule_ids_.end());
    }

    /**
     * @brief Get the cluster molecule containing the given atom block.
     *
     *  @param blk_id       The atom block to get the molecule of.
     */
    inline PackMoleculeId get_atom_molecule(AtomBlockId blk_id) const {
        // Safety debug to ensure the blk is valid and has a molecule entry.
        VTR_ASSERT_SAFE(blk_id.is_valid() && (size_t)blk_id < atom_molecule_.size());
        // Safety debug to ensure the molecule is valid
        VTR_ASSERT_DEBUG(atom_molecule_[blk_id].is_valid());
        return atom_molecule_[blk_id];
    }

    /**
     * @brief Get the expected lowest cost physical block graph node for the
     *        given atom block.
     *
     *  @param blk_id       The atom block to get the pb graph node of.
     */
    inline t_pb_graph_node* get_expected_lowest_cost_pb_gnode(AtomBlockId blk_id) const {
        // Safety debug to ensure the blk is valid and has an entry.
        VTR_ASSERT_SAFE(blk_id.is_valid() && (size_t)blk_id < expected_lowest_cost_pb_gnode.size());
        // Ensure the entry is valid.
        VTR_ASSERT(expected_lowest_cost_pb_gnode[blk_id] != nullptr);
        return expected_lowest_cost_pb_gnode[blk_id];
    }

    /**
     * @brief Search through all primitives and return the lowest cost primitive that fits this atom block.
     *
     * @param blk_id               Atom to find a primitive for.
     * @param logical_block_types  Block types to search through.
     */
    t_pb_graph_node* get_expected_lowest_cost_primitive_for_atom_block(const AtomBlockId blk_id,
                                                                       const std::vector<t_logical_block_type>& logical_block_types) const;

    /*
     * @brief Calculates molecule statistics for a single molecule.
     */
    t_molecule_stats calc_molecule_stats(PackMoleculeId molecule_id,
                                         const AtomNetlist& atom_netlist,
                                         const LogicalModels& models) const;

    /**
     * @brief Calculates maximum molecule statistics across all molecules,
     */
    t_molecule_stats calc_max_molecule_stats(const AtomNetlist& netlist,
                                             const LogicalModels& models) const;

    /**
     * @brief Returns the sets of nets that cross the boundary of a molecule,
     *        partitioned into clock, input, and output nets.
     *
     * @param molecule_id  Molecule to compute external nets for.
     * @param atom_nlist   Used to iterate pins and look up net drivers.
     *
     * Note: Complexity is O(A * P * S) in the worst case, where A is the number of
     *       atoms in the molecule, P is the maximum number of pins per atom, and S
     *       is the maximum fanout of an output net.
     */
    t_molecule_external_nets calc_molecule_external_nets(PackMoleculeId molecule_id,
                                                         const AtomNetlist& atom_nlist) const;

    /**
     * @brief Gets the largest number of blocks (atoms) that any molecule contains.
     */
    inline size_t get_max_molecule_size() const {
        size_t max_molecule_size = 1;
        for (const t_pack_molecule& molecule : pack_molecules_) {
            max_molecule_size = std::max<size_t>(max_molecule_size, molecule.atom_block_ids.size());
        }
        return max_molecule_size;
    }

    /**
     * @brief Get information about the molecule associated with the given ID.
     */
    inline const t_pack_molecule& get_molecule(PackMoleculeId molecule_id) const {
        VTR_ASSERT(molecule_id.is_valid());
        return pack_molecules_[molecule_id];
    }

    /**
     * @brief Get the root atom of this molecule.
     */
    inline AtomBlockId get_molecule_root_atom(PackMoleculeId molecule_id) const {
        VTR_ASSERT_SAFE_MSG(molecule_id.is_valid(), "Invalid molecule ID");
        const t_pack_molecule& mol = get_molecule(molecule_id);
        return mol.atom_block_ids[mol.root];
    }

    /**
     * @brief Get information about the chain associated with the given ID.
     */
    inline const t_chain_info& get_molecule_chain_info(MoleculeChainId chain_id) const {
        VTR_ASSERT(chain_id.is_valid());
        return chain_info_[chain_id];
    }

    /**
     * @brief Get the number of unique molecule chains from the prepacker.
     */
    inline size_t get_num_molecule_chains() const {
        return chain_info_.size();
    }

    /**
     * @brief Get a list of all the pack patterns in the architecture.
     */
    inline const std::vector<t_pack_patterns>& get_all_pack_patterns() const {
        return list_of_pack_patterns;
    }

  private:
    /**
     * Pre-pack atoms in netlist to molecules
     * 1.  Single atoms are by definition a molecule.
     * 2.  Forced pack molecules are groupings of atoms that matches a t_pack_pattern definition.
     * 3.  Chained molecules are molecules that follow a carry-chain style pattern,
     *     ie. a single linear chain that can be split across multiple complex blocks
     */
    void alloc_and_load_pack_molecules(std::multimap<AtomBlockId, PackMoleculeId>& atom_molecules_multimap,
                                       const AtomNetlist& atom_nlist,
                                       const LogicalModels& models,
                                       const std::vector<t_logical_block_type>& logical_block_types);

    /**
     * Given a pattern and an atom block to serve as the root block, determine if
     * the candidate atom block serving as the root node matches the pattern.
     * If yes, return the molecule with this atom block as the root, if not, return NULL
     *
     * Limitations: Currently assumes that forced pack nets must be single-fanout as
     *              this covers all the reasonable architectures we wanted. More complicated
     *              structures should probably be handled either downstream (general packing)
     *              or upstream (in tech mapping).
     *              If this limitation is too constraining, code is designed so that this limitation can be removed
     *
     * Side Effect: If successful, link atom to molecule
     */
    PackMoleculeId try_create_molecule(const int pack_pattern_index,
                                       AtomBlockId blk_id,
                                       std::multimap<AtomBlockId, PackMoleculeId>& atom_molecules_multimap,
                                       const AtomNetlist& atom_nlist);

  private:
    /**
     * @brief Collection of all molecule IDs. If an entry in this map is invalid
     *        it means that the molecule should be destroyed.
     */
    vtr::vector_map<PackMoleculeId, PackMoleculeId> pack_molecule_ids_;

    /**
     * @brief Lookup between each molecule ID and the information associated with
     *        that molecule.
     */
    vtr::vector_map<PackMoleculeId, t_pack_molecule> pack_molecules_;

    /**
     * @brief The molecules associated with each atom block.
     *
     * This vector is loaded in the init method and cleared in the reset method.
     * The pointers in this vector are shared with list_of_pack_molecules.
     */
    vtr::vector<AtomBlockId, PackMoleculeId> atom_molecule_;

    /// @brief A vector of the expected lowest cost physical block graph node.
    vtr::vector<AtomBlockId, t_pb_graph_node*> expected_lowest_cost_pb_gnode;

    /// @brief A list of the pack patterns used for prepacking. I think the
    ///        molecules keep pointers to this vector, so this needs to remain
    ///        for the lifetime of the molecules.
    std::vector<t_pack_patterns> list_of_pack_patterns;

    /**
     * @brief Lookup between each chain ID and the information associated with
     *        that chain.
     */
    vtr::vector<MoleculeChainId, t_chain_info> chain_info_;
};
