using MKL


include("../PartsHelpers/src/PartsHelpers.jl");
include("../zoo/PES.jl");
include("parse_commandline.jl")


using .PartsHelpers
using Distributions: NegativeBinomial
using Flux
using PlotlyJS
using Statistics



function test_callback(data::Union{AbstractVecOrMat{T}, Tuple, NamedTuple}, m::PES{M}) where {M, T}
    horizon, nskus = size(data[1]);
    explanatories = data[2];

    scalers = m[:scale];
    invscalers = m[:rescale][:apply_factors];
    cell = m[:smooth].layers[:cell];
    alpha = Ses.alpha(cell);
    theta = Ses.theta(cell);

    num_samples = 200;
    yhat = zeros(Float32, (nskus, horizon));
    trajectories = zeros(Float32, (nskus, horizon, num_samples));


    for nsample in 1:num_samples
        _, state = cell.state;

        for tidx in 1:horizon
            texplanatory = explanatories[tidx];

            tmpr = state .* inv.(theta);
            r = softplus.(tmpr);
            tmpp = theta ./ (1.0f0 .+ theta);
            p = sigmoid.(tmpp);
            label = rand.(NegativeBinomial.(r, p));
            for scaler in scalers
                factors = Ses.get_factors(scaler);
                sidx = Ses.filter_indexes(scaler, texplanatory);
                label = Ses.scale(scaler.mode, label, factors[sidx]);
            end
            tmpsample = smooth(cell.ses_op, label, state, alpha);
            sample = tmpsample;
            for invscaler in invscalers
                factors = Ses.get_factors(invscaler.scaler);
                sidx = Ses.filter_indexes(invscaler, texplanatory);
                sample = Ses.rescale(invscaler.scaler.mode, sample, factors[sidx]);
            end
            state = tmpsample;
            trajectories[:, tidx, nsample] = sample;
        end
    end

    yhat = sum(trajectories, dims = 3) ./ num_samples;

    return trajectories, yhat;
end



function test_loop(test_data::AbstractVecOrMat{T}, m::PES,
                   loss::Function, test_callback::Function) where {T}
    labels = test_data[2];
    explanatories = test_data[3];
    horizon = size(labels)[1];

    _, yhat = test_callback((labels, explanatories), m);
    ls = zeros(Float32, (horizon, 1));
    theta = Ses.theta(m[:smooth].layers[:cell]);
    for idx in 1:horizon
        ls[idx] = sum(loss(yhat[:, idx], labels[idx, :], theta = theta));
    end

    return sum(ls);
end



function train_loop!(data::AbstractVecOrMat{T}, model::PES{M},
                     loss::Function, opt::AbstractOptimiser;
                     logger::Union{L, Nothing} = nothing, verbose = false) where {L, M, T}
    skus = data[1];
    labels = data[2];
    explanatories = data[3];
    train_len = size(labels)[1];

    ps = params(model);

    model((labels[1, :], explanatories[1]));
    model((labels[2, :], explanatories[2]));

    agg_loss = 0.0f0;
    for idx in 3:train_len
        train_loss, back = pullback(ps) do
            labelsi = labels[idx, :];
            explanatoryi = explanatories[idx];
            yhat = model((labelsi, explanatoryi));
            tht = Ses.theta(model[:smooth].layers[:cell]);
            sum(loss(yhat, labelsi, theta=tht));
        end    
        Ses._check_nan(train_loss);

        agg_loss += train_loss;

        Ses.update!(train_loss, ps, back, opt, verbose = verbose);
    end

    if !isnothing(logger)
        logging(model, logger, skus);
    end

    return agg_loss; 
end



