import java.awt.Color
import java.awt.Dimension
import java.awt.Graphics
import java.awt.Graphics2D
import java.awt.image.BufferedImage
import javax.swing.JFrame
import javax.swing.JPanel
import javax.swing.WindowConstants
import scala.language.implicitConversions    

// complex numbers
case class Complex(val re: Double, val im: Double) { 
  // represents the complex number re + im * i
  def +(that: Complex) = Complex(this.re + that.re, this.im + that.im)
  def -(that: Complex) = Complex(this.re - that.re, this.im - that.im)
  def *(that: Complex) = Complex(this.re * that.re - this.im * that.im,
                                 this.re * that.im + that.re * this.im)
  def *(that: Double) = Complex(this.re * that, this.im * that)
  def abs() = Math.sqrt(this.re * this.re + this.im * this.im)
}

object i extends Complex(0, 1)

implicit def double2complex(re: Double): Complex = Complex(re, 0)

// some customn colours for the "sliding effect"
val colours = List(
  new Color(66, 30, 15),    new Color(25, 7, 26),
  new Color(9, 1, 47),      new Color(4, 4, 73),
  new Color(0, 7, 100),     new Color(12, 44, 138),
  new Color(24, 82, 177),   new Color(57, 125, 209),
  new Color(134, 181, 229), new Color(211, 236, 248),
  new Color(241, 233, 191), new Color(248, 201, 95),
  new Color(255, 170, 0),   new Color(204, 128, 0),
  new Color(153, 87, 0),    new Color(106, 52, 3))

// the viewer panel with a canvas
class Viewer(width: Int, height: Int) extends JPanel {
  val canvas = new BufferedImage(width, height, BufferedImage.TYPE_INT_ARGB)
  
  override def paintComponent(g: Graphics) = 
    g.asInstanceOf[Graphics2D].drawImage(canvas, null, null)
  
  override def getPreferredSize() = 
    new Dimension(width, height)

  def clearCanvas(color: Color) = {
    for (x <- 0 to width - 1; y <- 0 to height - 1) 
      canvas.setRGB(x, y, color.getRGB())
    repaint()
  }  
}

// initialising the viewer
def openViewer(width: Int, height: Int) : Viewer = {
  val frame = new JFrame("XYPlane")
  val viewer = new Viewer(width, height)
  frame.add(viewer)
  frame.pack()
  frame.setVisible(true)
  frame.setResizable(false)
  frame.setDefaultCloseOperation(WindowConstants.EXIT_ON_CLOSE)
  viewer
}

// some hardcoded data
val W = 900   // width
val H = 800   // height
val black = Color.black
val viewer = openViewer(W, H)

// drawing a pixel on the canvas
def pixel(x: Int, y: Int, color: Color) = 
  viewer.canvas.setRGB(x, y, color.getRGB())


// calculating the number of iterations using lazy streams
//   the iteration goes on for a maximum of max steps,
//   but might leave early when the pred is satisfied
def iterations(c: Complex, max: Int) : Int = {
  def next(z: Complex) = z * z + c    
  def pred(z: Complex) = z.abs < 2    // exit condition
  Stream.iterate(0.0 * i, max)(next).takeWhile(pred).size
}

// main function 
def mandelbrot(start: Complex, end: Complex, max: Int) : Unit = {
  viewer.clearCanvas(black)
  
  // deltas for each grid step 
  val d_x = (end.re - start.re) / W
  val d_y = (end.im - start.im) / H
   
  for (y <- (0 until H).par) {
    for (x <- (0 until W).par) {
    
     val c = start + 
      (x * d_x + y * d_y * i)
     val iters = iterations(c, max) 
     val col = 
       if (iters == max) black 
       else colours(iters % 16)

     pixel(x, y, col)
    }
    viewer.updateUI()
  }   
}

// Examples
//==========

//for measuring time
def time_needed[T](code: => T) = {
  val start = System.nanoTime()
  code
  val end = System.nanoTime()
  (end - start) / 1.0e9
}


// example 1
val exa1 = -2.0 + -1.5 * i
val exa2 =  1.0 +  1.5 * i

time_needed(mandelbrot(exa1, exa2, 1000))

// example 2
val exb1 = -0.37465401 + 0.659227668 * i
val exb2 = -0.37332410 + 0.66020767 * i

//time_needed(mandelbrot(exb1, exb2, 1000))

// example 3
val exc1 = 0.435396403 + 0.367981352 * i
val exc2 = 0.451687191 + 0.380210061 * i

//time_needed(mandelbrot(exc1, exc2, 1000))

// some more computations with example 3

val delta = (exc2 - exc1) * 0.0333


time_needed(
  for (n <- (0 to 12)) 
     mandelbrot(exc1 + delta * n, 
                exc2 - delta * n, 100)) 

/*
time_needed(
  for (n <- (0 to 12)) 
     mandelbrot(exc1 + delta * n, 
                exc2 - delta * n, 1000))
*/