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

package org.apache.spark.ml.feature

import org.apache.spark.SparkException
import org.apache.spark.ml.Transformer
import org.apache.spark.ml.attribute.{Attribute, AttributeGroup}
import org.apache.spark.ml.linalg.{Vector, VectorUDT, Vectors}
import org.apache.spark.ml.param.ParamMap
import org.apache.spark.ml.param.shared._
import org.apache.spark.ml.util._
import org.apache.spark.sql.functions._
import org.apache.spark.sql.types._
import org.apache.spark.sql.{DataFrame, Dataset, Row}

import scala.collection.mutable

/** A fast vector assembler.  The columns given must be ordered such that categorical columns come first
  * (otherwise spark learners will give categorical attributes to the wrong index).
  * Does not keep spurious numeric data which can significantly slow down computations when there are
  * millions of columns.
  */
class FastVectorAssembler (override val uid: String)
  extends Transformer with HasInputCols with HasOutputCol with DefaultParamsWritable {

  def this() = this(Identifiable.randomUID("FastVectorAssembler"))

  /** @group setParam */
  def setInputCols(value: Array[String]): this.type = set(inputCols, value)

  /** @group setParam */
  def setOutputCol(value: String): this.type = set(outputCol, value)

  //scalastyle:off cyclomatic.complexity
  override def transform(dataset: Dataset[_]): DataFrame = {  //scalastyle:ignore method.length
    // Schema transformation.
    val schema = dataset.schema
    var addedNumericField = false

    // Propagate only nominal (categorical) attributes (others only slow down the code)
    val attrs: Array[Attribute] = $(inputCols).flatMap { c =>
      val field = schema(c)
      field.dataType match {
        case _: NumericType | BooleanType =>
          val attr = Attribute.fromStructField(field)
          if (attr.isNominal) {
            if (addedNumericField) {
              throw new SparkException("Categorical columns must precede all others, column out of order: " + c)
            }
            Some(attr.withName(c))
          } else {
            addedNumericField = true
            None
          }
        case _: VectorUDT =>
          val group = AttributeGroup.fromStructField(field)
          if (group.attributes.isDefined) {
            // If attributes are defined, copy them with updated names.
            group.attributes.get.zipWithIndex.map { case (attr, i) =>
              if (attr.isNominal && attr.name.isDefined) {
                if (addedNumericField) {
                  throw new SparkException("Categorical columns must precede all others, column out of order: " + c)
                }
                attr.withName(c + "_" + attr.name.get)
              } else if (attr.isNominal) {
                if (addedNumericField) {
                  throw new SparkException("Categorical columns must precede all others, column out of order: " + c)
                }
                attr.withName(c + "_" + i)
              } else {
                addedNumericField = true
                null  //scalastyle:ignore null
              }
            }.filter(attr => attr != null)
          } else {
            addedNumericField = true
            None
          }
        case otherType =>
          throw new SparkException(s"FastVectorAssembler does not support the $otherType type")
      }
    }
    val metadata = new AttributeGroup($(outputCol), attrs).toMetadata()

    // Data transformation.
    val assembleFunc = udf { r: Row =>
      FastVectorAssembler.assemble(r.toSeq: _*)
    }
    val args = $(inputCols).map { c =>
      schema(c).dataType match {
        case DoubleType => dataset(c)
        case _: VectorUDT => dataset(c)
        case _: NumericType | BooleanType => dataset(c).cast(DoubleType).as(s"${c}_double_$uid")
      }
    }

    dataset.select(col("*"), assembleFunc(struct(args: _*)).as($(outputCol), metadata))
  }
  //scalastyle:on cyclomatic.complexity

  override def transformSchema(schema: StructType): StructType = {
    val inputColNames = $(inputCols)
    val outputColName = $(outputCol)
    val inputDataTypes = inputColNames.map(name => schema(name).dataType)
    inputDataTypes.foreach {
      case _: NumericType | BooleanType =>
      case t if t.isInstanceOf[VectorUDT] =>
      case other =>
        throw new IllegalArgumentException(s"Data type $other is not supported.")
    }
    if (schema.fieldNames.contains(outputColName)) {
      throw new IllegalArgumentException(s"Output column $outputColName already exists.")
    }
    StructType(schema.fields :+ StructField(outputColName, new VectorUDT, nullable = true))
  }

  override def copy(extra: ParamMap): FastVectorAssembler = defaultCopy(extra)

}

object FastVectorAssembler extends DefaultParamsReadable[FastVectorAssembler] {

  override def load(path: String): FastVectorAssembler = super.load(path)

  private[feature] def assemble(vv: Any*): Vector = {
    val indices = mutable.ArrayBuilder.make[Int]
    val values = mutable.ArrayBuilder.make[Double]
    var cur = 0
    vv.foreach {
      case v: Double =>
        if (v != 0.0) {
          indices += cur
          values += v
        }
        cur += 1
      case vec: Vector =>
        vec.foreachActive { case (i, v) =>
          if (v != 0.0) {
            indices += cur + i
            values += v
            ()
          }
        }
        cur += vec.size
      case null =>  //scalastyle:ignore null
        throw new SparkException("Values to assemble cannot be null.")
      case o =>
        throw new SparkException(s"$o of type ${o.getClass.getName} is not supported.")
    }
    Vectors.sparse(cur, indices.result(), values.result()).compressed
  }

}
