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

import org.apache.commons.math3.linear.{Array2DRowRealMatrix, ArrayRealVector, LUDecomposition}
import pl.edu.pw.mini.sg.rng.RandomGenerator
import org.log4s._
import pl.edu.pw.mini.sg.game.deterministic.PlayerVisibleState
import pl.edu.pw.mini.sg.mixed.IterationInfo
import pl.edu.pw.mini.sg.mixed.adjuster.{AssesmentBasedAdjuster, StatefulAdjuster}
import pl.edu.pw.mini.sg.mixed.tree.{InnerResultNode, MixedStrategyNode}

abstract class BisectRotateAdjusterBase(
  val fallbackAdjuster: AssesmentBasedAdjuster,
  val edgeDist: Double
) extends StatefulAdjuster {
  import BisectRotateAdjusterBase.orthogonalUpdate
  protected val logger = getLogger
  override type AdjusterState = BisectAdjusterWithOldUpdateState[fallbackAdjuster.AdjusterState]

  implicit val ost = new BisectAdjusterWithOldUpdateStateType(fallbackAdjuster.ost)

  def normalizationAfterRotate[DM, DS <: PlayerVisibleState[DM]](
    mixedStrategy: MixedStrategyNode[DM,DS, AdjusterState],
    desiredAttackerPayoff: InnerResultNode[DM,DS],
    optimalAttackerPayoff: InnerResultNode[DM,DS]
  )(implicit iterationInfo: IterationInfo) : Unit

  def inTreeAdjust[DM, DS <: PlayerVisibleState[DM]](
    mixedStrategy: MixedStrategyNode[DM,DS, AdjusterState],
    desiredAttackerPayoff: InnerResultNode[DM,DS],
    optimalAttackerPayoff: InnerResultNode[DM,DS]
  )(
    implicit rng: RandomGenerator, iterationInfo: IterationInfo) : Boolean = {
    val diff = desiredAttackerPayoff.payoff.distortedAttacker.attacker - optimalAttackerPayoff.payoff.distortedAttacker.attacker
    def diffm(move: DM) = desiredAttackerPayoff.payoff(move).distortedAttacker.attacker - optimalAttackerPayoff.payoff(move).distortedAttacker.attacker
    lazy val assesments = mixedStrategy.probabilities.view.map(p => diffm(p._1) - diff).toVector
    logger.debug(s"Strategy before adjust ${mixedStrategy.probabilities.map(_._2).toVector}")
    logger.debug(s"Payoffs before update DAP:${{desiredAttackerPayoff.updateValues();desiredAttackerPayoff.payoff}} OAP:${{optimalAttackerPayoff.updateValues();optimalAttackerPayoff.payoff}}")
    mixedStrategy.optimizerState match {
      case found : BisectAlreadyFoundState => {
        found.initialNewPosition match {
          case Some(initialPos) if found.lastGoodPosition.zip(mixedStrategy.probabilities)
              .map{case (a,(_,b)) => (a-b)*(a-b)}.sum < edgeDist*edgeDist =>
            val feasibilityDir = assesments
            val payoffDir = initialPos.zip(found.lastGoodPosition).map{case (a,b) => a-b}
            logger.debug(s"Projected update length ${Math.sqrt(payoffDir.map(x => x*x).sum)}")
            mixedStrategy.probabilities.view.zipWithIndex.zip(orthogonalUpdate(feasibilityDir, payoffDir)).zip(found.lastGoodPosition)
              .foreach{case ((((dm, _), idx), dir), lgp) =>
              mixedStrategy.updateProb(idx, Math.max(0,(dir+lgp)))
              }
            normalizationAfterRotate(mixedStrategy, desiredAttackerPayoff, optimalAttackerPayoff)
            logger.debug(s"Strategy after rotate adjust ${mixedStrategy.probabilities.map(_._2).toVector}")
            logger.debug(s"Payoffs after update DAP:${{desiredAttackerPayoff.updateValues();desiredAttackerPayoff.payoff}} OAP:${{optimalAttackerPayoff.updateValues();optimalAttackerPayoff.payoff}}")
            mixedStrategy.optimizerState=found.registerFeasible(found.lastGoodPosition)
          case _ => {
            mixedStrategy.probabilities.view.zipWithIndex.zip(found.lastGoodPosition).foreach{case (((dm, pr), idx), oldpr) =>
              mixedStrategy.updateProb(idx, (pr+oldpr)/2)
            }
            logger.debug(s"Last good position ${found.lastGoodPosition}")
            logger.debug(s"Strategy after bisect ${mixedStrategy.probabilities.map(_._2).toVector}")
            mixedStrategy.optimizerState=found.registerInfeasible(mixedStrategy.probabilities.map(_._2).toVector)
          }
        }
        true
      }
      case fallback : BisectFallbackState[fallbackAdjuster.AdjusterState] => {
        val newFS = fallbackAdjuster.updateState(fallback.fallback, assesments)
        mixedStrategy.optimizerState = new BisectFallbackState(newFS)
        val r = fallbackAdjuster.updateStrategy(mixedStrategy, newFS, assesments)
        logger.debug(s"Strategy after fallback ${mixedStrategy.probabilities.map(_._2).toVector}")
        r
      }
    }
  }
  
  def outOfTreeAdjust[DM, DS <: PlayerVisibleState[DM]](
    mixedStrategy: MixedStrategyNode[DM,DS,AdjusterState],
    optimalAttackerPayoff: InnerResultNode[DM,DS]
  )(
    implicit rng: RandomGenerator, iterationInfo: IterationInfo) : Boolean = {
    val state = mixedStrategy.optimizerState match {
      case s : BisectFallbackState[fallbackAdjuster.AdjusterState] => s
      case _ => throw new IllegalStateException(s"Bad node state ${mixedStrategy.optimizerState}")
    }
    val assesments = mixedStrategy.probabilities.view.map(p =>
      optimalAttackerPayoff.payoff.distortedAttacker.attacker - optimalAttackerPayoff.payoff(p._1).distortedAttacker.attacker).toVector
    val fState = fallbackAdjuster.updateState(
      state.fallback,
      assesments
    )

    mixedStrategy.optimizerState = new BisectFallbackState(fState)
    fallbackAdjuster.updateStrategy(mixedStrategy, fState, assesments)
  }
  def positiveAdjust[DM, DS <: PlayerVisibleState[DM]](
    mixedStrategy: MixedStrategyNode[DM,DS,AdjusterState],
    gamePayoff: InnerResultNode[DM,DS]
  )(
    implicit rng: RandomGenerator, iterationInfo: IterationInfo) : Boolean = {
    logger.debug(s"Feasible payoff: ${gamePayoff.payoff}")
    val state = mixedStrategy.optimizerState
    val multiplier = 1.0 /* state match {
      case s : BisectFallbackState[fallbackAdjuster.AdjusterState] => 1
      case s : BisectAlreadyFoundState => Math.pow(Math.sqrt(s.lastGoodPosition.zipAll(mixedStrategy.strategy.map(_._2), 0.0, 0.0)
        .map{case (a,b) => (a-b)*(a-b)}.sum),1.0/20.0)
    }*/

//    logger.debug(s"Positive multiplier: ${multiplier}")
    mixedStrategy.optimizerState = state.registerFeasible(mixedStrategy.probabilities.view.map(_._2).toVector)
    val assesments = mixedStrategy.probabilities.view.map(p =>
     gamePayoff.payoff(p._1).real.defender - gamePayoff.payoff.real.defender).toVector

    mixedStrategy.probabilities.view.zipWithIndex.zip(assesments).foreach{
      case (((dm, prob), idx), as) => {
        mixedStrategy.updateProb(idx, Math.max(0, prob+as*multiplier))
      }}

    mixedStrategy.normalizeOrEqual()
    logger.debug(s"Strategy after positive pass ${mixedStrategy.probabilities.map(_._2).toVector}")
    true
  }
}

