import torch
import os

# deepFM settings below
RANDOM_SEED = 2019
NUM_SPLITS = 1

YELP_CATEGORICAL_COLS = ['attributes.RestaurantsDelivery', 'attributes.Open24Hours',
                         'attributes.DogsAllowed', 'attributes.CoatCheck',
                         'attributes.GoodForKids', 'attributes.BYOB',
                         'attributes.RestaurantsTableService',
                         'attributes.RestaurantsCounterService', 'attributes.Corkage',
                         'attributes.RestaurantsGoodForGroups',
                         'attributes.BusinessAcceptsBitcoin', 'attributes.HappyHour',
                         'attributes.WheelchairAccessible', 'attributes.Caters',
                         'attributes.ByAppointmentOnly',
                         'attributes.BusinessAcceptsCreditCards', 'attributes.DriveThru',
                         'attributes.HasTV', 'attributes.GoodForDancing',
                         'attributes.AcceptsInsurance',
                         'attributes.RestaurantsReservations',
                         'attributes.RestaurantsTakeOut', 'attributes.BikeParking',
                         'attributes.OutdoorSeating', 'is_open']

YELP_NUMERIC_COLS = ['yelping_since', 'useful', 'compliment_photos', 'compliment_list',
                     'compliment_funny', 'compliment_plain', 'review_count_x', 'elite',
                     'fans', 'compliment_note', 'funny', 'compliment_writer',
                     'compliment_cute', 'average_stars', 'compliment_more', 'friends',
                     'compliment_hot', 'cool', 'compliment_profile', 'compliment_cool',
                     'attributes.Smoking', 'attributes.AgesAllowed',
                     'attributes.Alcohol', 'attributes.RestaurantsPriceRange2',
                     'attributes.BYOBCorkage', 'attributes.NoiseLevel',
                     'attributes.WiFi', 'attributes.RestaurantsAttire', 'stars',
                     'review_count_y', 'categories']

AMAZON_NUMERIC_COLS = ['price', 'categories', 'also_bought_num', 'also_viewed', 'bought_together',
                       'rank', 'review_num', 'help_num', 'nohelp_num']

AMAZON_CATEGORICAL_COLS = ['brand']

IGNORE_COLS = ['user', 'item', 'rating']


# NCF settings below
def save_checkpoint(model, model_dir):
    torch.save(model.state_dict(), model_dir)

def resume_checkpoint(model, model_dir, device_id):
    state_dict = torch.load(model_dir, map_location=lambda storage, loc: storage.cuda(device=device_id)) # ensure all storage are on gpu
    model.load_state_dict(state_dict)

def use_cuda(enabled, device_id=0):
    if enabled:
        assert torch.cuda.is_available(), 'CUDA is not available'
        torch.cuda.set_device(device_id)

def user_optimizer(network, params):
    if params['optimizer'] == 'sgd':
        optimizer = torch.optim.SGD(network.parameters(),
                                    lr=params['sgd_lr'], 
                                    momentum=params['sgd_momentum'], 
                                    weight_decay=params['l2_regularization'])
    elif params['optimizer'] == 'adam':
        optimizer = torch.optim.Adam(network.parameters(), 
                                     lr=params['adam_lr'],
                                     weight_decay=params['l2_regularization'])
    elif params['optimizer'] == 'rmsprop':
        optimizer = torch.optim.RMSprop(network.parameters(), 
                                        lr=params['rmsprop_lr'], 
                                        alpha=params['rmsprop_alpha'], 
                                        momentum=params['rmsprop_momentum'])
    return optimizer
