package sim

import model.HHCTypes
import scala.reflect.ClassTag
import scala.collection.mutable.{ArrayBuffer, ListBuffer}
import scala.collection.immutable.ArraySeq
import model.HHCEdge2
import util.Util
import model.Customer
import scala.io.BufferedSource
import model.Coord
import scala.util.Random
import com.typesafe.scalalogging.LazyLogging
import cats.implicits
import scala.util.chaining._

class HHCSim2(
    epsilonMax: Int = 0,
    iterations: Int = 0,
    WLDMax: Float = 0,
    customers: String Either ArraySeq[Customer]
) extends LazyLogging {
  import HHCSim2._
  import Ordering.Double.IeeeOrdering

  private var adjMatrix: Option[GraphMatrix[Double]] = None
  def go() = {
    customers
      .flatMap(checkEmpty)
      .map(formAdjMatrix)
      .tap(logGraph)
      .tap(_.foreach(e => adjMatrix = Some(e)))
      .map(e => mstUsingPrims(e))
      .tap(logMST)
      .map { case (mst, epsilon) => removeEdges(mst)(epsilon) }
      .tap(logMST2)
      .map { case (mstUpdated, removed) => findClusters(mstUpdated)(removed) }
      .map {
        case (mstUpdated, clusters, edgeMappings, removedEdges) =>
          printClusters(clusters)
          val e =
            groupClusters(mstUpdated, clusters, edgeMappings, removedEdges)
          printGroups(e._5)
          e
      }
      .flatMap(value => {
        val (mst, clusters, edgeMappings, removed, clusterGroups) = value
        balanceWorkload(clusters, clusterGroups)
      })
  }
  // def go2() = {
  //   val res = go().flatMap(value => {
  //     val (mst, clusters, edgeMappings, removed, clusterGroups) = value
  //     balanceWorkload(mst, clusters, clusterGroups)
  //   })
  //   res
  // }

  def printGraph(graph: GraphMatrix[Double]) = graph.foreach { e =>
    e.foreach { d =>
      print(f"$d%.2f, ")
    }
    println
  }

  def printGraph(graph: Graph[Double]) = {
    for (i <- 0 until graph.size) {
      print(s"$i: ")
      for (edge <- graph(i)) {
        print(f"(${edge.toNode}, ${edge.weight}%.2f), ")
      }
      println
    }
  }

  def logGraph(maybeGraph: Either[String, GraphMatrix[Double]]): Unit =
    maybeGraph.foreach(
      graph => {
        logger.whenDebugEnabled {
          logger.debug("Graph: ")
          printGraph(graph)
        }
      }
    )

  def logMST(maybeMST: Either[String, (Graph[Double], Double)]): Unit =
    maybeMST.foreach(
      t => {
        logger.whenDebugEnabled {
          logger.debug(s"Epsilon = ${t._2}")
          logger.debug("MST: ")
          printGraph(graph = t._1)
        }
      }
    )

  def logMST2(it: Either[String, (Graph[Double], RemovedEdges[Double])]): Unit =
    it.foreach(
      t => {
        logger.whenDebugEnabled {
          logger.debug(s"Removed Edges = ${t._2}")
          logger.debug("MST: ")
          printGraph(graph = t._1)
        }
      }
    )

  def printClusters(clusters: Clusters) = {
    logger.whenDebugEnabled {
      logger.debug("Clusters:")
      var i = 0
      clusters.foreach(e => {
        print(s"Cluster-$i: ")
        i += 1
        e.foreach(f => print(s"$f "))
        println
      })
    }
  }

  def printClusters(clusters: Array[List[Int]]) = {
    logger.whenDebugEnabled {
      logger.debug("Clusters:")
      var i = 0
      clusters.foreach(e => {
        print(s"Cluster-$i: ")
        i += 1
        e.foreach(f => print(s"$f "))
        println
      })
    }
  }

  def printGroups(groups: ArraySeq[List[(Int, Double)]]) = {
    logger.whenDebugEnabled {
      logger.debug("Neighbours: ")
      var i = 0
      groups.foreach(e => {
        print(s"Cluster-$i: ")
        i += 1
        e.foreach(f => { print(f"(${f._1}%d, ${f._2}%.2f), ") })
        println
      })
    }
  }

  def balanceWorkload[N: Numeric](
      clusters: Clusters,
      groups: ArraySeq[List[(Int, N)]]
  ) = {
    customers.map(c => {
      val mutClusters =
        clusters.toArray
      val clusterOrder = (0 until clusters.length).toArray
      logger.debug(s"Initial order: $clusterOrder")
      val k = mutClusters.size
      for (_ <- 0 until iterations) {
        val WL = mutClusters.view
          .map(nodes => {
            val customers = nodes.map(node => c(node))
            val wl = customers.foldLeft(0f)(_ + _.workload)
            wl
          })
          .to(ArraySeq)
        logger.debug(s"$WL")
        clusterOrder.sortInPlaceWith((x, y) => WL(x) < WL(y))
        logger.debug(s"Current order $clusterOrder")
        val (minWL, maxWL) = WL.foldLeft((99999f, 0f))((acc, wl) => {
          val (currMin, currMax) = acc
          if (wl < currMin) (wl, currMax)
          else if (wl > currMax) (currMin, wl)
          else acc
        })
        val WLD = maxWL - minWL

        logger.debug(s"maxWL = $maxWL, minWL = $minWL")
        logger.debug(s"WLD = $WLD")
        if (WLD > WLDMax) {
          val randNum = Random.between(((k / 2) + 1), k)
          logger.debug(s"Rand Num=$randNum")
          val effectiveNum = clusterOrder(randNum)
          logger.debug(s"Chosen cluster = $effectiveNum")
          val (s, _) = groups(effectiveNum)(0)
          logger.debug(s"Nearest neighbour = $s")
          // val avg = averagePairwiseDistance(mst(s))
          val avg = adjMatrix
            .map(m => averagePairwiseDistance(m)(mutClusters(s)))
            .getOrElse(9999d)
          //  if (avg < epsilonMax) true else false

          logger.debug(s"Pairwise average = $avg")

          if (avg < epsilonMax) {
            logger.debug("Average less than epsilonMax. Updating Clusters.")
            if (!mutClusters(s).contains(effectiveNum)) {
              mutClusters(s) = effectiveNum :: mutClusters(s)
            }
            // mutClusters.foreach(println)
            printClusters(mutClusters)
          }
        }
      }
      ArraySeq.unsafeWrapArray(mutClusters)
    })
  }

  def averagePairwiseDistance(graph: GraphMatrix[Double])(
      cluster: List[Int]
  ) = {
    val (sum, size) = cluster
      .combinations(2)
      .foldLeft((0d, 0))((acc, pair) => {
        val (sum, size) = acc
        val (i, j) = pair(0) -> pair(1)
        logger.trace(s"$i $j - ")
        logger.trace(s"weight = ${graph(i)(j)}")
        (sum + graph(i)(j), size + 1)
      })

    val n = if (size == 0) 1 else size
    sum / n
  }
}

