from collections import defaultdict

import numpy as np
import scipy as sp
import pandas as pd
from ipypb import track

from polara.evaluation import evaluation_engine as ee
from polara.evaluation.pipelines import random_grid, find_optimal_config
from polara.recommender.coldstart.models import ItemColdStartEvaluationMixin
from polara.recommender.external.turi.turiwrapper import (TuriFactorizationRecommender,
                                                          ColdStartRecommendationsMixin)

from data_preprocessing import (get_yahoo_music_data,
                                get_similarity_data,
                                prepare_data_model,
                                prepare_cold_start_data_model)
from utils import (report_results, save_results,
                   apply_config, print_data_stats,
                   save_training_time, save_cv_training_time)

%matplotlib inline

seed = 42

experiment_name = 'fm'

data_labels = ['YaMus']

# according to https://apple.github.io/turicreate/docs/api/generated/turicreate.recommender.ranking_factorization_recommender.RankingFactorizationRecommender.html
init_config = dict(with_data_feedback = False, # implicit case
                   ranking_optimization = True,
                   solver = 'adagrad',
                   sgd_step_size = 0, # let Turi autotune it
                   seed = seed,
                   max_iterations = 25,
                   other_tc_params = {}
                   )
fm_init_config = dict.fromkeys(data_labels, {'FM': init_config, # standard scenario
                                             'FM(cs)': init_config}) # cold start

params = {
          'regularization': [1e-10, 3e-10, 1e-9, 3e-9, 1e-8, 3e-8, 1e-7, 3e-7, 1e-6, 3e-6],
          'linear_regularization': [1e-10, 3e-10, 1e-9, 3e-9, 1e-8, 3e-8, 1e-7, 3e-7, 1e-6, 3e-6],
          'rank': [100] # for initial tuning (exploration)
         }

if init_config['solver'] == 'adagrad':
    params.update({
                   'adagrad_momentum_weighting': [0.9, 0.95, 0.99]
                  })

ranks_grid = [1, 10, 100, 500, 1000, 2000, 3000]

fm_ranks = {'YaMus': ranks_grid}

topk_values = [1, 3, 10, 20, 30]

target_metric = 'mrr'

data_dict = dict.fromkeys(data_labels)
meta_dict = dict.fromkeys(data_labels)
similarities = dict.fromkeys(data_labels)
sim_indices = dict.fromkeys(data_labels)
feature_idx = dict.fromkeys(data_labels)

all_data = [data_dict, similarities, sim_indices, meta_dict]

lbl = 'YaMus'

data_dict[lbl], meta_dict[lbl] = get_yahoo_music_data('/data/recsys/yahoo_music/yamus_train0_rating5.gz',
                                                      meta_path='/data/recsys/yahoo_music/yamus_attrs.gz',
                                                      implicit=True,
                                                      pcore=5,
                                                      filter_data={'genreid': [0]}, # filter unknown genre
                                                      filter_no_meta=True)

similarities[lbl], sim_indices[lbl], feature_idx[lbl] = get_similarity_data(meta_dict[lbl])

# conform with Turi input requirement for non-existent features
meta_dict[lbl].loc[:, 'genreid'] = meta_dict[lbl]['genreid'].apply(lambda x: {i:1 for i in x} if x else {})

def prepare_recommender_models(data_label, data_models, config):
    data_model = data_models[data_label]
    fm = TuriFactorizationRecommender(data_model, item_side_info=meta_dict[data_label])
    fm.method = 'FM'
    models = [fm]
    apply_config(models, config, data_label)
    return models

def fine_tune_fm(model, params, label, ntrials=60, record_time_as=None):
    param_grid, param_names = random_grid(params, n=ntrials)
    best_fm_config, fm_scores = find_optimal_config(model, param_grid, param_names,
                                                    target_metric,
                                                    return_scores=True,
                                                    force_build=True,
                                                    iterator=lambda x: track(x, label=label))
    model_config = {model.method: dict(zip(param_names, best_fm_config))}
    model_scores = {model.method: fm_scores}
    try:
        if record_time_as:
            save_training_time(f'{experiment_name}_{record_time_as}', model, fm_scores.index, label)
    finally:
        return model_config, model_scores

config = {}
scores = {}
data_models = {}

save_results(f'{experiment_name}_param', config=config, tuning=scores)

rank_config = {}
rank_scores = {}

save_results(f'{experiment_name}_rank', config=rank_config, tuning=rank_scores)

save_results(experiment_name, cv=result)

class TuriRecommenderColdStart(ItemColdStartEvaluationMixin,
                               ColdStartRecommendationsMixin,
                               TuriFactorizationRecommender): pass

def prepare_cold_start_recommender_models(data_label, data_models, config):
    data_model = data_models[data_label]
    fm = TuriRecommenderColdStart(data_model, item_side_info=meta_dict[data_label])
    fm.method = 'FM(cs)'
    models = [fm]
    apply_config(models, config, data_label)
    return models

config_cold = {}
scores_cold = {}
data_models_cold = {}

for label in track(data_labels):
    data_models_cold[label] = prepare_cold_start_data_model(label, *all_data, seed)
    model, =  prepare_cold_start_recommender_models(label, data_models_cold, fm_init_config)
    config_cold[label], scores_cold[label] = fine_tune_fm(model, params, label, ntrials=30)
del model

report_results('tuning', scores_cold);

config_cold

save_results(f'{experiment_name}_coldstart_param', config=config_cold, tuning=scores_cold)

rank_config_cold = {}
rank_scores_cold = {}

for label in track(data_labels):
    model, = prepare_cold_start_recommender_models(label, data_models_cold,
                                                   [fm_init_config, config_cold]) # initiate with optimal config
    rank_config_cold[label], rank_scores_cold[label] = fine_tune_fm(model, {'rank': fm_ranks[label]},
                                                                    label, ntrials=0)
del model

report_results('rank', {lbl: v.sort_index() for lbl, scr in rank_scores_cold.items() for k, v in scr.items()});

rank_config_cold

save_results(f'{experiment_name}_coldstart_rank', config=rank_config_cold, tuning=rank_scores_cold)

result_cold = {}
for label in track(data_labels):
    models_cold = prepare_cold_start_recommender_models(label, data_models_cold,
                                                        [fm_init_config, config_cold, rank_config_cold])
    result_cold[label] = ee.run_cv_experiment(models_cold,
                                              fold_experiment=ee.topk_test,
                                              topk_list=topk_values,
                                              ignore_feedback=True,
                                              iterator=lambda x: track(x, label=label))

report_results('topn', result_cold, target_metric);

report_results('topn', result_cold, 'coverage');

save_results(f'{experiment_name}_coldstart', cv=result_cold)