from pathlib import Path

import numpy as np
from sklearn.metrics import mean_squared_error
from sklearn.preprocessing import StandardScaler

DATA_DIR = Path(__file__).resolve().parent / "data"
if not DATA_DIR.exists():
    DATA_DIR = Path.cwd() / "data"

MSD_TRAIN_PATH = DATA_DIR / "YearPredictionMSD_train.txt"
MSD_TEST_PATH = DATA_DIR / "YearPredictionMSD_test.txt"

OFFSET = 1922
LAYER_1_SIZE = 20
LAYER_2_SIZE = 10
OUTPUT_SIZE = 1
DEMO_N_TRACKS = 5


class NeuralNet:
    def __init__(self, input_size: int):
        """
        Requires: Nothing
        Modifies: self.W1, self.b1, self.W2, self.b2, self.W3, self.b3
        Effects: Creates and initializes the weights and biases for the
                 neural network. The seed is set to 183.
        """
        np.random.seed(183)
        self._init_weights(input_size)

    def _init_weights(self, input_size: int) -> None:
        """
        Requires: Nothing
        Modifies: self.W1, self.b1, self.W2, self.b2, self.W3, self.b3
        Effects: Initializes the weights and biases to small random numbers.
        """
        self.W1 = np.random.randn(input_size, LAYER_1_SIZE) * 0.01
        self.b1 = np.zeros((1, LAYER_1_SIZE))
        self.W2 = np.random.randn(LAYER_1_SIZE, LAYER_2_SIZE) * 0.01
        self.b2 = np.zeros((1, LAYER_2_SIZE))
        self.W3 = np.random.randn(LAYER_2_SIZE, OUTPUT_SIZE) * 0.01
        self.b3 = np.zeros((1, OUTPUT_SIZE))

    def relu(self, x: np.ndarray) -> np.ndarray:
        """
        Requires: Nothing
        Modifies: Nothing
        Effects: Returns the ReLU activation of the input.
                 Formula: ReLU(x) = x if x > 0 else 0 for
                 each data point in the input array x. Use
                 vectorization - do not use a for loop.
        """
        return np.zeros_like(x)

    def relu_derivative(self, x: np.ndarray) -> np.ndarray:
        """
        Requires: Nothing
        Modifies: Nothing
        Effects: Returns the derivative of the ReLU activation of the input.
                 Formula: ReLU_derivative(x) = 1 if x > 0 else 0 for
                 each data point in the input array x. Use
                 vectorization - do not use a for loop. Cast the result
                 to an array of floats.
        """
        return np.zeros_like(x)

    def sigmoid(self, x: np.ndarray) -> np.ndarray:
        """
        Requires: Nothing
        Modifies: Nothing
        Effects: Returns the sigmoid activation of the input.
                 Formula: Sigmoid(x) = 1 / (1 + e^-x) for
                 each data point in the input array x. However,
                 before doing the exponentiation (e^-x), clip the
                 input x to be between -500 and 500 to prevent the
                 denominator from being too small. Use
                 vectorization - do not use a for loop.
        """
        return np.zeros_like(x)

    def sigmoid_derivative(self, x: np.ndarray) -> np.ndarray:
        """
        Requires: Nothing
        Modifies: Nothing
        Effects: Returns the derivative of the sigmoid activation of the input.
                 Formula: Sigmoid_derivative(x) = sigmoid(x) * (1 - sigmoid(x)).
                 You should call sigmoid in this function.
        """
        return np.zeros_like(x)

    def fc(self, X: np.ndarray, W: np.ndarray, b: np.ndarray) -> np.ndarray:
        """
        Requires: Nothing
        Modifies: Nothing
        Effects: Returns the fully connected output of the input.
                 Formula: FC(X, W, b) = X * W + b for
                 each data point in the input array X. Use
                 vectorization - do not use a for loop. Remember
                 that in this context, X * W is the dot product of X and W.
        """
        return np.zeros_like(X)

    def forward(self, X: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
        """
        Requires: Nothing
        Modifies: Nothing
        Effects: Runs the full forward pass for input X:
                 Input -> FC -> ReLU -> FC -> Sigmoid -> FC.
                 Returns (fc1_output, relu_output, fc2_output, sigmoid_output,
                 pred) where pred is the final FC output. Used by train() and
                 predict().
        """
        return (np.zeros_like(X), np.zeros_like(X), np.zeros_like(X), np.zeros_like(X), np.zeros_like(X))

    def predict(self, X: np.ndarray) -> np.ndarray:
        """
        Requires: Nothing
        Modifies: Nothing
        Effects: Returns the predicted output for each row of X by running
                 the same forward pass as training (via forward()).
        """
        return np.zeros_like(X)

    def train(self, X: np.ndarray, y: np.ndarray, iterations: int = 5, lr: float = 1e-4) -> None:
        """
        Requires: X.shape[0] == y.shape[0]
        Modifies: self.W1, self.b1, self.W2, self.b2, self.W3, self.b3
        Effects: Trains the neural network on the input data.
                 The neural network is structured as follows:
                 Input -> FC -> ReLU -> FC -> Sigmoid -> FC
                 (FC means Fully Connected). Remember, the output of a 
                 layer is the direct input to the next layer.
                 This method uses stochastic gradient descent (SGD) and 
                 trains one input sample at a time.
        """
        for epoch in range(iterations):
            squared_error = 0
            indices = np.arange(len(X))
            np.random.shuffle(indices)

            for i in indices:
                data_point = X[i:i+1]

                # TODO: perform forward pass (store results in fc1_output,
                # relu_output, fc2_output, sigmoid_output, pred)

                # Pred is a 2D numpy array of shape 1x1. This gets the element as a float.
                pred = pred.item()

                # TODO:
                # Compute the error (pred - y[i]) and store in a variable called error.
                # Compute the squared error
                # Add the current squared error to the squared_error variable

                # Backward Pass
                grad_W3 = np.dot(sigmoid_output.T, error)
                delta2 = np.dot(error, self.W3.T) * self.sigmoid_derivative(
                    fc2_output
                )
                grad_W2 = np.dot(relu_output.T, delta2)
                delta1 = (
                    np.dot(delta2, self.W2.T)
                    * self.relu_derivative(fc1_output)
                )
                grad_W1 = np.dot(data_point.T, delta1)

                # SGD update
                self.W1 -= lr * grad_W1
                self.b1 -= lr * np.sum(delta1, axis=0, keepdims=True)
                self.W2 -= lr * grad_W2
                self.b2 -= lr * np.sum(delta2, axis=0, keepdims=True)
                self.W3 -= lr * grad_W3
                self.b3 -= lr * error

            print(f"Epoch {epoch}: RMSE {np.sqrt(squared_error / len(X)):.2f}")

def load_msd_data():
    # Load train and test data from the text files under DATA_DIR
    train_data = np.loadtxt(MSD_TRAIN_PATH, delimiter=",")
    test_data = np.loadtxt(MSD_TEST_PATH, delimiter=",")

    X_train_raw = train_data[:, 1:]
    X_test_raw = test_data[:, 1:]

    # Get the ground-truth labels for training and testing sets
    # subtract so the years are relative to 1922 - this makes
    # the data easier for the neural network to learn
    y_train = train_data[:, 0] - OFFSET
    y_test = test_data[:, 0] - OFFSET

    # Normalize the data
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train_raw)
    X_test_scaled = scaler.transform(X_test_raw)

    return X_train_scaled, X_test_scaled, y_train, y_test


