package main

import (
	"fmt"
	"math"
	"runtime"
	"sync"
	"unsafe"

	"github.com/veandco/go-sdl2/sdl"
)

type PointArgs struct {
	resX         int
	resY         int
	imageX       int
	imageY       int
	zoom         float64
	maxIteration int
	threadNumber int
}

type Piece struct {
	pixels []int
	initX  int
	initY  int
	width  int
	height int
}

func getX(linearPoint int, width int) int {
	return linearPoint % width
}

func getY(linearPoint int, height int) int {
	return int(math.Floor(float64(linearPoint / height)))
}

func mapXMandelbrot(x int, width int, zoom float64) float64 {
	return ((float64(x) / float64(width)) * (3.5 * zoom)) - (2.5 - (1.0 - zoom))
}

func mapYMandelbrot(y int, height int, zoom float64) float64 {
	return ((float64(y) / float64(height)) * (2.0 * zoom)) - (1.00001 - (1.0 - zoom))
}

func mandebrotPoint(resX int, resY int, imageX int, imageY int, zoom float64, maxIteration int) int {
	posX := mapXMandelbrot(imageX, resX, zoom)
	posY := mapYMandelbrot(imageY, resY, zoom)

	var x float64 = 0.0
	var y float64 = 0.0
	iteration := 0

	// Period-2 bulb check
	xTerm := posX + 1.0
	posY2 := posY * posY
	if xTerm*xTerm+posY2 < 0.0625 {
		return 0
	}

	// Cardioid check
	xTerm = posX - 0.25
	q := xTerm*xTerm + posY2
	q = q * (q + xTerm)
	if q < (0.25 * posY2) {
		return 0
	}

	var xtemp, xx, yy, xplusy float64

	for iteration < maxIteration {
		xx = x * x
		yy = y * y
		xplusy = x + y
		if (xx)+(yy) > (4.0) {
			break
		}
		y = xplusy*xplusy - xx - yy
		y = y + posY
		xtemp = xx - yy + posX

		x = xtemp
		iteration++
	}

	if iteration > maxIteration {
		return 0
	} else {
		return iteration
	}
}

func mandelbrotPiece(threads int, numProcess int, width int, height int, maxIteration int, zoom float64, res chan Piece, work *sync.WaitGroup) {
	defer work.Done()

	var split int
	if threads > 2 {
		split = int(math.Floor(math.Sqrt(float64(threads))))
	} else if threads == 2 {
		split = 2
	} else {
		split = 1
	}

	pieceX := 0
	pieceY := 0

	if numProcess > 0 {
		pieceX = numProcess % split
		pieceY = int(math.Floor(float64(numProcess) / float64(split)))
	}

	threadWidth := int(math.Floor(float64(width) / float64(split)))
	threadHeight := int(math.Floor(float64(height) / float64(split)))
	initX := threadWidth * pieceX
	initY := threadHeight * pieceY

	iterations := make([]int, threadWidth*threadHeight)
	for y := 0; y < threadHeight; y++ {
		for x := 0; x < threadWidth; x++ {
			iterations[x+y*threadWidth] = mandebrotPoint(width, height, x+initX, y+initY, zoom, maxIteration)
		}
	}

	// Return the portion we have calculated
	res <- Piece{
		pixels: iterations,
		initX:  initX,
		initY:  initY,
		width:  threadWidth,
		height: threadHeight,
	}
}

func stitchPiece(iterationPixels *[]uint32, piece Piece, resX int, maxIteration int, format sdl.PixelFormat, rank int) {
	itPixels := *iterationPixels
	patchWidth := piece.width + piece.initX
	patchHeight := piece.height + piece.initY

	for y := piece.initY; y < patchHeight; y++ {
		for x := piece.initX; x < patchWidth; x++ {
			pieceX := x - piece.initX
			pieceY := y - piece.initY

			var pixel uint32
			iteration := piece.pixels[pieceX+pieceY*piece.width]
			if (iteration < 128) && (iteration > 0) {
				value := uint8(20 + iteration)
				pixel = sdl.MapRGBA(&format, value, 0, value, 255)
			} else if (iteration >= 128) && (iteration < maxIteration) {
				pixel = sdl.MapRGBA(&format, uint8(iteration), 148, uint8(iteration), 255)
			} else {
				pixel = sdl.MapRGBA(&format, 0, 0, 0, 255)
			}
			itPixels[x+y*rank] = pixel
		}
	}
}

