package eu.dnetlib.dhp.contextpropagation

import eu.dnetlib.dhp.contextpropagation.model.{EnrichedEntries, PropagationStructure}
import org.apache.spark.sql.expressions.Aggregator
import org.apache.spark.sql.{Encoder, Encoders}

object PropagationAggregator {

  def getDatasetAggregator(): Aggregator[(String, PropagationStructure), PropagationStructure, PropagationStructure] = new Aggregator[(String, PropagationStructure), PropagationStructure, PropagationStructure]{

    override def zero: PropagationStructure = new PropagationStructure()

    override def reduce(b: PropagationStructure, a: (String, PropagationStructure)): PropagationStructure = {
      b.mergeFrom(a._2)

    }

    override def merge(wx: PropagationStructure, wy: PropagationStructure): PropagationStructure = {

      wx.mergeFrom(wy)

    }
    override def finish(reduction: PropagationStructure): PropagationStructure = reduction

    override def bufferEncoder: Encoder[PropagationStructure] =
      Encoders.kryo(classOf[PropagationStructure])

    override def outputEncoder: Encoder[PropagationStructure] =
      Encoders.kryo(classOf[PropagationStructure])

  }


  def mergeEnrichedEntries(): Aggregator[(String, EnrichedEntries), EnrichedEntries, EnrichedEntries] = new Aggregator[(String, EnrichedEntries), EnrichedEntries, EnrichedEntries]{

    override def zero: EnrichedEntries = new EnrichedEntries()

    override def reduce(b: EnrichedEntries, a: (String, EnrichedEntries)): EnrichedEntries = {

      b.mergeWith(a._2)

    }

    override def merge(wx: EnrichedEntries, wy: EnrichedEntries): EnrichedEntries = {

      wx.mergeWith(wy)

    }
    override def finish(reduction: EnrichedEntries): EnrichedEntries = reduction

    override def bufferEncoder: Encoder[EnrichedEntries] =
      Encoders.kryo(classOf[EnrichedEntries])

    override def outputEncoder: Encoder[EnrichedEntries] =
      Encoders.kryo(classOf[EnrichedEntries])

  }


}
