# coding=utf-8

import torch
import logging
from sklearn.metrics import *
import numpy as np
import torch.nn.functional as F
import os
import pandas as pd
from utils.rank_metrics import *


class BaseModel(torch.nn.Module):
    """
    Base model, the following methods need to be overridden.
    parse_model_args,
    __init__,
    _init_weights,
    predict,
    forward,
    """
    append_id = False
    include_id = True
    include_user_features = True
    include_item_features = True
    include_context_features = False

    @staticmethod
    def parse_model_args(parser, model_name='BaseModel'):
        """
        parse model related arguments
        :param parser: parser object, which is from main() method
        :param model_name: model name
        :return:
        """
        parser.add_argument('--model_path', type=str,
                            default='../model/%s/%s.pt' % (model_name, model_name),
                            help='Model save path.')
        return parser

    @staticmethod
    def evaluate_method(p, data, metrics):
        """
        calculate evaluation metrics
        :param p: prediction valuds，np.array，generated by runner.predict()
        :param data: data dict，generated by DataProcessor
        :param metrics: metrics list，generated by runner.metrics，for example ['rmse', 'auc']
        :return:
        """
        l = data['Y']
        evaluations = []
        for metric in metrics:
            if metric == 'rmse':
                evaluations.append(np.sqrt(mean_squared_error(l, p)))
            elif metric == 'mae':
                evaluations.append(mean_absolute_error(l, p))
            elif metric == 'auc':
                evaluations.append(roc_auc_score(l, p))
            elif metric == 'f1':
                evaluations.append(f1_score(l, p))
            elif metric == 'accuracy':
                evaluations.append(accuracy_score(l, p))
            elif metric == 'precision':
                evaluations.append(precision_score(l, p))
            elif metric == 'recall':
                evaluations.append(recall_score(l, p))
            else:
                k = int(metric.split('@')[-1])
                df = pd.DataFrame()
                df['uid'] = data['uid']
                df['p'] = p
                df['l'] = l
                df = df.sort_values(by='p', ascending=False)
                df_group = df.groupby('uid')
                if metric.startswith('ndcg@'):
                    ndcgs = []
                    for uid, group in df_group:
                        ndcgs.append(ndcg_at_k(group['l'].tolist()[:k], k=k, method=1))
                    evaluations.append(np.average(ndcgs))
                elif metric.startswith('hit@'):
                    hits = []
                    for uid, group in df_group:
                        hits.append(int(np.sum(group['l'][:k]) > 0))
                    evaluations.append(np.average(hits))
                elif metric.startswith('precision@'):
                    precisions = []
                    for uid, group in df_group:
                        precisions.append(precision_at_k(group['l'].tolist()[:k], k=k))
                    evaluations.append(np.average(precisions))
                elif metric.startswith('recall@'):
                    recalls = []
                    for uid, group in df_group:
                        recalls.append(1.0 * np.sum(group['l'][:k]) / np.sum(group['l']))
                    evaluations.append(np.average(recalls))
        return evaluations

    @staticmethod
    def init_paras(m):
        """
        Initialize model parameters
        :param m: model parameters
        :return:
        """
        if type(m) == torch.nn.Linear:
            torch.nn.init.normal_(m.weight, mean=0.0, std=0.01)
            if m.bias is not None:
                torch.nn.init.normal_(m.bias, mean=0.0, std=0.01)
        elif type(m) == torch.nn.Embedding:
            torch.nn.init.normal_(m.weight, mean=0.0, std=0.01)

    def __init__(self, label_min, label_max, feature_num, random_seed=2018, model_path='../model/Model/Model.pt'):
        super(BaseModel, self).__init__()
        self.label_min = label_min
        self.label_max = label_max
        self.feature_num = feature_num
        self.random_seed = random_seed
        torch.manual_seed(self.random_seed)
        torch.cuda.manual_seed(self.random_seed)
        self.model_path = model_path

        self._init_weights()
        logging.debug(list(self.parameters()))

        self.total_parameters = self.count_variables()
        logging.info('# of params: %d' % self.total_parameters)

        # optimizer 由runner生成并赋值
        self.optimizer = None

    def _init_weights(self):
        self.x_bn = torch.nn.BatchNorm1d(self.feature_num)
        self.prediction = torch.nn.Linear(self.feature_num, 1)

    def count_variables(self):
        """
        count number of parameters in the model
        :return:
        """
        total_parameters = sum(p.numel() for p in self.parameters() if p.requires_grad)
        return total_parameters

    def l2(self):
        """
        calculate l2 regularization
        :return:
        """
        l2 = 0
        for p in self.parameters():
            l2 += (p ** 2).sum()
        return l2

    def predict(self, feed_dict):
        """
        prediction only. No loss calculation
        :param feed_dict: input data dictionary
        :return: output dictionary，prediction: contains predicted values; check: holds info
        to be output in the training process (for monitoring).
        """
        check_list = []
        x = self.x_bn(feed_dict['X'].float())
        x = torch.nn.Dropout(p=feed_dict['dropout'])(x)
        prediction = F.relu(self.prediction(x)).view([-1])
        out_dict = {'prediction': prediction,
                    'check': check_list}
        return out_dict

    def forward(self, feed_dict):
        """
        calculate loss
        :param feed_dict: input dictionary
        :return: output dictionary, which add loss info to the output dictionary of predict() method.
        """
        out_dict = self.predict(feed_dict)
        if feed_dict['rank'] == 1:
            # 计算topn推荐的loss，batch前一半是正例，后一半是负例
            batch_size = int(feed_dict['Y'].shape[0] / 2)
            pos, neg = out_dict['prediction'][:batch_size], out_dict['prediction'][batch_size:]
            loss = -(pos - neg).sigmoid().log().sum()
        else:
            # 计算rating/clicking预测的loss，默认使用mse
            loss = torch.nn.MSELoss()(out_dict['prediction'], feed_dict['Y'])
        out_dict['loss'] = loss
        return out_dict

    def lrp(self):
        pass

    def save_model(self, model_path=None):
        """
        save model
        """
        if model_path is None:
            model_path = self.model_path
        dir_path = os.path.dirname(model_path)
        if not os.path.exists(dir_path):
            os.mkdir(dir_path)
        torch.save(self.state_dict(), model_path)
        logging.info('Save model to ' + model_path)

    def load_model(self, model_path=None):
        """
        load model
        """
        if model_path is None:
            model_path = self.model_path
        self.load_state_dict(torch.load(model_path))
        self.eval()
        logging.info('Load model from ' + model_path)