func main() {
	fmt.Println("Starting fractal")
	if err := sdl.Init(sdl.INIT_EVERYTHING); err != nil {
		panic(err)
	}
	defer sdl.Quit()

	resX := 1024
	resY := 768

	window, err := sdl.CreateWindow("fractal", sdl.WINDOWPOS_CENTERED, sdl.WINDOWPOS_CENTERED,
		int32(resX), int32(resY), sdl.WINDOW_SHOWN)

	if err != nil {
		panic(err)
	}
	defer window.Destroy()

	renderer, err := sdl.CreateRenderer(window, -1, 0)
	if err != nil {
		panic(err)
	}

	// Blank the Window
	renderer.SetDrawColor(0, 0, 0, 255)
	renderer.Clear()
	renderer.Present()

	textureScreen, err := renderer.CreateTexture(sdl.PIXELFORMAT_ARGB8888, sdl.TEXTUREACCESS_STREAMING, int32(resX), int32(resY))

	if err != nil {
		panic(err)
	}

	screen, err := sdl.CreateRGBSurface(0, int32(resX), int32(resY), 32, 0, 0, 0, 0)

	if err != nil {
		panic(err)
	}

	numCPUs := runtime.NumCPU() // Get how many CPUs the system reports
	fmt.Println("Using ", numCPUs, " threads")

	fmt.Println("Rendering and zooming")

	zoom := 1.0
	stopPoint := 0.00001

	sizeOfInt32 := int32(unsafe.Sizeof(int32(0)))

	// Here we write the pixels to be sent to the screen
	iterationPixels := make([]uint32, resX*resY)

	pieces := make(chan Piece, numCPUs)

	running := true
	for (running) && (zoom > stopPoint) {
		// Let the user quit any time
		if event := sdl.PollEvent(); event != nil {
			switch event.(type) {
			case *sdl.QuitEvent:
				fmt.Println("Quit by user")
				running = false
				break
			}
		}
		// Draw the fractal
		// Set the iteration to a nice number to not overdo it
		maxIteration := 170
		if (zoom < -0.02) && (zoom > -1.0) {
			maxIteration = 100
		}

		var wg sync.WaitGroup
		wg.Add(numCPUs)

		// Launch the processes and wait for the results
		for threadCount := 0; threadCount < numCPUs; threadCount++ {
			go mandelbrotPiece(numCPUs, threadCount, resX, resY, maxIteration, zoom, pieces, &wg)
		}

		wg.Wait()

		pieceCount := 0
		for piece := range pieces {
			stitchPiece(&iterationPixels, piece, resX, maxIteration, *screen.Format, int(screen.Pitch/sizeOfInt32))
			pieceCount++
			if pieceCount >= numCPUs {
				break
			}
		}

		zoom = zoom * 0.99
		screenPixels, _, err := textureScreen.Lock(nil)

		if err != nil {
			panic(err)
		}

		count := 0
		for _, pixel := range iterationPixels {
			for i := uint32(0); i < 4; i++ {
				screenPixels[count] = byte((pixel >> (8 * i)) & 0xff)
				count++
			}
		}

		textureScreen.Unlock()

		if err != nil {
			panic(err)
		}

		// textureScreen.Update(nil, unsafe.Pointer(&iterationPixels[0]), int(resX*sizeOfInt32))
		renderer.Clear()
		renderer.Copy(textureScreen, nil, nil)
		renderer.Present()
	}
}