// NewtonChaos.java
// Robert Sedgewick and Kevin Wayne
// Modified by Andrew Davison, Nov. 2022

// https://introcs.cs.princeton.edu/java/32class/NewtonChaos.java.html

/* Plot an n-by-n grid of points showing which of the 
   four roots that NewtonChaos's method converges to in the 
   2-by-2 box centered at (0, 0).

  https://www.dreamincode.net/forums/topic/297478-chaos-in-newtons-method/

  Uses Complex.java

   Usage:
     > compile NewtonChaos.java
     > run NewtonChaos
*/

import java.awt.*;


public class NewtonChaos
{
  private static final double CON_VALUE = 0.001;   //convergence value

  public static void main(String[] args)
  {
    int size = 200;
    double min = -1.0;
    double max = 1.0;

    // constants
    Complex newOne = new Complex(1,0);
    Complex newFour = new Complex(4,0);

    // root coordinates for colors
    Complex whiteRoot = new Complex(1,0);
    Complex redRoot = new Complex(-1,0);
    Complex blueRoot = new Complex(0,1);
    Complex greenRoot = new Complex(0,-1);

    StdDraw.setCanvasSize(size, size);
    StdDraw.setScale(min, max);

    Complex z, func, derivFun;
    for (double y = -1.0; y < 1.0; y += 2.0/size) {
      for (double x = -1.0; x < 1.0; x += 2.0/size) {
        z = new Complex(x,y);

        for(int i=0; i < 100; i++) {
          func = z.power(4).minus(newOne);
          derivFun = newFour.times(z.power(3));
          z = z.minus(func.divides(derivFun));
        }

        if (z.minus(whiteRoot).abs() < CON_VALUE)
          StdDraw.setPenColor(StdDraw.WHITE);
        else if (z.minus(redRoot).abs() < CON_VALUE)
          StdDraw.setPenColor(StdDraw.RED);
        else if (z.minus(blueRoot).abs() < CON_VALUE)
          StdDraw.setPenColor(StdDraw.BLUE);
        else if (z.minus(greenRoot).abs() < CON_VALUE)
          StdDraw.setPenColor(StdDraw.GREEN);
        else
          StdDraw.setPenColor(StdDraw.BLACK);

        StdDraw.point(x, y);
      }
    }
    System.out.println("Drawing finished");
  } /* main */


}