// Copyright (C) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License. See LICENSE in project root for information.

package com.microsoft.azure.synapse.ml.services.face

import com.microsoft.azure.synapse.ml.services._
import com.microsoft.azure.synapse.ml.core.test.fuzzing.{TestObject, TransformerFuzzing}
import org.apache.spark.ml.util.MLReadable
import org.apache.spark.sql.functions.{col, explode, lit}
import org.apache.spark.sql.{DataFrame, Row}
import org.scalactic.Equality

import java.time.LocalDateTime
import java.time.format.{DateTimeFormatterBuilder, DateTimeParseException, SignStyle}
import java.time.temporal.ChronoField
import scala.util.matching.Regex

class DetectFaceSuite extends TransformerFuzzing[DetectFace] with CognitiveKey {

  import spark.implicits._

  lazy val df: DataFrame = Seq(
    "https://mmlspark.blob.core.windows.net/datasets/DSIR/test2.jpg"
  ).toDF("url")

  lazy val face = new DetectFace()
    .setSubscriptionKey(cognitiveKey)
    .setLocation("eastus")
    .setImageUrlCol("url")
    .setOutputCol("face")
    .setReturnFaceId(true)
    .setReturnFaceLandmarks(true)
    .setReturnFaceAttributes(Seq("exposure"))

  override def assertDFEq(df1: DataFrame, df2: DataFrame)(implicit eq: Equality[DataFrame]): Unit = {
    def prep(df: DataFrame) = df.select(explode(col("face"))).select("col.*").drop("faceId")
    super.assertDFEq(prep(df1), prep(df2))(eq)
  }

  test("Basic Usage") {
    face.transform(df)
    val results = face.transform(df)
    val fromRow = Face.makeFromRowConverter

    val f1 = fromRow(results.select("face").collect().head.getSeq[Row](0).head)
    assert(f1.faceAttributes.get.exposure.get.value != 0.0)

    results.show(truncate = false)
  }

  override def testObjects(): Seq[TestObject[DetectFace]] =
    Seq(new TestObject(face, df))

  override def reader: MLReadable[_] = DetectFace
}

class FindSimilarFaceSuite extends TransformerFuzzing[FindSimilarFace] with CognitiveKey {

  import spark.implicits._

  lazy val df: DataFrame = Seq(
    "https://mmlspark.blob.core.windows.net/datasets/DSIR/test1.jpg",
    "https://mmlspark.blob.core.windows.net/datasets/DSIR/test2.jpg",
    "https://mmlspark.blob.core.windows.net/datasets/DSIR/test3.jpg"
  ).toDF("url")

  lazy val detector = new DetectFace()
    .setSubscriptionKey(cognitiveKey)
    .setLocation("eastus")
    .setImageUrlCol("url")
    .setOutputCol("detected_faces")
    .setReturnFaceId(true)
    .setReturnFaceLandmarks(false)
    .setReturnFaceAttributes(Seq())

  lazy val fromRow = Face.makeFromRowConverter

  lazy val faceIdDF = detector.transform(df)
    .select(
      col("detected_faces").getItem(0).getItem("faceId").alias("id"))
    .cache()

  lazy val faceIds = faceIdDF.collect().map(row =>
    row.getAs[String]("id"))

  lazy val findSimilar = new FindSimilarFace()
    .setSubscriptionKey(cognitiveKey)
    .setLocation("eastus")
    .setOutputCol("similar")
    .setFaceIdCol("id")
    .setFaceIds(faceIds)

  test("Basic Usage") {
    val similarFaceDF = findSimilar.transform(faceIdDF)
    val numMatches = similarFaceDF.select("similar").collect().map(row =>
      row.getSeq[Row](0).length
    )
    assert(numMatches === List(1, 2, 2))
  }

  test("Throw errors if required fields not set") {
    val caught = intercept[AssertionError] {
      new FindSimilarFace()
        .setSubscriptionKey(cognitiveKey)
        .setLocation("eastus")
        .setOutputCol("similar")
        .transform(faceIdDF).collect()
    }
    assert(caught.getMessage.contains("Missing required params"))
    assert(caught.getMessage.contains("faceId"))
  }

  override def testObjects(): Seq[TestObject[FindSimilarFace]] =
    Seq(new TestObject(findSimilar, faceIdDF))

  override def reader: MLReadable[_] = FindSimilarFace
}

class GroupFacesSuite extends TransformerFuzzing[GroupFaces] with CognitiveKey {

  import spark.implicits._

