/*
* Face recognition:
* 1. For each photo, extract all faces. Store face rectangles.
*    face_id unique
*    photo_id foreign key
*    top left bottom right
*    identity_id
*    distance (0 == truth; manually assigned identity)
* 2. For each face_id, create:
*    /${picturesPath}face-data/${face_id % 100}/
*      ${face_id}-normalized
*      ${face_id}-original
*      ${face_id}-data
*/

"use strict";

process.env.TZ = "Etc/GMT";

console.log("Loading face-recognizer");

require('@tensorflow/tfjs-node');

const config = require("config"),
  Promise = require("bluebird"),
  { exists, mkdir, unlink } = require("./lib/util"),
  faceapi = require("face-api.js"),
  fs = require("fs"),
  canvas = require("canvas");

const { createCanvas, Canvas, Image, ImageData } = canvas;

faceapi.env.monkeyPatch({ Canvas, Image, ImageData });

const maxConcurrency = require("os").cpus().length;

require("./console-line.js"); /* Monkey-patch console.log with line numbers */

const picturesPath = config.get("picturesPath").replace(/\/$/, "") + "/",
  facesPath = config.get("facesPath").replace(/\/$/, "") + "/";

let photoDB = null;

console.log("Loading pictures out of: " + picturesPath);

function alignFromLandmarks(image, landmarks) {
  const faceMargin = 0.3,
    width = 256, height = 256,
    dY = landmarks._positions[45]._y - landmarks._positions[36]._y,
    dX = landmarks._positions[45]._x - landmarks._positions[36]._x,
    mid = {
      x: landmarks._positions[36]._x + 0.5 * dX,
      y: landmarks._positions[36]._y + 0.5 * dY
    },
    rotation = -Math.atan2(dY, dX),
    cosRotation = Math.cos(rotation),
    sinRotation = Math.sin(rotation),
    eyeDistance = Math.sqrt(dY * dY + dX * dX),
    scale = width * (1.0 - 2. * faceMargin) / eyeDistance,
    canvas = createCanvas(width, height),
    ctx = canvas.getContext("2d");

  const prime = {
    x: mid.x * cosRotation - mid.y * sinRotation,
    y: mid.y * cosRotation + mid.x * sinRotation
  };
  
  mid.x = prime.x;
  mid.y = prime.y;

  ctx.translate(
    0.5 * width - mid.x * scale,
    0.5 * height - (height * (0.5 - faceMargin)) - mid.y * scale);
  ctx.rotate(rotation);
  ctx.scale(scale, scale);
  ctx.drawImage(image, 0, 0);
  /*
  ctx.strokeStyle = "red";
  ctx.strokeWidth = "1";
  ctx.beginPath();
  landmarks._positions.forEach((point, index) => {
    if (index == 0) {
      ctx.moveTo(point._x, point._y);
    } else {
      ctx.lineTo(point._x, point._y);
    }
  });
  ctx.stroke();
  */
  return canvas;
}

