function _check_nan(loss::Real)
    if isnan(loss)
        @error "Loss is NaN";
        stop();
    end
end



function fit!(data::AbstractVecOrMat{D}, models::Union{M, AbstractVector{M}},
              loss::Function, opt::AbstractOptimiser;
              train_loop::Function = train_loop!,
              test_loop::Function = (::Vararg) -> nothing,
              logger::Union{L, Nothing} = nothing,
              test_data::Union{AbstractVecOrMat{TD}, Nothing} = nothing,
              test_callback::Function = (::Vararg) -> nothing,
              reset_callback!::Function = (::Vararg) -> nothing,
              verbose = false) where {D, L<:AbstractLogger, M, TD}
    train_loss = train_loop(data, models, loss, opt, logger = logger, verbose = verbose);
    @info "Train loss: $train_loss";

    if isnothing(test_data)
        if !isnothing(logger)
            log_loss_and_increase!(train_loss, logger, "loss/train");
        end
    else
        test_loss = test_loop(test_data, models, loss, test_callback);
        @info "Test loss: $test_loss";
        if !isnothing(logger)
            log_loss(train_loss, logger, "loss/train");
            log_loss_and_increase!(test_loss, logger, "loss/test");
        end
    end

    reset_callback!(models);
end



# function loop(data::AbstractVector{D},
#               loss::Function; opt::Union{AbstractOptimiser, Nothing} = nothing,
#               logger::Union{L, Nothing} = nothing, verbose = false,
#               test_mode = false, test_callback::Function = () -> ()) where {D<:Union{Tuple, NamedTuple},
#                                                                             L<:AbstractLogger}
#     local agg_loss = 0.0;

#     for elem in data
#         sku = elem.Id;

#         model = elem.Model;
#         labels = elem.Labels;

#         sample = labels;
#         if :TExplanatory in propertynames(first(data))
#             explanatory = elem.TExplanatory;
#             sample = (labels, explanatory);
#         end

#         if !test_mode
#             if isnothing(opt)
#                 throw(ArgumentError("`opt` has to be denifed in training mode!"));
#             end

#             ps = params(model);

#             train_loss, back = pullback(ps) do
#                 yhat = model(sample);
#                 loss(yhat, labels);
#             end

#             _check_nan(train_loss);
#             agg_loss += train_loss;
#             update!(train_loss, ps, back, opt, verbose = verbose);
#             if !isnothing(logger)
#                 logging(model, logger, "$sku");
#             end
#         else
#             yhat = test_callback(model, elem);
#             test_loss = loss(yhat, labels);

#             _check_nan(test_loss);
#             agg_loss += test_loss;
#         end
#     end

#     return agg_loss;
# end
function restart!(models::M; kwargs...) where {M}
    for models in models
        restart!(model; kwargs...);
    end
end



function test_loop(data::AbstractVector{D}, models::M,
                   loss::Function, test_callback::Function = (::Vararg) -> nothing) where {D<:Union{Tuple, NamedTuple},
                                                                                           M<:AbstractVector}
    local agg_loss = 0.0;

    for (elem, model) in zip(data, models)

        labels = elem.Labels;

        sample = labels;
        if :TExplanatory in propertynames(first(data))
            explanatory = elem.TExplanatory;
            sample = (labels, explanatory);
        end

        yhat = test_callback(model, elem);
        test_loss = loss(yhat, labels);

        _check_nan(test_loss);
        agg_loss += test_loss;
    end

    return agg_loss;
end
function train_loop!(data::AbstractVector{D}, models::M,
                     loss::Function, opt::AbstractOptimiser = nothing;
                     logger::Union{L, Nothing} = nothing, verbose = false) where {D<:Union{Tuple, NamedTuple},
                                                                                  L<:AbstractLogger,
                                                                                  M<:AbstractVector}
    local agg_loss = 0.0;

    for (elem, model) in zip(data, models)
        sku = elem.Id;

        labels = elem.Labels;

        sample = labels;
        if :TExplanatory in propertynames(first(data))
            explanatory = elem.TExplanatory;
            sample = (labels, explanatory);
        end

        if isnothing(opt)
            throw(ArgumentError("`opt` has to be denifed in training mode!"));
        end

        ps = params(model);

        train_loss, back = pullback(ps) do
            yhat = model(sample);
            loss(yhat, labels);
        end

        _check_nan(train_loss);
        agg_loss += train_loss;
        update!(train_loss, ps, back, opt, verbose = verbose);
        if !isnothing(logger)
            logging(model, logger, "$sku");
        end
    end

    return agg_loss;
end



