text: # !/usr/bin/env python3
  # -*- coding: utf-8 -*-
  """
  Created on Thu Nov 11 12:01:05 2021
  
  @author: alexxcollins
  """
  
  # %% import packages
  from numpy.random import default_rng
  from abc import ABC, abstractmethod
  import numpy as np
  from sklearn.linear_model import LinearRegression
  import matplotlib.pyplot as plt
  import scipy.stats as sps
  import seaborn as sns
  
  sns.set(style="ticks")
  
  
  # %% GnerateData
  class GenerateData(ABC):
      """
      Base Class for X and y data, regressions and visualisations.
  
      The base class will generate random Y given parameters for sample size
      beta (coefficients for intercept plus independent variables), variance of
      epsilon and optional X.
      The X generated by this base class will match dimension of beta and will
      be generated from uniform distribution along each dimension. Subclasses can
      be used to generate or supply difference Xs, with outliers, co-depence etc.
      """
  
      def __init__(self, N=1000, beta=(1, 2, 3), noise_var=1, random_seed=None):
          """
          Initialise regression object.
  
          N : integer
              DSample size. The default is 1000.
          beta : array or array-like, optional
              The values of the coefficients. beta = (b0, b1, b2) will give
              y = b0 + b1*x1 + b2*x2 + e. Can by any length.
              beta = (1,1) will give univariate regression; beta=(1,1,1,1) will
              be consistent with X of dim 3.
          noise_var : float, optional
              variance of epsilon. The default is 1.
  
          """
          # not sure where to set random seed. Should we do this in Jupyter
          # notebook? and pass into the class?
          self.rng = default_rng(random_seed)
          self.N = N
          self._p = len(beta) - 1
          self._p1 = len(beta)
          self.beta = np.array(beta).reshape((self._p1, 1))
          # we often use p and p' in regression formulas, so I'm defining both
          # at the moment discourage changing with underscore rather than setters
          self.e_var = noise_var
  
      def _generate_epsilon(self, noise_var):
          return self.rng.normal(0, noise_var ** (1 / 2), size=(self.N, 1))
  
      @abstractmethod
      def generate_X(self, low=None, high=None):
          """Generate X data using uniform distribution between low and high."""
          pass
  
      def generate_y(self):
          """Generate y using X, noise_var and beta. y = b0 + b @ X + e."""
          self.e = self._generate_epsilon(self.e_var)
          self.y = np.matmul(self.X1, self.beta) + self.e
  
      def train_test_split(self, test_size=0.25, shuffle=True):
          """
          Split data into test and train datasets.
  
          will shuffle dataset if shuffle=True. Test size should be between
          0 and 1.0 and is the proportion of the sample held back for testing.
  
          """
          if shuffle:
              index = np.arange(self.N)
              self.rng.shuffle(index)
              self.X = self.X[index]
              self.y = self.y[index]
  
          test_N = np.round(self.N * test_size).astype(int)
          self.X_test = self.X[:test_N]
          self.y_test = self.y[:test_N]
          self.X_train = self.X[test_N:]
          self.y_train = self.y[test_N:]
  
      def generate_dataset(self, test_size=0.25, shuffle=True, **kwargs):
          """Generate test and train data with option to shuffle."""
          self.generate_X(**kwargs)
          self.generate_y()
          self.train_test_split(test_size=test_size, shuffle=True)
  
      def fit(self, fit_intercept=True):
          """
          Use sklearn to fit the linear regression.
  
          This function has to be run after .train_test_split method.
          It calculates:
              predicted y, for the X_test data
              score - which is R^2 as in sklearn LinearRegression
              intercept, coef from sklearn LinearRegression
              b_pred which is concatenation of intercept and coe and same shape
                  as self.beta
          Results set as object attributes.
  
          """
          reg = LinearRegression(fit_intercept=fit_intercept)
          reg = reg.fit(self.X_train, self.y_train)
          self.y_pred = reg.predict(self.X_test)
          self.residuals = self.y_test - self.y_pred
          self.score = reg.score(self.X_test, self.y_test)
          # intercept and coef attributes are equivalent to scikit learn
          # attributes
          self.intercept = reg.intercept_
          self.coef = reg.coef_
          # we also set b_pred attribute. This is equivalent to self.beta: it
          # includes intercept and other coeffs and is same shape as self.beta
          b = np.concatenate((self.intercept, np.resize(self.coef, (self.coef.size))))
          self.b_pred = b[:, np.newaxis]
  
      def plot2D(self, i=None, fitted_line=True, true_beta_line=True):
          """
          Plot the scatter of y and X_i.
  
          The user can select the column of X to plot by setting i. If i is None
          then y is plotted vs each dimension of X in turn.
  
          take the beta coefficient relating to X_i to caluculate slope. The
          intercept is b_0 + the dot product of betas for X_j (j != i) multiplied
          by X_j.mean().
  
          The scatter of X_i and y is plotted. The fitted line, and the "true"
          corresponding to the original beta coefficients can also be plotted.
  
          ###### ToDo
          options to print title for chart and equation for line of best fit
          ######
  
          Parameters
          ----------
          i : int
              optional parameter to select on column from X matrix. i indexed at
              1, so choose 1 for X_1 and p for final column of X
  
          """
          num_plots = 1
          if i is None:
              num_plots = self._p
  
          cols = min(2, num_plots)
          rows = (num_plots + 1) // 2
          fig, axs = plt.subplots(rows, cols, figsize=(cols * 5, rows * 4))
          fig.suptitle("y vs one dimension of X. Fitted line, true line and " "scatter")
          for j in range(num_plots):
              ax = axs.flatten()[j] if num_plots > 1 else axs
              if i is not None:
                  j = i - 1
  
              X_i = self.X[:, j, np.newaxis]
              X = np.linspace(X_i.min(), X_i.max(), 100)
              y_true_intercept = self.beta[0] + np.dot(
                  self.beta[1:].reshape(-1), self.X.mean(axis=0)
              )
              y_beta = y_true_intercept + self.beta[j + 1] * (X - self.X.mean(axis=0)[j])
              y_fitted_intercept = self.b_pred[0] + np.dot(
                  self.coef.reshape(-1), self.X.mean(axis=0)
              )
              y_fitted = y_fitted_intercept + self.b_pred[j + 1] * (
                  X - self.X.mean(axis=0)[j]
              )
  
              ax.scatter(
                  X_i, self.y, color="b", alpha=0.2, label="y vs X_{}".format(j + 1)
              )
              if true_beta_line:
                  ax.plot(X, y_beta, color="r", label="true beta")
              if fitted_line:
                  ax.plot(X, y_fitted, color="g", label="fitted beta")
              ax.legend(loc="upper left")
              ax.text(
                  X.max(),
                  self.y.min(),
                  r"true: $y = $"
                  + "{:.2f}".format(y_true_intercept[0])
                  + r"$ + $"
                  + r"{:.2f}".format(self.beta[j + 1][0])
                  + r"$*(X-\bar{X})$"
                  + "\n"
                  + r"fitted: $y = $"
                  + "{:.2f}".format(y_fitted_intercept[0])
                  + r"$ + $"
                  + r"{:.2f}".format(self.b_pred[j + 1][0])
                  + r"$*(X-\bar{X})$",
                  fontsize=8,
                  bbox={"facecolor": "white", "alpha": 0.5, "pad": 2},
                  verticalalignment="bottom",
                  horizontalalignment="right",
              )
  
          plt.show()
  
      def plot_residuals(self, i=None, n_bins=20):
          """Plot histogram of residuals and residuals against each X_i.
  
          Option to select a dimension of X to plot residuals against.
          n_bins - integer number of bins for histogram plot
  
          code to plot the normal pdf was adapted from stack overflow answer:
          https://stackoverflow.com/questions/11315641/python-plotting-a-histogram-with-a-function-line-on-top
          """
          num_plots = 2
          if i is None:
              num_plots = self._p + 1
  
          cols = min(2, num_plots)
          rows = (num_plots + 1) // 2
          fig, axs = plt.subplots(rows, cols, figsize=(cols * 5, rows * 4))
          fig.suptitle("residuals: histogram and against features")
  
          # plot a histogram of residuals
          ax = axs.flatten()[0]
          # parameterise the normal distribution to display
          n = sps.norm(0, np.std(self.residuals))
          # add histogram showing individual components
          ax.hist(
              self.residuals,
              bins=n_bins,
              histtype="bar",
              density=True,
              alpha=0.4,
              edgecolor="none",
          )
          # get X limits and fix them
          mn, mx = ax.get_xlim()
          ax.set_xlim(mn, mx)
          # add our distributions to figure
          x = np.linspace(mn, mx, 301)
          ax.plot(x, n.pdf(x), color="C1", ls="--", label="normal pdf")
          # estimate Kernel Density and plot
          kde = sps.gaussian_kde(self.residuals.reshape(-1))
          ax.plot(x, kde.pdf(x), color="C0", label="residual sample pdf")
          # finish up
          ax.legend()
          ax.set_ylabel("Probability density")
          sns.despine()
  
          # plot the residuals vs X (features)
          for j in range(1, num_plots):
              ax = axs.flatten()[j]
              if i is not None:
                  j = i - 1
  
              # try:
              X_i = self.X_test[:, j - 1, np.newaxis]
  
              ax.scatter(
                  X_i,
                  self.residuals,
                  color="b",
                  alpha=0.2,
                  label="residual vs X_{}".format(j),
              )
              ax.legend(loc="upper left")
          # if there are an even number of features there will be an
          # IndexError as the last plot will attempt to plot for non-existant
          # feature.
          # except IndexError:
          #     pass
  
          plt.show()
  
      # next two functions used to generate integer range around (a, b)
      # use round_down(a) - returns integer below a, and works if a is +ve or -ve
      # similarly for round_up(b)
      # rounds number "up". 9.8 -> 10; -9.8 -> -9
      def round_up(self, x):
          """Round x up to nearest integer."""
          return int(x) + (x % 1 > 0) * (x > 0)
  
      # rounds number "down". 9.8 -> 9; -9.8 -> -10
      def round_down(self, x):
          """Round x down to nearest integer."""
          return int(x) - (x % 1 > 0) * (x < 0)
  
      def plot3D(self):
          """
          Designed to work for a regression with two input features.
  
          takes the beta (intercept + x coefficients) of the original data to
          create a visualisation of the line showing the relationship between
          x and y, as well as a scatter plot of the data.
  
          ###### ToDo
          put line of best fit from regression results
  
          options to print title for chart and equation for line of best fit
          ######
          """
          if self._p != 2:
              raise Exception(
                  "dimension of X must be 2. X had dimension {}".format(self._p)
              )
  
          fig = plt.figure()
          ax = fig.add_subplot(projection="3d")
  
          X1 = np.linspace(
              self.round_down(self.X[:, 0].min()), self.round_up(self.X[:, 0].max()), 2
          )
          X2 = np.linspace(
              self.round_down(self.X[:, 1].min()), self.round_up(self.X[:, 1].max()), 2
          )
          X1, X2 = np.meshgrid(X1, X2)
          y = self.beta[0] + self.beta[1] * X1 + self.beta[2] * X2
  
          ax.scatter(self.X[:, 0], self.X[:, 1], self.y, alpha=0.2)
          ax.plot_surface(X1, X2, y, alpha=0.2, color="r")
  
  
  # %% UniformX
  class UniformX(GenerateData):
      """Generate X data from uniform distribution."""
  
      def generate_X(self, low=-10, high=10):
          """Generate X data from uniform distribution."""
          self.X = self.rng.uniform(low=low, high=high, size=(self.N, self._p))
          self.X0 = np.ones((self.N, 1))
          self.X1 = np.concatenate([self.X0, self.X], axis=1)
  