#! /usr/bin/env python3
#
def faithful_kmeans2 ( ):

#*****************************************************************************80
#
## faithful_kmeans2() does a simple clustering exercise using scipy kmeans2().
#
#  Discussion:
#
#    Clustering data.
#
#  Licensing:
#
#    This code is distributed under the GNU LGPL license.
#
#  Modified:
#
#    19 September 2019
#
#  Author:
#
#    John Burkardt
#
  import matplotlib.pyplot as plt
  import numpy as np
  import platform
  from scipy.cluster.vq import kmeans2

  print ( '' )
  print ( 'faithful_kmeans2():' )
  print ( '  Python version: %s' % ( platform.python_version ( ) ) )
  print ( '  Cluster Old Faithful data using scipy kmeans2().' )
#
#  Read the data.
#
  data = np.loadtxt ( 'faithful_normalized.txt' )
#
#  Create x and y.
#
  x = data[:,0]
  y = data[:,1]
#
#  Call kmeans2()
#
  k = 2
  Z, C = kmeans2 ( data, k )
  print ( '  Kmeans2 cluster centers Z:' )
  print ( Z )

  bd = ( x - Z[0,0] )**2 + ( y - Z[0,1] )**2
  rd = ( x - Z[1,0] )**2 + ( y - Z[1,1] )**2
  bc = np.where ( bd < rd )
  rc = np.where ( rd < bd )
  e0 = sum ( bd[bc] )
  e1 = sum ( rd[rc] )

  bn = len ( bd[bc] )
  cn = len ( rd[rc] )

  e = e0 + e1
  print ( '  Cluster energy E = ',  e, '=', e0, '+', e1 )
  print ( '  Cluster size = ', bn + cn, '=', bn, '+', cn )

  plt.clf ( )
  plt.plot ( x[C==0], y[C==0], 'c.', markersize = 10 )
  plt.plot ( x[C==1], y[C==1], 'm.', markersize = 10 )
  plt.plot ( Z[0,0], Z[0,1], 'bo', markersize = 15 )
  plt.plot ( Z[1,0], Z[1,1], 'ro', markersize = 15 )
  plt.xlabel ( '<-- Duration -->', fontsize = 16 )
  plt.ylabel ( '<-- Wait -->', fontsize = 16 )
  plt.title ( 'Clusters using kmeans2()', fontsize = 16 )
  plt.grid ( True )
  filename = 'faithful_kmeans2.png'
  plt.savefig ( filename )
  plt.show ( )

  print ( '' )
  print ( '  Graphics saved as "%s"' % ( filename ) )
#
#  Terminate.
#
  print ( '' )
  print ( 'faithful_kmeans2():' )
  print ( '  Normal end of execution.' )

  return

if ( __name__ == '__main__' ):
  faithful_kmeans2 ( )