// See LICENSE for license details.

package firrtl.transforms.clockfinder

import firrtl.RenameMap
import firrtl.analyses.ConnectionGraph
import firrtl.annotations.{Annotation, IsMember, IsModule, ReferenceTarget}

object ClockSources {
  def buildFromRefs(signalToClockRefs: Map[ReferenceTarget, Set[ReferenceTarget]]): ClockSources = {
    val signalToClockSources = signalToClockRefs.map { case (signal, sources) =>
      signal -> sources.map {
        case c: ReferenceTarget if ConnectionGraph.isAsClock(c) => (c.pathTarget: IsMember, Some(c.ref))
        case c: ReferenceTarget => (c, None)
      }
    }
    ClockSources(signalToClockSources)
  }
}

/** Generated by [[ClockFinderTransform]], maps a given clock source to signals that are synchronized with it
  */
case class ClockSources(signalToClocks: Map[ReferenceTarget, Set[(IsMember, Option[String])]]) extends Annotation {
  override def update(renames: RenameMap): Seq[Annotation] = {
    val newSignalToClocks = signalToClocks.flatMap { case (signal, sources) =>
      val newSignals = renames(signal).collect {
        case x: ReferenceTarget => x
        case other => sys.error("Don't make my signals not a ReferenceTarget!")
      }

      val newClockSources = sources.map { case (clockSource, isAsClock) =>
        renames.get(clockSource) match {
          case Some(Seq(x: IsMember)) => (x, isAsClock)
          case None => (clockSource, isAsClock)
          case other => sys.error("Don't split or delete my clock source!!")
        }
      }
      newSignals.map(s => s -> newClockSources)
    }
    Seq(ClockSources(newSignalToClocks))
  }

  lazy val clockToSignals: Map[Option[(IsMember, Option[String])], Seq[ReferenceTarget]] =
    signalToClocks.foldLeft(Map.empty[Option[(IsMember, Option[String])], Seq[ReferenceTarget]]){
      case (map, (signal, clocks)) if clocks.nonEmpty => clocks.foldLeft(map){ (m, clock) =>
        m + (Some(clock) -> (signal +: m.getOrElse(Some(clock), Nil)))
      }
      case (map, (signal, clocks)) if clocks.isEmpty =>
        map + (None -> (signal +: map.getOrElse(None, Nil)))
    }

  lazy val signalsToClockRefs: Map[ReferenceTarget, Set[ReferenceTarget]] =
    signalToClocks.map {
      case (signal, seq) => signal -> seq.map {
        case (ref: ReferenceTarget, None) => ref
        case (mod: IsModule, Some(x: String)) => mod.ref(x)
      }
    }

  lazy val clockRefsToSignals: Map[Set[ReferenceTarget], Seq[ReferenceTarget]] =
    signalsToClockRefs.foldLeft(Map.empty[Set[ReferenceTarget], Seq[ReferenceTarget]]){
      case (map, (signal, clocks)) => map + (clocks -> (signal +: map.getOrElse(clocks, Nil)))
    }

  def prettyPrint: String = {

    clockRefsToSignals.map {
      case (clks, signals) =>
        "Clock Domain(\n\t" +
          clks.toSeq.map(_.toString).sorted.mkString("\n\t") +
          "\n) with signals (\n\t" +
          signals.map(_.toString).sorted.mkString("\n\t") +
          "\n)\n"
    }.mkString("\n")
  }

}