import numpy as np

# 원본 행렬 R 생성
# 분해 행렬 P, Q 초기화
# 잠재 요인 차원 K는 3으로 설정
R = np.array([[4, np.NaN, np.NaN, 2, np.NaN],
    [np.NaN, 5, np.NaN, 3, 1],
    [np.NaN, np.NaN, 3, 4, 4],
    [5, 2, 1, 2, np.NaN]])

num_users, num_items = R.shape
K=3

# P와 Q 행렬의 크기를 지정하고 정규 분포를 가진 임의의 값으로 입력
np.random.seed(1)

# np.random.normal 함수는 정규분포를 띄며 scale은 표준편차를 의미함.
P = np.random.normal(scale=1./K, size=(num_users, K))
Q = np.random.normal(scale=1./K, size=(num_items, K))

print(P)
print(Q)

arr = np.array([[1,2,3],[4,5,6]])

print(f"arr: \n{arr} \n")
print(f"arr[[0, 1], [1,2]]: \n{arr[[0,1 ], [1,2]]}")

from sklearn.metrics import mean_squared_error

def get_rmse(R, P, Q, non_zeros):
    error = 0
    # 두 개의 분해된 행렬 P와 Q.T의 내적으로 예측 R 행렬 생성
    full_pred_matrix = np.dot(P, Q.T)

    # 실제 R 행렬에서 널이 아닌 값의 위치 인덱스 추출해 실제 R 행렬과 예측 행렬의 RMSE 추출
    x_non_zero_ind = [non_zero[0] for non_zero in non_zeros]
    y_non_zero_ind = [non_zero[1] for non_zero in non_zeros]
    R_non_zeros = R[x_non_zero_ind, y_non_zero_ind]
    full_pred_matrix_non_zeros = full_pred_matrix[x_non_zero_ind, y_non_zero_ind]       # 이런식으로 x좌표를 row 자리에, y좌표를 column 자리에 배치하여 여러 값들에 인덱싱이 가능하다
    mse = mean_squared_error(R_non_zeros, full_pred_matrix_non_zeros)
    rmse = np.sqrt(mse)

    return rmse

# R > 0인 행 위치, 열 위치, 값을 non_zeros 리스트에 저장
non_zeros = [(i, j, R[i, j])for i in range(num_users) for j in range(num_items) if R[i, j] > 0]

steps = 1000
learning_rate = 0.01
r_lambda = 0.01

# SGD 기법으로 P, Q 매트릭스 계속 업데이트.
for step in range(steps):
    for i, j, r in non_zeros:
        # 실제 값과 예측 값의 차이인 오류 값 구함
        # P, Q 행렬은 정규분포로 초기화 해둔 상태여서 Null값이 없음
        eij = r - np.dot(P[i, :], Q[j, :].T)        # Q에 대해 전치를 하기때문에 j 변수가 row 인덱스 쪽으로 온다
        # Regularization을 반영한 SGD 업데이트 공식 적용
        P[i, :] = P[i, :] + learning_rate * (eij * Q[j, :] - r_lambda * P[i, :])
        Q[j, :] = Q[j, :] + learning_rate * (eij * P[i, :] - r_lambda * Q[j, :])

        rmse = get_rmse(R, P, Q, non_zeros)
        if (step%50) == 0:
            print("### Iteration step: ", step, "rmse: ", rmse)

# P, Q.T 내적 결과 확인.
print(np.dot(P, Q.T))