package eu.dnetlib.dhp.sx.graph

import eu.dnetlib.dhp.application.ArgumentApplicationParser
import eu.dnetlib.dhp.schema.oaf.Relation
import eu.dnetlib.dhp.schema.sx.scholix.Scholix
import eu.dnetlib.dhp.schema.sx.summary.ScholixSummary
import eu.dnetlib.dhp.sx.graph.scholix.ScholixUtils
import eu.dnetlib.dhp.sx.graph.scholix.ScholixUtils.RelatedEntities
import org.apache.commons.io.IOUtils
import org.apache.spark.SparkConf
import org.apache.spark.sql.functions.count
import org.apache.spark.sql._
import org.slf4j.{Logger, LoggerFactory}

object SparkCreateScholix {

  def main(args: Array[String]): Unit = {
    val log: Logger = LoggerFactory.getLogger(getClass)
    val conf: SparkConf = new SparkConf()
    val parser = new ArgumentApplicationParser(IOUtils.toString(getClass.getResourceAsStream("/eu/dnetlib/dhp/sx/graph/create_scholix_params.json")))
    parser.parseArgument(args)
    val spark: SparkSession =
      SparkSession
        .builder()
        .config(conf)
        .appName(getClass.getSimpleName)
        .master(parser.get("master")).getOrCreate()

    val relationPath = parser.get("relationPath")
    log.info(s"relationPath  -> $relationPath")
    val summaryPath = parser.get("summaryPath")
    log.info(s"summaryPath  -> $summaryPath")
    val targetPath = parser.get("targetPath")
    log.info(s"targetPath  -> $targetPath")


    implicit val relEncoder: Encoder[Relation] = Encoders.kryo[Relation]
    implicit val summaryEncoder: Encoder[ScholixSummary] = Encoders.kryo[ScholixSummary]
    implicit val scholixEncoder: Encoder[Scholix] = Encoders.kryo[Scholix]

    import spark.implicits._


    val relationDS: Dataset[(String, Relation)] = spark.read.load(relationPath).as[Relation]
      .filter(r => (r.getDataInfo == null || r.getDataInfo.getDeletedbyinference == false) && !r.getRelClass.toLowerCase.contains("merge"))
      .map(r => (r.getSource, r))(Encoders.tuple(Encoders.STRING, relEncoder))

    val summaryDS: Dataset[(String, ScholixSummary)] = spark.read.load(summaryPath).as[ScholixSummary]
      .map(r => (r.getId, r))(Encoders.tuple(Encoders.STRING, summaryEncoder))


    relationDS.joinWith(summaryDS, relationDS("_1").equalTo(summaryDS("_1")), "left")
      .map { input: ((String, Relation), (String, ScholixSummary)) =>
        if (input._1 != null && input._2 != null) {
          val rel: Relation = input._1._2
          val source: ScholixSummary = input._2._2
          (rel.getTarget, ScholixUtils.scholixFromSource(rel, source))
        }
        else null
      }(Encoders.tuple(Encoders.STRING, scholixEncoder))
      .filter(r => r != null)
      .write.mode(SaveMode.Overwrite).save(s"$targetPath/scholix_from_source")

    val scholixSource: Dataset[(String, Scholix)] = spark.read.load(s"$targetPath/scholix_from_source").as[(String, Scholix)](Encoders.tuple(Encoders.STRING, scholixEncoder))

    scholixSource.joinWith(summaryDS, scholixSource("_1").equalTo(summaryDS("_1")), "left")
      .map { input: ((String, Scholix), (String, ScholixSummary)) =>
        if (input._2 == null) {
          null
        } else {
          val s: Scholix = input._1._2
          val target: ScholixSummary = input._2._2
          ScholixUtils.generateCompleteScholix(s, target)
        }
      }.filter(s => s != null).write.mode(SaveMode.Overwrite).save(s"$targetPath/scholix_one_verse")


    val scholix_o_v: Dataset[Scholix] = spark.read.load(s"$targetPath/scholix_one_verse").as[Scholix]

    scholix_o_v.flatMap(s => List(s, ScholixUtils.createInverseScholixRelation(s))).as[Scholix]
      .map(s => (s.getIdentifier, s))(Encoders.tuple(Encoders.STRING, scholixEncoder))
      .groupByKey(_._1)
      .agg(ScholixUtils.scholixAggregator.toColumn)
      .map(s => s._2)
      .write.mode(SaveMode.Overwrite).save(s"$targetPath/scholix")

    val scholix_final: Dataset[Scholix] = spark.read.load(s"$targetPath/scholix").as[Scholix]

    val stats: Dataset[(String, String, Long)] = scholix_final.map(s => (s.getSource.getDnetIdentifier, s.getTarget.getObjectType)).groupBy("_1", "_2").agg(count("_1")).as[(String, String, Long)]


    stats
      .map(s => RelatedEntities(s._1, if ("dataset".equalsIgnoreCase(s._2)) s._3 else 0, if ("publication".equalsIgnoreCase(s._2)) s._3 else 0))
      .groupByKey(_.id)
      .reduceGroups((a, b) => RelatedEntities(a.id, a.relatedDataset + b.relatedDataset, a.relatedPublication + b.relatedPublication))
      .map(_._2)
      .write.mode(SaveMode.Overwrite).save(s"$targetPath/related_entities")

    val relatedEntitiesDS: Dataset[RelatedEntities] = spark.read.load(s"$targetPath/related_entities").as[RelatedEntities].filter(r => r.relatedPublication > 0 || r.relatedDataset > 0)

    relatedEntitiesDS.joinWith(summaryDS, relatedEntitiesDS("id").equalTo(summaryDS("_1")), "inner").map { i =>
      val re = i._1
      val sum = i._2._2

      sum.setRelatedDatasets(re.relatedDataset)
      sum.setRelatedPublications(re.relatedPublication)
      sum
    }.write.mode(SaveMode.Overwrite).save(s"${summaryPath}_filtered")

  }
}