#include "EigenKalmanBBoxTrack.h"

namespace ByteTrackEigen {
    // Initializing the static shared Kalman filter
    KalmanFilter KalmanBBoxTrack::shared_kalman;

    /**
     * @brief Default constructor for KalmanBBoxTrack.
     *        Initializes member variables to their default state.
     */
    KalmanBBoxTrack::KalmanBBoxTrack() : BaseTrack(),
        _tlwh(Eigen::Vector4d::Zero()), // Initializes the bounding box to a zero vector
        kalman_filter(KalmanFilter()),  // Default initialization of the Kalman filter
        mean(Eigen::VectorXd()),        // Initializes the mean state vector to default
        covariance(Eigen::MatrixXd()),  // Initializes the covariance matrix to default
        tracklet_len(0)                 // Initializes the tracklet length to zero
    { }

    /**
     * @brief Constructor for KalmanBBoxTrack with initial bounding box and detection score.
     *        Throws an exception if the tlwh vector does not contain exactly 4 values.
     *
     * @param tlwh Initial bounding box in top-left width-height format.
     * @param score Detection score associated with the bounding box.
     */

    KalmanBBoxTrack::KalmanBBoxTrack(const std::vector<float> tlwh, float score, int class_id_,std::string object_id_):
        BaseTrack(score),
        _tlwh(tlwh.size() == 4 ? Eigen::Vector4d(tlwh[0], tlwh[1], tlwh[2], tlwh[3]) : Eigen::Vector4d::Zero()), // Ensures the bounding box has 4 elements
        kalman_filter(KalmanFilter()), // Initializes the Kalman filter
        mean(Eigen::VectorXd()),       // Initializes the mean state vector
        covariance(Eigen::MatrixXd()), // Initializes the covariance matrix
        tracklet_len(0),                // Initializes the tracklet length to zero
        class_id(class_id_),
        object_id(object_id_),
        detection_count_(1),
        class_id_locked_(false)
    {
        // Validation of tlwh size
        if (tlwh.size() != 4) {
            throw std::invalid_argument("tlwh vector must contain exactly 4 values.");
        }
        class_id_scores_[class_id_] = score;
    }

    void KalmanBBoxTrack::voteClassId(int new_class_id, float score)
    {
        if (class_id_locked_) return;

        class_id_scores_[new_class_id] += score;
        detection_count_++;

        int best_id = class_id;
        float best_score = 0.0f;
        for (const auto& entry : class_id_scores_)
        {
            if (entry.second > best_score)
            {
                best_score = entry.second;
                best_id = entry.first;
            }
        }
        class_id = best_id;

        if (detection_count_ >= CLASS_ID_LOCK_FRAMES)
        {
            class_id_locked_ = true;
        }
    }

    /**
     * @brief Static method to perform the prediction step on multiple KalmanBBoxTrack objects.
     *        Updates the mean and covariance of each track based on the shared Kalman filter.
     *
     * @param tracks Vector of shared pointers to KalmanBBoxTrack instances.
     */
    void KalmanBBoxTrack::multi_predict(std::vector<std::shared_ptr<KalmanBBoxTrack>>& tracks) {
        if (tracks.empty()) {
            return; // Early exit if no tracks to process
        }

        // Extract mean and covariance for each track
        std::vector<Eigen::VectorXd> multi_means;
        std::vector<Eigen::MatrixXd> multi_covariances;

        for (const auto track : tracks) {
            multi_means.push_back(track->mean);  // Eigen performs deep copy by default
            multi_covariances.push_back(track->covariance);
        }

        // For each track, set velocity to 0 if it's not in Tracked state
        for (size_t i = 0; i < tracks.size(); ++i) {
            if (tracks[i]->get_state() != TrackState::Tracked) {
                multi_means[i](7) = 0;  // Zero out the velocity if not tracked
            }
        }

        // Convert vectors to Eigen matrices for Kalman filter processing
        Eigen::MatrixXd means_matrix(multi_means.size(), multi_means[0].size());
        for (size_t i = 0; i < multi_means.size(); ++i) {
            means_matrix.row(i) = multi_means[i];
        }
        int n = (int)multi_covariances[0].rows();  // Assuming square matrices for covariance
        Eigen::MatrixXd covariances_matrix(n, n * multi_covariances.size());
        for (size_t i = 0; i < multi_covariances.size(); ++i) {
            covariances_matrix.middleCols(i * n, n) = multi_covariances[i];
        }

        // Use shared kalman filter for prediction
        Eigen::MatrixXd predicted_means, predicted_covariances;
        std::tie(predicted_means, predicted_covariances) = shared_kalman.multi_predict(means_matrix, covariances_matrix);

        // Update each track with the predicted mean and covariance
        for (size_t i = 0; i < tracks.size(); ++i) {
            tracks[i]->mean = predicted_means.col(i);
        }

        int pcr = (int)predicted_covariances.rows();  // Assuming the block matrix is composed of square matrices
        size_t num_matrices = predicted_covariances.cols() / pcr;

        for (size_t i = 0; i < num_matrices; ++i) {
            Eigen::MatrixXd block = predicted_covariances.middleCols(i * pcr, pcr);
            tracks[i]->covariance = block;
        }
    }

