# -*- coding: utf-8 -*-
#https://www.kaggle.com/dasolmar/xgb-with-whq-jaccard/code/code

#==============================================================================
# Ensemble learning creation:
# base_algorithms = [logistic_regression, decision_tree_classification, ...] #for classification
#
# stacking_train_dataset = matrix(row_length=len(target), column_length=len(algorithms))
# stacking_test_dataset = matrix(row_length=len(test), column_length=len(algorithms))
#
# for i,base_algorithm in enumerate(base_algorithms):
#     for trainix, testix in split(train, k=10): #you may use sklearn.cross_validation.KFold of sklearn library
#         stacking_train_dataset[testix,i] = base_algorithm.fit(train[trainix], target[trainix]).predict(train[testix])
#     stacking_test_dataset[,i] = base_algorithm.fit(train).predict(test)
#
#
# final_predictions = combiner_algorithm.fit(stacking_train_dataset, target).predict(stacking_test_dataset)
#==============================================================================

import numpy as np
import pandas as pd
import xgboost as xgb
import operator
from sklearn.model_selection import train_test_split
from collections import Counter
from nltk.corpus import stopwords
import matplotlib.pyplot as plt
from pylab import plot, show, subplot, specgram, imshow, savefig

print("Started")
RS = 12357
np.random.seed(RS)

input_folder = 'F:/DS-main/Kaggle-main/Quora Question Pairs - inputs/data/'
df_train = pd.read_csv(input_folder + 'train.csv')
df_test  = pd.read_csv(input_folder + 'test.csv')

#==============================================================================
# #rnn 1 lstm
# train_rnn1Lstm = pd.read_csv('F:/DS-main/Kaggle-main/Quora Question Pairs - inputs/data/nn 1 lstm best results/0.2647_lstm_300_200_0.30_0.30Train.csv', header=0)
# test_rnn1Lstm = pd.read_csv('F:/DS-main/Kaggle-main/Quora Question Pairs - inputs/data/nn 1 lstm best results/0.2647_lstm_300_200_0.30_0.30.csv', header=0)
# df_rnn1Lstm = pd.concat([train_rnn1Lstm, test_rnn1Lstm])
#
# #rnn 1 GRU
# train_rnn1GRU = pd.read_csv('F:/DS-main/Kaggle-main/Quora Question Pairs - inputs/data/nn 1 gru/0.3143_lstm_300_200_0.50_0.50_gru_Train.csv', header=0)
# test_rnn1GRU = pd.read_csv('F:/DS-main/Kaggle-main/Quora Question Pairs - inputs/data/nn 1 gru/0.3143_lstm_300_200_0.50_0.50_gru.csv', header=0)
# df_rnn1GRU = pd.concat([train_rnn1GRU, test_rnn1GRU])
#==============================================================================

#adding final x features
x = pd.read_csv('F:/DS-main/Kaggle-main/Quora Question Pairs - inputs/data/x_final_features.csv', header=0)

#==============================================================================
# ##rnn 1 lstm
# print("rnn features")
# x['rnn1Lstm'] = df_rnn1Lstm['is_duplicate']
# ##rnn 1 gru
# print("rnn gru features")
# x['rnn1GRU'] = df_rnn1GRU['is_duplicate']
#==============================================================================

print(x.columns)
print(x.describe())

feature_names = list(x.columns.values)
print("Features: {}".format(feature_names))

x_train = x[:df_train.shape[0]]
x_test  = x[df_train.shape[0]:]
y_train = df_train['is_duplicate'].values
x_train_for_ensemble =  x_train
y_train_for_ensemble =  y_train
del x, df_train
#params['scale_pos_weight'] = 0.360

if 1: # Now we oversample the negative class - on your own risk of overfitting!
	pos_train = x_train[y_train == 1]
	neg_train = x_train[y_train == 0]

	print("Oversampling started for proportion: {}".format(len(pos_train) / (len(pos_train) + len(neg_train))))
	p = 0.165
	scale = ((len(pos_train) / (len(pos_train) + len(neg_train))) / p) - 1
	while scale > 1:
		neg_train = pd.concat([neg_train, neg_train])
		scale -=1
	neg_train = pd.concat([neg_train, neg_train[:int(scale * len(neg_train))]])
	print("Oversampling done, new proportion: {}".format(len(pos_train) / (len(pos_train) + len(neg_train))))

	x_train = pd.concat([pos_train, neg_train])
	y_train = (np.zeros(len(pos_train)) + 1).tolist() + np.zeros(len(neg_train)).tolist()
	del pos_train, neg_train

####################################################################
# XGBoost
####################################################################
ratio = float(np.sum(y_train == 0)) / np.sum(y_train==1)

print("Training data: X_train: {}, Y_train: {}, X_test: {}".format(x_train.shape, len(y_train), x_test.shape))
ROUNDS = 1000
#training
def create_feature_map(features):
	outfile = open('xgb.fmap', 'w')
	i = 0
	for feat in features:
		outfile.write('{0}\t{1}\tq\n'.format(i, feat))
		i = i + 1
	outfile.close()
create_feature_map(feature_names)

print("Will train XGB for {} rounds, RandomSeed: {}".format(ROUNDS, RS))
params = {}
#?params['scale_pos_weight'] = ratio
params['objective'] = 'binary:logistic'
params['eval_metric'] = 'logloss'
params['eta'] = 0.11
params['max_depth'] = 5
params['silent'] = 1
params['seed'] = RS
X_training, X_val, y_training, y_val = train_test_split(x_train, y_train, test_size=0.2, random_state=RS)
xg_train = xgb.DMatrix(X_training, label=y_training)
xg_val = xgb.DMatrix(X_val, label=y_val)
watchlist  = [(xg_train,'train'), (xg_val,'eval')]

clr = xgb.train(params, xg_train, ROUNDS, watchlist, early_stopping_rounds=50)

preds = clr.predict(xgb.DMatrix(x_test))

print("Writing output...")
sub = pd.DataFrame()
sub['test_id'] = df_test['test_id']
sub['is_duplicate'] = preds *.75
sub.to_csv("xgb1_seed{}_n{}.csv".format(RS, ROUNDS), index=False)

print("Features importances...")
importance = clr.get_fscore(fmap='xgb.fmap')
importance = sorted(importance.items(), key=operator.itemgetter(1))
ft = pd.DataFrame(importance, columns=['feature', 'fscore'])
ft.plot(kind='barh', x='feature', y='fscore', legend=False, figsize=(10, 25))
plt.gcf().savefig('features_importance.png')

print("Done.")

#--- pred training for ensemble

print("pred training for ensemble")
preds = clr.predict(xgb.DMatrix(x_train_for_ensemble))

print("Writing output...")
sub = pd.DataFrame()
sub['is_duplicate'] = preds *.75
sub.to_csv("xgb_seed{}_n{}training.csv".format(RS, ROUNDS), index=False)

####################################################################
# ANN
####################################################################
# Importing the Keras libraries and packages
#==============================================================================
# import keras
# from keras.models import Sequential
# from keras.layers import Dense, PReLU
# from keras.layers import Dropout
# from keras.layers.normalization import BatchNormalization
# from keras.callbacks import EarlyStopping, ModelCheckpoint
#
#
# # Initialising the ANN
# classifier = Sequential()
# # Adding the input layer and the first hidden layer
# classifier.add(Dense(units = 200, kernel_initializer = 'truncated_normal', activation='relu', input_dim = 71))
# classifier.add(Dropout(rate = 0.1))
#
# classifier.add(Dense(200, kernel_initializer = 'truncated_normal', activation='relu'))
# classifier.add(Dropout(rate = 0.1))
#
# classifier.add(Dense(200, kernel_initializer = 'truncated_normal', activation='relu'))
# classifier.add(Dropout(rate = 0.1))
#
# classifier.add(Dense(units = 1, kernel_initializer = 'truncated_normal', activation = 'sigmoid'))
#
# early_stopping =EarlyStopping(monitor='val_loss', patience = 10)
# model_checkpoint = ModelCheckpoint('weightsANN_X_final_Features.h5', save_best_only=True, save_weights_only=True, verbose=1)
#
# classifier.compile(optimizer = 'nadam', loss = 'binary_crossentropy', metrics = ['accuracy'])
# classifier.fit(x_train_for_ensemble.values, np.array(y_train_for_ensemble), batch_size=2048, epochs=300, verbose=1, validation_split=0.1, shuffle=True, callbacks=[early_stopping, model_checkpoint])
#
# classifier.load_weights('weightsANN_X_final_Features.h5')
#
# y_pred = classifier.predict(x_test.values, batch_size=2048, verbose=1)
# submissionTest = pd.DataFrame({'is_duplicate':y_pred.ravel()})
# submissionTest.to_csv('denseANN_Test.csv', index=False)
#
# y_pred_training = classifier.predict(x_train_for_ensemble.values, batch_size=2048, verbose=1)
# submissionTraining = pd.DataFrame({'is_duplicate':y_pred_training.ravel()})
# submissionTraining.to_csv('denseANN_Training.csv', index=False)
#==============================================================================

####################################################################
# K-nn
####################################################################


try:
    from onitama.connect.card import LI
except:
    from card import LI
from matplotlib import pyplot as plt
import matplotlib.image as img
import numpy
import os

try:
    os.mkdir(r'C:\Users\吕\Desktop\supergame\static\onitama\card')
except:
    pass

plt.rcParams['font.sans-serif'] = ['KaiTi']

for card in LI:
    print(card.name)
    fig = plt.figure(figsize=(5, 5))
    # plt.title(card.name)

    plt.gca().xaxis.set_major_locator(plt.NullLocator())
    plt.gca().yaxis.set_major_locator(plt.NullLocator())
    plt.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=1, wspace=1)
    # # plt.margins(0, 0)

    plt.axis([0.3, 5.7, 0.3, 5.7])
    plt.axis('off')

    bgimg = img.imread(r'C:\Users\吕\Desktop\supergame\static\onitama\other\yangpi.jpg')
    plt.figimage(bgimg)

    tick = numpy.asarray((0, 1, 2, 3, 4, 5)) + 0.5
    plt.xticks(tick, [])
    plt.yticks(tick, [])
    for _ in tick:
        plt.plot(tick, numpy.ones_like(tick)*_, c='black')
        plt.plot(numpy.ones_like(tick)*_, tick, c='black')

    old = numpy.asarray((3, 3))
    new = numpy.asarray(card.method) + old
    plt.scatter(new[:, 0], new[:, 1], c='black', s=2400, marker='*')
    plt.scatter(old[0], old[1], c=card.seq, s=1800, marker='o')

    plt.text(3, 3, s=card.name, c=card.seq, size=150, alpha=0.3,
             va="center", ha="center")

    # plt.savefig(r'C:\Users\吕\Desktop\supergame\static\onitama\card\%s.jpg' % card.name)
    plt.show()
    break


import xml.etree.ElementTree as ET
from typing import List

import pandas as pd
import matplotlib.pyplot as plt

from mplsoccer import Radar, FontManager, grid

