zipfun.scala

trait MapSeqZip[A,B] { self =>
  def focus   :(A,B)
  def back    :Map[A,Seq[B]]
  def forward :Map[A,Seq[B]]

  lazy val index = back.map(_._2.size).sum + back.size - 1
  
  lazy val underlying :Map[A,Seq[B]] = back.get(focus._1) match {
    case Some( seq ) =>
      back + ((focus._1, seq :+ focus._2 )) ++ forward
    case None => forward.get(focus._1) match {
      case Some(seq) =>
        back ++ ( forward + (( focus._1, focus._2 +: seq )) )
      case None =>
        back ++ Map(focus._1 -> Seq(focus._2)) ++ forward
    }   
  }
  
  def map[C]( f : (A, B) => C ) = new MapSeqZip[A,C] {
    lazy val focus   = (self.focus._1, f(self.focus._1, self.focus._2))
    lazy val back    = self.back.map    { case (a, seq) => (a, seq.map( b => f(a,b) )) }
    lazy val forward = self.forward.map { case (a, seq) => (a, seq.map( b => f(a,b) )) }
  }

  def next :Option[MapSeqZip[A,B]] = forward.get(focus._1) match {
    case None if forward.isEmpty => None
    case None =>
      val a  = forward.head._1
      val b  = forward.head._2.head
      val bk = back.get(focus._1) match {
        case Some( seq ) => back + ((focus._1, seq :+ focus._2))
        case None => back + (( focus._1, Seq( focus._2 ) ))
      }

      Some( MapSeqZip(bk,(a,b),forward.tail + ((forward.head._1 -> forward.head._2.tail)) ) )

    case Some(seq) =>
      val a  = focus._1
      val b  = seq.head
      val bk = back.get(focus._1) match {
        case Some( bseq ) => back + ((focus._1, bseq :+ focus._2))
        case None => back + (( focus._1, Seq( focus._2 ) ))
      } 
      Some( MapSeqZip(bk,(a,b), forward + ((a,seq.tail)) ) )
  }
}

object MapSeqZip {
  def apply[A,B]( m :Map[A,Seq[B]]) :Option[MapSeqZip[A,B]] = m.isEmpty match {
    case true => None
    case false => 
    Some( apply( 
        Map.empty[A,Seq[B]], 
        (m.head._1,m.head._2.head),
        Map( m.head._1 -> m.head._2.tail ) ) )
  }  
  
  def apply[A,B]( _back :Map[A,Seq[B]], _focus :(A,B),_forward :Map[A,Seq[B]] ) = new MapSeqZip[A,B] {
    val focus   = _focus
    val back    = _back
    val forward = _forward
  }
}