object HHCSim2 extends HHCTypes {
  import Ordering.Implicits._

  val composed = (formAdjMatrix _)
    .andThen(e => mstUsingPrims(e))
    .andThen { case (mst, epsilon) => removeEdges(mst)(epsilon) }
    .andThen { case (mstUpdated, removed) => findClusters(mstUpdated)(removed) }

  def getCustomers(
      infile: BufferedSource
  ): String Either ArraySeq[Customer] = {
    var customers: String Either ArraySeq[Customer] = Right(ArraySeq.empty)
    val it = infile.getLines
    @annotation.tailrec
    def loop(
        lst: ListBuffer[Customer],
        iter: Iterator[String]
    ): String Either ListBuffer[Customer] = {
      if (!iter.hasNext) Right(lst)
      else {
        val line = iter.next
        val arr = line.split(",").map(_.trim)
        arr match {
          case Array(latitude, longitude, workLoad) => {
            val cust =
              for {
                lat <- latitude.toDoubleOption
                lon <- longitude.toDoubleOption
                coord = Coord(lat, lon)
                wl <- workLoad.toFloatOption
              } yield (Customer(coord, wl))
            cust match {
              case Some(c) => loop(lst += c, iter)
              case None    => Left(s"Error reading customers at line - $line")
            }
          }
          case _ => {
            if (line.equals(" ") || line.contains("\n"))
              Left("Error newline")
            else {
              Left(
                "Error reading customers from" +
                  s" file at line - $line}"
              )
            }
          }
        }
      }
    }
    try {
      customers = loop(ListBuffer.empty, it).map(_.to(ArraySeq))
    } catch {
      case e: NumberFormatException =>
        customers = Left(
          s"Expected number but received string ${e.getMessage()}"
        )
      case e: NullPointerException =>
        customers = Left("Input was null")
    }
    customers
  }

