const fs = require('fs');
const path = require('path');
const ArgumentParser = require('argparse');
const cv = require('opencv');
const yaml = require('js-yaml');
const { Pose, POSE_LANDMARKS } = require('@mediapipe/pose');

const logging = console;
const warnings = console;

class Landmarker {
    static resizedHeight = 256;
    static resizedWidth = 256;

    constructor() {
        this.args = this.parseArgs();
        this.measurements = this.loadLandmarks();
        if (!this.args.frontImage) {
            throw new Error("Front image needs to be passed");
        }
        if (!this.args.sideImage) {
            throw new Error("Side image needs to be passed");
        }

        this.frontImage = cv.imread(this.args.frontImage);
        this.sideImage = cv.imread(this.args.sideImage);

        this.frontImageResized = this.frontImage.resize(Landmarker.resizedHeight, Landmarker.resizedWidth);
        this.sideImageResized = this.sideImage.resize(Landmarker.resizedHeight, Landmarker.resizedWidth);

        this.distances = {};

        this.personHeight = this.args.personHeight;
        this.pixelHeight = this.args.pixelHeight;

        this.pose = new Pose({
            locateFile: (file) => {
                return `https://cdn.jsdelivr.net/npm/@mediapipe/pose/${file}`;
            }
        });

        this.landmarksIndices = [
            POSE_LANDMARKS.LEFT_SHOULDER,
            POSE_LANDMARKS.RIGHT_SHOULDER,
            POSE_LANDMARKS.LEFT_ELBOW,
            POSE_LANDMARKS.RIGHT_ELBOW,
            POSE_LANDMARKS.LEFT_WRIST,
            POSE_LANDMARKS.RIGHT_WRIST,
            POSE_LANDMARKS.LEFT_HIP,
            POSE_LANDMARKS.RIGHT_HIP,
            POSE_LANDMARKS.LEFT_KNEE,
            POSE_LANDMARKS.RIGHT_KNEE,
            POSE_LANDMARKS.LEFT_ANKLE,
            POSE_LANDMARKS.RIGHT_ANKLE
        ];
    }

    loadLandmarks() {
        const file = fs.readFileSync(this.args.yamlFile, 'utf8');
        const landmarksData = yaml.load(file);
        const measurements = {};
        for (const measurement of landmarksData.measurements) {
            measurements[measurement.name] = measurement.landmarks;
        }
        return measurements;
    }

    parseArgs() {
        const parser = new ArgumentParser();
        parser.add_argument('--front', { dest: 'frontImage', type: 'str', help: 'Front image' });
        parser.add_argument('--side', { dest: 'sideImage', type: 'str', help: 'Side image' });
        parser.add_argument('--poseDetectionConfidence', { dest: 'poseDetectionConfidence', default: 0.5, type: 'float', help: 'Confidence score for pose detection' });
        parser.add_argument('--poseTrackingConfidence', { dest: 'poseTrackingConfidence', default: 0.5, type: 'float', help: 'Confidence score for pose tracking' });
        parser.add_argument('--personHeight', { dest: 'personHeight', type: 'int', help: 'Person height in cm' });
        parser.add_argument('--pixelHeight', { dest: 'pixelHeight', type: 'int', help: 'Pixel height of person' });
        parser.add_argument('--measurement', { dest: 'measurement', nargs: '+', type: 'str', help: 'Type of measurement' });
        parser.add_argument('--yamlFile', { dest: 'yamlFile', type: 'str', help: 'Path to the YAML file containing landmarks' });
        return parser.parse_args();
    }

    async run() {
        await this.pose.initialize();
        const { frontResults, sideResults } = await this.processImages();

        this.getCenterTopPoint(sideResults);

        const table = [];
        if (this.args.measurement) {
            for (const m of this.args.measurement) {
                if (!this.measurements[m]) {
                    throw new Error("Incorrect input (input not present in config.yml)");
                } else {
                    const distance = this.calculateDistanceBetweenLandmarks(frontResults, m);
                    table.push([m, distance]);
                }
            }
        } else {
            for (const m in this.measurements) {
                const distance = this.calculateDistanceBetweenLandmarks(frontResults, m);
                table.push([m, distance]);
            }
        }

        console.table(table);

        this.pose.close();
    }