from constants import UNDEFINED_INT
from player import Player
from team import Team

class MatchResults:

    """
    Class constructor that initializes instance variables and calls a function to parse Opta F9 xml file data

    Parameters:
    file_path (str): The path to the Opta F9 xml file that contains the data
    """
    def __init__(self, file_path):
        self.__teams = []
        self.__players = []
        self.__players_stats = []
        self.__all_stat_types_set = set()
        self.__parse_f9opta_file(file_path)

    def get_df_players_stats(self):
        df = pd.DataFrame.from_dict(self.__players_stats)
        return df

    def get_df_all_stats(self):
        return pd.DataFrame(self.__all_stat_types_set, columns = ['stat'])

    def get_df_players(self):
        return pd.DataFrame([{'player_id': p.player_id, 'team_id': p.team_id,
                            'match_position': p.match_position, 'first': p.first,
                            'last': p.last, 'known': p.known} for p in self.__players])

    def get_df_teams(self):
        """
        Returns a pandas DataFrame with information about the teams that played the match.

        Each row of the DataFrame corresponds to a team, and contains the team's
        identifier and name.

        Returns:
            A pandas DataFrame with columns 'team_id' and 'team_name'.
        """
        return pd.DataFrame([{'team_id': t.team_id, 'team_name': t.team_name} for t in self.__teams])

    def get_df_specific_stat(self, stat, player_id = UNDEFINED_INT, team_id = UNDEFINED_INT):
        min_columns = ['player_id', 'team_id', 'match_position', 'first', 'last', 'known',
                        'shirt_number', 'status', 'captain', 'mins_played', stat]
        df = self.get_df_players_stats().loc[:, min_columns]
        if player_id != UNDEFINED_INT:
            df = df.query(f'player_id == {player_id}').reset_index(drop=True)
        if team_id != UNDEFINED_INT:
            df = df.query(f'team_id == {team_id}').reset_index(drop=True)

        return df.sort_values(by=[stat], ascending=False)

    def create_player_radar(self, player_id, list_params,
                        min_minutes = 1, normalized_to_90 = False):
        # radar params
        params = self.__reduce_params(list_params)
        # "lower number is better" list
        lower_is_better_list = self.__get_keys_value_zero(list_params)

        # Player stats dataframe (players with more than {min_minutes} minutes played)
        df = self.get_df_players_stats()
        df_mins = df.query(f'mins_played >= {min_minutes}').reset_index(drop=True)

        if len(df_mins.query(f'player_id == {player_id}')) == 0:
            raise("The selected player did not play the defined minimum number of minutes")

        low_values, high_values = self.__calculate_low_high_values(df_mins, params,
                                            normalized_to_90 = normalized_to_90)
        player_values = self.__calculate_player_values(df_mins, params, player_id,
                                            normalized_to_90 = normalized_to_90)

        # creating the figure using the grid function from mplsoccer:
        fig, axs = grid(figheight=14, grid_height=0.915, title_height=0.06, endnote_height=0.025,
                        title_space=0, endnote_space=0, grid_key='radar', axis=False)

        radar = Radar(params, low_values, high_values,
              lower_is_better=lower_is_better_list,
              round_int=[False]*len(params),
              num_rings=4,
              ring_width=1, center_circle_radius=1)

        # Plot radar
        radar.setup_axis(ax=axs['radar'])
        rings_inner = radar.draw_circles(ax=axs['radar'], facecolor='#6fa8dc', edgecolor='#2c4358')
        radar_output = radar.draw_radar(player_values, ax=axs['radar'],
                                        kwargs_radar={'facecolor': '#dca36f'},
                                        kwargs_rings={'facecolor': '#66d8ba'})
        radar_poly, rings_outer, vertices = radar_output
        range_labels = radar.draw_range_labels(ax=axs['radar'], fontsize=12)
        param_labels = radar.draw_param_labels(ax=axs['radar'], fontsize=12)
        axs['radar'].scatter(vertices[:, 0], vertices[:, 1], c='#00f2c1', edgecolors='#6d6c6d',
                                marker='o', s=40, zorder=2)

        # Legends
        # Player info
        obj_player = self.__get_player(player_id)
        obj_team = self.__get_team(UNDEFINED_INT if obj_player is None else obj_player.team_id)
        player_name = "" if obj_player is None else f"{obj_player.first} {obj_player.last}"
        player_team_name = "" if obj_team is None else obj_team.team_name
        title1_text = axs['title'].text(0.04, 0.65, player_name, fontsize=16,
                                        ha='left', va='center')
        title2_text = axs['title'].text(0.90, 0.65, player_team_name, fontsize=14,
                                        ha='right', va='center')
        # Endnote
        min_minutes_text = f"Players who played at least {min_minutes} minutes."
        normalized_text = "Values normalized to 90 mins." if normalized_to_90 else "Absolute values."
        endnote_text = f"INFO: {min_minutes_text}{normalized_text}"
        endnote_text = axs['endnote'].text(0.04, 0.5, endnote_text, fontsize=10,
                                   ha='left', va='center')

        return plt

    def create_comparison_radar(self, player1_id, player2_id, list_params,
                        min_minutes = 1, normalized_to_90 = False):
        # radar params
        params = self.__reduce_params(list_params)
        # "lower number is better" list
        lower_is_better_list = self.__get_keys_value_zero(list_params)

        # Player stats dataframe (players with more than 0 minutes played)
        df = self.get_df_players_stats()
        df_mins = df.query(f'mins_played >= {min_minutes}').reset_index(drop=True)

        if len(df_mins.query(f'player_id == {player1_id}')) == 0 or \
            len(df_mins.query(f'player_id == {player2_id}')) == 0:
            raise("One of the selected players did not play the defined minimum number of minutes")
        low_values, high_values = self.__calculate_low_high_values(df_mins, params,
                                            normalized_to_90 = normalized_to_90)
        player1_values = self.__calculate_player_values(df_mins, params, player1_id,
                                            normalized_to_90 = normalized_to_90)
        player2_values = self.__calculate_player_values(df_mins, params, player2_id,
                                            normalized_to_90 = normalized_to_90)

        radar = Radar(params, low_values, high_values,
              lower_is_better=lower_is_better_list,
              round_int=[False]*len(params),
              num_rings=4,
              ring_width=1, center_circle_radius=1)

        # creating the figure using the grid function from mplsoccer:
        fig, axs = grid(figheight=14, grid_height=0.915, title_height=0.06, endnote_height=0.025,
                title_space=0, endnote_space=0, grid_key='radar', axis=False)

        # plot radar
        radar.setup_axis(ax=axs['radar'])  # format axis as a radar
        rings_inner = radar.draw_circles(ax=axs['radar'], facecolor='#ffb2b2', edgecolor='#fc5f5f')
        radar_output = radar.draw_radar_compare(player1_values, player2_values, ax=axs['radar'],
                                                kwargs_radar={'facecolor': '#00f2c1', 'alpha': 0.6},
                                                kwargs_compare={'facecolor': '#d80499', 'alpha': 0.6})
        radar_poly, radar_poly2, vertices1, vertices2 = radar_output
        range_labels = radar.draw_range_labels(ax=axs['radar'], fontsize=12)
        param_labels = radar.draw_param_labels(ax=axs['radar'], fontsize=12)
        axs['radar'].scatter(vertices1[:, 0], vertices1[:, 1],
                            c='#00f2c1', edgecolors='#6d6c6d', marker='o', s=40, zorder=2)
        axs['radar'].scatter(vertices2[:, 0], vertices2[:, 1],
                            c='#d80499', edgecolors='#6d6c6d', marker='o', s=40, zorder=2)
        # Legend
        obj_player1 = self.__get_player(player1_id)
        obj_player2 = self.__get_player(player2_id)
        obj_team1 = self.__get_team(UNDEFINED_INT if obj_player1 is None else obj_player1.team_id)
        obj_team2 = self.__get_team(UNDEFINED_INT if obj_player2 is None else obj_player2.team_id)
        player1_name = "" if obj_player1 is None else f"{obj_player1.first} {obj_player1.last}"
        player1_team_name = "" if obj_team1 is None else obj_team1.team_name
        player2_name = "" if obj_player2 is None else f"{obj_player2.first} {obj_player2.last}"
        player2_team_name = "" if obj_team2 is None else obj_team2.team_name
        title1_text = axs['title'].text(0.01, 0.65, player1_name, fontsize=16, color='#01c49d',
                                        ha='left', va='center')
        title2_text = axs['title'].text(0.01, 0.25, player1_team_name, fontsize=10,
                                        ha='left', va='center', color='#01c49d')
        title3_text = axs['title'].text(0.99, 0.65, player2_name, fontsize=16,
                                        ha='right', va='center', color='#d80499')
        title4_text = axs['title'].text(0.99, 0.25, player2_team_name, fontsize=10,
                                        ha='right', va='center', color='#d80499')
        return plt

    def __parse_f9opta_file(self, file_path):
        tree = ET.parse(file_path)
        root = tree.getroot()

        list_team_tag = root.findall('.//Team')
        for team_tag in list_team_tag:
            obj_team = self.__parse_team_tag(team_tag)
            self.__teams.append(obj_team)
            for player_tag in team_tag.iter('Player'):
                obj_player = self.__parse_player_tag(player_tag, obj_team.team_id)
                self.__players.append(obj_player)

        list_team_data_tag = root.findall('.//MatchData/TeamData')
        for team_data_tag in list_team_data_tag:
            player_lineup_tag = team_data_tag.find('PlayerLineUp')
            for match_player_tag in player_lineup_tag.iter('MatchPlayer'):
                obj_match_player = self.__parse_match_player_tag(match_player_tag)
                self.__players_stats.append(obj_match_player)
        #adjust player stats with all stats
        self.__adjust_player_stats()

    def __parse_team_tag(self, team_tag):
        id = self.__tag_value_to_int(team_tag.get('uID')[1:])
        name = self.__tag_value_to_str(team_tag.find('Name').text)
        return Team(team_id=id, team_name=name)

    def __parse_player_tag(self, player_tag, team_id):
        id = self.__tag_value_to_int(player_tag.get('uID')[1:])
        match_position = self.__tag_value_to_str(player_tag.get('Position'))
        first_tag = player_tag.find('.//First')
        last_tag = player_tag.find('.//Last')
        known_tag = player_tag.find('.//Known')
        first = self.__tag_value_to_str(None if first_tag is None else first_tag.text)
        last = self.__tag_value_to_str(None if last_tag is None else last_tag.text)
        known = self.__tag_value_to_str(None if known_tag is None else known_tag.text)
        return Player(player_id=id, team_id=team_id, match_position=match_position,
                        first=first, last=last, known=known)

    def __parse_match_player_tag(self, match_player_tag):
        id = self.__tag_value_to_int(match_player_tag.get('PlayerRef')[1:])
        match_position = self.__tag_value_to_str(match_player_tag.get('Position'))
        shirt_number = self.__tag_value_to_int(match_player_tag.get('ShirtNumber'))
        status = self.__tag_value_to_str(match_player_tag.get('Status'))
        captain = self.__tag_value_to_int(match_player_tag.get('Captain'), error_value=0)
        player = next(filter(lambda p: p.player_id == id, self.__players))
        dict_player = {'player_id':id, 'team_id': player.team_id, 'match_position': match_position,
                        'first': player.first, 'last': player.last, 'known': player.known,
                        'shirt_number': shirt_number, 'status': status, 'captain': captain}
        for stat_tag in match_player_tag.iter('Stat'):
            stat_type = self.__tag_value_to_str(stat_tag.get('Type'))
            stat_value = self.__tag_value_to_int(stat_tag.text, error_value=0)
            dict_player.update({stat_type:stat_value})
            self.__all_stat_types_set.add(stat_type)
        return dict_player

    def __adjust_player_stats(self):
        for player_stats in self.__players_stats:
            for stat_type in self.__all_stat_types_set:
                if stat_type not in player_stats:
                    player_stats.update({stat_type: 0})

    def __get_player(self, player_id):
        filtered = [pl for pl in self.__players if pl.player_id == player_id]
        if len(filtered) == 0:
            return None
        else:
            return filtered[0]

    def __get_team(self, team_id):
        filtered = [t for t in self.__teams if t.team_id == team_id]
        if len(filtered) == 0:
            return None
        else:
            return filtered[0]

    @staticmethod
    def __tag_value_to_int(value, error_value = UNDEFINED_INT):
        try:
            return int(value)
        except (ValueError, TypeError):
            return error_value

    @staticmethod
    def __tag_value_to_str(value):
        if value is None:
            return ''
        return str(value)

    @staticmethod
    def __reduce_params(list_params):
        return list(map(lambda d: list(d.keys())[0], list_params))

    @staticmethod
    def __get_keys_value_zero(list_params):
        return list(map(lambda d: list(d.keys())[0],
                    filter(lambda d: d[list(d.keys())[0]] == 0, list_params)))

    @staticmethod
    def __calculate_low_high_values(df, list_params, normalized_to_90 = False):
        df_copy = df.copy()
        low_values = []
        high_values = []
        for param in list_params:
            if normalized_to_90:
                df_copy[param] = df_copy.apply(lambda x: 90 * x[param] / x['mins_played'], axis=1)
            low_values.append(df_copy[param].min())
            high_values.append(df_copy[param].max())
        return low_values, high_values

    @staticmethod
    def __calculate_player_values(df, list_params, player_id, normalized_to_90 = False):
        df_copy = df.copy()
        player_values = []
        for param in list_params:
            if normalized_to_90:
                df_copy[param] = df_copy.apply(lambda x: 90 * x[param] / x['mins_played'], axis=1)
            player_values.append(df_copy.loc[df_copy['player_id'] == player_id, param].iloc[0])
        return player_values