function train!(train_vm::AbstractVecOrMat{D}, model::PES{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 main()
    default_args = get_commandline_defaults();
    parsed_args = parse_commandline();

    testing = parsed_args[:test];

    model_location = parsed_args[:model_location];
    model_name = parsed_args[:model_name];
    use_bson = parsed_args[:use_bson];

    if use_bson
        @eval using Flux
    end

    cell = Ses.available_cells[parsed_args[:cell]];
    horizon = parsed_args[:horizon];
    frequencies = parsed_args[:frequencies];
    modes = parsed_args[:modes];
    delimiter = parsed_args[:delimiter];


    
    parts = PartsHelpers.read_parts();
    gdf = groupby(parts, :Date);
    
    trainDates = gdf[[1:end - horizon;]];
    testDates = gdf[[end - horizon + 1:end;]];
    gdf = nothing;
    
    df = Ses.preprocess(trainDates, byDate = true);
    tdf = Ses.preprocess(testDates, byDate = true);
    Ses.add_time_explanatory!(df, modes, frequencies, delimiter);
    Ses.add_time_explanatory!(tdf, modes, frequencies, delimiter);
    trainDates = nothing;
    testDates = nothing;

    Ses._check_columns([:Id, :Qty, :TExplanatory, :TMode], df,
                       message = "[:Id, :Qty, :TExplanatory, :TMode] are mandatory to init PES model.")

    skus = df[1, :].Id;
    train_mat = permutedims(hcat(df.Labels...));
    test_mat = permutedims(hcat(tdf.Labels...));
    group_label = sum(train_mat, dims = 2)[:, 1];
    tgroup_label = sum(test_mat, dims = 2)[:, 1];


    train_texpl = df.TExplanatory;
    test_texpl = tdf.TExplanatory;
    train_data = [skus, train_mat, train_texpl];
    test_data = if testing 
        [skus, test_mat, test_texpl];
    else
        nothing;
    end


    if parsed_args[:train]
        state01 = zeros(Float32, (length(skus), 1));
        state02 = zeros(Float32, (length(skus), 1));
        model = PES(cell, train_texpl[1], df.TMode[1],
                           (state01, state02), length(skus));

        optimizer = ADAM(parsed_args[:learning_rate]);

        tstart = now();
        _, _, log_name = train!(
            train_data, model, train_loop!, loss = Ses.nb_loss,
            opt = optimizer, epochs = parsed_args[:epochs],
            test_vm = test_data, test_loop = test_loop, test_callback = test_callback,
            log = parsed_args[:log], log_prefix = parsed_args[:log_prefix], log_location = parsed_args[:log_location],
            verbose = parsed_args[:verbose]
        );
        tstop = now();
        if !parsed_args[:no_timeit_training]
            @info "Training ended in: $(Dates.canonicalize(tstop - tstart))"
        end

        model_name = if model_name != default_args[:model_name]
            model_name
        else 
            "$(log_name)"
        end;

        mkpath(model_location);
        model_path = joinpath(model_location, model_name * ".model");
        Ses.save_model(model_path, model, use_bson = use_bson);
    else
        model_path = joinpath(model_location, model_name * ".model");
        model = Ses.load_model(model_path, use_bson = use_bson);
    end

    restart!(model);
    szl = size(train_mat);
    train_hat = zeros(Float32, szl);
    train_hat[1, :] = model((train_mat[1, :], train_texpl[1]));
    train_hat[2, :] = model((train_mat[2, :], train_texpl[2]));
    for idx in 2:szl[1]
        train_hat[idx, :] = model((train_mat[idx, :], train_texpl[idx]))
    end
    test_labels = test_mat';

    
    trjs, yhat = test_callback((test_mat, test_texpl), model);
    nskus = size(trjs)[1];
    mquantile = zeros(Float32, (nskus, horizon));
    nquantile = zeros(Float32, (nskus, horizon));
    for sku in 1:nskus
        trajectoryi = trjs[sku, :, :];
        for tidx in 1:horizon
            mquantile[sku, tidx] = quantile(trajectoryi[tidx, :], 0.5, alpha = 0, beta = 0);
            nquantile[sku, tidx] = quantile(trajectoryi[tidx, :], 0.9, alpha = 0, beta = 0);
        end
    end

    min_p50 = 1e6;
    min_p90 = 1e6;
    max_p50 = 1e-16;
    max_p90 = 1e-16;
    min_p50_sku = 1;
    min_p90_sku = 1;
    max_p50_sku = 1;
    max_p90_sku = 1;
    for sku in 1:nskus
        labelsi = test_labels[sku, :];
        mquantilei = mquantile[sku, :];
        nquantilei = nquantile[sku, :];
        ssm_rl_c, ssm_sl_c = Ses.deepSSM_ql(labelsi, mquantilei, 0.5);
        ssm_rl_n, ssm_sl_n = Ses.deepSSM_ql(labelsi, nquantilei, 0.9);
        p50 = ssm_rl_c / ssm_sl_c;
        p90 = ssm_rl_n / ssm_sl_n;

        if p50 >= max_p50
            max_p50 = p50;
            max_p50_sku = sku;
        end

        if p90 >= max_p90
            max_p90 = p90;
            max_p90_sku = sku;
        end


        if p50 <= min_p50
            min_p50 = p50;
            min_p50_sku = sku;
        end
        if p90 <= min_p90
            min_p90 = p90;
            min_p90_sku = sku;
        end
    end

    ssm_rl_c, ssm_sl_c = Ses.deepSSM_ql(test_labels, mquantile, 0.5);
    ssm_rl_n, ssm_sl_n = Ses.deepSSM_ql(test_labels, nquantile, 0.9);
    println("Rho-risk (as presented in Rangapuram et al.)\np50 = $(ssm_rl_c / ssm_sl_c)\tp90 = $(ssm_rl_n / ssm_sl_n)")
    
    if parsed_args[:plot]
        fig_name = if isempty(parsed_args[:fig_name]) model_name else parsed_args[:fig_name] end;
        fig_location = parsed_args[:fig_location];
        mkpath(fig_location);

        
        skus = df[1, :].Id;
        train_dates = df.Date;
        test_dates = tdf.Date;


        minp50_trace = scatter(x=train_dates, y = train_mat[:, min_p50_sku], name="Train", mode="lines");
        test_minp50_trace = scatter(x=test_dates, y = test_labels[min_p50_sku, :], name="Test", mode="lines");
        test_trjs_trace = scatter(x=test_dates, y = mquantile[min_p50_sku, :], line_color="green", colorname="50% int.", mode="lines", dash="dash", fill="tozeroy");
    minp_fig = plot([minp50_trace, test_minp50_trace, test_trjs_trace], Layout(title = "Min. p50 error - $(skus[min_p50_sku])"));
        fig_path = joinpath(fig_location, fig_name * "_minerror_p50" * parsed_args[:fig_extension]);
        savefig(minp_fig, fig_path);


        minp90_trace = scatter(x=train_dates, y = train_mat[:, min_p90_sku], name="Train", mode="lines");
        test_minp90_trace = scatter(x=test_dates, y = test_labels[min_p90_sku, :], name="Test", mode="lines");
        test_trjs_trace = scatter(x=test_dates, y = nquantile[min_p90_sku, :], line_color="purple", name="90% int.", mode="lines", dash="dash", fill="tozeroy");
        minp_fig = plot([minp50_trace, test_minp90_trace, test_trjs_trace], Layout(title = "Min p90 error - $(skus[min_p90_sku])"));
        fig_path = joinpath(fig_location, fig_name * "_minerror_p90" * parsed_args[:fig_extension]);
        savefig(minp_fig, fig_path);


        maxp50_trace = scatter(x=train_dates, y = train_mat[:, max_p50_sku], name="Train", mode="lines");
        test_maxp50_trace = scatter(x=test_dates, y = test_labels[max_p50_sku, :], name="Test", mode="lines");
        test_trjs_trace = scatter(x=test_dates, y = mquantile[max_p50_sku, :], line_color="green", name="50% int.", mode="lines", dash="dash", fill="tozeroy");
        maxp_fig = plot([maxp50_trace, test_maxp50_trace, test_trjs_trace], Layout(title = "Max p50 error - $(skus[max_p50_sku])"));
        fig_path = joinpath(fig_location, fig_name * "_maxerror_p50" * parsed_args[:fig_extension]);
        savefig(maxp_fig, fig_path);


        maxp90_trace = scatter(x=train_dates, y = train_mat[:, max_p90_sku], name="Train", mode="lines");
        test_maxp90_trace = scatter(x=test_dates, y = test_labels[max_p90_sku, :], name="Test", mode="lines");
        test_trjs_trace = scatter(x=test_dates, y = nquantile[max_p90_sku, :], line_color="purple", name="90% int.", mode="lines", dash="dash", fill="tozeroy");
        maxp_fig = plot([maxp90_trace, test_maxp90_trace, test_trjs_trace], Layout(title = "Max p90 error - $(skus[max_p90_sku])"));
        fig_path = joinpath(fig_location, fig_name * "_maxerror_p90" * parsed_args[:fig_extension]);
        savefig(maxp_fig, fig_path);

        gtrain_hat = sum(train_hat, dims = 2)[:, 1];
        gtest_trjs = sum(yhat, dims = 1)[1, :];
        train_group_trace = scatter(x=train_dates, y = group_label[:, 1], name="Train", mode="lines");
        train_hat_trace = scatter(x=train_dates, y = gtrain_hat, name="FTrain", mode="lines", dash="dash");
        test_group_trace = scatter(x=test_dates, y = tgroup_label, name="Test", mode="lines");
        test_trjs_trace = scatter(x=test_dates, y = gtest_trjs[:, 1], name="Trjs", mode="lines", dash="dash");
        agg_fig = plot([train_group_trace, train_hat_trace,
                        test_group_trace, test_trjs_trace]);
        fig_path = joinpath(fig_location, fig_name * "_agg" * parsed_args[:fig_extension]);
        savefig(agg_fig, fig_path);
    end

    return nothing
end



main();