package eu.dnetlib.dhp.contextpropagation

import java.util

import com.fasterxml.jackson.databind.ObjectMapper
import eu.dnetlib.dhp.application.ArgumentApplicationParser
import eu.dnetlib.dhp.contextpropagation.model.{DatasetPropagationStructure, EnrichedEntries, Node, PropagationStructure, PropagationUse, RelationPropagation}
import eu.dnetlib.dhp.provision.SparkConvertDatasetToJson
import eu.dnetlib.dhp.provision.scholix.Scholix
import org.apache.commons.io.IOUtils
import org.apache.hadoop.io.compress.GzipCodec
import org.apache.spark.SparkConf
import org.apache.spark.sql.{Dataset, Encoder, Encoders, SaveMode, SparkSession}

import scala.collection.JavaConverters._


object SparkContextPropagation {

  implicit val relationEncoder: Encoder[RelationPropagation] = Encoders.kryo[RelationPropagation]
  implicit val tupleForRelation: Encoder[(String, RelationPropagation)] = Encoders.tuple(Encoders.STRING, relationEncoder)
  implicit val enrichedEntitiesEncoder: Encoder[EnrichedEntries] = Encoders.kryo[EnrichedEntries]
  implicit val tupleForEntities: Encoder[(String, EnrichedEntries)] = Encoders.tuple(Encoders.STRING, enrichedEntitiesEncoder)

  implicit val mapEncoderPub: Encoder[PropagationStructure] = Encoders.kryo[PropagationStructure]
  implicit val mapEncoderDats: Encoder[DatasetPropagationStructure] = Encoders.kryo[DatasetPropagationStructure]
  implicit val tupleForPropagation: Encoder[(String, PropagationStructure)] = Encoders.tuple(Encoders.STRING, mapEncoderPub)
  implicit val tupleForPropagationDars: Encoder[(String, DatasetPropagationStructure)] = Encoders.tuple(Encoders.STRING, mapEncoderDats)
  implicit val stringEncoder: Encoder[String] = Encoders.STRING

  def getEnrichedPublications (allowedRelations: Dataset[RelationPropagation], enrichedEntitiesPath: String, spark: SparkSession): Dataset[RelationPropagation] = {

    val startingPropagation = allowedRelations
      .filter(r => r.getSource.getId.startsWith("50")).map(r => (r.getSource.getId, r))

    val enrichedNodes = spark.read.load(enrichedEntitiesPath).as[EnrichedEntries]
      .map(e => (e.getScholixId, e))

    startingPropagation.joinWith(enrichedNodes, startingPropagation("_1").equalTo(enrichedNodes("_1"))).map(tuple => tuple._1._2)

  }

  def propagatePublicationDataset(pubs_rel : Dataset[RelationPropagation]): Dataset[(String, PropagationStructure)] ={

    val publication_dataset : Dataset[(String, PropagationStructure)] = pubs_rel.map(r => {
      val ps = new PropagationStructure

      val pv : List[PropagationUse] = List(PropagationUse.copyInstance(Costants.getPublicationValue(r.getSemantics)))
      ps.add(r.getSource.getId, pv.asJava)
      (r.getTarget.getId, ps)

    })
    publication_dataset.groupByKey(_._1)
      .agg(PropagationAggregator.getDatasetAggregator().toColumn)
  }


  def propagateDatasetDataset(propagation: Dataset[(String, PropagationStructure)], dataset_dataset : Dataset[(String, DatasetPropagationStructure)], count :Int): Dataset[(String, PropagationStructure)] = {

    val pl2_step1 : Dataset [(String, PropagationStructure)] = propagation.joinWith(dataset_dataset, propagation("_1").equalTo(dataset_dataset("_1")))
      .flatMap(PropagationUtils.propagateDataset)

    val pl2 : Dataset [(String, PropagationStructure)] = propagation.union(pl2_step1).groupByKey(_._1)
      .agg(PropagationAggregator.getDatasetAggregator().toColumn)

   pl2

  }



  def main(args: Array[String]): Unit = {
    val parser = new ArgumentApplicationParser(IOUtils.toString(SparkConvertDatasetToJson.getClass.getResourceAsStream("/eu/dnetlib/dhp/provision/dataset2Json.json")))
    parser.parseArgument(args)
    val conf = new SparkConf
    val spark = SparkSession.builder.config(conf).appName(SparkConvertDatasetToJson.getClass.getSimpleName).master(parser.get("master")).getOrCreate
    val propagationOutputPath = parser.get("propagationOutputPath")

    val allowedRelations = spark.read.load(parser.get("allowedRelationPath")).as[RelationPropagation]

    val dataset_dataset : Dataset[(String, DatasetPropagationStructure)] = allowedRelations.filter(r => r.getSource.getId.startsWith("60"))
      .map(r => {
        val ps = new DatasetPropagationStructure

        ps.add(r.getTarget.getId, PropagationUse.copyInstance(Costants.getDatasetValue(r.getSemantics)))
        (r.getSource.getId, ps)

      })


    val pl1 : Dataset[(String, PropagationStructure)] = propagatePublicationDataset(
      getEnrichedPublications(allowedRelations, parser.get("enrichedEntitiesPath"), spark ))


    pl1.write.mode(SaveMode.Overwrite).save(s"$propagationOutputPath/pl1")

    var propagation = pl1

    var count = 1

    do {

      count += 1
      propagation = propagateDatasetDataset(propagation, dataset_dataset, count )
      propagation.write.mode(SaveMode.Overwrite).save(s"$propagationOutputPath/pl${count}" )

    }
    while(propagation.count() > 0)


  }

}