from django.contrib import admin
from apps.suggestions.models import ProblemType, Problem, Location, ProblemImages

@admin.register(ProblemType)
class ProblemTypeAdmin(admin.ModelAdmin):
    list_display = ("name",)
    search_fields = ("name",)

@admin.register(ProblemImages)
class ProblemImagesAdmin(admin.ModelAdmin):
    list_display = ("id", "problem", "image")
    search_fields = ("id", "problem", "image")
    list_filter = ("problem",)
    readonly_fields = ("id", "problem", "image")

class ProblemImagesInline(admin.TabularInline):
    model = ProblemImages
    extra = 1

@admin.register(Problem)
class ProblemAdmin(admin.ModelAdmin):
    list_display = ("id", "user", "city", "district", "date", "status")
    search_fields = ("id", "user", "city", "district", "date", "status")
    list_filter = ("status",)
    readonly_fields = ("id", "city", "district", "date")
    inlines = [ProblemImagesInline]

@admin.register(Location)
class LocationAdmin(admin.ModelAdmin):
    list_display = ['id', 'lon', 'lat']


N,L = map(int,input().split())
street = [list(map(int,input().split())) for _ in range(N)]
cnt = 0

def check(line):
    for i in range(1,N):
        #1보다 크면 패스
        if abs(line[i] - line[i-1]) > 1:
            return False
        #앞이 클때
        if line[i] < line[i-1]:
            for j in range(L):
                if i+j >= N or line[i] != line[i+j] or slope[i+j]:
                    return False
                if line[i] == line[i+j]:
                    slope[i+j] = True
        #뒤가 클때
        elif line[i] > line[i-1]:
            for j in range(L):
                if i-j-1<0 or line[i-1] != line[i-j-1] or slope[i-j-1]:
                    return False
                if line[i-1] == line[i-j-1]:
                    slope[i-j-1] = True
    return True

for i in range(N):
    slope = [False] * N
    if check([street[i][j] for j in range(N)]):
        cnt += 1

for j in range(N):
    slope = [False] * N
    if check([street[i][j] for i in range(N)]):
        cnt += 1

print(cnt)


def binary_search(arr, key):
    start = 0
    end = len(arr) - 1

    while start <= end:

        mid = (start + end) // 2

        # Check if key == mid
        if arr[mid] == key:
            return mid

        # If key is greater, ignore left half
        elif arr[mid] < key:
            start = mid + 1

        # If key is smaller, ignore right half
        elif arr[mid] > key:
            end = mid - 1

    return -1

# arr = [1,4,6,8,14,45]
# key = 11
# print(binary_search(arr, key))


from math import *

def sang(a):
    a[0], a[1] = False, False
    for i in range(2, isqrt(10**6)+1):
        if a[i]:
            for j in range(i*i, 10**6+1, i):
                a[j] = False

if __name__ == '__main__':
    a = [True] * (10 ** 6 + 1)
    sang(a)
    n, x = map(int, input().split())
    m = n
    pri = []
    for i in range(0, 10**6+1):
        if a[i] and m > 0:
            m -= 1
            pri.append(i)
    # print(pri)
    print(x, end = ' ')
    for i in range(n):
        v = x + pri[i]
        print(v, end = ' ')
        x += pri[i]


from random import randint

t = 100
print(t)
for _ in range(t):
	n = randint(1, 12)
	print(n)
	arr = [1]
	for i in range(n-1):
		r = (arr[-1] + randint(-1, 1))%10
		arr += [r]
	print(*arr, sep='')


from trainer import GPTTrainer, GPTDataset, activation_functions_dict
from preprocessing.tokenizer import Tokenizer
import warnings
import os

warnings.filterwarnings('ignore')

def train_model(
        manifest_path: str,
        tokenizer_path: str,
        n: int,
        d_model: int,
        heads: int,
        activation: str,
        dropout_rate: float,
        eps: float,
        device: str,
        learning_rate: float,
        epochs: int,
        batch_size: int,
        checkpoint: str,
        tracking_config_path: str,
        val_path: str,
        validation_config_path: str
):
    assert os.path.exists(tokenizer_path) and os.path.exists(manifest_path)

    tokenizer = Tokenizer(tokenizer_path)

    train_dataset = GPTDataset(manifest_path, tokenizer=tokenizer)

    val_dataset = None
    if val_path is not None:
        if os.path.exists(val_path) == True:
            val_dataset = GPTDataset(val_path, tokenizer=tokenizer)

    trainer = GPTTrainer(
        token_size=len(tokenizer.dictionary),
        n=n,
        d_model=d_model,
        heads=heads,
        activation=activation_functions_dict[activation],
        dropout_rate=dropout_rate,
        eps=eps,
        device=device,
        init_learning_rate=learning_rate,
        checkpoint=checkpoint
    )

    trainer.fit(
        train_dataset=train_dataset,
        val_dataset=val_dataset,
        epochs=epochs,
        batch_size=batch_size,
        tracking_config_path=tracking_config_path,
        validation_config_path=validation_config_path
    )

if __name__ == '__main__':
    from argparse import ArgumentParser

    parser = ArgumentParser()

    # Pre-config
    parser.add_argument("--manifest_path", type=str)
    parser.add_argument("--tokenizer_path", type=str)
    parser.add_argument("--val_path", type=str, default=None)

    # Model Hyper-parameters
    parser.add_argument("--n", type=int, default=12)
    parser.add_argument("--d_model", type=int, default=768)
    parser.add_argument("--heads", type=int, default=12)
    parser.add_argument("--dropout_rate", type=float, default=0.1)
    parser.add_argument("--eps", type=float, default=0.02)
    parser.add_argument("--activation", type=str, default='gelu')

    # Training
    parser.add_argument("--device", type=str, default="cpu")
    parser.add_argument("--learning_rate", type=float, default=3e-5)
    parser.add_argument("--epochs", type=int, default=1)
    parser.add_argument("--batch_size", type=int, default=1)
    parser.add_argument("--checkpoint", type=str, default="./gpt.pt")

    # Config
    parser.add_argument("--tracking_config_path", type=str, default=None)
    parser.add_argument("--validation_config_path", type=str, default=None)

    # Main Handle
    args = parser.parse_args()

    train_model(
        manifest_path=args.manifest_path,
        tokenizer_path=args.tokenizer_path,
        n=args.n,
        d_model=args.d_model,
        heads=args.heads,
        activation=args.activation,
        dropout_rate=args.dropout_rate,
        eps=args.eps,
        device=args.device,
        learning_rate=args.learning_rate,
        epochs=args.epochs,
        batch_size=args.batch_size,
        checkpoint=args.checkpoint,
        tracking_config_path=args.tracking_config_path,
        validation_config_path=args.validation_config_path,
        val_path=args.val_path
    )


from detection import detect_people
from scipy.spatial import distance as dist
import numpy as np
import streamlit as st
import cv2

def video_detect(video_name, MIN_CONF):

    cuda = 'False'

    labelsPath = "yolo-coco/coco.names"
    LABELS = open(labelsPath).read().strip().split("\n")

    USE_GPU = bool(cuda)
    MIN_DISTANCE = 50
    file_url = 'https://pjreddie.com/media/files/yolov3.weights'
    file_name = wget.download(file_url)
    #weightsPath = "yolov3.weights"
    weightsPath = file_name
    configPath = "yolo-coco/yolov3.cfg"

    net = cv2.dnn.readNetFromDarknet(configPath, weightsPath)

    if USE_GPU:
        net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
        net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)

    ln = net.getLayerNames()
    ln = [ln[i[0] - 1] for i in net.getUnconnectedOutLayers()]

    vs = cv2.VideoCapture(video_name)
    writer = None
    image_placeholder = st.empty()
    while True:

        (grabbed, frame) = vs.read()

        if not grabbed:
            break

        frame = cv2.resize(frame, (700,700))
        results = detect_people(frame, net, ln,
                                personIdx=LABELS.index("person"))

        violate = set()

        if len(results) >= 2:

            centroids = np.array([r[2] for r in results])
            D = dist.cdist(centroids, centroids, metric="euclidean")

            for i in range(0, D.shape[0]):
                for j in range(i + 1, D.shape[1]):

                    if D[i, j] < MIN_DISTANCE:

                        violate.add(i)
                        violate.add(j)

        for (i, (prob, bbox, centroid)) in enumerate(results):

            (startX, startY, endX, endY) = bbox
            (cX, cY) = centroid
            color = (0, 255, 0)

            if i in violate:
                color = (0, 0, 255)

            cv2.rectangle(frame, (startX, startY), (endX, endY), color, 2)
            #cv2.circle(frame, (cX, cY), 5, color, 1)

        font = cv2.FONT_HERSHEY_SIMPLEX
        #datet = str(datetime.datetime.now())
        #frame = cv2.putText(frame, datet, (0, 35), font, 1,
         #                   (0, 255, 255), 2, cv2.LINE_AA)
        text = "Social Distancing Violations: {}".format(len(violate))
        cv2.putText(frame, text, (10, frame.shape[0] - 25),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.85, (0, 0, 255), 3)

        display = 1
        if display > 0:

            image_placeholder.image(
                frame, caption='Live Social Distancing Monitor Running..!', channels="BGR")

        if writer is not None:
            writer.write(frame)