  def checkEmpty(
      customers: IndexedSeq[Customer]
  ): Either[String, IndexedSeq[Customer]] =
    customers.length match {
      case 0 => Left("Error input was empty")
      case _ => Right(customers)
    }

  def formAdjMatrix(customers: IndexedSeq[Customer]): GraphMatrix[Double] = {
    val n = customers.length
    val edges: Array[Array[Double]] = Array.ofDim(n, n)
    for (i <- 0 until n) {
      for (j <- i until n) {
        val weight =
          Util.getHaversineDistance(
            customers(i).location,
            customers(j).location
          )
        edges(i)(j) = weight
        edges(j)(i) = weight
      }
    }
    edges
  }

  def mstUsingPrims[T: Numeric](
      edges: GraphMatrix[T]
  ): (Graph[T], T) = {
    val num = implicitly[Numeric[T]]
    val n = edges.length
    val selected: Array[Boolean] = Array.fill(n)(false)

    selected(0) = true

    val adjList =
      Array.fill(n)(ListBuffer.empty[HHCEdge2[T]])

    var sum = 0
    var count = 0

    for (_ <- 0 until n - 1) {
      var min = 999999
      var x = 0
      var y = 0
      for (i <- 0 until n) {
        if (selected(i) == true) {
          for (j <- 0 until n) {
            if (selected(j) == false && edges(i)(j) != 0) {
              if (min > num.toInt(edges(i)(j))) {
                min = num.toInt(edges(i)(j))
                x = i
                y = j
              }
            }
          }
        }
      }
      sum += num.toInt(edges(x)(y))
      adjList(x) += HHCEdge2(x, y, edges(x)(y))
      adjList(y) += HHCEdge2(y, x, edges(x)(y))
      selected(y) = true
    }
    val adjList2 = adjList.map(l => {
      count += l.size
      l.toList
    })
    // adjList2.foreach(println)
    (ArraySeq.unsafeWrapArray(adjList2), num.fromInt(sum / (count / 2)))
  }

  def removeEdges[N: Ordering](
      mst: Graph[N]
  )(epsilon: N): (Graph[N], RemovedEdges[N]) = {

    val removed = ListBuffer.empty[HHCEdge2[N]]

    val result = ArraySeq.tabulate(mst.length) { i =>
      val (filtered, rm) = mst(i)
      // .view
      // .filter(e => (e.fromNode <= e.toNode))
        .partition(_.weight <= epsilon)

      // val rm2 = rm.filter(e => (e.fromNode <= e.toNode))

      removed ++= rm

      filtered
    }
    println
    // result.foreach(println)
    // println
    (result, removed.toList)
  }

