using OpenCVForUnity.CoreModule;
using OpenCVForUnity.DnnModule;
using OpenCVForUnity.ImgprocModule;
using OpenCVForUnity.UnityUtils;
using System;
using System.Collections.Generic;
using UnityEngine;
using UnityEngine.UIElements;
using Yoga;

public class OpenPoseEsimater : Estimator
{
    KeypointsModel _openPoseModel;
    private Net _net;

    private YOLOv7ObjectDetector _objectDetector;
    public int inpWidth = 416;
    public int inpHeight = 416;
    public float confThreshold = 0.35f;
    public float nmsThreshold = 0.6f;
    public int topK = 1000;

    private double threshold = 0.5;
    public override void InitModel()
    {
        _objectDetector = new YOLOv7ObjectDetector(
            Utils.getFilePath("OpenCVForUnity/dnn/yolov7-tiny.weights"),
            Utils.getFilePath("OpenCVForUnity/dnn/yolov7-tiny.cfg"),
            Utils.getFilePath("OpenCVForUnity/dnn/coco.names"),
            new Size(inpWidth, inpHeight), confThreshold, nmsThreshold/*, topK*/);

        _net = null;
        var modelFilePath = Utils.getFilePath(YogaConfig.MODEL_PATHS[ModelType.OpenPose]);
        if (string.IsNullOrEmpty(modelFilePath))
        {
            Debug.LogError("modelFilePath is empty. Please copy from “OpenCVForUnity/StreamingAssets/” to “Assets/StreamingAssets/” folder. ");
            return;
        }

        _net = Dnn.readNet(modelFilePath);

        _openPoseModel = new KeypointsModel(_net);
        _openPoseModel.setInputScale(new Scalar(YogaConfig.InScale));
        _openPoseModel.setInputSize(new Size(YogaConfig.InWidth, YogaConfig.InHeight));
        _openPoseModel.setInputMean(new Scalar(YogaConfig.InMean));
        _openPoseModel.setInputSwapRB(false);
        _openPoseModel.setInputCrop(false);
    }

    public override bool Esitmate(Mat bgrMat, Mat rgbaMat, out List<Point> points)
    {
        points = null;
        if (_objectDetector == null)
        {
            Debug.LogWarning("ObjectDetector is not ready. ");
            return false; //重新检测
        }

        //volo 人物检测
        //获取数据
        Mat results = _objectDetector.infer(bgrMat);

        var voloResultBox = new List<float[]>();

        bool hasValidObject = false;
        List<List<float[]>> boxList = new();

        for (int i = results.rows() - 1; i >= 0; --i)
        {
            float[] box = new float[4];
            results.get(i, 0, box); //方框
            float[] conf = new float[1];
            results.get(i, 4, conf); //检测数据
            float[] cls = new float[1];
            results.get(i, 5, cls); //类别

            if (!IsObjectValid(box, conf, cls, rgbaMat)) //不符合规则的，跳过
                continue;
            
            hasValidObject = true;
            List<float[]> item= new List<float[]>();
            item.Add(box);
            item.Add(conf);
            item.Add(cls);
            boxList.Add(item);
        }

        float area = 0;
        List<float[]> maxSizeBox = new List<float[]>();
        //选择最大的一个
        foreach(var box in boxList)
        {
            var rect = box[0];
            var tmp = (rect[2] - rect[0]) * (rect[3] - rect[1]);
            area = Math.Max(area, tmp);
            if (area == tmp)
            {
                maxSizeBox = box;
            }
        }
        voloResultBox.Clear();
        voloResultBox = maxSizeBox;

        if (!hasValidObject) //没有检测到人体
        {
            Debug.Log("No person block found. Re-Estimation.");
            return false; //重新检测
        }

        //更新人物框检测结果
        YogaManager.Instance.VoloResult = voloResultBox;

        points = null;

        if (_openPoseModel == null)
            throw new ArgumentException("CVEstimator.Init: args[1] is not KeypointsModel")
;
        OpenCVForUnity.CoreModule.Rect roiRect = new OpenCVForUnity.CoreModule.Rect(
                (int)voloResultBox[0][0],
                (int)voloResultBox[0][1],
                Math.Abs((int)(voloResultBox[0][2] - voloResultBox[0][0])),
                Math.Abs((int)(voloResultBox[0][3] - voloResultBox[0][1])));

        if (roiRect.y < 0 || //0 <= _rowRange.start
            (roiRect.y + roiRect.height) < roiRect.y || // _rowRange.start <= _rowRange.end
            bgrMat.rows() < (roiRect.y + roiRect.height)) //_rowRange.end <= m.rows 
            return false; //重新检测

        Mat personRectImg = new Mat(bgrMat, roiRect);//获取人体区域
        points = _openPoseModel.estimate(personRectImg, (float)threshold).toList();
        //将人体区域的坐标转换为原图坐标
        for (int j = 0; j < points.Count; j++)
        {
            if (points[j] == null ||
                (points[j].x == -1 && points[j].y == -1)) //没找到的点，跳过
                continue;

            points[j].x += roiRect.x;
            points[j].y += roiRect.y;
        }

        if (!YogaManager.Instance.ActionCheckPoints(points))
        {
            Debug.Log("ActionCheckPoints failed. Re-Estimation.");
            return false; //重新检测
        }
        return true;
    }
    public override void DisposeModel()
    {
        if (_net != null)
            _net.Dispose();
    }


    public override void Check(ref Mat img)
    {
        if (_net == null)
        {
            Imgproc.putText(img, "model file is not loaded.", new Point(5, img.rows() - 30), Imgproc.FONT_HERSHEY_SIMPLEX, 0.7, new Scalar(255, 255, 255), 2, Imgproc.LINE_AA, false);
            Imgproc.putText(img, "Please read console message.", new Point(5, img.rows() - 10), Imgproc.FONT_HERSHEY_SIMPLEX, 0.7, new Scalar(255, 255, 255), 2, Imgproc.LINE_AA, false);
        }
    }

    public override void DebugPrint(Mat image, bool isRGB = false)
    {
        if (YogaManager.Instance.VoloResult.Count > 2)
            return;

        var box = YogaManager.Instance.VoloResult[0];
        var conf = YogaManager.Instance.VoloResult[1];
        var cls = YogaManager.Instance.VoloResult[2];

        float left = box[0];
        float top = box[1];
        float right = box[2];
        float bottom = box[3];
        int classId = (int)cls[0];

        Scalar c = _objectDetector.palette[classId % _objectDetector.palette.Count];
        Scalar color = isRGB ? c : new Scalar(c.val[2], c.val[1], c.val[0], c.val[3]);

        Imgproc.rectangle(image, new Point(left, top), new Point(right, bottom), color, 2);

        string label = String.Format("{0:0.00}", conf[0]);
        if (_objectDetector.classNames != null && _objectDetector.classNames.Count != 0)
        {
            if (classId < (int)_objectDetector.classNames.Count)
            {
                label = _objectDetector.classNames[classId] + " " + label;
            }
        }

        int[] baseLine = new int[1];
        Size labelSize = Imgproc.getTextSize(label, Imgproc.FONT_HERSHEY_SIMPLEX, 0.5, 1, baseLine);

        top = Mathf.Max((float)top, (float)labelSize.height);
        Imgproc.rectangle(image, new Point(left, top - labelSize.height),
            new Point(left + labelSize.width, top + baseLine[0]), color, Core.FILLED);
        Imgproc.putText(image, label, new Point(left, top), Imgproc.FONT_HERSHEY_SIMPLEX, 0.5, Scalar.all(255), 1, Imgproc.LINE_AA);
    }
}