    /**
     * @brief Converts bounding box from tlwh format to xyah format (center x, center y, aspect ratio, height).
     *
     * @param tlwh Bounding box in tlwh format.
     * @return Eigen::VectorXd Bounding box in xyah format.
     */
    Eigen::VectorXd KalmanBBoxTrack::tlwh_to_xyah(Eigen::VectorXd tlwh) {
        Eigen::VectorXd ret = tlwh;
        ret.head(2) += ret.segment(2, 2) / 2.0; // Adjust x, y to be the center of the bounding box
        if (ret(3) > 0)
            ret(2) /= ret(3); // Update aspect ratio
        else
            ret(2) = 0;
        return ret;
    }

    /**
     * @brief Activates the track with a given Kalman filter and frame ID.
     *        Initializes the track ID and sets up the Kalman filter with the bounding box information.
     *
     * @param kalman_filter Kalman filter to be used for this track.
     * @param frame_id Frame ID at which this track is activated.
     */
    void KalmanBBoxTrack::activate(KalmanFilter& kalman_filter, int frame_id) {
        // Link to the provided Kalman filter
        this->kalman_filter = kalman_filter;

        // Initialize track ID
        this->track_id = BaseTrack::next_id();

        // Convert bounding box to the xyah format and initiate the Kalman filter
        std::tie(this->mean, this->covariance) = this->kalman_filter.initiate(this->tlwh_to_xyah(this->_tlwh));

        this->tracklet_len = 0;
        this->state = TrackState::Tracked;

        // Check if track is activated in the first frame
        this->is_activated = (frame_id == 1);
        this->frame_id = frame_id;
        this->start_frame = frame_id;
    }

    /**
     * @brief Updates the track with a new detection, potentially assigning a new ID.
     *
     * @param new_track The new detection represented as a KalmanBBoxTrack object.
     * @param frame_id Frame ID of the new detection.
     * @param new_id Flag indicating whether to assign a new ID to this track.
     */
    void KalmanBBoxTrack::update_track(const KalmanBBoxTrack& new_track, int frame_id, bool new_id) {
        // Track update logic
        this->frame_id = frame_id;
        this->tracklet_len++;

        std::tie(this->mean, this->covariance) = this->kalman_filter.update(this->mean, this->covariance, this->tlwh_to_xyah(new_track.tlwh()));

        this->state = TrackState::Tracked;
        this->is_activated = true;

        // Assigning a new ID if necessary
        if (new_id) {
            this->track_id = BaseTrack::next_id();
        }

        // Update the score and vote on class_id
        this->score = new_track.get_score();
        voteClassId(new_track.class_id, new_track.get_score());
        this->object_id = new_track.object_id;
    }

    /**
     * @brief Re-activates the track with a new detection, possibly assigning a new ID.
     *        Essentially a wrapper for the update_track method.
     *
     * @param new_track The new detection to reactivate the track with.
     * @param frame_id Frame ID of the reactivation.
     * @param new_id Flag indicating whether to assign a new ID.
     */
    void KalmanBBoxTrack::re_activate(const KalmanBBoxTrack& new_track, int frame_id, bool new_id) {
        update_track(new_track, frame_id, new_id);
    }

    /**
     * @brief Updates the track with a new detection without changing its ID.
     *        Wrapper for the update_track method without the new_id flag.
     *
     * @param new_track The new detection to update the track with.
     * @param frame_id Frame ID of the update.
     */
    void KalmanBBoxTrack::update(const KalmanBBoxTrack& new_track, int frame_id) {
        update_track(new_track, frame_id);
    }

    /**
     * @brief Converts tlwh bounding box to top-left bottom-right (tlbr) format.
     *
     * @param tlwh Bounding box in tlwh format.
     * @return Eigen::Vector4d Bounding box in tlbr format.
     */
    Eigen::Vector4d KalmanBBoxTrack::tlwh_to_tlbr(const Eigen::Vector4d tlwh) {
        Eigen::Vector4d ret = tlwh;
        ret.tail<2>() += ret.head<2>();
        return ret;
    }

    /**
     * @brief Converts tlbr bounding box to tlwh format.
     *
     * @param tlbr Bounding box in tlbr format.
     * @return Eigen::Vector4d Bounding box in tlwh format.
     */
    Eigen::Vector4d KalmanBBoxTrack::tlbr_to_tlwh(const Eigen::Vector4d tlbr) {
        Eigen::Vector4d ret = tlbr;
        ret.tail<2>() -= ret.head<2>();
        return ret;
    }

    /**
     * @brief Get current bounding box in tlwh format
     *
     * @return Eigen::Vector4d Bounding box in tlwh format.
     */
    Eigen::Vector4d KalmanBBoxTrack::tlwh() const {
        if (mean.isZero(0)) { // Checking if 'mean' is uninitialized or zero
            return _tlwh;
        }

        Eigen::Vector4d ret = mean.head(4);
        ret[2] *= ret[3];
        ret[0] -= ret[2] / 2.0;
        ret[1] -= ret[3] / 2.0;

        return ret;
    }

    /**
     * @brief Returns the bounding box in tlbr format.
     *
     * @return Eigen::Vector4d Bounding box in tlbr format.
     */
    Eigen::Vector4d KalmanBBoxTrack::tlbr() const {
        return this->tlwh_to_tlbr(this->tlwh());
    }
}