import numpy as np

# =============================================================================
# TODO: Implement the Evaluator class from scratch.
#
# Requirements:
#   1. Define a class named Evaluator.
#   2. Define an __init__ method (no attributes needed).
#   3. Implement the accuracy and mean_squared_error methods inside the class.
#
# Note: You can copy and paste the function headers provided in the comment
# below into your new Evaluator class. Make sure they are indented correctly!
# =============================================================================

# Write your Evaluator class here:


'''
# You may copy/paste these headers into your Evaluator class:

    def accuracy(self, y_true: np.ndarray, y_pred: np.ndarray) -> float:
        """
        Requires: y_true and y_pred are 1D numpy arrays.
        Modifies: Nothing.
        Effects: Returns the accuracy of the predictions (exact matches).
                 See the lab spec for intuition.
        """
        # TODO: Implement this method
        pass

    def mean_squared_error(
            self, y_true: np.ndarray, y_pred: np.ndarray
            ) -> float:
        """
        Requires: y_true and y_pred are 1D numpy arrays.
        Modifies: Nothing.
        Effects: Returns the mean squared error (MSE) of the predictions.
                 Algorithm:
                 1. Subtract y_pred from y_true to get the error.
                 2. Square the error.
                 3. Take the mean of the squared errors.
        """
        # TODO: Implement this method
        pass
'''


class LinearRegressor:
    def __init__(self) -> None:
        self.w = None

    def fit(self, X: np.ndarray, y: np.ndarray) -> None:
        """
        Requires: Nothing
        Modifies: self.w
        Effects: Uses the closed form solution of linear regression
                 to fit the weights of the linear regression model.
                 The weights are stored in self.w. The closed form solution is:
                 w = (X^T X)^-1 X^T y
                 See the spec for hints and more details.
        """
        # TODO: Implement this function
        pass

    def predict(self, X: np.ndarray) -> np.ndarray:
        """
        Requires: Nothing
        Modifies: Nothing
        Effects: Uses the weights in self.w to predict the labels for
                 the testing dataset X. The prediction equation is:
                 y_pred = Xw
                 See the spec for hints and more details. Note that
                 the predictions may be floats, but we are predicting the
                 release year, which is an integer, so you should use
                 .astype(int) to return the predictions as integers.
        """
        # TODO: Implement this function
        return np.zeros(X.shape[0])


def add_intercept(X: np.ndarray) -> np.ndarray:
    """
    Requires: Nothing
    Modifies: Nothing
    Effects: Adds a column of ones to the front of the matrix X.
    """
    X = np.asarray(X, dtype=float)
    ones = np.ones((X.shape[0], 1), dtype=float)
    return np.hstack([ones, X])


def main():
    print("=== Loading dataset into numpy array ===")
    # Loading dataset into numpy array
    X_train_data = np.loadtxt("YearPredictionMSD_train.txt", delimiter=",")
    X_test_data = np.loadtxt("YearPredictionMSD_test.txt", delimiter=",")
    print()

    # Get the features for training and testing sets
    X_train_raw = X_train_data[:, 1:]
    X_test_raw = X_test_data[:, 1:]

    # Get the ground-truth labels for training and testing sets
    y_train = X_train_data[:, 0].astype(int)
    y_test = X_test_data[:, 0].astype(int)

    # Add intercept column to training and testing sets
    X_train = add_intercept(X_train_raw)
    X_test = add_intercept(X_test_raw)

    # Train and test linear regression
    print("=== Linear Regression: predict release year ===")
    regressor = LinearRegressor()
    regressor.fit(X_train, y_train)
    y_pred = regressor.predict(X_test)

    # Evaluate the regressor
    evaluator = Evaluator()
    mse = evaluator.mean_squared_error(y_test, y_pred)
    rmse = float(np.sqrt(mse))
    accuracy = evaluator.accuracy(y_test, y_pred) * 100
    print(f"Test MSE:  {mse:.3f}")
    print(f"Test RMSE: {rmse:.3f}")
    print(f"Test accuracy: {accuracy:.2f}%")
    print()

    # Print a few example predictions.
    print("Sample predictions (predicted year, actual year):")
    for i in range(10):
        print(f"  {y_pred[i]}, {y_test[i]}")


if __name__ == "__main__":
    main()
