import org.bouncycastle.jcajce.provider.digest.SHA3
import org.bouncycastle.util.encoders.Hex

object Ethash {
  
  val CACHE_BYTES_INIT:Long   = Math.pow(2, 24).toLong
  val CACHE_BYTES_GROWTH:Long = Math.pow(2, 17).toLong
  val EPOCH_LENGTH:Long       = 30000
  val HASH_BYTES:Int          = 64
  val CACHE_ROUNDS:Int        = 3
  
  def cacheSize(blockNumber:Long):Long = {
    val linearSize:Long = CACHE_BYTES_INIT + CACHE_BYTES_GROWTH * (blockNumber / EPOCH_LENGTH) - HASH_BYTES
    Stream.from(0).map(x => linearSize - (2*x*HASH_BYTES)).filter(x => Utilities.isPrime(x/HASH_BYTES)).head
  }
  
  def makeCache(size:Long, seed:Array[Byte]):Array[Byte] = {
    import Utilities.sha3_512
    val n:Int = (size/HASH_BYTES).toInt
    
    //Sequentially produce initial dataset
    val ans = Stream.iterate(seed){x => Utilities.sha3_512(x)}.drop(1).flatten.take(size.toInt).toArray
    val hashes:Array[Sha3Hash] = ans.sliding(64, 64).map { x => new Sha3Hash(x) }.toArray
    
    //Use a low-round version of randmemohash
    0.to(CACHE_ROUNDS-1).foreach{ _ =>
      0.to(n - 1).foreach{ i =>
        val v = (hashes(i).bytes.head & 0xFF) % n
        hashes(i) = new Sha3Hash(sha3_512( hashes((i-1+n) % n).xor(hashes(v)) ))
      }
    }
    
    hashes.map { x => x.bytes.toSeq }.toSeq.flatten.toArray
  }
}

object Utilities {
  
  def isPrime(x:Long):Boolean = 
    !2L.to(Math.sqrt(x).toLong).exists{factor => x%factor == 0}
 
  def sha3_512(x:Array[Byte]):Array[Byte] = new SHA3.Digest512().digest(x)
  
}

class Sha3Hash(val bytes:Array[Byte]) {
  if (bytes.size != 64) throw new Exception(s"Byte array for a Sha3Hash must be 64 bytes")

  def xor(other:Sha3Hash):Array[Byte] = bytes.zip(other.bytes).map(pr => (pr._1 ^ pr._2).toByte).toArray
  
  override def toString = Hex.toHexString(bytes)
  
}