package pl.edu.pw.mini.sg.mixed.adjuster.momentum

import pl.edu.pw.mini.sg.rng.RandomGenerator
import org.log4s._
import pl.edu.pw.mini.sg.mixed.IterationInfo
import pl.edu.pw.mini.sg.mixed.adjuster.AssesmentBasedAdjuster
import pl.edu.pw.mini.sg.mixed.tree.MixedStrategyNode



final case class SquaredNormalizerMomentumAdjuster(exponent: Double) extends AssesmentBasedAdjuster {
  private val logger = getLogger

  override type AdjusterState = Momentum
  override val ost = MomentumOptimizerStateType

  override def updateState(
    oldState: AdjusterState,
    assesments: Seq[Double]
  )(implicit
    rng: RandomGenerator
  ): AdjusterState = {
    val ret = oldState.accumulateAssessments(assesments.toVector)
    val squaredNormalizer = Math.pow(ret.normalizer, exponent)
    logger.debug(s"Normalized Momentum value: ${ret.momentum.map(_/squaredNormalizer)}")
    logger.debug(s"Normalizer: ${ret.normalizer}")
    ret
  }


  override def updateStrategy(
    mixedStrategyNode: MixedStrategyNode[_,_,_],
    adjusterState: AdjusterState,
    assesments: Seq[Double]
  )(implicit
    rng: RandomGenerator,  iterationInfo: IterationInfo
  ) : Boolean = {
    if(adjusterState.normalizer > 0) {
      val squaredNormalizer = Math.pow(adjusterState.normalizer, exponent)
      mixedStrategyNode.updateBy(adjusterState.momentum.map(_/squaredNormalizer))
      true
    } else {
      false
    }   
  }
}
