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

#*****************************************************************************80
#
## basketball_model uses a least squares solver to model basketball data.
#
#  Discussion:
#
#    Least squares solver for optimal (b,m) to fit (x,y) data
#    y = b + m*x
#
#  Licensing:
#
#    This code is distributed under the MIT license.
#
#  Modified:
#
#    31 January 2022
#
#  Author:
#
#    John Burkardt
#
  import matplotlib.pyplot as plt
  import numpy as np

  print ( '' )
  print ( 'basketball_model:' )
  print ( '  Given age x and sponsorship y for basketball players,' )
  print ( '  seek (m,b) so that y=b+mx approximates the data.' )
  print ( '  Use least squares solver to estimate best b and m.' )
#
#  Get the (normalized) Ford data.
#
  data = np.loadtxt ( 'basketball_data.txt' )
  y = data[:,3]
  x = data[:,4]
  n = len ( y )
#
#  Create the matrix X:
#
  X = np.zeros ( [ n, 2 ] )
  X[:,0] = 1;
  X[:,1] = x;
#
#  Call least squares solver for rectangular linear system.
#
  c4 = np.linalg.lstsq ( X, y, rcond = None )[0]

  print ( '' )
  print ( '  Seek relationship sponsor = c[0] + c[1] * age' )
  print ( '  c[0] = %g' % ( c4[0] ) )
  print ( '  c[1] = %g' % ( c4[1] ) )
  mse4 = 1 / n * np.sum ( ( y - c4[0] - c4[1] * x )**2 )
  print ( '  MSE is %g' % ( mse4 ) )
#
#  Plot.
#
  plt.clf ( )
  plt.plot ( x, y, 'bo', markersize = 8 )
  m1 = np.min ( x )
  m2 = np.max ( x )
  p1 = c4[0] + c4[1] * m1
  p2 = c4[0] + c4[1] * m2
  plt.plot ( [m1,m2], [p1,p2], 'r-', linewidth = 3 )
  plt.grid ( True )
  plt.xlabel ( '<--- Age (years) --->' )
  plt.ylabel ( '<--- Sponsorship --->' )
  plt.title ( 'Sponsorship as function of age', fontsize = 16 )
  filename = 'basketball_model_fit.png'
  plt.savefig ( filename )
  print ( '' )
  print ( '  Graphics saved as "%s"' % ( filename ) )
  plt.show ( )
  plt.close ( )
#
#  Terminate.
#
  print ( '' )
  print ( 'basketball_model:' )
  print ( '  Normal end of execution.' )
  return

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