package pl.edu.pw.mini.sg.efgame.gen

import java.util.UUID

import scala.annotation.tailrec

import org.apache.commons.math3.distribution.EnumeratedIntegerDistribution
import org.apache.commons.math3.linear.Array2DRowRealMatrix
import org.apache.commons.math3.random.{CorrelatedRandomVectorGenerator, GaussianRandomGenerator, RandomDataGenerator, RandomGenerator}
import pl.edu.pw.mini.sg.efgame.{EfGameConfig, GameNode, InnerGameNode, LeafGameNode}

object Generator {
  private class CachedRangeList {
    import scala.collection.mutable.{Map => MMap}
    private val cache = MMap[Int, List[Int]]()

    def get(endExclusive: Int) : List[Int] =
      cache.getOrElseUpdate(endExclusive, {if(endExclusive == 1) List(0)
      else (endExclusive-1)::get(endExclusive-1)
        })
  }

  def leafProbability(n: Int, bf: Double) : Double = -(bf-1)/(bf*(Math.pow(1.0/bf, n.toDouble) -1))
  @tailrec
  def assignIs(stateCount: Int, minIs: Int, isSize: Vector[Int], eqBias: Double, currAssignment: Vector[Int])(implicit rng: RandomGenerator) : Vector[Int] = {
    val emptyIs = isSize.count(_==0)
    require(emptyIs <= stateCount, s"eis: $emptyIs, sc: $stateCount")
    if(isSize.length == 0) {
      Vector()
    } else if(stateCount == 0) {
      val rgd = new RandomDataGenerator(rng)
      val perm = rgd.nextPermutation(currAssignment.length, currAssignment.length)
      perm.map(currAssignment(_)).toVector
    } else {
      val assignment = if(emptyIs == stateCount) {
        isSize.zipWithIndex.find(_._1 == 0).get._2
      } else {
        val equalProb = 1/isSize.length.toDouble
        val assignedCount = Math.max(1,isSize.sum)
        val weights = isSize.map(ss => eqBias*(assignedCount-ss).toDouble/assignedCount+equalProb)
        val sum = weights.sum
        val probs = weights.map(_/sum)
        val ee = new EnumeratedIntegerDistribution(rng, Range(0, isSize.length).toArray, probs.toArray)
        ee.sample()
      }
      assignIs(stateCount-1, minIs, isSize.updated(assignment, isSize(assignment)+1),
        eqBias, currAssignment:+(assignment+minIs))
    }
  }

  def theFun(x: Double, y: Double)= Math.pow(Math.abs((x/2)+.5),(Math.tan(y*3.14)))

  def isCount(nodes: Int, lambda: Double)(implicit rng: RandomGenerator) : Int = {
    if(lambda == 0.0) 1
    else if(lambda == 1.0 || nodes==1 || nodes==0) nodes
    else {
      val y = if(lambda <=0.5) 0.5-lambda else lambda-0.5
      val probs = Range.inclusive(1, nodes).map(nn => theFun(nn.toDouble/(nodes), y)).toArray
      require(probs.sum > 0 && probs.forall(x => !x.isInfinite()), s"Bad probs: ${probs.toVector}, lambda: $lambda, y: $y, nodes: $nodes")
      val samples = (if(lambda <= 0.5) Range.inclusive(1,nodes).reverse else Range.inclusive(1, nodes)).toArray
      val gen = new EnumeratedIntegerDistribution(rng, samples, probs)
      gen.sample()
    }
  }

  private sealed trait NodeDescription {
    val defenderHistory: Vector[(Int, Int)]
    val attackerHistory: Vector[(Int, Int)]
  }
  private final case class InnerDescription(defenderIs: Int, attackerIs: Int, payoff: Double,
    defenderHistory: Vector[(Int, Int)], attackerHistory: Vector[(Int, Int)]) extends NodeDescription
  private final case class LeafDescription(defenderPayoff: Double, attackerPayoff: Double, defenderHistory: Vector[(Int, Int)], attackerHistory: Vector[(Int, Int)]) extends NodeDescription

