package com.google.ar.sceneform.samples.hellosceneform;

import android.content.Context;
import android.graphics.Bitmap;
import android.media.Image;

import com.google.ar.core.CameraIntrinsics;
import com.google.ar.core.Frame;
import com.google.ar.core.PointCloud;
import com.google.ar.core.Pose;
import com.google.ar.core.exceptions.NotYetAvailableException;

import org.json.JSONArray;
import org.json.JSONException;
import org.json.JSONObject;

import java.nio.FloatBuffer;
import java.text.DateFormat;
import java.text.SimpleDateFormat;
import java.util.ArrayList;
import java.util.Date;
import java.util.Random;

import static java.lang.Math.sqrt;

public class Helper {

    public static JSONObject getPointCloudFromFrame(Frame frame, int hashCode, Pose anchorPos){
        JSONObject pointsJson = new JSONObject();
        PointCloud pointCloud = frame.acquirePointCloud();
        FloatBuffer pointsBuffer = pointCloud.getPoints();
        float[] points = new float[pointsBuffer.remaining()];
        pointsBuffer.get(points);
        for (int i = 0; i < points.length; i += 4){
            float x = points[i];
            float y = points[i+1];
            float z = points[i+2];
            float[] pCoord = {x, y, z};
            float[] ndcCoord = anchorPos.transformPoint(pCoord);
            points[i] = ndcCoord[0];
            points[i+1] = ndcCoord[1];
            points[i+2] = ndcCoord[2];
        }
        try {
            JSONArray pointsArr = new JSONArray(points);
            pointsJson.put("pointcloud", pointsArr);
            pointsJson.put("hashCode", hashCode);
        } catch (JSONException e) {
            e.printStackTrace();
        }
        return pointsJson;
    }

    public static JSONObject get2DImageFromFrame(Frame frame, CameraIntrinsics imgIntrinsic, Context context, int hashCode, int requestId){
        float[] focal = imgIntrinsic.getFocalLength();
        float[] pp = imgIntrinsic.getPrincipalPoint();
        DateFormat dateFormat = new SimpleDateFormat("yyyy_MM_dd_HH_mm_ss");
        Date date = new Date();
        JSONObject imgData = new JSONObject();
        try {
            Image img = frame.acquireCameraImage();
            Bitmap bmp = HelperYuv.YUV_420_888_toRGBIntrinsics(img, context);
            String imgStr = HelperYuv.BitmapToString(bmp);
            Random rand = new Random();
            String imgName = "camera_image" + dateFormat.format(date) + "_" + rand.nextInt(100);
            imgData.put("name", imgName);
            imgData.put("image", imgStr);
            JSONArray focalJSON = new JSONArray(focal);
            imgData.put("focal", focalJSON);
            JSONArray ppJSON = new JSONArray(pp);
            imgData.put("pp", ppJSON);
            imgData.put("hashCode", hashCode);
            imgData.put("requestId", requestId);
        } catch (NotYetAvailableException | JSONException e) {
            e.printStackTrace();
        }
        return imgData;
    }

    public static JSONObject getPoseDataFromFrame(float[] matrix, float[] projectionMatrix, int hashCode){
        JSONObject data = new JSONObject();
        try {
            JSONArray matrixJson = new JSONArray(matrix);
            JSONArray matrixPJson = new JSONArray(projectionMatrix);
            data.put("matrix", matrixJson);
            data.put("matrixProjection", matrixPJson);
            data.put("hashCode", hashCode);
        } catch (JSONException e) {
            e.printStackTrace();
        }
        return data;
    }

    public static float[] JSONArrayToFloatArray(JSONArray jsonArray){
        float[] fData = new float[jsonArray.length()];
        for (int i = 0; i < jsonArray.length(); i++) {
            try {
                fData[i] = Float.parseFloat(jsonArray.getString(i));
            } catch (JSONException e) {
                e.printStackTrace();
            }
        }
        return fData;
    }

    public static Pose matrixToPose(float[] matrix){
        float[] translation = new float[3];
        float[] rotationQuaternion = Helper.rotationToQuaternion(matrix);
        translation[0] = matrix[12];
        translation[1] = matrix[13];
        translation[2] = matrix[14];
        return new Pose(translation, rotationQuaternion);
    }
    // http://www.euclideanspace.com/maths/geometry/rotations/conversions/matrixToQuaternion/
    private static float[] rotationToQuaternion(float[] matrix) {
//        0 4 8 12
//        1 5 9 13
//        2 6 10 14
//        3 7 11 15
        float m00 = matrix[0];
        float m10 = matrix[1];
        float m20 = matrix[2];
        float m01 = matrix[4];
        float m11 = matrix[5];
        float m21 = matrix[6];
        float m02 = matrix[8];
        float m12 = matrix[9];
        float m22 = matrix[10];
        float tr = m00 + m11 + m22;
        float qx, qy, qz, qw;

        if (tr > 0) {
            float S = (float) (sqrt(tr+1.0) * 2); // S=4*qw
            qw = (float) (0.25 * S);
            qx = (m21 - m12) / S;
            qy = (m02 - m20) / S;
            qz = (m10 - m01) / S;
        } else if ((m00 > m11)&(m00 > m22)) {
            float S = (float) (sqrt(1.0 + m00 - m11 - m22) * 2); // S=4*qx
            qw = (m21 - m12) / S;
            qx = (float) (0.25 * S);
            qy = (m01 + m10) / S;
            qz = (m02 + m20) / S;
        } else if (m11 > m22) {
            float S = (float) (sqrt(1.0 + m11 - m00 - m22) * 2); // S=4*qy
            qw = (m02 - m20) / S;
            qx = (m01 + m10) / S;
            qy = (float) (0.25 * S);
            qz = (m12 + m21) / S;
        } else {
            float S = (float) (sqrt(1.0 + m22 - m00 - m11) * 2); // S=4*qz
            qw = (m10 - m01) / S;
            qx = (m02 + m20) / S;
            qy = (m12 + m21) / S;
            qz = (float) (0.25 * S);
        }
        return new float[]{qx, qy, qz, qw};
    }

    public static ArrayList<Pose> getAndroidPoses(JSONArray jsonArray){
        ArrayList<Pose> output = new ArrayList<>();
        for (int i = 0; i < jsonArray.length(); i++) {
            try {
                JSONArray matrixJSON = jsonArray.getJSONArray(i);
                float[] matrix =  Helper.JSONArrayToFloatArray(matrixJSON);
                Pose currPose = Helper.matrixToPose(matrix);
                output.add(currPose);
            } catch (JSONException e) {
                e.printStackTrace();
            }
        }
        return output;
    }
}
