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

import pl.edu.pw.mini.sg.rng.RandomGenerator
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 object MomentumUpdateSignLimitAdjuster extends AssesmentBasedAdjuster {
  override type AdjusterState = Momentum
  override val ost = MomentumOptimizerStateType

  override def updateState(
    oldState: AdjusterState,
    assesments: Seq[Double]
  )(implicit
    rng: RandomGenerator
  ): AdjusterState = {
    if(assesments.view.zip(oldState.momentum).map{case (a, b) => a*b}.sum <= 0.1) {
      oldState.accumulateAssessments(assesments.toVector)
    } else {
      oldState.extend(assesments.length)
    }
  }


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