from typing import Tuple
import numpy as np
import paddle
from paddle import Tensor
from paddle import nn
from paddle.nn import functional as F

def frame(x: Tensor,
          num_samples: Tensor,
          win_length: int,
          hop_length: int,
          clip: bool = True) -> Tuple[Tensor, Tensor]:
    """Extract frames from audio.

    Parameters
    ----------
    x : Tensor
        Shape (N, T), batched waveform.
    num_samples : Tensor
        Shape (N, ), number of samples of each waveform.
    win_length : int
        Window length.
    hop_length : int
        Number of samples shifted between ajancent frames.
    clip : bool, optional
        Whether to clip audio that does not fit into the last frame, by
        default True

    Returns
    -------
    frames : Tensor
        Shape (N, T', win_length).
    num_frames : Tensor
        Shape (N, ) number of valid frames
    """
    assert hop_length <= win_length
    num_frames = (num_samples - win_length) // hop_length
    padding = (0, 0)
    if not clip:
        num_frames += 1
        # NOTE: pad hop_length - 1 to the right to ensure that there is at most
        # one frame dangling to the righe edge
        padding = (0, hop_length - 1)

    weight = paddle.eye(win_length).unsqueeze(1)

    frames = F.conv1d(x.unsqueeze(1),
                      weight,
                      padding=padding,
                      stride=(hop_length, ))
    return frames, num_frames

class STFT(nn.Layer):
    """A module for computing stft transformation in a differentiable way.

    Parameters
    ------------
    n_fft : int
        Number of samples in a frame.

    hop_length : int
        Number of samples shifted between adjacent frames.

    win_length : int
        Length of the window.

    clip: bool
        Whether to clip audio is necesaary.
    """
    def __init__(self,
                 n_fft: int,
                 hop_length: int,
                 win_length: int,
                 window_type: str = None,
                 clip: bool = True):
        super().__init__()

        self.hop_length = hop_length
        self.n_bin = 1 + n_fft // 2
        self.n_fft = n_fft
        self.clip = clip

        # calculate window
        if window_type is None:
            window = np.ones(win_length)
        elif window_type == "hann":
            window = np.hanning(win_length)
        elif window_type == "hamming":
            window = np.hamming(win_length)
        else:
            raise ValueError("Not supported yet!")

        if win_length < n_fft:
            window = F.pad(window, (0, n_fft - win_length))
        elif win_length > n_fft:
            window = window[:n_fft]

        # (n_bins, n_fft) complex
        kernel_size = min(n_fft, win_length)
        weight = np.fft.fft(np.eye(n_fft))[:self.n_bin, :kernel_size]
        w_real = weight.real
        w_imag = weight.imag

        # (2 * n_bins, kernel_size)
        w = np.concatenate([w_real, w_imag], axis=0)
        w = w * window

        # (2 * n_bins, 1, kernel_size) # (C_out, C_in, kernel_size)
        w = np.expand_dims(w, 1)
        weight = paddle.cast(paddle.to_tensor(w), paddle.get_default_dtype())
        self.register_buffer("weight", weight)

    def forward(self, x: Tensor, num_samples: Tensor) -> Tuple[Tensor, Tensor]:
        """Compute the stft transform.
        Parameters
        ------------
        x : Tensor [shape=(B, T)]
            The input waveform.
        num_samples : Tensor
            Number of samples of each waveform.
        Returns
        ------------
        D : Tensor
            Shape(N, T', n_bins, 2) Spectrogram.

        num_frames: Tensor
            Shape (N,) number of samples of each spectrogram
        """
        num_frames = (num_samples - self.win_length) // self.hop_length
        padding = (0, 0)
        if not self.clip:
            num_frames += 1
            padding = (0, self.hop_length - 1)

        batch_size, _, _ = paddle.shape(x)
        x = x.unsqueeze(-1)
        D = F.conv1d(self.weight,
                     x,
                     stride=(self.hop_length, ),
                     padding=padding,
                     data_format="NLC")
        D = paddle.reshape(D, [batch_size, -1, self.n_bin, 2])
        return D, num_frames


import sys
import logging
from diro2c.data_generation.helper import *
import pandas as pd
import numpy as np
from diro2c.data_generation.neighborhood_generation import modified_gpdatagenerator
from diro2c.data_generation.distance_functions import simple_match_distance, normalized_euclidean_distance, mixed_distance
import _pickle as cPickle
from diro2c.enums.diff_classifier_method_type import diff_classifier_method_type
from diro2c.data_generation.neighborhood_generation.gpdatagenerator import calculate_feature_values
from sklearn.tree import DecisionTreeClassifier
from sklearn import metrics
from data.getdata import loaddata, prepare_df
from data.split3fold import split3fold
from scipy.cluster.hierarchy import linkage, fcluster
import pickle
from collections import defaultdict
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from datetime import date
from FromLocalToGlobal import getglobal, getexplainer, clusterbasedinstances, stratifiedrandominstances, randominstances

if __name__ == '__main__':

    data = sys.argv[1]
    methodin = sys.argv[2]

    if methodin == 'randominstances':
        method = randominstances
        methodname = '"Approach 1: Random sampling"'
        npointsin = [8,16,32,64]
        ts = [1]
    elif methodin == 'classstratifiedinstances':
        method = stratifiedrandominstances
        methodname = '"Approach 2: Class-stratified sampling"'
        npointsin = [8,16,32,64]
        ts = [1]
    elif methodin == 'clusterbasedinstances':
        method = clusterbasedinstances
        methodname = '"Approach 3: Cluster-stratified sampling"'

    if data == 'compas':
        ts = [{'2|0': 2.5, '0|1': 7.5,'2|1': 5, '1|2': 2,'0|2': 4, '1|0': 5},
              {'2|0': 4, '0|1': 15,'2|1': 10, '1|2': 7,'0|2': 12, '1|0': 15},
              {'2|0': 4, '0|1': 15,'2|1': 10, '1|2': 4.5,'0|2': 6, '1|0': 10}]
        npointsin = [1]
    elif data == 'bankmarketing':
        ts = [{'0|1': 25, '1|0': 10},
              {'0|1': 30, '1|0': 15},
              {'0|1': 20, '1|0': 7.5}]
        npointsin = [1]

    logging.basicConfig(filename='FromLocalToGlobalfit'+str(date.today())+'_'+data+'_'+methodin+'.log',
                        level=logging.DEBUG, format = '%(asctime)s %(message)s')
    performancefile = 'FromLocalToGlobalperformance'+str(date.today())+'_'+data+'_'+methodin+'.txt'

    with open(performancefile, 'a') as myfile:
        myfile.write('model data iteration npoints maxdepthexplainer depthexplainer leavesexplainer accuracy precisionmacro recallmacro\n')

    dataA, dataB, cols, discrete, continuous, le = loaddata(data)
    blackboxtrainA, trainA, testA = split3fold(dataA, 0.4, 0.2, random_state=1)
    blackboxtrainB, trainB, testB = split3fold(dataB, 0.4, 0.2, random_state=1)
    modelA = pickle.load(open('blackboxes/'+data+'A.sav', 'rb'))
    modelB = pickle.load(open('blackboxes/'+data+'B.sav', 'rb'))

    train = pd.concat([trainA, trainB])
    train['predA'] = modelA.predict(train[cols].values)
    train['predB'] = modelB.predict(train[cols].values)
    train['difference'] = train.apply(lambda row: str(int(row['predA'])) + '|' + str(int(row['predB'])), axis = 1)
    train.drop(columns=['predA', 'predB', 'y'], inplace=True, errors='ignore')
    train = train.reset_index(drop=True)

    test = pd.concat([testA, testB])
    test['predA'] = modelA.predict(test[cols].values)
    test['predB'] = modelB.predict(test[cols].values)
    test['difference'] = test.apply(lambda row: str(int(row['predA'])) + '|' + str(int(row['predB'])), axis = 1)
    test.drop(columns=['predA', 'predB', 'y'], inplace=True)
    test = test.reset_index(drop=True)

    discrete.append('difference')
    populationsize = 1000

    if methodin == 'advancedstratifiedrandominstances':
        for i in range(5):
            logging.info('start processing round: ' + str(i))
            for t in ts:
                globaldata = getglobal(train, classname='difference', npoints=1, populationsize=populationsize,
                                       method=method, random_state=i, discrete = discrete, continuous=continuous, t = t)
                nchoseninstances = len(np.unique(globaldata.index.get_level_values(0)))
                filename = 'neighborhoods/' + methodin + '_round' + str(i) + '_' + str(nchoseninstances) + '_' + data + '.npy'
                with open(filename, 'wb') as f:
                    np.save(f, globaldata)
                for j in [3,5,7,None]:
                    globalexplainer = getexplainer(globaldata.values, modelA, modelB, max_depth = j)
                    pred = globalexplainer.predict(test[cols].values)
                    with open(performancefile, 'a') as f:
                        line = ' '.join([methodname,
                                         data,
                                         str(i),
                                         str(nchoseninstances),
                                         str(j),
                                         str(globalexplainer.get_depth()),
                                         str(globalexplainer.get_n_leaves()),
                                         str(metrics.accuracy_score(test.difference, pred)),
                                         str(metrics.precision_score(test.difference, pred, average='macro')),
                                         str(metrics.recall_score(test.difference, pred, average='macro'))
                                         ])
                        f.write(line + '\n')
                logging.info('    finished nchoseninstances: ' + str(nchoseninstances))
    else:
        for i in range(5):
            logging.info('start processing round: ' + str(i))
            for npoints in npointsin:
                globaldata = getglobal(train, classname='difference', npoints=npoints, populationsize=populationsize,
                                       method=method, random_state=i, discrete = discrete, continuous=continuous)
                filename = 'neighborhoods/' + methodin + '_round' + str(i) + '_' + str(npoints) + '_' + data +'.npy'
                with open(filename, 'wb') as f:
                    np.save(f, globaldata)
                for j in [3,5,7,None]:
                    globalexplainer = getexplainer(globaldata.values, modelA, modelB, max_depth = j)
                    pred = globalexplainer.predict(test[cols].values)
                    with open(performancefile, 'a') as f:
                        line = ' '.join([methodname,
                                         data,
                                         str(i),
                                         str(npoints),
                                         str(j),
                                         str(globalexplainer.get_depth()),
                                         str(globalexplainer.get_n_leaves()),
                                         str(metrics.accuracy_score(test.difference, pred)),
                                         str(metrics.precision_score(test.difference, pred, average='macro')),
                                         str(metrics.recall_score(test.difference, pred, average='macro'))
                                         ])
                        f.write(line + '\n')
                logging.info('    finished npoints: ' + str(npoints))


