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

#*****************************************************************************80
#
## gopher() classifies gophers as species 0 or 1.
#
#  Licensing:
#
#    This code is distributed under the GNU LGPL license.
#
#  Modified:
#
#    14 February 2022
#
#  Author:
#
#    John Burkardt
#
  from sklearn.linear_model import LogisticRegression
  import matplotlib.pyplot as plt
  import numpy as np

  print ( 'gopher():' )
#
#  Read the data file.
#
  data = np.loadtxt ( 'gopher_data.txt' )
#
#  skull width, skull length, species (-1 or +1 )
#
  swi = data[:,0]
  sle = data[:,1]
  spe = data[:,2]
  n = len ( swi )

  X = np.zeros ( [ n, 3 ] )
  X[:,0] = 1
  X[:,1] = swi[:]
  X[:,2] = sle[:]

  y = ( spe == 1 ).astype(int)
#
#  Fit the logistic regression model to the data.
#
  classifier = LogisticRegression(random_state=0).fit ( X, y )
#
#  Evaluate the model on the original data.
#
  yp = classifier.predict ( X )
#
#  Compute MSE
#
  mse = ( 1.0 / n ) * sum ( ( yp - y )**2 )
  print ( '  MSE = ', mse )
#
#  Construct the coefficient array W.
#
  w = classifier.coef_[0].copy()
  w[0] = classifier.intercept_

  print ( '' )
  print ( '  Computed weights W = ', w )
#
#  Display data.
#
  swimin = np.min ( swi )
  swimax = np.max ( swi )
  slemin = ( - w[0] - w[1] * swimin ) / w[2]
  slemax = ( - w[0] - w[1] * swimax ) / w[2]

  plt.plot ( X[y==0,1], X[y==0,2], 'r.', markersize = 15 )
  plt.plot ( X[y==1,1], X[y==1,2], 'b.', markersize = 15 )
  plt.plot ( [swimin,swimax], [slemin,slemax], 'k-', linewidth = 3 )
  plt.xlabel ( '<-- Skull width -->' )
  plt.ylabel ( '<-- Skull length -->' )
  plt.title ( 'Blue = Species A, Red = Species B', fontsize = 16 )
  plt.grid ( True )
  filename = 'gopher.png'
  plt.savefig ( filename )
  plt.show ( )
  print ( '' )
  print ( '  Graphics saved as "%s"' % ( filename ) )
#
#  Evaluate the species for new data.
#
  X2 = np.array ( [ \
    [ 1.0, 6.5, 10.8 ], \
    [ 1.0, 6.0, 10.0 ], \
    [ 1.0, 5.6,  9.3 ] ] )
  print ( '  New gopher data:' )
  print ( X2 )
  y2 = classifier.predict ( X2 )
  print ( '  Gopher is species A? (0/1)' )
  print ( y2 )
  p2 = classifier.predict_proba ( X2 )
  print ( '  Strength of decision (0/1)' )
  print ( p2 )
#
#  Terminate.
#
  print ( '' )
  print ( 'gopher():' )
  print ( '  Normal end of execution.' )

  return

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