package com.qualcomm.qti.snpe.imageclassifiers.thread;

import com.qualcomm.qti.snpe.imageclassifiers.detector.Bbox;

import java.util.Collections;
import java.util.concurrent.CopyOnWriteArrayList;

class UnionMergeProcess extends PostProcessPrediction{
    CopyOnWriteArrayList<Bbox> sourceObjectPredictions;
    CopyOnWriteArrayList<Bbox> newSourceObjectPredictions;
    Bbox selectedObjectPredictions;
    CopyOnWriteArrayList<Bbox> selectedObjectPredictionsList = new CopyOnWriteArrayList<Bbox>();
    public UnionMergeProcess(float matchThreshold_, String matchMetric_, boolean classAgnostic_, CopyOnWriteArrayList<Bbox> sourceObjectPredictions_) {
        super(matchThreshold_, matchMetric_, classAgnostic_);
        this.sourceObjectPredictions = sourceObjectPredictions_;
    }

    public CopyOnWriteArrayList<Bbox>  UnionMerging(){
        while(this.sourceObjectPredictions.size() > 0){
            Collections.sort(sourceObjectPredictions, Collections.reverseOrder());
            selectedObjectPredictions = sourceObjectPredictions.get(0);
            sourceObjectPredictions.remove(0);
            newSourceObjectPredictions = new CopyOnWriteArrayList<Bbox>();
            for (int i = 0; i<sourceObjectPredictions.size();i++){
                if (this.hasMatch(selectedObjectPredictions,sourceObjectPredictions.get(i))){
                    selectedObjectPredictions = this.mergeObjectPredictionPair(selectedObjectPredictions,sourceObjectPredictions.get(i));
                }
                else{
                    newSourceObjectPredictions.add(sourceObjectPredictions.get(i));
                }
            }
            sourceObjectPredictions = newSourceObjectPredictions;
            selectedObjectPredictionsList.add(selectedObjectPredictions);
        }

        return  selectedObjectPredictionsList;
    }

    public Bbox mergeObjectPredictionPair(Bbox mPred1, Bbox mPred2){
        int[] mergedBoundingBox = this.getMergedBBox(mPred1,mPred2);
        float mergedScore = this.getMergedScore(mPred1,mPred2);
        String mergedCategory = this.getMergedCategory(mPred1,mPred2);
        return new Bbox(mergedBoundingBox[0],mergedBoundingBox[1],mergedBoundingBox[2],mergedBoundingBox[3],mergedScore,mergedCategory);
    }

    public int[] getMergedBBox(Bbox mPred1, Bbox mPred2){
        Bbox Bbox1 = mPred1;
        Bbox Bbox2 = mPred2;
        int[] BboxParams = this.calculateBoxUnion(Bbox1,Bbox2);
        return BboxParams;
    }

    public float getMergedScore(Bbox mPred1, Bbox mPred2){
        return Math.max(mPred1.conf,mPred2.conf);
    }

    public String getMergedCategory (Bbox mPred1, Bbox mPred2) {
        if (mPred1.conf > mPred2.conf){
            return mPred1.label;
        }
        else{
            return mPred2.label;
        }
    }

    public int[] calculateBoxUnion (Bbox Bbox1, Bbox Bbox2)
    {
        int[] Bbox = new int[]{(int) Math.min(Bbox1.x1,Bbox2.x1), (int) Math.min(Bbox1.y1,Bbox2.y1),
                (int) Math.max(Bbox1.x2,Bbox2.x2), (int) Math.max(Bbox1.y2,Bbox2.y2)
        };
        return Bbox;

    }


}