    async processImages() {
        const frontResults = await this.pose.estimatePoses(this.frontImageResized);
        const sideResults = await this.pose.estimatePoses(this.sideImageResized);

        this.sideImageKeypoints = this.sideImageResized.copy();
        this.frontImageKeypoints = this.frontImageResized.copy();

        if (frontResults[0].landmarks) {
            this.drawLandmarks(this.frontImageKeypoints, frontResults[0].landmarks, this.landmarksIndices);
        }
        if (sideResults[0].landmarks) {
            this.drawLandmarks(this.sideImageKeypoints, sideResults[0].landmarks, this.landmarksIndices);
        }
        return {
            frontResults: frontResults[0],
            sideResults: sideResults[0]
        };
    }

    pixelToMetricRatio() {
        this.pixelHeight = this.pixelDistance * 2;
        const pixelToMetricRatio = this.personHeight / this.pixelHeight;
        logging.debug("pixelToMetricRatio %s", pixelToMetricRatio);
        return pixelToMetricRatio;
    }

    drawLandmarks(image, landmarks, indices) {
        for (const idx of indices) {
            const landmark = landmarks[idx];
            const h = image.rows;
            const w = image.cols;
            const cx = Math.round(landmark.x * w);
            const cy = Math.round(landmark.y * h);
            this.circle(image, cx, cy);
        }
    }

    circle(image, cx, cy) {
        return image.drawCircle(new cv.Point(cx, cy), 2, new cv.Vec(255, 0, 0), -1);
    }

    calculateDistanceBetweenLandmarks(frontResults, measurementName) {
        if (!frontResults.landmarks) {
            return;
        }

        const landmarks = frontResults.landmarks;
        const landmarkNames = this.measurements[measurementName];

        let totalDistance = 0;
        for (let i = 0; i < landmarkNames.length - 1; i++) {
            const current = landmarks[landmarkNames[i]];
            const next = landmarks[landmarkNames[i + 1]];
            const pixelDistance = this.euclideanDistance(
                current.x * Landmarker.resizedWidth,
                current.y * Landmarker.resizedHeight,
                next.x * Landmarker.resizedWidth,
                next.y * Landmarker.resizedHeight
            );
            const realDistance = pixelDistance * this.pixelToMetricRatio();
            totalDistance += realDistance;
        }
        return totalDistance;
    }

    euclideanDistance(x1, y1, x2, y2) {
        return Math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2);
    }

    getCenterTopPoint(sideResults) {
        const grayImage = this.sideImageKeypoints.cvtColor(cv.COLOR_BGR2GRAY);
        const blurredImage = grayImage.gaussianBlur(new cv.Size(5, 5), 0);
        const roi = blurredImage.getRegion(new cv.Rect(0, 0, this.sideImageResized.cols, Math.floor(this.sideImageResized.rows / 2)));
        this.edges = roi.canny(50, 150);
        const contours = this.edges.findContours(cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE);

        let xt, yt;
        this.topmostPoint = null;

        if (contours.length > 0) {
            const largestContour = contours.sort((a, b) => b.area - a.area)[0];
            this.topmostPoint = largestContour.minBy(pt => pt.y);
            xt = this.topmostPoint.x;
            yt = this.topmostPoint.y;

            this.circle(this.sideImageKeypoints, xt, yt);
        }

        let xc, yc;
        if (sideResults.landmarks) {
            const leftHip = sideResults.landmarks[POSE_LANDMARKS.LEFT_HIP];
            const rightHip = sideResults.landmarks[POSE_LANDMARKS.RIGHT_HIP];
            const centerPoint = [
                (leftHip.x + rightHip.x) / 2,
                (leftHip.y + rightHip.y) / 2
            ];
            xc = Math.round(centerPoint[0] * this.sideImageResized.cols);
            yc = Math.round(centerPoint[1] * this.sideImageResized.rows);
            this.circle(this.sideImageKeypoints, xc, yc);
        }

        this.pixelDistance = this.euclideanDistance(xt, yt, xc, yc);
        this.pixelDistance *= 2;
    }
}

(async () => {
    const landmarker = new Landmarker();
    await landmarker.run();
})();