import os
import urllib.request
from zipfile import ZipFile
from io import StringIO

import numpy as np
import pandas as pd
import scipy.sparse as sp


def globally_normalize_bipartite_adjacency(adjacencies, symmetric=True):
 """ Globally Normalizes set of bipartite adjacency matrices """

 print('{} normalizing bipartite adj'.format(
 ['Asymmetrically', 'Symmetrically'][symmetric]))

 adj_tot = np.sum([adj for adj in adjacencies])
 degree_u = np.asarray(adj_tot.sum(1)).flatten()
 degree_v = np.asarray(adj_tot.sum(0)).flatten()

 # set zeros to inf to avoid dividing by zero
 degree_u[degree_u == 0.] = np.inf
 degree_v[degree_v == 0.] = np.inf

 degree_u_inv_sqrt = 1. / np.sqrt(degree_u)
 degree_v_inv_sqrt = 1. / np.sqrt(degree_v)
 degree_u_inv_sqrt_mat = sp.diags([degree_u_inv_sqrt], [0])
 degree_v_inv_sqrt_mat = sp.diags([degree_v_inv_sqrt], [0])

 degree_u_inv = degree_u_inv_sqrt_mat.dot(degree_u_inv_sqrt_mat)

 if symmetric:
 adj_norm = [degree_u_inv_sqrt_mat.dot(adj).dot(
 degree_v_inv_sqrt_mat) for adj in adjacencies]

 else:
 adj_norm = [degree_u_inv.dot(adj) for adj in adjacencies]

 return adj_norm


def get_adjacency(edge_df, num_user, num_movie, symmetric_normalization):
 user2movie_adjacencies = []
 movie2user_adjacencies = []
 train_edge_df = edge_df.loc[edge_df['usage'] == 'train']
 for i in range(5):
 edge_index = train_edge_df.loc[train_edge_df.ratings == i, [
 'user_node_id', 'movie_node_id']].to_numpy()
 support = sp.csr_matrix((np.ones(len(edge_index)), (edge_index[:, 0], edge_index[:, 1])),
 shape=(num_user, num_movie), dtype=np.float32)
 user2movie_adjacencies.append(support)
 movie2user_adjacencies.append(support.T)

 user2movie_adjacencies = globally_normalize_bipartite_adjacency(user2movie_adjacencies,
 symmetric=symmetric_normalization)
 movie2user_adjacencies = globally_normalize_bipartite_adjacency(movie2user_adjacencies,
 symmetric=symmetric_normalization)

 return user2movie_adjacencies, movie2user_adjacencies


def get_node_identity_feature(num_user, num_movie):
 """one-hot encoding for nodes"""
 identity_feature = np.identity(num_user + num_movie, dtype=np.float32)
 user_identity_feature, movie_indentity_feature = identity_feature[
 :num_user], identity_feature[num_user:]

 return user_identity_feature, movie_indentity_feature


def get_user_side_feature(node_user: pd.DataFrame):
 """用户节点属性特征,包括年龄,性别,职业"""
 age = node_user['age'].to_numpy().astype('float32')
 age /= age.max()
 age = age.reshape((-1, 1))
 gender_arr, gender_index = pd.factorize(node_user['gender'])
 gender_arr = np.reshape(gender_arr, (-1, 1))
 occupation_arr = pd.get_dummies(node_user['occupation']).to_numpy()

 user_feature = np.concatenate([age, gender_arr, occupation_arr], axis=1)

 return user_feature


def get_movie_side_feature(node_movie: pd.DataFrame):
 """电影节点属性特征,主要是电影类型"""
 movie_genre_cols = ['Action', 'Adventure', 'Animation',
 'Childrens', 'Comedy', 'Crime', 'Documentary', 'Drama', 'Fantasy',
 'Film-Noir', 'Horror', 'Musical', 'Mystery', 'Romance', 'Sci-Fi',
 'Thriller', 'War', 'Western']
 movie_genre_arr = node_movie.loc[:,
 movie_genre_cols].to_numpy().astype('float32')
 return movie_genre_arr


def convert_to_homogeneous(user_feature: np.ndarray, movie_feature: np.ndarray):
 """通过补零将用户和电影的属性特征对齐到同一维度"""
 num_user, user_feature_dim = user_feature.shape
 num_movie, movie_feature_dim = movie_feature.shape
 user_feature = np.concatenate(
 [user_feature, np.zeros((num_user, movie_feature_dim))], axis=1)
 movie_feature = np.concatenate(
 [np.zeros((num_movie, user_feature_dim)), movie_feature], axis=1)

 return user_feature, movie_feature


def normalize_feature(feature):
 row_sum = feature.sum(1)
 row_sum[row_sum == 0] = np.inf
 normalized_feat = feature / row_sum.reshape(-1, 1)
 return normalized_feat


