import warnings
from collections import defaultdict

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

warnings.filterwarnings(action='ignore')
torch.set_printoptions(sci_mode=True)

import argparse

def get_config(jupyter_notebook=False):
    parser = argparse.ArgumentParser()
    parser.add_argument('--valid_samples', default=10, type=int, help="검증에 사용할 sample 수")
    parser.add_argument('--seed', default=2024, type=int)
    return parser.parse_args(args=[]) if jupyter_notebook else parser.parse_args()

config = get_config(jupyter_notebook=True)

device = 'cuda' if torch.cuda.is_available() else 'cpu'
config

import lightning as L
from mmrec.utils.common import get_root_path

ROOT = get_root_path()
L.seed_everything(config.seed)

class MakeMatrixDataSet():
    """
    MatrixDataSet 생성
    """
    def __init__(self, config):
        self.config = config
        df = pd.read_parquet(f"{ROOT}/data/ml-1m/logs.parquet")
        # valid sample 수 만큼 이력이 있는 warm 유저만 필터링하여 학습/평가
        count_user_hist = df.groupby('user_id').item_id.apply(len)
        warm_user_set = set(count_user_hist[count_user_hist > config.valid_samples].index.tolist())
        self.df = df[df.user_id.isin(warm_user_set)]
        
        self.item_encoder, self.item_decoder = self.generate_encoder_decoder('item_id')
        self.user_encoder, self.user_decoder = self.generate_encoder_decoder('user_id')
        self.num_item, self.num_user = len(self.item_encoder), len(self.user_encoder)

        self.df['item_idx'] = self.df['item_id'].apply(lambda x : self.item_encoder[x])
        self.df['user_idx'] = self.df['user_id'].apply(lambda x : self.user_encoder[x])

        self.user_train, self.user_valid = self.generate_sequence_data()

    def generate_encoder_decoder(self, col : str) -> dict:
        """
        encoder, decoder 생성

        Args:
            col (str): 생성할 columns 명
        Returns:
            dict: 생성된 user encoder, decoder
        """

        encoder = {}
        decoder = {}
        ids = self.df[col].unique()

        for idx, _id in enumerate(ids):
            encoder[_id] = idx
            decoder[idx] = _id

        return encoder, decoder
    
    def generate_sequence_data(self) -> dict:
        """
        sequence_data 생성

        Returns:
            dict: train user sequence / valid user sequence
        """
        users = defaultdict(list)
        user_train = {}
        user_valid = {}
        for user, item in zip(self.df['user_idx'], self.df['item_idx']):
            users[user].append(item)
        
        for user in users:

            user_total = users[user]
            valid = np.random.choice(user_total, size = self.config.valid_samples, replace = False).tolist()
            train = list(set(user_total) - set(valid))

            user_train[user] = train
            user_valid[user] = valid # valid_samples 개수 만큼 검증에 활용 (현재 Task와 가장 유사하게)

        return user_train, user_valid
    
    def get_train_valid_data(self):
        return self.user_train, self.user_valid

    def make_matrix(self, user_list, train = True):
        """
        user_item_dict를 바탕으로 행렬 생성
        """
        mat = torch.zeros(size = (user_list.size(0), self.num_item))
        for idx, user in enumerate(user_list):
            if train:
                mat[idx, self.user_train[user.item()]] = 1
            else:
                mat[idx, self.user_train[user.item()] + self.user_valid[user.item()]] = 1
        return mat

    def make_sparse_matrix(self):
        X = sp.dok_matrix((self.num_user, self.num_item), dtype=np.float32)
        for user in self.user_train.keys():
            item_list = self.user_train[user]
            X[user, item_list] = 1.0
                
        return X.tocsr()

class EASE():
    def __init__(self, X, reg):
        self.X = self._convert_sp_mat_to_sp_tensor(X)
        self.reg = reg
    
    def _convert_sp_mat_to_sp_tensor(self, X):
        """
        Convert scipy sparse matrix to PyTorch sparse matrix

        Arguments:
        ----------
        X = Adjacency matrix, scipy sparse matrix
        """
        coo = X.tocoo().astype(np.float32)
        i = torch.LongTensor(np.mat([coo.row, coo.col]))
        v = torch.FloatTensor(coo.data)
        res = torch.sparse.FloatTensor(i, v, coo.shape).to(device)
        return res
    
    def fit(self):
        '''

        진짜 정말 간단한 식으로 모델을 만듬

        '''
        G = self.X.to_dense().t() @ self.X.to_dense()
        diagIndices = torch.eye(G.shape[0]) == 1
        G[diagIndices] += self.reg

        P = G.inverse()
        B = P / (-1 * P.diag())
        B[diagIndices] = 0

        self.pred = self.X.to_dense() @ B

def get_ndcg(pred_list, true_list):
    idcg = sum((1 / np.log2(rank + 2) for rank in range(1, len(pred_list))))
    dcg = 0
    for rank, pred in enumerate(pred_list):
        if pred in true_list:
            dcg += 1 / np.log2(rank + 2)
    ndcg = dcg / idcg
    return ndcg

# hit == recall == precision
def get_hit(pred_list, true_list):
    hit_list = set(true_list) & set(pred_list)
    hit = len(hit_list) / len(true_list)
    return hit

def evaluate(model, X, user_train, user_valid):

    mat = torch.from_numpy(X)

    NDCG = 0.0 # NDCG@10
    HIT = 0.0 # HIT@10

    recon_mat1 = model.pred.cpu()
    recon_mat1[mat == 1] = -np.inf  # 이미 interaction이 있는 아이템 제거
    rec_list1 = recon_mat1.argsort(dim = 1)

    for user, rec1 in enumerate(rec_list1):
        uv = user_valid[user]

        # ranking
        up = rec1[-10:].cpu().numpy().tolist()[::-1]

        NDCG += get_ndcg(pred_list = up, true_list = uv)
        HIT += get_hit(pred_list = up, true_list = uv)

    NDCG /= len(user_train)
    HIT /= len(user_train)

    return NDCG, HIT

make_matrix_data_set = MakeMatrixDataSet(config = config)
user_train, user_valid = make_matrix_data_set.get_train_valid_data()
X = make_matrix_data_set.make_sparse_matrix()

print(X.shape)

make_matrix_data_set.df

for reg in [1000, 100, 10, 1, 0.1, 0.01]:
    model = EASE(X = X, reg = reg)
    model.fit()
    ndcg, hit = evaluate(model = model, X = X.todense(), user_train = user_train, user_valid = user_valid)
    print(f'reg: {reg}| NDCG@10: {ndcg:.5f}| HIT@10: {hit:.5f}')