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

class HHCSim2(
    private val epsilonMax: Int = 0,
    private val iterations: Int = 0,
    private val WLDMax: Float = 0,
    private val customers: String Either ArraySeq[Customer]
) {
  import HHCSim2._
  import Ordering.Double.IeeeOrdering
  def go() = {
    customers
      .flatMap(checkEmpty)
      .map(formAdjMatrix)
      .map(e => mstUsingPrims(e))
      .map { case (mst, epsilon) => removeEdges(mst)(epsilon) }
      .map { case (mstUpdated, removed) => findClusters(mstUpdated)(removed) }
      .map {
        case (clusters, edgeMappings, removedEdges) =>
          groupClusters(clusters, edgeMappings, removedEdges)
      }
  }
  def go2() = {
    val res = go().flatMap(value => {
      val (clusters, edgeMappings, removed, clusterGroups) = value
      workloadBalance(clusters, edgeMappings, removed, clusterGroups)
    })
    res
  }

  def workloadBalance[N: Numeric](
      clusters: Clusters,
      edgeMappings: ArraySeq[List[HHCEdge2[N]]],
      removedEdges: RemovedEdges[N],
      groups: ArraySeq[(Int, List[(Int, N)])]
  ) = {
    customers.map { c =>
      val k = clusters.size
      val mutClusters = clusters.toArray
      val worlkLoads =
        mutClusters.map(_.map(j => c(j)).foldLeft(0f)((x, y) => x + y.workload))
      val sortedClusters =
        mutClusters.map(_.sortWith((x, y) => c(x).workload < c(y).workload))
      val minWL =
        c.foldLeft(c(0).workload)(_ min _.workload)
      val maxWL =
        c.foldLeft(0f)(_ max _.workload)
      val WLD = maxWL - minWL
      if (WLD > WLDMax) {
        var t = Random.between(((k / 2) + 1), k)
        while (!mutClusters(t).isEmpty) t = Random.between(((k / 2) + 1), k)
        val Ds = groups(t)._2(0)
        mutClusters.updated(Ds._1, t)
      }
      for (_ <- 0 until iterations) {
        // val workloads = clusters.flatMap(_.map(j => c(j)))

      }
    }
  }

  def checkAveragePairwiseDistance[T](
      lst: List[(Int, T)]
  )(implicit num: Numeric[T]) = {
    val sum = lst.foldLeft(0)((x, y) => x + num.toInt(y._2))
    val avg = sum / lst.length
    if (avg < epsilonMax) true else false
  }
}

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 =
              Customer(
                Coord(latitude.toDouble, longitude.toDouble),
                workLoad.toFloat
              )
            loop(lst += cust, iter)
          }
          case _ => {
            if (arr.mkString.equals(" ") || arr.mkString.contains("\n"))
              Left("Error newline")
            else {
              Left(
                "Error reading customers from" +
                  s" file - ${arr.mkString(", ")}"
              )
            }
          }
        }
      }
    }
    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 += 1
      l.toList
    })
    adjList2.foreach(println)
    (ArraySeq.unsafeWrapArray(adjList2), num.fromInt(sum / count))
  }

  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]
  ): (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 result = ArraySeq.tabulate(mstUpdated.length) { i =>
      {
        val buf = ListBuffer[Int]()
        // mstUpdated(i).isEmpty match {
        //   case true => { lst += i }
        //   case false => {
        //     if (!visited(i)) {
        //       lst ++= DFS(i)(mstUpdated)
        //       for (j <- lst) visited(j) = true
        //     }
        //   }
        // }
        visited(i) match {
          case true => {
            // val (nds, rms) = assignEdges(List(i), removedEdges2)
            // egdeMappings(i) = nds
            // removedEdges2 = rms
          }
          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}")

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

  def assignEdges[T: Ordering](
      nodes: List[Int],
      removedEdges: RemovedEdges[T]
  ) = {
    val it = nodes.iterator
    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)
        // removedEdges.partition(e => e.toNode == node || e.fromNode == node)
        loop(nodes, rm, iter, filt)
      }
    }
    // val filtered = ListBuffer.empty[HHCEdge2[T]]
    // val updated = ListBuffer.empty[HHCEdge2[T]]
    // for (node <- nodes) {
    //   val (filt, rm) =
    //     removedEdges.partition(e => e.toNode == node || e.fromNode == node)
    //   if (!filtered.isEmpty) {
    //     // return (filt.toList, rm.toList)
    //   } else {
    //     filtered ++= filt
    //     updated ++= rm
    //   }
    // }

    // print(filtered)
    // (filtered.toList, updated.toList)
    val (edges, removedEdgesUpdated) =
      loop(nodes, removedEdges, it, List.empty[HHCEdge2[T]])
    // val edges2 = edges.sortWith {
    //   case (HHCEdge2(_, _, weight1), HHCEdge2(_, _, weight2)) =>
    //     weight1 < weight2
    // }
    (edges, removedEdgesUpdated)
  }

  def groupClusters[N: Ordering](
      clusters: Clusters,
      edgeMappings: ArraySeq[List[HHCEdge2[N]]],
      removed: RemovedEdges[N]
  ) = {
    var k = -1
    val x = ArraySeq.tabulate(edgeMappings.size)(i => {
      val buf = ListBuffer.empty[(Int, N)]
      val lst = edgeMappings(i)
      val y = lst.foreach(e => {
        if (edgeMappings(e.toNode).isEmpty) {
          // var buf = ListBuffer.empty[(Int, N)]
          for (j <- 0 until clusters.size) {
            for (edge <- clusters(j)) {
              if (edge == e.toNode)
                buf += ((j, e.weight))
            }
          }
          // Left(buf.toList)
        } else {
          buf += ((e.toNode, e.weight))
        }
        // Right(lst.map {
        //   case HHCEdge2(fromNode, toNode, weight) => (toNode, weight)
        // })
      })
      if (!clusters(i).isEmpty) k += 1
      // val lst2 = lst.map {
      //   case HHCEdge2(fromNode, toNode, weight) => (toNode, weight)
      // }
      // if (buf.isEmpty) lst2 else buf.toList ::: lst2
      // if (lt.isEmpty) rt else lt
      (k, buf.toList.sortWith {
        case ((_, weight1), (_, weight2)) =>
          weight1 < weight2
      })
    })
    (clusters, edgeMappings, removed, x)
  }
}