  lazy val df: DataFrame = Seq(
    "https://mmlspark.blob.core.windows.net/datasets/DSIR/test1.jpg",
    "https://mmlspark.blob.core.windows.net/datasets/DSIR/test2.jpg",
    "https://mmlspark.blob.core.windows.net/datasets/DSIR/test3.jpg"
  ).toDF("url")

  lazy val detector = new DetectFace()
    .setSubscriptionKey(cognitiveKey)
    .setLocation("eastus")
    .setImageUrlCol("url")
    .setOutputCol("detected_faces")
    .setReturnFaceId(true)
    .setReturnFaceLandmarks(false)
    .setReturnFaceAttributes(Seq())

  lazy val fromRow = Face.makeFromRowConverter

  lazy val faceIdDF = detector.transform(df)
    .select(
      col("detected_faces").getItem(0).getItem("faceId").alias("id"))
    .cache()

  lazy val faceIds: Array[String] = faceIdDF.collect().map(row =>
    row.getAs[String]("id"))

  lazy val group = new GroupFaces()
    .setSubscriptionKey(cognitiveKey)
    .setLocation("eastus")
    .setOutputCol("grouping")
    .setFaceIds(faceIds)

  test("Basic Usage") {
    val groupDF = group.transform(faceIdDF)
    val numMatches = groupDF.select(col("grouping.groups").getItem(0)).collect().map(row =>
      row.getSeq[Row](0).length
    )
    assert(numMatches === List(2, 2, 2))
  }

  test("Throw errors if required fields not set") {
    val caught = intercept[AssertionError] {
      new GroupFaces()
        .setSubscriptionKey(cognitiveKey)
        .setLocation("eastus")
        .setOutputCol("grouping")
        .transform(faceIdDF).collect()
    }
    assert(caught.getMessage.contains("Missing required params"))
    assert(caught.getMessage.contains("faceIds"))
  }

  override def testObjects(): Seq[TestObject[GroupFaces]] =
    Seq(new TestObject(group, faceIdDF))

  override def reader: MLReadable[_] = GroupFaces
}

class IdentifyFacesSuite extends TransformerFuzzing[IdentifyFaces] with CognitiveKey {

  import spark.implicits._

  lazy val satyaFaces = Seq(
    "https://mmlspark.blob.core.windows.net/datasets/DSIR/test1.jpg"
  )

  lazy val bradFaces = Seq(
    "https://mmlspark.blob.core.windows.net/datasets/DSIR/test2.jpg",
    "https://mmlspark.blob.core.windows.net/datasets/DSIR/test3.jpg"
  )

  // When a date pattern starts with 'yyyy' and has no separator following, the parser can sometimes decide
  // to take the whole string to match the year, which results in an exception. The following is a hackaround.
  val formatter = new DateTimeFormatterBuilder()
    .appendValue(ChronoField.YEAR_OF_ERA, 4, 4, SignStyle.EXCEEDS_PAD)
    .appendPattern("MMddHHmmssSSS").toFormatter()

  lazy val pgName = "group" + formatter.format(LocalDateTime.now())

  lazy val pgId = {
    cleanOldGroups()
    PersonGroup.create(pgName, pgName)
    Thread.sleep(500) // A little insurance
    PersonGroup.list().find(_.name == pgName).get.personGroupId
  }

  lazy val satyaId = Person.create("satya", pgId)
  lazy val bradId = Person.create("brad", pgId)

  lazy val satyaFaceIds = satyaFaces.map(Person.addFace(_, pgId, satyaId))
  lazy val bradFaceIds = bradFaces.map(Person.addFace(_, pgId, bradId))

  lazy val detector = new DetectFace()
    .setSubscriptionKey(cognitiveKey)
    .setLocation("eastus")
    .setImageUrlCol("url")
    .setOutputCol("detected_faces")
    .setReturnFaceId(true)
    .setReturnFaceLandmarks(false)
    .setReturnFaceAttributes(Seq())

  lazy val otherFaceIds = detector.transform((satyaFaces ++bradFaces).toDF("url"))
    .select(col("detected_faces").getItem(0).getItem("faceId"))
    .collect()
    .map(r => r.getString(0)).toSeq

  override def beforeAll(): Unit = {
    super.beforeAll()
    println(satyaFaceIds ++ bradFaceIds)
    PersonGroup.train(pgId)
    tryWithRetries() { () =>
      assert(PersonGroup.getTrainingStatus(pgId).status === "succeeded")
      ()
    }
    println("done training face group")
  }

  override def afterAll(): Unit = {
    PersonGroup.list().find(_.name == pgName).foreach { pgi =>
      PersonGroup.delete(pgi.personGroupId)
      println("deleted group")
    }

    super.afterAll()
  }