object BisectRotateAdjusterBase {
  def orthogonalUpdate(feasibilityDirection: Vector[Double], defenderPayoffDirection: Vector[Double]) : Vector[Double] = {
    require(feasibilityDirection.length >= defenderPayoffDirection.length)
    val feasibilityDirectionS = feasibilityDirection.take(defenderPayoffDirection.length)
    // n+1 zmiennych, n pierwszych -- współrzędne szukanego wektora, 1, współczynnik t
    // x wsp. szukanego wektora
    // f - wsp feasibility wektora
    // d - wsp def wektora
    // x_1*f_1+x_2f_2+++x_nf_n = 0
    // \forall i x_i-f_i*t=a_i
    val n = feasibilityDirectionS.length
    val hyperplaneCoefficients : Vector[Double] = (feasibilityDirectionS:+0.0)
    val hypeplaneLhs : Double = 0
    val zeroVector : Vector[Double] = Vector.fill(n+1)(0)
    def parallelCoeff(pos: Int) : Vector[Double] = zeroVector.updated(pos, 1.0)
      .updated(n, -feasibilityDirectionS(pos))
    def parallelCoeffLhs(pos: Int) : Double = defenderPayoffDirection(pos)

    val cm : Seq[Array[Double]] = hyperplaneCoefficients.toArray+:Range(0,n).map(parallelCoeff(_).toArray[Double])

    val coefficients = new Array2DRowRealMatrix(cm.toArray, false)
    val constants = new ArrayRealVector((hypeplaneLhs+:Range(0,n).map(x => parallelCoeffLhs(x))).toArray)
    val solver = new LUDecomposition(coefficients).getSolver()
    val solution = solver.solve(constants)
    Vector.tabulate(n)(solution.getEntry(_))
  }
}