import json
import warnings
from inspect import signature
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
from tqdm import trange
from torch import nn
from torchmetrics import Metric as TorchMetric
from torch.utils.data import DataLoader

from pytorch_widedeep.losses import ZILNLoss
from pytorch_widedeep.metrics import Metric
from pytorch_widedeep.wdtypes import (
    Dict,
    List,
    Tuple,
    Union,
    Tensor,
    Literal,
    Optional,
    WideDeep,
    Optimizer,
    Transforms,
    LRScheduler,
)
from pytorch_widedeep.callbacks import Callback
from pytorch_widedeep.dataloaders import DataLoaderDefault
from pytorch_widedeep.initializers import Initializer
from pytorch_widedeep.training._finetune import FineTune
from pytorch_widedeep.utils.general_utils import Alias
from pytorch_widedeep.training._wd_dataset import WideDeepDataset
from pytorch_widedeep.training._base_trainer import BaseTrainer
from pytorch_widedeep.training._trainer_utils import (
    save_epoch_logs,
    wd_train_val_split,
    print_loss_and_metric,
)
from pytorch_widedeep.training._feature_importance import (
    Explainer,
    FeatureImportance,
)

class Trainer(BaseTrainer):
    r"""Class to set the of attributes that will be used during the
    training process.

    Parameters
    ----------
    model: `WideDeep`
        An object of class `WideDeep`
    objective: str
        Defines the objective, loss or cost function.

        Param aliases: `loss_function`, `loss_fn`, `loss`,
        `cost_function`, `cost_fn`, `cost`.

        Possible values are:

        - `binary`, aliases: `logistic`, `binary_logloss`, `binary_cross_entropy`

        - `binary_focal_loss`

        - `multiclass`, aliases: `multi_logloss`, `cross_entropy`, `categorical_cross_entropy`,

        - `multiclass_focal_loss`

        - `regression`, aliases: `mse`, `l2`, `mean_squared_error`

        - `mean_absolute_error`, aliases: `mae`, `l1`

        - `mean_squared_log_error`, aliases: `msle`

        - `root_mean_squared_error`, aliases:  `rmse`

        - `root_mean_squared_log_error`, aliases: `rmsle`

        - `zero_inflated_lognormal`, aliases: `ziln`

        - `quantile`

        - `tweedie`
    custom_loss_function: `nn.Module`. Optional, default = None
        It is possible to pass a custom loss function. See for example
        `pytorch_widedeep.losses.FocalLoss` for the required structure of the
        object or the Examples section in this documentation or in the repo.
        Note that if `custom_loss_function` is not `None`, `objective` must
        be _'binary'_, _'multiclass'_ or _'regression'_, consistent with the
        loss function
    optimizers: `Optimzer` or dict. Optional, default= None
        - An instance of Pytorch's `Optimizer` object
          (e.g. `torch.optim.Adam()`) or
        - a dictionary where there keys are the model components (i.e.
          _'wide'_, _'deeptabular'_, _'deeptext'_, _'deepimage'_ and/or _'deephead'_)  and
          the values are the corresponding optimizers. If multiple optimizers are used
          the  dictionary **MUST** contain an optimizer per model component.

        if no optimizers are passed it will default to `Adam` for all
        model components
    lr_schedulers: `LRScheduler` or dict. Optional, default=None
        - An instance of Pytorch's `LRScheduler` object (e.g
          `torch.optim.lr_scheduler.StepLR(opt, step_size=5)`) or
        - a dictionary where there keys are the model componenst (i.e. _'wide'_,
          _'deeptabular'_, _'deeptext'_, _'deepimage'_ and/or _'deephead'_) and the
          values are the corresponding learning rate schedulers.
    initializers: `Initializer` or dict. Optional, default=None
        - An instance of an `Initializer` object see `pytorch-widedeep.initializers` or
        - a dictionary where there keys are the model components (i.e. _'wide'_,
          _'deeptabular'_, _'deeptext'_, _'deepimage'_ and/or _'deephead'_)
          and the values are the corresponding initializers.
    transforms: List. Optional, default=None
        List with `torchvision.transforms` to be applied to the image
        component of the model (i.e. `deepimage`) See
        [torchvision transforms](https://pytorch.org/docs/stable/torchvision/transforms.html).
    callbacks: List. Optional, default=None
        List with `Callback` objects. The three callbacks available in
        `pytorch-widedeep` are: `LRHistory`, `ModelCheckpoint` and
        `EarlyStopping`. The `History` and the `LRShedulerCallback` callbacks
        are used by default. This can also be a custom callback as long as
        the object of type `Callback`. See
        `pytorch_widedeep.callbacks.Callback` or the examples folder in the
        repo.
    metrics: List. Optional, default=None
        - List of objects of type `Metric`. Metrics available are:
          `Accuracy`, `Precision`, `Recall`, `FBetaScore`,
          `F1Score` and `R2Score`. This can also be a custom metric as long
          as it is an object of type `Metric`. See
          `pytorch_widedeep.metrics.Metric` or the examples folder in the
          repo
        - List of objects of type `torchmetrics.Metric`. This can be any
          metric from torchmetrics library
          [Examples](https://torchmetrics.readthedocs.io/en/latest/).
          This can also be a custom metric as long as
          it is an object of type `Metric`. See
          [the instructions](https://torchmetrics.readthedocs.io/en/latest/).
    verbose: int, default=1
        Verbosity level. If set to 0 nothing will be printed during training
    seed: int, default=1
        Random seed to be used internally for train/test split

    Other Parameters
    ----------------
    **kwargs: dict
        Other infrequently used arguments that can also be passed as kwargs are:

        - **device**: `str`

            string indicating the device. One of _'cpu'_ or _'gpu'_

        - **num_workers**: `int`

            number of workers to be used internally by the data loaders

        - **lambda_sparse**: `float`

            lambda sparse parameter in case the `deeptabular` component is `TabNet`

        - **class_weight**: `List[float]`

            This is the `weight` or `pos_weight` parameter in
            `CrossEntropyLoss` and `BCEWithLogitsLoss`, depending on whether
        - **reducelronplateau_criterion**: `str`
            This sets the criterion that will be used by the lr scheduler to
            take a step: One of _'loss'_ or _'metric'_. The ReduceLROnPlateau
            learning rate is a bit particular.

    Attributes
    ----------
    cyclic_lr: bool
        Attribute that indicates if any of the lr_schedulers is cyclic_lr
        (i.e. `CyclicLR` or
        `OneCycleLR`). See [Pytorch schedulers](https://pytorch.org/docs/stable/optim.html).
    feature_importance: dict
        dict where the keys are the column names and the values are the
        corresponding feature importances. This attribute will only exist
        if the `deeptabular` component is a Tabnet model.

    Examples
    --------
    >>> import torch
    >>> from torchvision.transforms import ToTensor
    >>>
    >>> # wide deep imports
    >>> from pytorch_widedeep.callbacks import EarlyStopping, LRHistory
    >>> from pytorch_widedeep.initializers import KaimingNormal, KaimingUniform, Normal, Uniform
    >>> from pytorch_widedeep.models import TabResnet, Vision, BasicRNN, Wide, WideDeep
    >>> from pytorch_widedeep import Trainer
    >>>
    >>> embed_input = [(u, i, j) for u, i, j in zip(["a", "b", "c"][:4], [4] * 3, [8] * 3)]
    >>> column_idx = {k: v for v, k in enumerate(["a", "b", "c"])}
    >>> wide = Wide(10, 1)
    >>>
    >>> # build the model
    >>> deeptabular = TabResnet(blocks_dims=[8, 4], column_idx=column_idx, cat_embed_input=embed_input)
    >>> deeptext = BasicRNN(vocab_size=10, embed_dim=4, padding_idx=0)
    >>> deepimage = Vision()
    >>> model = WideDeep(wide=wide, deeptabular=deeptabular, deeptext=deeptext, deepimage=deepimage)
    >>>
    >>> # set optimizers and schedulers
    >>> wide_opt = torch.optim.Adam(model.wide.parameters())
    >>> deep_opt = torch.optim.AdamW(model.deeptabular.parameters())
    >>> text_opt = torch.optim.Adam(model.deeptext.parameters())
    >>> img_opt = torch.optim.AdamW(model.deepimage.parameters())
    >>>
    >>> wide_sch = torch.optim.lr_scheduler.StepLR(wide_opt, step_size=5)
    >>> deep_sch = torch.optim.lr_scheduler.StepLR(deep_opt, step_size=3)
    >>> text_sch = torch.optim.lr_scheduler.StepLR(text_opt, step_size=5)
    >>> img_sch = torch.optim.lr_scheduler.StepLR(img_opt, step_size=3)
    >>>
    >>> optimizers = {"wide": wide_opt, "deeptabular": deep_opt, "deeptext": text_opt, "deepimage": img_opt}
    >>> schedulers = {"wide": wide_sch, "deeptabular": deep_sch, "deeptext": text_sch, "deepimage": img_sch}
    >>>
    >>> # set initializers and callbacks
    >>> initializers = {"wide": Uniform, "deeptabular": Normal, "deeptext": KaimingNormal, "deepimage": KaimingUniform}
    >>> transforms = [ToTensor]
    >>> callbacks = [LRHistory(n_epochs=4), EarlyStopping]
    >>>
    >>> # set the trainer
    >>> trainer = Trainer(model, objective="regression", initializers=initializers, optimizers=optimizers,
    ... lr_schedulers=schedulers, callbacks=callbacks, transforms=transforms)
    """

    @Alias(  # noqa: C901
        "objective",
        ["loss_function", "loss_fn", "loss", "cost_function", "cost_fn", "cost"],
    )
    def __init__(
        self,
        model: WideDeep,
        objective: str,
        custom_loss_function: Optional[nn.Module] = None,
        optimizers: Optional[Union[Optimizer, Dict[str, Optimizer]]] = None,
        lr_schedulers: Optional[Union[LRScheduler, Dict[str, LRScheduler]]] = None,
        initializers: Optional[Union[Initializer, Dict[str, Initializer]]] = None,
        transforms: Optional[List[Transforms]] = None,
        callbacks: Optional[List[Callback]] = None,
        metrics: Optional[Union[List[Metric], List[TorchMetric]]] = None,
        verbose: int = 1,
        seed: int = 1,
        **kwargs,
    ):
        super().__init__(
            model=model,
            objective=objective,
            custom_loss_function=custom_loss_function,
            optimizers=optimizers,
            lr_schedulers=lr_schedulers,
            initializers=initializers,
            transforms=transforms,
            callbacks=callbacks,
            metrics=metrics,
            verbose=verbose,
            seed=seed,
            **kwargs,
        )

    @Alias("finetune", "warmup")
    def fit(  # noqa: C901
        self,
        X_wide: Optional[np.ndarray] = None,
        X_tab: Optional[np.ndarray] = None,
        X_text: Optional[np.ndarray] = None,
        X_img: Optional[np.ndarray] = None,
        X_train: Optional[Dict[str, np.ndarray]] = None,
        X_val: Optional[Dict[str, np.ndarray]] = None,
        val_split: Optional[float] = None,
        target: Optional[np.ndarray] = None,
        n_epochs: int = 1,
        validation_freq: int = 1,
        batch_size: int = 32,
        custom_dataloader: Optional[DataLoader] = None,
        feature_importance_sample_size: Optional[int] = None,
        finetune: bool = False,
        with_lds: bool = False,
        **kwargs,
    ):
        r"""Fit method.

        The input datasets can be passed either directly via numpy arrays
        (`X_wide`, `X_tab`, `X_text` or `X_img`) or alternatively, in
        dictionaries (`X_train` or `X_val`).

        Parameters
        ----------
        X_wide: np.ndarray, Optional. default=None
            Input for the `wide` model component.
            See `pytorch_widedeep.preprocessing.WidePreprocessor`
        X_tab: np.ndarray, Optional. default=None
            Input for the `deeptabular` model component.
            See `pytorch_widedeep.preprocessing.TabPreprocessor`
        X_text: np.ndarray, Optional. default=None
            Input for the `deeptext` model component.
            See `pytorch_widedeep.preprocessing.TextPreprocessor`
        X_img: np.ndarray, Optional. default=None
            Input for the `deepimage` model component.
            See `pytorch_widedeep.preprocessing.ImagePreprocessor`
        X_train: Dict, Optional. default=None
            The training dataset can also be passed in a dictionary. Keys are
            _'X_wide'_, _'X_tab'_, _'X_text'_, _'X_img'_ and _'target'_. Values
            are the corresponding matrices.
        X_val: Dict, Optional. default=None
            The validation dataset can also be passed in a dictionary. Keys
            are _'X_wide'_, _'X_tab'_, _'X_text'_, _'X_img'_ and _'target'_.
            Values are the corresponding matrices.
        val_split: float, Optional. default=None
            train/val split fraction
        target: np.ndarray, Optional. default=None
            target values
        n_epochs: int, default=1
            number of epochs
        validation_freq: int, default=1
            epochs validation frequency
        batch_size: int, default=32
            batch size
        custom_dataloader: `DataLoader`, Optional, default=None
            object of class `torch.utils.data.DataLoader`. Available
            predefined dataloaders are in `pytorch-widedeep.dataloaders`.If
            `None`, a standard torch `DataLoader` is used.
        finetune: bool, default=False
            fine-tune individual model components. This functionality can also
            be used to 'warm-up' (and hence the alias `warmup`) individual
            components before the joined training starts, and hence its
            alias. See the Examples folder in the repo for more details

            `pytorch_widedeep` implements 3 fine-tune routines.

            - fine-tune all trainable layers at once. This routine is
              inspired by the work of Howard & Sebastian Ruder 2018 in their
              [ULMfit paper](https://arxiv.org/abs/1801.06146). Using a
              Slanted Triangular learing (see
              [Leslie N. Smith paper](https://arxiv.org/pdf/1506.01186.pdf) ) ,
              the process is the following: *i*) the learning rate will
              gradually increase for 10% of the training steps from max_lr/10
              to max_lr. *ii*) It will then gradually decrease to max_lr/10
              for the remaining 90% of the steps. The optimizer used in the
              process is `Adam`.

            and two gradual fine-tune routines, where only certain layers are
            trained at a time.

            - The so called `Felbo` gradual fine-tune rourine, based on the the
              Felbo et al., 2017 [DeepEmoji paper](https://arxiv.org/abs/1708.00524).
            - The `Howard` routine based on the work of Howard & Sebastian Ruder 2018 in their
              [ULMfit paper](https://arxiv.org/abs/1801.06146>).

            For details on how these routines work, please see the Examples
            section in this documentation and the Examples folder in the repo.

            Param Alias: `warmup`
        with_lds: bool, default=False
            Boolean indicating if Label Distribution Smoothing will be used.

            information_source: **NOTE**: We consider this feature absolutely
            experimental and we recommend the user to not use it unless the
            corresponding [publication](https://arxiv.org/abs/2102.09554) is
            well understood

        Other Parameters
        ----------------
        **kwargs : dict
            Other keyword arguments are:

            - **DataLoader related parameters**:

                For example,  `sampler`, `batch_sampler`, `collate_fn`, etc.
                Please, see the pytorch
                [DataLoader docs](https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader)
                for details.

            - **Label Distribution Smoothing related parameters**:

                - lds_kernel (`Literal['gaussian', 'triang', 'laplace']`):
                    choice of kernel for Label Distribution Smoothing
                - lds_ks (`int`):
                    LDS kernel window size
                - lds_sigma (`float`):
                    standard deviation of ['gaussian','laplace'] kernel for LDS
                - lds_granularity (`int`):
                    number of bins in histogram used in LDS to count occurence of sample values
                - lds_reweight (`bool`):
                    option to reweight bin frequency counts in LDS
                - lds_y_max (`Optional[float]`):
                    option to restrict LDS bins by upper label limit
                - lds_y_min (`Optional[float]`):
                    option to restrict LDS bins by lower label limit

                See `pytorch_widedeep.trainer._wd_dataset` for more details on
                the implications of these parameters

            - **Finetune related parameters**:

                see the source code at `pytorch_widedeep._finetune`. Namely, these are:

                - `finetune_epochs` (`int`):
                    number of epochs use for fine tuning
                - `finetune_max_lr` (`float`):
                   max lr during fine tuning
                - `routine` (`str`):
                   one of _'howard'_ or _'felbo'_
                - `deeptabular_gradual` (`bool`):
                   boolean indicating if the `deeptabular` component will be fine tuned gradually
                - `deeptabular_layers` (`List[nn.Module]`):
                   List of pytorch modules indicating the layers of the
                   `deeptabular` that will be fine tuned
                - `deeptabular_max_lr` (`float`):
                   max lr for the `deeptabular` componet during fine tuning
                - `deeptext_gradual` (`bool`):
                   same as `deeptabular_gradual` but for the `deeptext` component
                - `deeptext_layers` (`List[nn.Module]`):
                   same as `deeptabular_gradual` but for the `deeptext` component
                - `deeptext_max_lr` (`float`):
                   same as `deeptabular_gradual` but for the `deeptext` component
                - `deepimage_gradual` (`bool`):
                   same as `deeptabular_gradual` but for the `deepimage` component
                - `deepimage_layers` (`List[nn.Module]`):
                   same as `deeptabular_gradual` but for the `deepimage` component
                - `deepimage_max_lr` (`float`):
                    same as `deeptabular_gradual` but for the `deepimage` component

        Examples
        --------

        For a series of comprehensive examples on how to use the `fit` method, please see the
        [Examples](https://github.com/jrzaurin/pytorch-widedeep/tree/master/examples)
        folder in the repo
        """

        lds_args, dataloader_args, finetune_args = self._extract_kwargs(kwargs)
        lds_args["with_lds"] = with_lds
        self.with_lds = with_lds

        self.batch_size = batch_size
        train_set, eval_set = wd_train_val_split(
            self.seed,
            self.method,
            X_wide,
            X_tab,
            X_text,
            X_img,
            X_train,
            X_val,
            val_split,
            target,
            **lds_args,
        )
        if isinstance(custom_dataloader, type):
            if issubclass(custom_dataloader, DataLoader):
                train_loader = custom_dataloader(  # type: ignore[misc]
                    dataset=train_set,
                    batch_size=batch_size,
                    num_workers=self.num_workers,
                    **dataloader_args,
                )
            else:
                NotImplementedError(
                    "Custom DataLoader must be a subclass of "
                    "torch.utils.data.DataLoader, please see the "
                    "pytorch documentation or examples in "
                    "pytorch_widedeep.dataloaders"
                )
        else:
            train_loader = DataLoaderDefault(
                dataset=train_set,
                batch_size=batch_size,
                num_workers=self.num_workers,
                **dataloader_args,
            )
        train_steps = len(train_loader)
        if eval_set is not None:
            eval_loader = DataLoader(
                dataset=eval_set,
                batch_size=batch_size,
                num_workers=self.num_workers,
                shuffle=False,
            )
            eval_steps = len(eval_loader)

        if finetune:
            self._finetune(train_loader, **finetune_args)
            if self.verbose:
                print(
                    "Fine-tuning (or warmup) of individual components completed. "
                    "Training the whole model for {} epochs".format(n_epochs)
                )

        self.callback_container.on_train_begin(
            {"batch_size": batch_size, "train_steps": train_steps, "n_epochs": n_epochs}
        )
        for epoch in range(n_epochs):
            epoch_logs: Dict[str, float] = {}
            self.callback_container.on_epoch_begin(epoch, logs=epoch_logs)

            self.train_running_loss = 0.0
            with trange(train_steps, disable=self.verbose != 1) as t:
                for batch_idx, (data, targett, lds_weightt) in zip(t, train_loader):
                    t.set_description("epoch %i" % (epoch + 1))
                    train_score, train_loss = self._train_step(
                        data, targett, batch_idx, epoch, lds_weightt
                    )
                    print_loss_and_metric(t, train_loss, train_score)
                    self.callback_container.on_batch_end(batch=batch_idx)
            epoch_logs = save_epoch_logs(epoch_logs, train_loss, train_score, "train")

            on_epoch_end_metric = None
            if eval_set is not None and epoch % validation_freq == (
                validation_freq - 1
            ):
                self.callback_container.on_eval_begin()
                self.valid_running_loss = 0.0
                with trange(eval_steps, disable=self.verbose != 1) as v:
                    for i, (data, targett) in zip(v, eval_loader):
                        v.set_description("valid")
                        val_score, val_loss = self._eval_step(data, targett, i)
                        print_loss_and_metric(v, val_loss, val_score)
                epoch_logs = save_epoch_logs(epoch_logs, val_loss, val_score, "val")

                if self.reducelronplateau:
                    if self.reducelronplateau_criterion == "loss":
                        on_epoch_end_metric = val_loss
                    else:
                        on_epoch_end_metric = val_score[
                            self.reducelronplateau_criterion
                        ]
            else:
                if self.reducelronplateau:
                    raise NotImplementedError(
                        "ReduceLROnPlateau scheduler can be used only with validation data."
                    )
            self.callback_container.on_epoch_end(epoch, epoch_logs, on_epoch_end_metric)

            if self.early_stop:
                # self.callback_container.on_train_end(epoch_logs)
                break

            if self.model.with_fds:
                self._update_fds_stats(train_loader, epoch)

        self.callback_container.on_train_end(epoch_logs)

        if feature_importance_sample_size is not None:
            self.feature_importance = FeatureImportance(
                self.device, feature_importance_sample_size
            ).feature_importance(train_loader, self.model)
        self._restore_best_weights()
        self.model.train()

    def predict(  # type: ignore[return]
        self,
        X_wide: Optional[np.ndarray] = None,
        X_tab: Optional[np.ndarray] = None,
        X_text: Optional[np.ndarray] = None,
        X_img: Optional[np.ndarray] = None,
        X_test: Optional[Dict[str, np.ndarray]] = None,
        batch_size: int = 256,
    ) -> np.ndarray:
        r"""Returns the predictions

        The input datasets can be passed either directly via numpy arrays
        (`X_wide`, `X_tab`, `X_text` or `X_img`) or alternatively, in
        a dictionary (`X_test`)

        Parameters
        ----------
        X_wide: np.ndarray, Optional. default=None
            Input for the `wide` model component.
            See `pytorch_widedeep.preprocessing.WidePreprocessor`
        X_tab: np.ndarray, Optional. default=None
            Input for the `deeptabular` model component.
            See `pytorch_widedeep.preprocessing.TabPreprocessor`
        X_text: np.ndarray, Optional. default=None
            Input for the `deeptext` model component.
            See `pytorch_widedeep.preprocessing.TextPreprocessor`
        X_img: np.ndarray, Optional. default=None
            Input for the `deepimage` model component.
            See `pytorch_widedeep.preprocessing.ImagePreprocessor`
        X_test: Dict, Optional. default=None
            The test dataset can also be passed in a dictionary. Keys are
            `X_wide`, _'X_tab'_, _'X_text'_, _'X_img'_ and _'target'_. Values
            are the corresponding matrices.
        batch_size: int, default = 256
            If a trainer is used to predict after having trained a model, the
            `batch_size` needs to be defined as it will not be defined as
            the `Trainer` is instantiated

        Returns
        -------
        np.ndarray:
            array with the predictions
        """
        preds_l = self._predict(X_wide, X_tab, X_text, X_img, X_test, batch_size)
        if self.method == "regression":
            return np.vstack(preds_l).squeeze(1)
        if self.method == "binary":
            preds = np.vstack(preds_l).squeeze(1)
            return (preds > 0.5).astype("int")
        if self.method == "qregression":
            return np.vstack(preds_l)
        if self.method == "multiclass":
            preds = np.vstack(preds_l)
            return np.argmax(preds, 1)  # type: ignore[return-value]

    def predict_uncertainty(  # type: ignore[return]
        self,
        X_wide: Optional[np.ndarray] = None,
        X_tab: Optional[np.ndarray] = None,
        X_text: Optional[np.ndarray] = None,
        X_img: Optional[np.ndarray] = None,
        X_test: Optional[Dict[str, np.ndarray]] = None,
        batch_size: int = 256,
        uncertainty_granularity=1000,
    ) -> np.ndarray:
        r"""Returns the predicted ucnertainty of the model for the test dataset
        using a Monte Carlo method during which dropout layers are activated
        in the evaluation/prediction phase and each sample is predicted N
        times (`uncertainty_granularity` times).

        This is based on
        [Dropout as a Bayesian Approximation: Representing
        Model Uncertainty in Deep Learning](https://arxiv.org/abs/1506.02142?context=stat).

        Parameters
        ----------
        X_wide: np.ndarray, Optional. default=None
            Input for the `wide` model component.
            See `pytorch_widedeep.preprocessing.WidePreprocessor`
        X_tab: np.ndarray, Optional. default=None
            Input for the `deeptabular` model component.
            See `pytorch_widedeep.preprocessing.TabPreprocessor`
        X_text: np.ndarray, Optional. default=None
            Input for the `deeptext` model component.
            See `pytorch_widedeep.preprocessing.TextPreprocessor`
        X_img: np.ndarray, Optional. default=None
            Input for the `deepimage` model component.
            See `pytorch_widedeep.preprocessing.ImagePreprocessor`
        X_test: Dict, Optional. default=None
            The test dataset can also be passed in a dictionary. Keys are
            _'X_wide'_, _'X_tab'_, _'X_text'_, _'X_img'_ and _'target'_. Values
            are the corresponding matrices.
        batch_size: int, default = 256
            If a trainer is used to predict after having trained a model, the
            `batch_size` needs to be defined as it will not be defined as
            the `Trainer` is instantiated
        uncertainty_granularity: int default = 1000
            number of times the model does prediction for each sample

        Returns
        -------
        np.ndarray:
            - if `method = regression`, it will return an array with `(max, min, mean, stdev)`
              values for each sample.
            - if `method = binary` it will return an array with
              `(mean_cls_0_prob, mean_cls_1_prob, predicted_cls)` for each sample.
            - if `method = multiclass` it will return an array with
              `(mean_cls_0_prob, mean_cls_1_prob, mean_cls_2_prob, ... , predicted_cls)`
              values for each sample.

        """
        preds_l = self._predict(
            X_wide,
            X_tab,
            X_text,
            X_img,
            X_test,
            batch_size,
            uncertainty_granularity,
            uncertainty=True,
        )
        preds = np.vstack(preds_l)
        samples_num = int(preds.shape[0] / uncertainty_granularity)
        if self.method == "regression":
            preds = preds.squeeze(1)
            preds = preds.reshape((uncertainty_granularity, samples_num))
            return np.array(
                (
                    preds.max(axis=0),
                    preds.min(axis=0),
                    preds.mean(axis=0),
                    preds.std(axis=0),
                )
            ).T
        if self.method == "qregression":
            raise ValueError(
                "Currently predict_uncertainty is not supported for qregression method"
            )
        if self.method == "binary":
            preds = preds.squeeze(1)
            preds = preds.reshape((uncertainty_granularity, samples_num))
            preds = preds.mean(axis=0)
            probs = np.zeros([preds.shape[0], 3])
            probs[:, 0] = 1 - preds
            probs[:, 1] = preds
            return probs
        if self.method == "multiclass":
            preds = preds.reshape(uncertainty_granularity, samples_num, preds.shape[1])
            preds = preds.mean(axis=0)
            preds = np.hstack((preds, np.vstack(np.argmax(preds, 1))))
            return preds

    def predict_proba(  # type: ignore[return]
        self,
        X_wide: Optional[np.ndarray] = None,
        X_tab: Optional[np.ndarray] = None,
        X_text: Optional[np.ndarray] = None,
        X_img: Optional[np.ndarray] = None,
        X_test: Optional[Dict[str, np.ndarray]] = None,
        batch_size: int = 256,
    ) -> np.ndarray:
        r"""Returns the predicted probabilities for the test dataset for  binary
        and multiclass methods

        The input datasets can be passed either directly via numpy arrays
        (`X_wide`, `X_tab`, `X_text` or `X_img`) or alternatively, in
        a dictionary (`X_test`)

        Parameters
        ----------
        X_wide: np.ndarray, Optional. default=None
            Input for the `wide` model component.
            See `pytorch_widedeep.preprocessing.WidePreprocessor`
        X_tab: np.ndarray, Optional. default=None
            Input for the `deeptabular` model component.
            See `pytorch_widedeep.preprocessing.TabPreprocessor`
        X_text: np.ndarray, Optional. default=None
            Input for the `deeptext` model component.
            See `pytorch_widedeep.preprocessing.TextPreprocessor`
        X_img: np.ndarray, Optional. default=None
            Input for the `deepimage` model component.
            See `pytorch_widedeep.preprocessing.ImagePreprocessor`
        X_test: Dict, Optional. default=None
            The test dataset can also be passed in a dictionary. Keys are
            `X_wide`, _'X_tab'_, _'X_text'_, _'X_img'_ and _'target'_. Values
            are the corresponding matrices.
        batch_size: int, default = 256
            If a trainer is used to predict after having trained a model, the
            `batch_size` needs to be defined as it will not be defined as
            the `Trainer` is instantiated

        Returns
        -------
        np.ndarray
            array with the probabilities per class
        """

        preds_l = self._predict(X_wide, X_tab, X_text, X_img, X_test, batch_size)
        if self.method == "binary":
            preds = np.vstack(preds_l).squeeze(1)
            probs = np.zeros([preds.shape[0], 2])
            probs[:, 0] = 1 - preds
            probs[:, 1] = preds
            return probs
        if self.method == "multiclass":
            return np.vstack(preds_l)

    def explain(self, X_tab: np.ndarray, save_step_masks: Optional[bool] = None):
        # TO DO: Add docs to this, to the feat imp parameter and the all
        # related classes
        explainer = Explainer(self.device)

        res = explainer.explain(
            self.model, X_tab, self.num_workers, self.batch_size, save_step_masks
        )

        return res

    def save(
        self,
        path: str,
        save_state_dict: bool = False,
        model_filename: str = "wd_model.pt",
    ):
        r"""Saves the model, training and evaluation history, and the
        `feature_importance` attribute (if the `deeptabular` component is a
        Tabnet model) to disk

        The `Trainer` class is built so that it 'just' trains a model. With
        that in mind, all the torch related parameters (such as optimizers,
        learning rate schedulers, initializers, etc) have to be defined
        externally and then passed to the `Trainer`. As a result, the
        `Trainer` does not generate any attribute or additional data
        products that need to be saved other than the `model` object itself,
        which can be saved as any other torch model (e.g. `torch.save(model,
        path)`).

        The exception is Tabnet. If the `deeptabular` component is a Tabnet
        model, an attribute (a dict) called `feature_importance` will be
        created at the end of the training process. Therefore, a `save`
        method was created that will save the feature importance dictionary
        to a json file and, since we are here, the model weights, training
        history and learning rate history.

        Parameters
        ----------
        path: str
            path to the directory where the model and the feature importance
            attribute will be saved.
        save_state_dict: bool, default = False
            Boolean indicating whether to save directly the model or the
            model's state dictionary
        model_filename: str, Optional, default = "wd_model.pt"
            filename where the model weights will be store
        """

        save_dir = Path(path)
        history_dir = save_dir / "history"
        history_dir.mkdir(exist_ok=True, parents=True)

        # the trainer is run with the History Callback by default
        with open(history_dir / "train_eval_history.json", "w") as teh:
            json.dump(self.history, teh)  # type: ignore[attr-defined]

        has_lr_history = any(
            [clbk.__class__.__name__ == "LRHistory" for clbk in self.callbacks]
        )
        if self.lr_scheduler is not None and has_lr_history:
            with open(history_dir / "lr_history.json", "w") as lrh:
                json.dump(self.lr_history, lrh)  # type: ignore[attr-defined]

        model_path = save_dir / model_filename
        if save_state_dict:
            torch.save(self.model.state_dict(), model_path)
        else:
            torch.save(self.model, model_path)

        if self.model.is_tabnet:
            with open(save_dir / "feature_importance.json", "w") as fi:
                json.dump(self.feature_importance, fi)

    @Alias("n_epochs", ["finetune_epochs", "warmup_epochs"])
    @Alias("max_lr", ["finetune_max_lr", "warmup_max_lr"])
    def _finetune(
        self,
        loader: DataLoader,
        n_epochs: int = 5,
        max_lr: float = 0.01,
        routine: Literal["howard", "felbo"] = "howard",
        deeptabular_gradual: bool = False,
        deeptabular_layers: Optional[List[nn.Module]] = None,
        deeptabular_max_lr: float = 0.01,
        deeptext_gradual: bool = False,
        deeptext_layers: Optional[List[nn.Module]] = None,
        deeptext_max_lr: float = 0.01,
        deepimage_gradual: bool = False,
        deepimage_layers: Optional[List[nn.Module]] = None,
        deepimage_max_lr: float = 0.01,
    ):
        r"""
        Simple wrap-up to individually fine-tune model components
        """
        if self.model.deephead is not None:
            raise ValueError(
                "Currently warming up is only supported without a fully connected 'DeepHead'"
            )

        finetuner = FineTune(self.loss_fn, self.metric, self.method, self.verbose)
        if self.model.wide:
            finetuner.finetune_all(self.model.wide, "wide", loader, n_epochs, max_lr)

        if self.model.deeptabular:
            if deeptabular_gradual:
                finetuner.finetune_gradual(
                    self.model.deeptabular,
                    "deeptabular",
                    loader,
                    deeptabular_max_lr,
                    deeptabular_layers,
                    routine,
                )
            else:
                finetuner.finetune_all(
                    self.model.deeptabular, "deeptabular", loader, n_epochs, max_lr
                )

        if self.model.deeptext:
            if deeptext_gradual:
                finetuner.finetune_gradual(
                    self.model.deeptext,
                    "deeptext",
                    loader,
                    deeptext_max_lr,
                    deeptext_layers,
                    routine,
                )
            else:
                finetuner.finetune_all(
                    self.model.deeptext, "deeptext", loader, n_epochs, max_lr
                )

        if self.model.deepimage:
            if deepimage_gradual:
                finetuner.finetune_gradual(
                    self.model.deepimage,
                    "deepimage",
                    loader,
                    deepimage_max_lr,
                    deepimage_layers,
                    routine,
                )
            else:
                finetuner.finetune_all(
                    self.model.deepimage, "deepimage", loader, n_epochs, max_lr
                )

    def _train_step(
        self,
        data: Dict[str, Tensor],
        target: Tensor,
        batch_idx: int,
        epoch: int,
        lds_weightt: Tensor,
    ):
        lds_weight = (
            None
            if torch.all(lds_weightt == 0)
            else lds_weightt.view(-1, 1).to(self.device)
        )
        if (
            self.with_lds
            and lds_weight is not None
            and "lds_weight" not in signature(self.loss_fn.forward).parameters
        ):
            warnings.warn(
                """LDS weights are not None but the loss function used does not"
                " support LDS weightening. For loss functions that support LDS"
                " weightening please read the docs""",
                UserWarning,
            )

        self.model.train()
        X = {k: v.to(self.device) for k, v in data.items()}
        y = (
            target.view(-1, 1).float()
            if self.method not in ["multiclass", "qregression"]
            else target
        )
        y = y.to(self.device)

        self.optimizer.zero_grad()

        if self.model.with_fds:
            _, y_pred = self.model(X, y, epoch)
        else:
            y_pred = self.model(X)

        if self.model.is_tabnet:
            loss = self.loss_fn(y_pred[0], y) - self.lambda_sparse * y_pred[1]
            score = self._get_score(y_pred[0], y)
        else:
            loss = (
                self.loss_fn(y_pred, y)
                if not self.with_lds
                else self.loss_fn(y_pred, y, lds_weight=lds_weight)
            )
            score = self._get_score(y_pred, y)

        loss.backward()
        self.optimizer.step()

        self.train_running_loss += loss.item()
        avg_loss = self.train_running_loss / (batch_idx + 1)

        return score, avg_loss

    def _eval_step(self, data: Dict[str, Tensor], target: Tensor, batch_idx: int):
        self.model.eval()
        with torch.no_grad():
            X = {k: v.to(self.device) for k, v in data.items()}
            y = (
                target.view(-1, 1).float()
                if self.method not in ["multiclass", "qregression"]
                else target
            )
            y = y.to(self.device)

            y_pred = self.model(X)
            if self.model.is_tabnet:
                loss = self.loss_fn(y_pred[0], y) - self.lambda_sparse * y_pred[1]
                score = self._get_score(y_pred[0], y)
            else:
                score = self._get_score(y_pred, y)
                loss = self.loss_fn(y_pred, y)

            self.valid_running_loss += loss.item()
            avg_loss = self.valid_running_loss / (batch_idx + 1)

        return score, avg_loss

    def _get_score(self, y_pred, y):
        if self.metric is not None:
            if self.method == "regression":
                score = self.metric(y_pred, y)
            if self.method == "binary":
                score = self.metric(torch.sigmoid(y_pred), y)
            if self.method == "qregression":
                score = self.metric(y_pred, y)
            if self.method == "multiclass":
                score = self.metric(F.softmax(y_pred, dim=1), y)
            return score
        else:
            return None

    def _fds_step(
        self,
        data: Dict[str, Tensor],
        target: Tensor,
        epoch: int,
    ) -> Tuple[Tensor, Tensor]:
        self.model.train()
        X = {k: v.to(self.device) for k, v in data.items()}
        y = target.view(-1, 1).float().to(self.device)
        smoothed_features, _ = self.model(X, y, epoch)
        return smoothed_features, y

    def _update_fds_stats(self, train_loader: DataLoader, epoch: int):
        train_steps = len(train_loader)
        features_l, y_pred_l = [], []
        with torch.no_grad():
            with trange(train_steps, disable=self.verbose != 1) as t:
                for idx, (data, targett, lds_weight) in zip(t, train_loader):
                    t.set_description("FDS update")
                    deeptab_features, deeptab_preds = self._fds_step(
                        data,
                        targett,
                        epoch,
                    )
                    features_l.append(deeptab_features)
                    y_pred_l.append(deeptab_preds)
        features = torch.cat(features_l)
        y_pred = torch.cat(y_pred_l)
        self.model.fds_layer.update_last_epoch_stats(epoch)
        self.model.fds_layer.update_running_stats(features, y_pred, epoch)

    def _predict(  # noqa: C901
        self,
        X_wide: Optional[np.ndarray] = None,
        X_tab: Optional[np.ndarray] = None,
        X_text: Optional[np.ndarray] = None,
        X_img: Optional[np.ndarray] = None,
        X_test: Optional[Dict[str, np.ndarray]] = None,
        batch_size: int = 256,
        uncertainty_granularity=1000,
        uncertainty: bool = False,
        quantiles: bool = False,
    ) -> List:
        r"""Private method to avoid code repetition in predict and
        predict_proba. For parameter information, please, see the .predict()
        method documentation
        """
        if X_test is not None:
            test_set = WideDeepDataset(**X_test)  # type: ignore[arg-type]
        else:
            load_dict = {}
            if X_wide is not None:
                load_dict = {"X_wide": X_wide}
            if X_tab is not None:
                load_dict.update({"X_tab": X_tab})
            if X_text is not None:
                load_dict.update({"X_text": X_text})
            if X_img is not None:
                load_dict.update({"X_img": X_img})
            test_set = WideDeepDataset(**load_dict)  # type: ignore[arg-type]

        if not hasattr(self, "batch_size"):
            self.batch_size = batch_size

        test_loader = DataLoader(
            dataset=test_set,
            batch_size=self.batch_size,
            num_workers=self.num_workers,
            shuffle=False,
        )
        test_steps = (len(test_loader.dataset) // test_loader.batch_size) + 1  # type: ignore[arg-type]

        self.model.eval()
        preds_l = []

        if uncertainty:
            for m in self.model.modules():
                if m.__class__.__name__.startswith("Dropout"):
                    m.train()
            prediction_iters = uncertainty_granularity
        else:
            prediction_iters = 1

        with torch.no_grad():
            with trange(uncertainty_granularity, disable=uncertainty is False) as t:
                for i, k in zip(t, range(prediction_iters)):
                    t.set_description("predict_UncertaintyIter")

                    with trange(
                        test_steps, disable=self.verbose != 1 or uncertainty is True
                    ) as tt:
                        for j, data in zip(tt, test_loader):
                            tt.set_description("predict")
                            X = {k: v.to(self.device) for k, v in data.items()}
                            preds = (
                                self.model(X)
                                if not self.model.is_tabnet
                                else self.model(X)[0]
                            )
                            if self.method == "binary":
                                preds = torch.sigmoid(preds)
                            if self.method == "multiclass":
                                preds = F.softmax(preds, dim=1)
                            if self.method == "regression" and isinstance(
                                self.loss_fn, ZILNLoss
                            ):
                                preds = self._predict_ziln(preds)
                            preds = preds.cpu().data.numpy()
                            preds_l.append(preds)
        self.model.train()
        return preds_l

    @staticmethod
    def _predict_ziln(preds: Tensor) -> Tensor:
        """Calculates predicted mean of zero inflated lognormal logits.

        Adjusted implementaion of `code
        <https://github.com/google/lifetime_value/blob/master/lifetime_value/zero_inflated_lognormal.py>`

        Arguments:
            preds: [batch_size, 3] tensor of logits.
        Returns:
            ziln_preds: [batch_size, 1] tensor of predicted mean.
        """
        positive_probs = torch.sigmoid(preds[..., :1])
        loc = preds[..., 1:2]
        scale = F.softplus(preds[..., 2:])
        ziln_preds = positive_probs * torch.exp(loc + 0.5 * torch.square(scale))
        return ziln_preds

    @staticmethod
    def _extract_kwargs(kwargs):
        dataloader_params = [
            "shuffle",
            "sampler",
            "batch_sampler",
            "num_workers",
            "collate_fn",
            "pin_memory",
            "drop_last",
            "timeout",
            "worker_init_fn",
            "generator",
            "prefetch_factor",
            "persistent_workers",
            "oversample_mul",
        ]
        lds_params = [
            "lds_kernel",
            "lds_ks",
            "lds_sigma",
            "lds_granularity",
            "lds_reweight",
            "lds_y_max",
            "lds_y_min",
        ]
        finetune_params = [
            "n_epochs",
            "finetune_epochs",
            "warmup_epochs",
            "max_lr",
            "finetune_max_lr",
            "warmup_max_lr",
            "routine",
            "deeptabular_gradual",
            "deeptabular_layers",
            "deeptabular_max_lr",
            "deeptext_gradual",
            "deeptext_layers",
            "deeptext_max_lr",
            "deepimage_gradual",
            "deepimage_layers",
            "deepimage_max_lr",
        ]

        lds_args, dataloader_args, finetune_args = {}, {}, {}
        for k, v in kwargs.items():
            if k in lds_params:
                lds_args[k] = v
            if k in dataloader_params:
                dataloader_args[k] = v
            if k in finetune_params:
                finetune_args[k] = v

        return lds_args, dataloader_args, finetune_args