class MovielensDataset(object):
 url = "http://files.grouplens.org/datasets/movielens/ml-100k.zip"

 def __init__(self, data_root="data"):
 self.data_root = data_root
 self.maybe_download()

 @staticmethod
 def build_graph(edge_df: pd.DataFrame, user_df: pd.DataFrame,
 movie_df: pd.DataFrame, symmetric_normalization=False):
 node_user = edge_df[['user_node']
 ].drop_duplicates().sort_values('user_node')
 node_movie = edge_df[['movie_node']
 ].drop_duplicates().sort_values('movie_node')
 node_user.loc[:, 'user_node_id'] = range(len(node_user))
 node_movie.loc[:, 'movie_node_id'] = range(len(node_movie))

 edge_df = edge_df.merge(node_user, on='user_node', how='left')\
 .merge(node_movie, on='movie_node', how='left')

 node_user = node_user.merge(user_df, on='user_node', how='left')
 node_movie = node_movie.merge(movie_df, on='movie_node', how='left')
 num_user = len(node_user)
 num_movie = len(node_movie)

 # adjacency
 user2movie_adjacencies, movie2user_adjacencies = get_adjacency(edge_df, num_user, num_movie,
 symmetric_normalization)

 # node property feature
 user_side_feature = get_user_side_feature(node_user)
 movie_side_feature = get_movie_side_feature(node_movie)
 user_side_feature = normalize_feature(user_side_feature)
 movie_side_feature = normalize_feature(movie_side_feature)
 user_side_feature, movie_side_feature = convert_to_homogeneous(user_side_feature,
 movie_side_feature)

 # one-hot encoding for nodes
 user_identity_feature, movie_indentity_feature = get_node_identity_feature(
 num_user, num_movie)

 # user_indices, movie_indices, labels, train_mask
 user_indices, movie_indices, labels = edge_df[[
 'user_node_id', 'movie_node_id', 'ratings']].to_numpy().T
 train_mask = (edge_df['usage'] == 'train').to_numpy()

 return user2movie_adjacencies, movie2user_adjacencies, \
 user_side_feature, movie_side_feature, \
 user_identity_feature, movie_indentity_feature, \
 user_indices, movie_indices, labels, train_mask

 def read_data(self):
 data_dir = os.path.join(self.data_root, "ml-100k")
 # edge data
 edge_train = pd.read_csv(os.path.join(data_dir, 'u1.base'), sep='\t',
 header=None, names=['user_node', 'movie_node', 'ratings', 'timestamp'])
 edge_train.loc[:, 'usage'] = 'train'
 edge_test = pd.read_csv(os.path.join(data_dir, 'u1.test'), sep='\t',
 header=None, names=['user_node', 'movie_node', 'ratings', 'timestamp'])
 edge_test.loc[:, 'usage'] = 'test'
 edge_df = pd.concat((edge_train, edge_test),
 axis=0).drop(columns='timestamp')
 edge_df.loc[:, 'ratings'] -= 1
 # item feature
 sep = r'|'
 movie_file = os.path.join(data_dir, 'u.item')
 movie_headers = ['movie_node', 'movie_title', 'release_date', 'video_release_date',
 'IMDb_URL', 'unknown', 'Action', 'Adventure', 'Animation',
 'Childrens', 'Comedy', 'Crime', 'Documentary', 'Drama', 'Fantasy',
 'Film-Noir', 'Horror', 'Musical', 'Mystery', 'Romance', 'Sci-Fi',
 'Thriller', 'War', 'Western']
 movie_df = pd.read_csv(movie_file, sep=sep, header=None,
 names=movie_headers, encoding='latin1')
 # user feature
 users_file = os.path.join(data_dir, 'u.user')
 users_headers = ['user_node', 'age',
 'gender', 'occupation', 'zip_code']
 users_df = pd.read_csv(users_file, sep=sep, header=None,
 names=users_headers, encoding='latin1')
 return edge_df, users_df, movie_df

 def maybe_download(self):
 save_path = os.path.join(self.data_root)
 if not os.path.exists(save_path):
 self.download_data(self.url, save_path)
 if not os.path.exists(os.path.join(self.data_root, "ml-100k")):
 zipfilename = os.path.join(self.data_root, "ml-100k.zip")
 with ZipFile(zipfilename, "r") as zipobj:
 zipobj.extractall(os.path.join(self.data_root))
 print("Extracting data from {}".format(zipfilename))

 @staticmethod
 def download_data(url, save_path):
 """数据下载工具,当原始数据不存在时将会进行下载"""
 print("Downloading data from {}".format(url))
 if not os.path.exists(save_path):
 os.makedirs(save_path)
 request = urllib.request.urlopen(url)
 filename = os.path.basename(url)
 with open(os.path.join(save_path, filename), 'wb') as f:
 f.write(request.read())
 return True


if __name__ == "__main__":
 data = MovielensDataset()
 user2movie_adjacencies, movie2user_adjacencies, \
 user_side_feature, movie_side_feature, \
 user_identity_feature, movie_indentity_feature, \
 user_indices, movie_indices, labels, train_mask = data.build_graph(
 *data.read_data())