process.stdout.write("Loading DB.");
require("./db/photos").then(function(db) {
  process.stdout.write("done\n");
  photoDB = db;
}).then(() => {
  console.log("DB connected.");
  process.stdout.write("Loading models.");
  return faceapi.nets.ssdMobilenetv1.loadFromDisk('./models');
}).then(() => {
  process.stdout.write(".");
  return faceapi.nets.faceLandmark68Net.loadFromDisk('./models');
}).then(() => {
  process.stdout.write(".");
  return faceapi.nets.faceRecognitionNet.loadFromDisk('./models');
}).then(() => {
  console.log("Beginning face detection scanning.");
  return photoDB.sequelize.query("SELECT photos.id,photos.filename,photos.width,photos.height,albums.path " +
    "FROM photos " +
    "LEFT JOIN albums ON (albums.id=photos.albumId) " +
    "WHERE faces=-1 AND photos.duplicate=0 AND photos.deleted=0 ORDER BY albums.path,photos.filename", {
    type: photoDB.sequelize.QueryTypes.SELECT,
    raw: true
  });
}).then((needToScan) => {
  const total = needToScan.length;
  let remaining = total,
    processed = 0,
    lastStatus = Date.now();

  console.log(`${needToScan.length} photos have not had faces scanned.`);

  return Promise.map(needToScan, (photo) => {
    const photoPath = photo.path + photo.filename;

    console.log(`Processing ${photoPath}...`);

    /* Remove any existing face data for this photo */
    return photoDB.sequelize.query("SELECT id FROM faces WHERE photoId=:id", {
      replacements: photo,
      type: photoDB.sequelize.QueryTypes.SELECT,
      raw: true
    }).then((faces) => {
      /* For each face-id, remove any face-data files, and then remove all the entries
        * from the DB */
      return Promise.map(faces, (face) => {
        return Promise.map([ "-data.json", "-original.png" ], (suffix) => {
          const id = face.id,
            dataPath = facesPath + (id % 100) + "/" + id + suffix;
          return exists(dataPath).then((result) => {
            if (result) {
              console.log(`...removing ${dataPath}`);
              return unlink(dataPath);
            }
          });
        });
      }).then(() => {
        return photoDB.sequelize.query("DELETE FROM faces WHERE photoId=:id", {
          replacements: photo,
        });
      });
    }).then(async () => {
      /* Process image for faces data */
      const image = await canvas.loadImage(picturesPath + photoPath);
      const detections = await faceapi.detectAllFaces(image, 
          new faceapi.SsdMobilenetv1Options({
            minConfidence: 0.9
          })
      ).withFaceLandmarks();
      
      if (detections.length > 0) {
        console.log(`...${detections.length} faces identified in ${photoPath}.`);
      }

      return Promise.map(detections, async (face) => {
        const detection = face.detection,
          canvas = alignFromLandmarks(image, face.landmarks);
        face.descriptor = await faceapi.computeFaceDescriptor(canvas);
       
        const width = detection._box._width,
          height = detection._box._height,
          replacements = {
            id: photo.id,
            top: detection._box._y / detection._imageDims.height,
            left: detection._box._x / detection._imageDims.width,
            bottom: (detection._box._y + height) / detection._imageDims.height,
            right: (detection._box._x + width) / detection._imageDims.width,
            faceConfidence: detection._score
          };

        return photoDB.sequelize.query("INSERT INTO faces (photoId,top,left,bottom,right,faceConfidence) " +
          "VALUES (:id,:top,:left,:bottom,:right,:faceConfidence)", {
          replacements: replacements
        }).then(([ results, metadata ]) => {
          return metadata.lastID;
        }).then((id) => {
          const path = facesPath + (id % 100);
          return mkdir(path).then(() => {
            const dataPath = `${path}/${id}-data.json`, data = [];
            console.log(`...writing descriptor data to ${dataPath}...`);
            /* Confert from sparse object to dense array */
            for (let i = 0; i < 128; i++) {
              data.push(face.descriptor[i]);
            }
            fs.writeFileSync(dataPath, JSON.stringify(data));
          }).then(() => {
            const target = `${path}/${id}-original.png`;
            console.log(`...writing aligned face crop to ${target}.`);
            fs.writeFileSync(target, canvas.toBuffer("image/png", {
              quality: 0.95,
              chromaSubsampling: false
            }));
          }).catch((error) => {
            console.error(error);
            process.exit(-1);
          });
        });
      }).then(() => {
        return photoDB.sequelize.query("UPDATE photos SET faces=:faces WHERE id=:id", {
          replacements: {
            id: photo.id,
            faces: detections.length
          },
        });
      });
    }).catch((error) => {
      console.log(error);
      console.warn("Skipping out on image " + photoPath + " and marking to 0 faces to prevent future scanning.");
      return photoDB.sequelize.query("UPDATE photos SET faces=:faces WHERE id=:id", {
        replacements: {
          id: photo.id,
          faces: 0
        },
      });
    }).then(() => {
      processed++;
      const now = Date.now();
      if (now - lastStatus > 5000) {
        const rate = Math.round(10000 * (remaining - (total - processed)) / (now - lastStatus)) / 10,
          eta = Math.round((total - processed) / rate);
        lastStatus = now;
        remaining = total - processed;
        console.log(`Processing ${rate} images per second. ${remaining} images to be processed. ETA: ${eta}s`);
      }
    });
  }, {
    concurrency: maxConcurrency
  });
}).then(() => {
  console.log("Looking for face distances that need to be updated...");
  let maxId;

  return photoDB.sequelize.query("SELECT faces.id FROM faces ORDER BY faces.id DESC LIMIT 1", {
    type: photoDB.sequelize.QueryTypes.SELECT,
    raw: true
  }).then((results) => {
    if (!results.length) {
      console.log("...no faces exist yet to generate distances.");
      maxId = 0;
      return [];
    }
    maxId = results[0].id;
    return photoDB.sequelize.query(
      "SELECT faces.id,faces.lastComparedId " +
      "FROM faces INNER JOIN photos ON photos.duplicate=0 AND photos.deleted=0 AND photos.id=faces.photoId " +
      "WHERE faces.lastComparedId<:maxId OR faces.lastComparedId IS NULL " +
      "ORDER BY faces.id ASC", {
      replacements: {
        maxId: maxId
      },
      type: photoDB.sequelize.QueryTypes.SELECT,
      raw: true
    });
  }).then((facesToUpdate) => {
    console.log(`...${facesToUpdate.length} faces need distances updated.`);
    console.log("---- run scanner/scanner !! ---");
    return [];
    if (facesToUpdate.length == 0) {
      return facesToUpdate;
    }

    const descriptors = {};

    return photoDB.sequelize.query(
      "SELECT id FROM faces ORDER BY id ASC", {
      type: photoDB.sequelize.QueryTypes.SELECT,
      raw: true
    }).then((allFaces) => {
      console.log(`...reading ${allFaces.length} descriptors...`);
      return Promise.map(allFaces, (face) => {
        const id = face.id,
          dataPath = facesPath + "/" + (id % 100) + "/" + id + "-data.json";

        if (id in descriptors) {
          return;
        }

        return exists(dataPath).then((doesExist) => {
          if (!doesExist) {
            console.warn(`${dataPath} is missing!`);
            return;
          }

          descriptors[id] = JSON.parse(fs.readFileSync(dataPath));
        });
      });
    }).then(() => {
      const total = facesToUpdate.length;
      let remaining = total,
        processed = 0,
        lastStatus = Date.now(),
        targets = [];

      for (let target in descriptors) {
        targets.push({ id: target, descriptor: descriptors[target] });
      }

      return Promise.mapSeries(facesToUpdate, (face) => {
        if (!(face.id in descriptors)) {
          console.warn(`...attempt to compare distance with no descriptor for ${face.id}`);
          return;
        }

        const faceDescriptor = descriptors[face.id];

        return photoDB.sequelize.transaction((transaction) => {
          return photoDB.sequelize.query(
            "SELECT distance,face1Id,face2Id " +
            "FROM facedistances " +
            "WHERE face1Id=:id OR face2Id=:id " +
            "ORDER BY face1Id ASC", {
            replacements: {
              id: face.id
            },
            type: photoDB.sequelize.QueryTypes.SELECT,
            raw: true,
            transaction: transaction
          }).then((distances) => {
            return Promise.map(targets, (target) => {
              /* Skip comparing to self */
              if (target.id == face.id) {
                return;
              }

              /* Only compare against newer faces */
              if (face.lastComparedId && target.id <= face.lastComparedId) {
                return;
              }

              const index = distances.findIndex((distance) => {
                return distance.face1Id == target.id || distance.face2Id == target.id
              });

              if (index != -1) {
                /* A distance has already been calculated between face and target */
                return;
              }

              const distance = faceapi.euclideanDistance(faceDescriptor, target.descriptor);

              /* If the distance > 0.6, we don't want to store this in the DB */
              if (distance > 0.6) {
                return;
              }

              if (distance < 0.4) {
                process.stdout.write(".");
//                console.log(`Face ${face.id} and ${target.id} have a distance of: ${distance}`);
              }

              return photoDB.sequelize.query(
                "INSERT INTO facedistances (face1Id,face2Id,distance) " +
                "VALUES (:first,:second,:distance)", {
                replacements: {
                  first: Math.min(face.id, target.id),
                  second: Math.max(face.id, target.id),
                  distance: distance
                },
                transaction: transaction
              });
            }, {
              concurrency: maxConcurrency
            });
          }).then(() => {
            return photoDB.sequelize.query(
              "UPDATE faces SET lastComparedId=:lastId WHERE id=:id", {
              replacements: {
                lastId: maxId,
                id: face.id
              },
              transaction: transaction
            });
          });
        }).then(() => {
          processed++;
          const now = Date.now();
          if (now - lastStatus > 5000) {
            const rate = Math.round(10000 * (remaining - (total - processed)) / (now - lastStatus)) / 10,
              eta = Math.round((total - processed) / rate);
            lastStatus = now;
            remaining = total - processed;
            console.log(`\nProcessing ${rate} faces per second. ${remaining} faces to be processed. ETA: ${eta}s`);
          }
        });
      });
    });
  });
}).then(() => {
  console.log("Face detection scanning completed.");
}).catch((error) => {
  console.error(error);
  process.exit(-1);
});