//   Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// The code is based on:
// https://github.com/CnybTseng/JDE/blob/master/platforms/common/trajectory.h
// Ths copyright of CnybTseng/JDE is as follows:
// MIT License

#pragma once

#include <vector>
#include <opencv2/opencv.hpp>

namespace PaddleDetection {

typedef enum
{
    New = 0,
    Tracked = 1,
    Lost = 2,
    Removed = 3
} TrajectoryState;

class Trajectory;
typedef std::vector<Trajectory> TrajectoryPool;
typedef std::vector<Trajectory>::iterator TrajectoryPoolIterator;
typedef std::vector<Trajectory *>TrajectoryPtrPool;
typedef std::vector<Trajectory *>::iterator TrajectoryPtrPoolIterator;

class TKalmanFilter : public cv::KalmanFilter
{
public:
    TKalmanFilter(void);
    virtual ~TKalmanFilter(void) {}
    virtual void init(const cv::Mat &measurement);
    virtual const cv::Mat &predict();
    virtual const cv::Mat &correct(const cv::Mat &measurement);
    virtual void project(cv::Mat &mean, cv::Mat &covariance) const;
private:
    float std_weight_position;
    float std_weight_velocity;
};

inline TKalmanFilter::TKalmanFilter(void) : cv::KalmanFilter(8, 4)
{
    cv::KalmanFilter::transitionMatrix = cv::Mat::eye(8, 8, CV_32F);
    for (int i = 0; i < 4; ++i)
        cv::KalmanFilter::transitionMatrix.at<float>(i, i + 4) = 1;
    cv::KalmanFilter::measurementMatrix = cv::Mat::eye(4, 8, CV_32F);
    std_weight_position = 1/20.f;
    std_weight_velocity = 1/160.f;
}

class Trajectory : public TKalmanFilter
{
public:
    Trajectory();
    Trajectory(cv::Vec4f &ltrb, float score, const cv::Mat &embedding);
    Trajectory(const Trajectory &other);
    Trajectory &operator=(const Trajectory &rhs);
    virtual ~Trajectory(void) {};

    static int next_id();
    virtual const cv::Mat &predict(void);
    virtual void update(Trajectory &traj, int timestamp, bool update_embedding=true);
    virtual void activate(int timestamp);
    virtual void reactivate(Trajectory &traj, int timestamp, bool newid=false);
    virtual void mark_lost(void);
    virtual void mark_removed(void);

    friend TrajectoryPool operator+(const TrajectoryPool &a, const TrajectoryPool &b);
    friend TrajectoryPool operator+(const TrajectoryPool &a, const TrajectoryPtrPool &b);
    friend TrajectoryPool &operator+=(TrajectoryPool &a, const TrajectoryPtrPool &b);
    friend TrajectoryPool operator-(const TrajectoryPool &a, const TrajectoryPool &b);
    friend TrajectoryPool &operator-=(TrajectoryPool &a, const TrajectoryPool &b);
    friend TrajectoryPtrPool operator+(const TrajectoryPtrPool &a, const TrajectoryPtrPool &b);
    friend TrajectoryPtrPool operator+(const TrajectoryPtrPool &a, TrajectoryPool &b);
    friend TrajectoryPtrPool operator-(const TrajectoryPtrPool &a, const TrajectoryPtrPool &b);

    friend cv::Mat embedding_distance(const TrajectoryPool &a, const TrajectoryPool &b);
    friend cv::Mat embedding_distance(const TrajectoryPtrPool &a, const TrajectoryPtrPool &b);
    friend cv::Mat embedding_distance(const TrajectoryPtrPool &a, const TrajectoryPool &b);

    friend cv::Mat mahalanobis_distance(const TrajectoryPool &a, const TrajectoryPool &b);
    friend cv::Mat mahalanobis_distance(const TrajectoryPtrPool &a, const TrajectoryPtrPool &b);
    friend cv::Mat mahalanobis_distance(const TrajectoryPtrPool &a, const TrajectoryPool &b);