function train(trainDF::Union{DF, GroupedDataFrame{DF}};
               cell = "ses", with_seasonality = false, modes = "mul:mul:mul:mul", frequencies = "moy:woy:doy:dow", delimiter = ":",
               loss::Function = mse, opt::AbstractOptimiser = ADAM(0.1, (0.9, 0.999)), epochs::Number = 500,
               testDF::Union{Union{DF, GroupedDataFrame{DF}}, Nothing} = nothing,
               test_callback::Function = (::Vararg) -> nothing,
               log_prefix = "log", log_location = default_log_location, verbose = false,
               kwargs...) where {DF<:AbstractDataFrame}
    chosen_cell = lowercase(cell);
    lwr_modes = lowercase(modes);
    lwr_frequencies = lowercase(frequencies);

    c = available_cells[chosen_cell];
    expected_cols = [:Id, :Date, :Qty];
    select_cols = [:Id, :Labels, :Model];

    _check_columns(expected_cols, trainDF,
                   message = "$(expected_cols) expected in training DataFrame!");

    # ----[ START - Initialize Models ]----
    df = preprocess(trainDF);
    if with_seasonality
        add_time_explanatory!(df, lwr_modes, lwr_frequencies, delimiter);
        insert!(select_cols, 2, :TExplanatory);
        default_model!(df, c, modes = lwr_modes, frequencies = lwr_frequencies, delimiter = delimiter);
    else
        default_model!(df, c);
    end

    sdf = select(df, select_cols);
    sdf = [NamedTuple(row) for row in eachrow(sdf)];
    test_sdf = nothing;

    if !isnothing(testDF)
        _check_columns(expected_cols, testDF,
                       message = "$(expected_cols) expected in test DataFrame!");

        test_df = preprocess(testDF);
        if with_seasonality
            add_time_explanatory!(test_df, lwr_modes, lwr_frequencies, delimiter);
        end
        test_sdf = select(test_df, select_cols[1:(end - 1)]);
        test_sdf = [NamedTuple(row) for row in eachrow(test_sdf)];

    end
    # ----[ END - Initialize Models ]----

    logger, log_name = init_log_file(sdf[1].Model, joinpath(log_location, "Parts"), log_prefix)
    try
        @epochs epochs fit!(sdf, sdf.Model,
                            loss, opt,
                            test_loop = test_loop,
                            logger = logger,
                            test_data = test_sdf,
                            test_callback = test_callback,
                            reset_callback! = restart!,
                            verbose = verbose);
    finally
        close(logger.file);
    end

    @info "Logdir: $(logger.logdir)"
    store_dct = Dict([(elem.Id, elem.Model) for elem in sdf])
    save_model(joinpath(logger.logdir, log_name * ".bson"), store_dct);
end
# function train!(train_vm::AbstractVecOrMat{D}, model::M,
#                 train_loop!::Function;
#                 loss::Function = mse, opt::AbstractOptimiser = ADAM(0.1, (0.9, 0.999)), epochs::Number = 500,
#                 test_vm::Union{AbstractVecOrMat{TD}, Nothing} = nothing,
#                 test_loop::Function = (::Vararg) -> nothing, test_callback::Function = (::Vararg) -> nothing,
#                 reset_callback!::Function = (::Vararg) -> nothing, 
#                 log = false, log_prefix = "log", log_location = joinpath(@__DIR__, "./Log"), verbose = false,
#                 kwargs...) where {D, M, TD}
#     log_name = log_prefix;
#     log_path = log_location;
#     logger = nothing;
#     if log
#         logger, log_name = init_log_file(model, log_location, log_name);
#     else
#         td = today();
#         nw = now();
#         sec = second(nw);
#         min = minute(nw);
#         hr = hour(nw);
#         log_name = "$(log_prefix)_$(td)_$(hr)h$(min)m$(sec)s";
#         log_path = joinpath(log_location, log_name);
#         mkpath(log_path);
#     end

#     try
#         @epochs epochs Ses.fit!(train_vm, model,
#                                 loss, opt,
#                                 train_loop = train_loop!,
#                                 test_loop = test_loop,
#                                 logger = logger,
#                                 test_data = test_vm,
#                                 test_callback = test_callback,
#                                 reset_callback! = reset_callback!,
#                                 verbose = verbose);
#     finally
#         if !isnothing(logger)
#             close(logger.file);
#         end
#     end

#     return logger, log_path, log_name
# end


function update!(loss::Real, params::Params, back::Function, opt::AbstractOptimiser; verbose::Bool = false)
    # println("[Update] Loss: $loss")
    # println("[Update] 1vLoss: $(one(loss))")
    # println("[Update] Gs: $(back(one(loss)))")
    gs = back(one(loss));
    if verbose
        for (param, g) in pairs(gs)
            @info "Param: $param";
            @info "Gradient: $g";
        end
    end
    update!(opt, params, gs); 
end