  def randomGame(config: GeneratorConfig)(implicit rng: RandomGenerator) : EfGameConfig = {
    val bf= (config.defenderMovesMax + config.defenderMovesMin.toDouble)*(config.attackerMovesMin + config.attackerMovesMax.toDouble)/4
    val covarianceMatrix = new Array2DRowRealMatrix(Array(
      Array(config.defenderPayoffVariance, config.playersPayoffCovaciance),
      Array(config.playersPayoffCovaciance, config.attackerPayoffVariance)
    ))

    def buildNextTreeLevel(depth: Int, prevLevel: Vector[NodeDescription], currentDIS: Vector[Int], currentAIS: Vector[Int]) : (Vector[NodeDescription], Vector[Int], Vector[Int]) = {
      val expandableNodes: Vector[InnerDescription] = prevLevel.collect{case (x : InnerDescription) => x}
      def expandSingleNode(id: InnerDescription) = {
        val defenderMoves = Range(0, currentDIS(id.defenderIs))
        val attackerMoves = Range(0, currentAIS(id.attackerIs))
        defenderMoves.flatMap(dm => attackerMoves.map(am => InnerDescription(-1, -1, id.payoff+rng.nextGaussian()*config.payoffDriftVariance,
          id.defenderHistory:+((id.defenderIs, dm)), id.attackerHistory:+((id.attackerIs, am)))))
      }

      val expandedPhase1 = expandableNodes.flatMap(expandSingleNode).toVector
      val lp = if(depth >= config.gameDepthMin) leafProbability(config.gameDepthMax+1-depth, bf) else 0.0
      val leafToBe : Vector[Boolean] = Vector.fill(expandedPhase1.size)(rng.nextDouble() < lp)
      val distinguishableDefenderNodes = Range(0, expandedPhase1.length).filter(leafToBe.andThen(!_)).groupBy(idx => expandedPhase1(idx).defenderHistory)
      val distinguishableAttackerNodes = Range(0, expandedPhase1.length).filter(leafToBe.andThen(!_)).groupBy(idx => expandedPhase1(idx).attackerHistory)

      val (defenderIsAssignment, newDIS) =  distinguishableDefenderNodes.foldLeft((Map[Int, Int](), currentDIS)){ case ((assignment, dis), (_, nodes)) => {
        val disc = isCount(nodes.length, config.defenderLambda)
        val dassignment: Vector[Int] = assignIs(nodes.length, dis.length, Vector.fill(disc)(0), config.defenderEqualIsBias, Vector())
        val newIs = Vector.fill(disc)(rng.nextInt(config.defenderMovesMax-config.defenderMovesMin)+config.defenderMovesMin)
        (nodes.zip(dassignment).foldLeft(assignment)((oldmap, pair) => oldmap+pair), dis++newIs)
      }}

      val (attackerIsAssignment, newAIS) =  distinguishableAttackerNodes.foldLeft((Map[Int, Int](), currentAIS)){ case ((assignment, ais), (_, nodes)) => {
        val aisc = isCount(nodes.length, config.attackerLambda)
        val aassignment: Vector[Int] = assignIs(nodes.length, ais.length, Vector.fill(aisc)(0), config.attackerEqualIsBias, Vector())
        val newIs = Vector.fill(aisc)(rng.nextInt(config.attackerMovesMax-config.attackerMovesMin)+config.attackerMovesMin)
        (nodes.zip(aassignment).foldLeft(assignment)((oldmap, pair) => oldmap+pair), ais++newIs)
      }}

      (expandedPhase1.zipWithIndex.map{case (id, idx) => {
        if(leafToBe(idx)) {
          val gen = new CorrelatedRandomVectorGenerator(Array(id.payoff, id.payoff), covarianceMatrix, 1e-6,
            new GaussianRandomGenerator(rng))
          val rec = gen.nextVector()
          require(rec.size == 2)
          LeafDescription(Math.round(rec(0)*100)/100.0, Math.round(rec(1)*100)/100.0, id.defenderHistory, id.attackerHistory)
        }
        else {
          id.copy(defenderIs = defenderIsAssignment(idx), attackerIs = attackerIsAssignment(idx))
        }
      }}, newDIS, newAIS)
    }

    val firstLevel : (Vector[NodeDescription], Vector[Int], Vector[Int]) = {
      val defMoves = rng.nextInt(config.defenderMovesMax-config.defenderMovesMin)+config.defenderMovesMin
      val attMoves = rng.nextInt(config.attackerMovesMax-config.attackerMovesMin)+config.attackerMovesMin
      (
        Vector(InnerDescription(0, 0, rng.nextGaussian()+0.5, Vector(), Vector())),
        Vector(defMoves), Vector(attMoves)
      )
    }

    val allLevels = Range.inclusive(1, config.gameDepthMax).toVector.scanLeft(firstLevel){case ((vec, defIs, atIs), depth) => buildNextTreeLevel(depth, vec, defIs, atIs)}
    val dIs = allLevels.last._2
    val aIs = allLevels.last._3
    require(allLevels.last._1.forall(n => n.getClass == classOf[LeafDescription]))

    val lastLevel : Map[(Vector[(Int, Int)], Vector[(Int, Int)]), GameNode] = allLevels.last._1.view
      .collect{case (ld: LeafDescription) => ((ld.defenderHistory, ld.attackerHistory),LeafGameNode(ld.defenderPayoff, ld.attackerPayoff))}.toMap

    val tree = allLevels.init.foldRight(lastLevel) {case ((theLevel, _, _), children) => {
      theLevel.view.map(nd => nd match {
        case LeafDescription(dp, ap, dh, ah) => ((dh, ah), LeafGameNode(dp, ap))
        case InnerDescription(dis, ais, _, dh, ah) => ((dh, ah), InnerGameNode(dis, ais,
          Range(0, dIs(dis)).map(dm => (dm, Range(0, aIs(ais)).map(am => (am, children((dh:+((dis, dm)), ah:+((ais, am)))))).toMap)).toMap
        ))
      }).toMap
    }}

    require(tree.size == 1)
    val crl = new CachedRangeList()
    EfGameConfig(tree.head._2.asInstanceOf[InnerGameNode],
      aIs.map(crl.get).zipWithIndex.map(_.swap).toMap,
      dIs.map(crl.get).zipWithIndex.map(_.swap).toMap,
      UUID.randomUUID())
  }
}