  lazy val id = new IdentifyFaces()
    .setSubscriptionKey(cognitiveKey)
    .setLocation("eastus")
    .setFaceIdsCol("faces")
    .setPersonGroupId(pgId)
    .setOutputCol("identified_faces")

  lazy val df = otherFaceIds.map(Seq[String](_)).toDF("faces")

  test("Basic Usage") {
    Person.list(pgId).foreach(println)
    val matches = id.transform(df).select(col("identified_faces")
      .getItem(0).getItem("candidates").getItem(0).getItem("personId"))
      .collect().map(_.getString(0))
    assert(matches === List(satyaId, bradId, bradId))
  }

  test("Throw errors if required fields not set") {
    val caught = intercept[AssertionError] {
      new IdentifyFaces()
        .setSubscriptionKey(cognitiveKey)
        .setLocation("eastus")
        .setPersonGroupId(pgId)
        .setOutputCol("identified_faces")
        .transform(df).collect()
    }
    assert(caught.getMessage.contains("Missing required params"))
    assert(caught.getMessage.contains("faceIds"))
  }
  override def testObjects(): Seq[TestObject[IdentifyFaces]] = Seq(new TestObject(id, df))

  def cleanOldGroups(): Unit = {
    val twoDaysAgo = LocalDateTime.now().minusDays(2)
    var groupDeleted: Boolean = false
    val uuidPattern = new Regex("[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}")

    // scalastyle:off while
    do {
      groupDeleted = false
      PersonGroup.list(top = Some("500")).foreach { pgi =>
        try {
          val pgDateString = pgi.personGroupId.replaceFirst("group", "")
          val pgDate = LocalDateTime.parse(pgDateString, formatter)
          if (pgDate.isBefore(twoDaysAgo)) {
            PersonGroup.delete(pgi.personGroupId)
            println(s"Deleted group $pgi")
            groupDeleted = true
          }
        } catch {
          // for uuid-based names
          // TODO: delete this after we can be assured that everyone has updated
          case _: DateTimeParseException => {
            if ((uuidPattern findFirstIn pgi.personGroupId).isDefined) {
              PersonGroup.delete(pgi.personGroupId)
              println(s"Deleted group $pgi")
              groupDeleted = true
            }
          }
          case t: Throwable => throw t
        }
      }
    } while (groupDeleted)
    // scalastyle:on while
  }

  def cleanAllGroups(): Unit = {
    var groupDeleted: Boolean = false
    // scalastyle:off while
    do {
      groupDeleted = false
      PersonGroup.list(top=Some("500")).foreach { pgi =>
        PersonGroup.delete(pgi.personGroupId)
        println(s"Deleted group $pgi")
        groupDeleted = true
      }
    } while (groupDeleted)
    // scalastyle:on while
  }

  override def reader: MLReadable[_] = IdentifyFaces
}

class VerifyFacesSuite extends TransformerFuzzing[VerifyFaces] with CognitiveKey {

  import spark.implicits._

  lazy val df: DataFrame = Seq(
    "https://mmlspark.blob.core.windows.net/datasets/DSIR/test1.jpg",
    "https://mmlspark.blob.core.windows.net/datasets/DSIR/test2.jpg",
    "https://mmlspark.blob.core.windows.net/datasets/DSIR/test3.jpg"
  ).toDF("url")

  lazy val detector = new DetectFace()
    .setSubscriptionKey(cognitiveKey)
    .setLocation("eastus")
    .setImageUrlCol("url")
    .setOutputCol("detected_faces")
    .setReturnFaceId(true)
    .setReturnFaceLandmarks(false)
    .setReturnFaceAttributes(Seq())

  lazy val fromRow = Face.makeFromRowConverter

  lazy val faceIdDF = detector.transform(df)
    .select(col("detected_faces").getItem(0).getItem("faceId").alias("faceId1"))
    .cache()

  lazy val faceIdDF2 = faceIdDF.withColumn(
    "faceId2", lit(faceIdDF.take(1).head.getString(0)))

  lazy val verify = new VerifyFaces()
    .setSubscriptionKey(cognitiveKey)
    .setLocation("eastus")
    .setOutputCol("same")
    .setFaceId1Col("faceId1")
    .setFaceId2Col("faceId2")

  test("Basic Usage") {
    val verifyDF = verify.transform(faceIdDF2)
    assert(verifyDF
      .select("same.isIdentical")
      .collect().map(_.getBoolean(0)) === List(true, false, false))
  }

  override def testObjects(): Seq[TestObject[VerifyFaces]] =
    Seq(new TestObject(verify, faceIdDF2))

  override def reader: MLReadable[_] = VerifyFaces

}
