#pragma once

#include <opencv2/opencv.hpp>
#include <cmath>
#include <iostream>
#include "UCMClapjv.h"
#include "UCMCKalman.h"

namespace UCMCKalman {

    void lapjv(const Eigen::MatrixXd& costMatrix, double costLimit,
        std::vector<int>& x, std::vector<int>& y) {
        //
        int n_rows = costMatrix.rows();
        int n_cols = costMatrix.cols();
        x.clear();
        y.clear();
        x.resize(n_rows);
        y.resize(n_cols);
        const int n = n_rows + n_cols;

        costLimit /= 2.;

        double** cost_ptr = new double* [n];
        for (int i = 0; i < n; i++) {
            cost_ptr[i] = new double[n];
            for (int j = 0; j < n; j++) {
                if (i < n_rows && j < n_cols) {
                    cost_ptr[i][j] = costMatrix(i, j);
                }
                else if (i >= n_rows && j >= n_cols) {
                    cost_ptr[i][j] = 0.;
                }
                else {
                    cost_ptr[i][j] = costLimit;
                }
            }
        }

        UCMC::int_t* x_c = new UCMC::int_t[n];
        UCMC::int_t* y_c = new UCMC::int_t[n];
        UCMC::lapjv_internal(n, cost_ptr, &x_c[0], &y_c[0]);

        for (int i = 0; i < n; i++) delete[] cost_ptr[i];
        delete[] cost_ptr;

        for (int i = 0; i < n_rows; i++) {
            x[i] = (x_c[i] >= n_cols) ? -1 : x_c[i];
        }

        for (int i = 0; i < n_cols; i++) {
            y[i] = (y_c[i] >= n_rows) ? -1 : y_c[i];
        }
        delete[] x_c;
        delete[] y_c;

    }


    void linearAssignment(const Eigen::MatrixXd& costMatrix, double thresh,
        std::vector<std::vector<int>>& matches,
        std::vector<int>& unmatched_a,
        std::vector<int>& unmatched_b) {

        matches.clear();
        unmatched_a.clear();
        unmatched_b.clear();

        int n = costMatrix.rows();
        int m = costMatrix.cols();

        if (n == 0 || m == 0) {
            for (int i = 0; i < n; i++) unmatched_a.push_back(i);
            for (int i = 0; i < m; i++) unmatched_b.push_back(i);
            return;
        }

        std::vector<int> x, y;

        lapjv(costMatrix, thresh, x, y);

        for (int i = 0; i < n; ++i) {
            if (x[i] < 0) {
                unmatched_a.push_back(i);
            }
            else {
                std::vector<int> this_match;
                this_match.push_back(i);
                this_match.push_back(x[i]);
                matches.push_back(this_match);
            }
        }
        for (int j = 0; j < m; ++j) {
            if (y[j] < 0) {
                unmatched_b.push_back(j);
            }
        }
    }


    Filter:: Filter(int dim_x, int dim_z, int dim_u)
            : dim_x(dim_x), dim_z(dim_z), dim_u(dim_u),
            x(dim_x, 1), P(dim_x, dim_x), Q(dim_x, dim_x),
            F(dim_x, dim_x), H(dim_z, dim_x), R(dim_z, dim_z),
            _alpha_sq(1.0), M(dim_z, dim_z), z(dim_z, 1),
            K(dim_x, dim_z), y(dim_z, 1), S(dim_z, dim_z),
            SI(dim_z, dim_z), _I(dim_x, dim_x),
            x_prior(dim_x, 1), P_prior(dim_x, dim_x),
            x_post(dim_x, 1), P_post(dim_x, dim_x) {
            x.setZero();
            P.setIdentity();
            Q.setIdentity();
            F.setIdentity();
            H.setZero();
            R.setIdentity();
            M.setZero();
            z.setOnes()* std::numeric_limits<float>::quiet_NaN();
            K.setZero();
            y.setZero();
            S.setZero();
            SI.setZero();
            _I.setIdentity();
            x_prior = x;
            P_prior = P;
            x_post = x;
            P_post = P;
        }
    void Filter::predict() {
            x = F * x;
            P = _alpha_sq * (F * P * F.transpose()) + Q;
            x_prior = x;
            P_prior = P;
        }
    void  Filter::update(const Eigen::VectorXd& z,
            const Eigen::MatrixXd& R) {
            if (R.rows() != dim_z || R.cols() != dim_z) {
                throw std::invalid_argument("R matrix has incorrect dimensions");
            }

            // Update state
            y = z - H * x;
            Eigen::MatrixXd PHT = P_prior * H.transpose();
            S = H * PHT + R;
            SI = S.inverse();
            K = PHT * S.inverse();
            x += K * y;

            // Update covariance
            Eigen::MatrixXd I_KH = _I - K * H;
            P = (I_KH * P_prior) * I_KH.transpose() + (K * R) * K.transpose();

            // Prepare for next prediction
            this->z = z;
            x_post = x;
            P_post = P;

        }

    Eigen::VectorXd  Filter::getState() const {
            return x;
        }
    Eigen::MatrixXd  Filter::getCovariance() const {
            return P;
        }



    void Tracker::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;
        }
    }

    Tracker::Tracker(Eigen::Vector2d y, Eigen::Matrix2d R, double wx, double wy,
        double vmax, double w, double h, double dt, int tracker_count)
        : id(tracker_count), age(0), death_count(0), birth_count(0), detidx(-1),
        w(w), h(h), status(TrackStatus::Tentative), kf(4, 2),
        class_id_(0), detection_count_(0), class_id_locked_(false) {
        kf.F = Eigen::Matrix4d::Identity();
        kf.F(0, 1) = dt;
        kf.F(2, 3) = dt;
        kf.H = Eigen::Matrix<double, 2, 4>::Zero();
        kf.H(0, 0) = 1;
        kf.H(1, 2) = 1;
        kf.R = R;
        kf.P.setZero(); //= Eigen::Matrix<double, 4, 4>::Zero();
        kf.P << 1, 0, 0, 0,
            0, vmax* vmax / 3.0, 0, 0,
            0, 0, 1, 0,
            0, 0, 0, vmax* vmax / 3.0;

        // INFO << "P:\n" << kf.P << ENDL;

        Eigen::Matrix<double, 4, 2> G;
        G << 0.5 * dt * dt, 0,
            dt, 0,
            0, 0.5 * dt * dt,
            0, dt;
        Eigen::Matrix<double, 2, 2> Q0;
        Q0.setZero();
        Q0(0, 0) = wx;
        Q0(1, 1) = wy;

        Eigen::MatrixXd _Q = (G * Q0) * G.transpose();
        kf.Q = _Q;

        kf.x << y(0), 0, y(1), 0;
    }

    void Tracker::update(const Eigen::Vector2d& y, Eigen::Matrix2d& R) {
        kf.update(y, R);
    }

    Eigen::Vector2d Tracker::predict() {
        kf.predict();
        this->age += 1;
        return kf.H * kf.x;
    }

    Eigen::Vector4d Tracker::getState() const {
        return kf.x;
    }

    double Tracker::distance(const Eigen::Vector2d& y, Eigen::Matrix2d& R, bool show = false) const {
        Eigen::Vector2d diff = y - (kf.H * kf.x);
        Eigen::Matrix2d S = kf.H * (kf.P * kf.H.transpose()) + R;
        Eigen::Matrix2d SI = S.inverse();
        auto mahalanobis = diff.transpose() * (SI * diff);
        double logdet = log(S.determinant());

        return mahalanobis(0, 0) + logdet;
    }
}