#! /usr/bin/env python
#
import numpy as np
import time

def mandelbrot_python ( ):

#*****************************************************************************80
#
## MANDELBROT_PYTHON computes an image of the Mandelbrot set.
#
#  Licensing:
#
#    This code is distributed under the GNU LGPL license. 
#
#  Modified:
#
#    15 May 2017
#
#  Author:
#
#    John Burkardt
#
#  Local parameters:
#
#    Local, integer XNUM, the number of pixels in the X direction.
#
#    Local, integer YNUM, the number of pixels in the Y direction.
#
#    Local, integer COUNT_MAX, the number of iterations.
#
  print ( '' )
  print ( 'MANDELBROT_PYTHON:' )
  print ( '  Python version' )
  print ( '  Compute the Mandelbrot set and save it in a graphics file.' )

  xmin = -1.0
  xmax = -0.6
  ymin = 0.0
  ymax = 0.4
  xnum = 1000
  ynum = 1000
  count_max = 200

  print ( '' )
  print ( '  X range: [ %g, %g ]' % ( xmin, xmax ) )
  print ( '  Y range: [ %g, %g ]' % ( ymin, ymax ) )
  print ( '  Xnum = %d x Ynum = %d = %d Pixels' % ( xnum, ynum, xnum * ynum ) )
  print ( '  Maximum number of iterations = %d' % ( count_max ) )

  mandelbrot_image ( xmin, xmax, ymin, ymax, xnum, ynum, count_max )
#
#  Terminate.
#
  print ( '' )
  print ( 'MANDELBROT_PYTHON:' )
  print ( '  Normal end of execution.' )
  return

def mandelbrot_image ( xmin, xmax, ymin, ymax, xnum, ynum, count_max ):

#*****************************************************************************80
#
## MANDELBROT_IMAGE creates an image of the Mandelbrot set.
#
#  Discussion:
#
#    Over the rectangle [XMIN,XMAX] x [YMIN,YMAX], determine the Mandelbrot
#    count for a grid of XNUMxYNUM points, using a particular value of
#    COUNT_MAX.
#
#  Licensing:
#
#    This code is distributed under the GNU LGPL license. 
#
#  Modified:
#
#    15 May 2017
#
#  Author:
#
#    John Burkardt
#
#  Parameters:
#
#    Input, real XMIN, XMAX, YMIN, YMAX, the physical limits of the rectangle.
#
#    Input, integer XNUM, YNUM, the number of points in X and Y directions.
#
#    Input, integer COUNT_MAX, the maximum number of iterations.
#
  time1 = time.time ( )
  clock1 = time.clock ( )

  count = mandelbrot_set ( xmin, xmax, ymin, ymax, xnum, ynum, count_max )

  clock2 = time.clock ( )
  time2 = time.time ( )

  print ( '' )
  print ( '  Compute time:' )
  print ( '  Wallclock: %.02f sec.' % ( time2 - time1 ) )
  print ( '        CPU: %.02f sec.' % ( clock2 - clock1 ) )
#
#  Save the data in a graphics file.
#
  file_name = 'mandelbrot_python.pgm'
  comment = 'PGM gray map created by mandelbrot_python.py'

  pgma_write ( file_name, comment, xnum, ynum, count_max, count )

  print ( '' )
  print ( '  Plot saved in file "%s"' % ( file_name ) )

  return

def mandelbrot_set ( xmin, xmax, ymin, ymax, xnum, ynum, count_max ):

#*****************************************************************************80
#
## MANDELBROT_SET computes the Mandelbrot count for a grid of points.
#
#  Discussion:
#
#    Over the rectangle [XMIN,XMAX] x [YMIN,YMAX], determine the Mandelbrot
#    count for a grid of XNUMxYNUM points, using a particular value of
#    COUNT_MAX.
#
#  Licensing:
#
#    This code is distributed under the GNU LGPL license. 
#
#  Modified:
#
#    15 May 2017
#
#  Author:
#
#    John Burkardt
#
#  Parameters:
#
#    Input, real XMIN, XMAX, YMIN, YMAX, the physical limits of the rectangle.
#
#    Input, integer XNUM, YNUM, the number of points in X and Y directions.
#
#    Input, integer COUNT_MAX, the maximum number of iterations.
#
#    Output, integer COUNT(XNUM,YNUM), the count for each point in the grid.
#
  count = np.zeros ( ( xnum, ynum ) )

  for i in range ( ynum ):

    y = ( float ( ynum - i - 1 ) * ymax   \
        + float (        i     ) * ymin ) \
        / float ( ynum     - 1 )

    for j in range ( xnum ):

      x = ( float ( xnum - j - 1 ) * xmin   \
          + float (        j     ) * xmax ) \
          / float ( xnum     - 1 );

      c = x + y * 1j

      count[i,j] = mandelbrot_count ( c, count_max )

  return count

def mandelbrot_count ( c, count_max ):

#*****************************************************************************80
#
## MANDELBROT_COUNT returns the Mandelbrot count for a single point.
#
#  Discussion:
#
#    Starting with the value 0, repeatedly square and add complex value C.
#
#    Return number of applications of this process at which the
#    value exceeds 2 in norm.
#
#    Repeat no more than COUNT_MAX times.
#
#  Licensing:
#
#    This code is distributed under the GNU LGPL license. 
#
#  Modified:
#
#    15 May 2017
#
#  Author:
#
#    John Burkardt
#
#  Parameters:
#
#    Input, complex C, the value added at each step.
#
#    Input, integer COUNT_MAX, the maximum number of iterations.
#
#    Output, integer MANDEL_COUNT, the number of operations at which
#    the iterate first exceeded 2 in norm.  If this never happens,
#    return COUNT_MAX.
#
  value = count_max

  z = 0.0 + 0.0 * 1j
  for k in range ( count_max ):
    if ( 2.0 <= abs ( z ) ):
      value = k
      break
    z = z * z + c

  return value

def pgma_write ( file_name, comment, width, height, maxval, gray ):

#*****************************************************************************80
#
## PGMA_WRITE writes an ASCII PGM graphics file.
#
#  Licensing:
#
#    This code is distributed under the GNU LGPL license. 
#
#  Modified:
#
#    13 May 2017
#
#  Author:
#
#    John Burkardt
#
#  Parameters:
#
#    Input, string FILE_NAME, the name of the file.
#
#    Input, string COMMENT, a comment, which may be empty ('');
#
#    Input, integer WIDTH, HEIGHT, the width and height of the graphics image.
#
#    Input, integer MAXVAL, the maximum allowed gray value.
#
#    Input, integer GRAY[WIDTH,HEIGHT], gray values between 0 and MAXVAL.
#
  file_type = 'P2'

  file_handle = open ( file_name, 'wt' )
  file_handle.write ( "%s\n" % ( file_type ) ) 
  file_handle.write ( "#%s\n" % ( comment ) ) 
  file_handle.write ( "%d %d\n" % ( width, height ) ) 
  file_handle.write ( "%d\n" % ( maxval ) )

  for i in range ( height ):
    for j in range ( width ):
      file_handle.write ( " %d" % ( gray[i,j] ) ) 
    file_handle.write ( "\n" )

  file_handle.close ( )

  return

if ( __name__ == '__main__' ):

  mandelbrot_python ( )