  def DFS[T: Numeric](start: Int)(graph: Graph[T]) = {
    val visited = Array.fill(graph.size)(false)
    val buf = ListBuffer[Int]()
    def loop(
        start: Int,
        graph: Graph[T],
        visited: Array[Boolean]
    ): Unit = {
      visited(start) = true
      buf += start
      // val iter = graph(start).iterator
      // while (iter.hasNext) {
      //   val edge = iter.next
      //   if (!visited(edge.toNode))
      //     loop(edge.toNode, graph, visited)
      // }
      for (edge <- graph(start)) {
        if (!visited(edge.toNode))
          loop(edge.toNode, graph, visited)
      }
    }

    loop(start, graph, visited)
    buf.toList
  }

  def findClusters[T: Numeric](mstUpdated: Graph[T])(
      removedEdges: RemovedEdges[T]
  ): (Graph[T], Clusters, ArraySeq[List[HHCEdge2[T]]], RemovedEdges[T]) = {
    val visited = Array.fill[Boolean](mstUpdated.length)(false)
    var removedEdges2 = removedEdges
    val egdeMappings =
      Array.fill(mstUpdated.length)(List.empty[HHCEdge2[T]])
    val nodeToClusterMappings = Array.fill(mstUpdated.length)(0)
    val result = ArraySeq.tabulate(mstUpdated.length) { i =>
      {
        val buf = ListBuffer[Int]()
        visited(i) match {
          case true =>
          case false if (!mstUpdated(i).isEmpty) => {
            val nodes = DFS(i)(mstUpdated)
            buf ++= nodes
            val (nds, rms) = assignEdges(nodes, removedEdges2)
            removedEdges2 = rms
            egdeMappings(i) = nds
            for (j <- nodes) visited(j) = true
          }
          case false => {
            buf += i
            val (nds, rms) = assignEdges(List(i), removedEdges2)
            egdeMappings(i) = nds
            removedEdges2 = rms
          }
        }
        buf.toList
      }
    }

    // println(s"Removed edges size: ${removedEdges2.size}")

    (
      mstUpdated,
      result.filterNot(_.isEmpty),
      ArraySeq.unsafeWrapArray(egdeMappings.filterNot(_.isEmpty)),
      removedEdges
    )
  }

  def assignEdges[T: Ordering](
      nodes: List[Int],
      removedEdges: RemovedEdges[T]
  ) = {
    val it = nodes.iterator
    @annotation.tailrec
    def loop(
        nodes: List[Int],
        removedEdges: RemovedEdges[T],
        iter: Iterator[Int],
        edges: List[HHCEdge2[T]]
    ): (List[HHCEdge2[T]], RemovedEdges[T]) = {
      if (!iter.hasNext) (edges, removedEdges)
      else if (!edges.isEmpty) (edges, removedEdges)
      else {
        val node = iter.next
        val (filt, rm) =
          removedEdges.partition(e => e.fromNode == node)
        loop(nodes, rm, iter, filt)
      }
    }
    val (edges, removedEdgesUpdated) =
      loop(nodes, removedEdges, it, List.empty[HHCEdge2[T]])
    (edges, removedEdgesUpdated)
  }

  def mapNodestoClusters[N: Ordering](mst: Graph[N], clusters: Clusters) = {
    val arr = Array.fill(mst.size)(0)
    for (i <- 0 until clusters.size) {
      val nodes = clusters(i)
      for (node <- nodes) {
        arr(node) = i
      }
    }
    ArraySeq.unsafeWrapArray(arr)
  }

  def groupClusters[N: Ordering](
      mst: Graph[N],
      clusters: Clusters,
      edgeMappings: ArraySeq[List[HHCEdge2[N]]],
      removed: RemovedEdges[N]
  ) = {
    val nodeMap = mapNodestoClusters(mst, clusters)
    val groups = ArraySeq.tabulate(edgeMappings.size)(i => {
      val buf = ListBuffer.empty[(Int, N)]
      val lst = edgeMappings(i)
      lst.foreach(e => {
        buf += nodeMap(e.toNode) -> e.weight
      })
      buf.sortWith {
        case ((_, weight1), (_, weight2)) =>
          weight1 < weight2
      }.toList
    })
    (mst, clusters, edgeMappings, removed, groups)
  }
}