    friend cv::Mat iou_distance(const TrajectoryPool &a, const TrajectoryPool &b);
    friend cv::Mat iou_distance(const TrajectoryPtrPool &a, const TrajectoryPtrPool &b);
    friend cv::Mat iou_distance(const TrajectoryPtrPool &a, const TrajectoryPool &b);
private:   
    void update_embedding(const cv::Mat &embedding);
public:
    TrajectoryState state;
    cv::Vec4f ltrb;
    cv::Mat smooth_embedding;
    int id;
    bool is_activated;
    int timestamp;
    int starttime;
    float score;
private:
    static int count;
    cv::Vec4f xyah;
    cv::Mat current_embedding;
    float eta;
    int length;
};

inline cv::Vec4f ltrb2xyah(cv::Vec4f &ltrb)
{
    cv::Vec4f xyah;
    xyah[0] = (ltrb[0] + ltrb[2]) * 0.5f;
    xyah[1] = (ltrb[1] + ltrb[3]) * 0.5f;
    xyah[3] =  ltrb[3] - ltrb[1];
    xyah[2] = (ltrb[2] - ltrb[0]) / xyah[3];
    return xyah;
}

inline Trajectory::Trajectory() :
    state(New), ltrb(cv::Vec4f()), smooth_embedding(cv::Mat()), id(0),
    is_activated(false), timestamp(0), starttime(0), score(0), eta(0.9), length(0)
{
}

inline Trajectory::Trajectory(cv::Vec4f &ltrb_, float score_, const cv::Mat &embedding) :
    state(New), ltrb(ltrb_), smooth_embedding(cv::Mat()), id(0),
    is_activated(false), timestamp(0), starttime(0), score(score_), eta(0.9), length(0)
{
    xyah = ltrb2xyah(ltrb);
    update_embedding(embedding);
}

inline Trajectory::Trajectory(const Trajectory &other):
    state(other.state), ltrb(other.ltrb), id(other.id), is_activated(other.is_activated),
    timestamp(other.timestamp), starttime(other.starttime), xyah(other.xyah),
    score(other.score), eta(other.eta), length(other.length)
{    
    other.smooth_embedding.copyTo(smooth_embedding);
    other.current_embedding.copyTo(current_embedding);
    // copy state in KalmanFilter
    
    other.statePre.copyTo(cv::KalmanFilter::statePre);
    other.statePost.copyTo(cv::KalmanFilter::statePost);
    other.errorCovPre.copyTo(cv::KalmanFilter::errorCovPre);
    other.errorCovPost.copyTo(cv::KalmanFilter::errorCovPost);
    
}

inline Trajectory &Trajectory::operator=(const Trajectory &rhs)
{
    this->state = rhs.state;
    this->ltrb = rhs.ltrb;
    rhs.smooth_embedding.copyTo(this->smooth_embedding);
    this->id = rhs.id;
    this->is_activated = rhs.is_activated;
    this->timestamp = rhs.timestamp;
    this->starttime = rhs.starttime; 
    this->xyah = rhs.xyah;
    this->score = rhs.score;
    rhs.current_embedding.copyTo(this->current_embedding);            
    this->eta = rhs.eta;    
    this->length = rhs.length;

    // copy state in KalmanFilter
    
    rhs.statePre.copyTo(cv::KalmanFilter::statePre);
    rhs.statePost.copyTo(cv::KalmanFilter::statePost);
    rhs.errorCovPre.copyTo(cv::KalmanFilter::errorCovPre);
    rhs.errorCovPost.copyTo(cv::KalmanFilter::errorCovPost);
    
    return *this;
}

inline int Trajectory::next_id()
{
    ++count;
    return count;
}

inline void Trajectory::mark_lost(void)
{
    state = Lost;
}

inline void Trajectory::mark_removed(void)
{
    state = Removed;
}

}   // namespace PaddleDetection