def print_demo_predictions(
    nn: NeuralNet, X_demo: np.ndarray, y_demo: np.ndarray
) -> None:
    """
    Requires: Nothing
    Modifies: Nothing
    Effects: Prints the actual and predicted release years for a few songs.
    """
    preds = nn.predict(X_demo)
    n = len(y_demo)
    print(f"{'#':>3} {'Actual year':>12} {'Predicted':>12} {'Error':>10}")
    total_abs_error = 0
    for i in range(n):
        actual_year = int(round(y_demo[i] + OFFSET))
        pred_year = int(round(preds[i, 0] + OFFSET))
        err = pred_year - actual_year
        total_abs_error += abs(err)
        print(f"{i + 1:>3} {actual_year:>12} {pred_year:>12} {err:>10}")


def main():
    X_train, X_test, y_train, y_test = load_msd_data()

    X_demo = X_test[:DEMO_N_TRACKS]
    y_demo = y_test[:DEMO_N_TRACKS]

    nn = NeuralNet(input_size=X_train.shape[1])
    nn.train(X_train, y_train, iterations=0)

    print("Sample predictions with untrained network:")
    print_demo_predictions(nn, X_demo, y_demo)
    print()

    print("Starting 3-Layer Neural Network Training...")
    nn.train(X_train, y_train, iterations=50, lr=5e-5)

    print()
    predictions = nn.predict(X_test)
    rmse = np.sqrt(mean_squared_error(y_test, predictions))

    print("Sample predictions with trained network:")
    print_demo_predictions(nn, X_demo, y_demo)
    print()

    print("-" * 35)
    print(f"Neural Network RMSE: {rmse:.2f} years")  
    print(f"Linear Regression RMSE: 7.30 years")
    print("-" * 35)

if __name__ == "__main__":
    main()