from esl_model.ch3.models import LeastSquareModel

# import prostate data set
from esl_model.datasets import ProstateDataSet

data = ProstateDataSet()

lsm = LeastSquareModel(train_x=data.train_x, train_y=data.train_y)
lsm.pre_processing()
lsm.train()

# after pre_processing and train, you can get the beta_hat
print(lsm.beta_hat)

# predict
y_hat = lsm.predict(data.test_x)

# get the test result
test_result = lsm.test(data.test_x, data.test_y)

# get the mean of square error
print(test_result.mse)

# get standard error
print(test_result.std_error)

class LeastSquareModel(LinearModel):
    def _pre_processing_x(self, X):
        X = self.standardize(X)
        X = np.insert(X, 0, 1, axis=1)
        return X

    def train(self):
        x = self.train_x
        y = self.train_y
        self.beta_hat = self.math.inv(x.T @ x) @ x.T @ y

    def predict(self, X):
        X = self._pre_processing_x(X)
        return X @ self.beta_hat