package com.gmei

import java.io.Serializable
import org.apache.spark.storage.StorageLevel
import org.apache.spark.mllib.linalg.Vectors
import org.apache.spark.mllib.linalg.distributed.{IndexedRow, IndexedRowMatrix}
import org.apache.spark.sql.{SaveMode,TiContext}
import org.apache.log4j.{Level, Logger}
import scopt.OptionParser
import com.gmei.lib.AbstractParams
import com.soundcloud.lsh.Lsh






object Main {

  Logger.getLogger("org.apache.spark").setLevel(Level.WARN)
  Logger.getLogger("org.apache.eclipse.jetty.server").setLevel(Level.OFF)


  case class Params(iter: Int = 10,
                    lr: Double = 0.025,
                    numPartition: Int = 10,
                    dim: Int = 128,
                    window: Int = 10,
                    walkLength: Int = 80,
                    numWalks: Int = 10,
                    p: Double = 1.0,
                    q: Double = 1.0,
                    weighted: Boolean = true,
                    directed: Boolean = false,
                    degree: Int = 30,
                    indexed: Boolean = false,
                    env: String = ENV.DEV,
                    nodePath: String = null
                   ) extends AbstractParams[Params] with Serializable
  val defaultParams = Params()

  val parser = new OptionParser[Params]("Node2Vec_Spark") {
    head("Main")
    opt[Int]("walkLength")
      .text(s"walkLength: ${defaultParams.walkLength}")
      .action((x, c) => c.copy(walkLength = x))
    opt[Int]("numWalks")
      .text(s"numWalks: ${defaultParams.numWalks}")
      .action((x, c) => c.copy(numWalks = x))
    opt[Double]("p")
      .text(s"return parameter p: ${defaultParams.p}")
      .action((x, c) => c.copy(p = x))
    opt[Double]("q")
      .text(s"in-out parameter q: ${defaultParams.q}")
      .action((x, c) => c.copy(q = x))
    opt[Boolean]("weighted")
      .text(s"weighted: ${defaultParams.weighted}")
      .action((x, c) => c.copy(weighted = x))
    opt[Boolean]("directed")
      .text(s"directed: ${defaultParams.directed}")
      .action((x, c) => c.copy(directed = x))
    opt[Int]("degree")
      .text(s"degree: ${defaultParams.degree}")
      .action((x, c) => c.copy(degree = x))
    opt[Boolean]("indexed")
      .text(s"Whether nodes are indexed or not: ${defaultParams.indexed}")
      .action((x, c) => c.copy(indexed = x))
    opt[String]("env")
      .text(s"the databases environment you used")
      .action((x,c) => c.copy(env = x))
    opt[String]("nodePath")
      .text("Input node2index file path: empty")
      .action((x, c) => c.copy(nodePath = x))
    note(
      """
        |For example, the following command runs this app on a synthetic dataset:
        |
        | bin/spark-submit --class com.nhn.sunny.vegapunk.ml.model.Node2vec \
      """.stripMargin +
        s"|   --lr ${defaultParams.lr}" +
        s"|   --iter ${defaultParams.iter}" +
        s"|   --numPartition ${defaultParams.numPartition}" +
        s"|   --dim ${defaultParams.dim}" +
        s"|   --window ${defaultParams.window}" +
        s"|   --node <nodeFilePath>" +
        s"|   --output <path>"
    )
  }


  def main(args: Array[String]):Unit = {

    parser.parse(args, defaultParams).map { param =>

      //1. get the input and node2vec
      GmeiConfig.setup(param)
      val context = GmeiConfig.getSparkSession()._1
      val sc = GmeiConfig.getSparkSession()._2

      val ti = new TiContext(sc)
      ti.tidbMapTable(dbName = GmeiConfig.config.getString("tidb.database"),tableName = "data_meigou_cid")

      val tidb_inupt = sc.sql(
        s"""
           |SELECT
           | service_id,cid
           |FROM data_meigou_cid
     """.stripMargin
      )

      Node2vec.setup(context, param)

      Node2vec.load(tidb_inupt)
        .initTransitionProb()
        .randomWalk()
        .embedding()

      val node2vector = context.parallelize(Word2vec.getVectors.toList)
        .map { case (nodeId, vector) =>
          (nodeId.toLong, vector.map(x => x.toDouble))
        }

      val id2Node = Node2vec.node2id.map{ case (strNode, index) =>
        (index, strNode)
      }
      println("get id2node")
      println(id2Node.first())

      val node2vec_2 = node2vector.join(id2Node)
        .map { case (nodeId, (vector, name)) => (name,vector) }
        .repartition(200)



      //2. compute similar cid and then take top k

      val storageLevel = StorageLevel.MEMORY_AND_DISK

      val indexed = node2vec_2.zipWithIndex.persist(storageLevel)

      // create indexed row matrix where every row represents one word
      val rows = indexed.map {
        case ((word, features), index) =>
          IndexedRow(index, Vectors.dense(features))
      }

      // store index for later re-mapping (index to word)
      val index = indexed.map {
        case ((word, features), index) =>
          (index, word)
      }.persist(storageLevel)

      // create an input matrix from all rows and run lsh on it
      val matrix = new IndexedRowMatrix(rows)
      val lsh = new Lsh(
        minCosineSimilarity = 0.5,
        dimensions = 20,
        numNeighbours = 200,
        numPermutations = 10,
        partitions = 200,
        storageLevel = storageLevel
      )
      val similarityMatrix = lsh.join(matrix)

      import sc.implicits._
      // remap both ids back to words
      val remapFirst = similarityMatrix.entries.keyBy(_.i).join(index).values
      val remapSecond = remapFirst.keyBy { case (entry, word1) => entry.j }.join(index).values.map {
        case ((entry, word1), word2) =>
          (word1, word2, entry.value)
      }

      remapSecond.take(20).foreach(println)

      val score_result = remapSecond.toDF("cid1","cid2","score")
      GmeiConfig.writeToJDBCTable(score_result, table="cid_pairs_cosine_distince", SaveMode.Overwrite)


      // group by neighbours to get a list of similar words and then take top k
      val result = remapSecond.groupBy(_._1).map {
        case (word1, similarWords) =>
          // sort by score desc. and take top 10 entries
          val similar = similarWords.toSeq.sortBy(-1 * _._3).filter(_._2.startsWith("diary")).take(10).map(_._2).mkString(",")
          (word1,s"$similar")
      }
      // print out the results for the first 10 words
      result.take(20).foreach(println)

      val similar_result = result.toDF("cid","similarity_cid")
      GmeiConfig.writeToJDBCTable(similar_result, table="cid_similarity_matrix", SaveMode.Overwrite)

    }
  } getOrElse {
    sys.exit(1)
  }
}

