package main

import (
	"fmt"
	"image"
	"image/png"
	"os"
	"log"
	"flag"
)

type Newton struct {
	dx, dy int
	z0, dz complex128
	n int
}

func (n *Newton) ColorModel() image.ColorModel {
	return image.RGBAColorModel
}

func (n *Newton) Width() int {
	return n.dx
}

func (n *Newton) Height() int {
	return n.dy
}

func (n *Newton) At(x, y int) image.Color {
	fx := float64(x) / float64(n.dx)
	fy := float64(y) / float64(n.dy)

	Z := n.z0 + cmplx(real(n.dz)*fx, imag(n.dz)*fy)
	r := [...]complex128{cmplx(1,0), cmplx(-1,1), cmplx(-1,-1)}
	c := [...]int{0,0,0}
	for i,v := range r {
		z := Z
		for c[i] = n.n; c[i] > 0; c[i]-- {
			z0 := z
			z -= (z*z*z - v)/(3*z*z)
			if z == z0 {
				break
			}
		}
	}

	adj := func(a int) uint8 { return uint8((0xff*a)/n.n) }
	return image.RGBAColor{adj(c[2]), adj(c[0]), adj(c[1]), 0xff}
}

func main() {
	var z0r *float64 = flag.Float64("z0r", -3, "z0 real part")
	var z0i *float64 = flag.Float64("z0i", -3, "z0 imaginary part")
	var z1r *float64 = flag.Float64("z1r", 3, "z1 real part")
	var z1i *float64 = flag.Float64("z1i", 3, "z1 imaginary part")

	flag.Parse()

	z0 := cmplx(*z0r, *z0i)
	z1 := cmplx(*z1r, *z1i)
	dz := z1-z0
	n := Newton{1024,1024,z0,dz,64}
	name := fmt.Sprintf("newton%+f%+fi%+f%+fi.png", real(z0), imag(z0), real(z1), imag(z1))
	if store, err := os.Open(name, os.O_CREATE|os.O_WRONLY, 0666); err == nil {
		defer store.Close()
		png.Encode(store, &n)
	} else {
		log.Stdout("newton: ", err)
	}
}