using HDF5
using Dates
using JLD2
using Distributed
using StatsBase
using Plots
using DPMMSubClusters
using Dates
using MultivariateStats



# DS_PATH = "/vildata/vlad/splitnet/experiments_data/"

# MNIST_PATH = DS_PATH * "mnist_dim_64.h5"
# FMNIST_PATH = DS_PATH * "fmnist_dim_64.h5"
# CIFAR_PATH = DS_PATH * "cifar10_dim_88.h5"
# CIFAR_64_PATH = DS_PATH * "cifar10_efficientnet_dim_64.h5"
# CIFAR100_PATH = DS_PATH * "cifar100.h5"
# SVHN_PATH = DS_PATH * "svhn_dim_64.h5"
# GAUSS_2D_PATH = DS_PATH * "gaussian2d.h5"
# GAUSS_3D_PATH = DS_PATH * "gaussian3d.h5"
# GAUSS_5D_PATH = DS_PATH * "gaussian5d.h5"
# GAUSS_10D_PATH = DS_PATH * "gaussian10d.h5"
# GAUSS_30D_PATH = DS_PATH * "gaussian30d.h5"
# GAUSS_100D_PATH = DS_PATH * "gaussian100d.h5"
# OTTO_PATH = DS_PATH * "otto.h5"
# FROGS_PATH = DS_PATH * "frogs.h5"
# WINE_PATH = DS_PATH * "wine.h5"
# YEAST_PATH = DS_PATH * "yeast.h5"
# BEES_PATH = DS_PATH * "bees.h5"

# IMAGENET_SWAV = DS_PATH * "imagenet_swav_D=128_K=1000.h5"

EXPERIMENTS_DIR = "/vildata/vlad/splitnet/experiments_results/"

# datasets = Dict(
#     :gauss2d => GAUSS_2D_PATH,
#     :gauss3d => GAUSS_3D_PATH,
#     :gauss5d => GAUSS_5D_PATH,
#     :gauss10d => GAUSS_10D_PATH,
#     :gauss30d => GAUSS_30D_PATH,
#     :gauss100d => GAUSS_100D_PATH,
#     :mnist_d64 => MNIST_PATH,
#     :fmnist_d64 => FMNIST_PATH,
#     :cifar_d88 => CIFAR_PATH,
#     :cifar_d64 => CIFAR_64_PATH,
#     :cifar100 => CIFAR100_PATH,
#     :svhn_d64 => SVHN_PATH,
#     :otto => OTTO_PATH,
#     :frogs => FROGS_PATH,
#     :wine => WINE_PATH,
#     :yeast => YEAST_PATH,
#     :bees => BEES_PATH,
#     :imagenet_swav_d128 => IMAGENET_SWAV,
# )

# function plot_dp_2d(pts, labels)
#     plt=Plots.plot()
#     Plots.plot!(pts[1,:],pts[2,:], seriestype=:scatter, color = Int64.(labels), markersize = 3, markerstrokewidth = 0.5)
#     return plt
# end

function load_data(path)
    data = HDF5.h5read(path, "data")
    gt = HDF5.h5read(path, "labels")
    return data, gt
end




function run_and_save_experiment_xy(
    data, 
    gt,
    alpha,
    init_type,
    burnout_period = 5,
    iterations = 1000,
    save = false,
    verbose = false;
    hyperparam = nothing,
)
    global datasets
    global EXPERIMENTS_DIR

    if hyperparam !== nothing
        labels, clusters, weights, timing, nmi, ari, ll, cluster_hist, sub_clstr = DPMMSubClusters.fit(
            Float32.(data),
            hyperparam,
            Float32(alpha),
            iters = iterations,
            burnout = burnout_period,
            init_type = init_type,
            gt = gt,
            verbose = verbose,
        )
    else
        labels, clusters, weights, timing, nmi, ari, ll, cluster_hist, sub_clstr = DPMMSubClusters.fit(
            Float32.(data),
            Float32(alpha),
            iters = iterations,
            burnout = burnout_period,
            init_type = init_type,
            gt = gt,
            verbose = verbose,
        )
    end

    new_path = EXPERIMENTS_DIR * "_temp_" * "/" *
               string(init_type)
    Base.Filesystem.mkpath(new_path)
    file_path = string(new_path * string(now()) * string(".h5"))
    ispath(file_path) && rm(file_path)
    if save
        @save file_path nmi labels ll timing
    end 

    result = Dict([
        "timing" => timing,
        "ll" => ll,
        "nmi" => nmi,
        "ari" => ari,
        "labels" => cluster_hist,
    ])
    # return nmi, ari, ll, timing, cluster_hist
    return result
end




function repeat_experiment_xy(
    data,
    gt,
    alpha,
    times,
    init_type,
    burnout_period = 5,
    iterations = 200,
    save = false,
    verbose = false;
    hyperparam = nothing,
)
    
    # hyperparam::DPMMSubClusters.distribution_hyper_params = nothing,

    result = run_and_save_experiment_xy(
        data,
        gt,
        alpha,
        init_type,
        burnout_period,
        iterations,
        save,
        verbose,
        hyperparam=hyperparam,
        )
        
        timing = result["timing"]
        nmi = result["nmi"]
        ari = result["ari"]
        ll = result["ll"]
        cluster_num = result["labels"]
        
        stds = []
    stds_nmi = []
    stds_ari = []
    cluster_nums = []

    push!(stds, ll)
    push!(stds_nmi, nmi)
    push!(stds_ari, ari)
    push!(cluster_nums, cluster_num)
    
    for j = 1:times-1
        result = run_and_save_experiment_xy(
            data,
            gt,
            alpha,
            init_type,
            burnout_period,
            iterations,
            save,
            verbose,
            hyperparam=hyperparam,
        )

        timing2 = result["timing"]
        nmi2 = result["nmi"]
        ari2 = result["ari"]
        ll2 = result["ll"]
        cluster_num2 = result["labels"]

        timing += timing2
        ll += ll2
        nmi += nmi2
        ari += ari2

        push!(stds, ll2)
        push!(stds_nmi, nmi2)
        push!(stds_ari, ari2)
        push!(cluster_nums, cluster_num2)

    end

        
    ll_mean = ll / (times)
    stds = [(x - ll_mean).^2 for x in stds]
    ll_std = sqrt.(mean(stds))

    nmi_mean = nmi / (times)
    stds_nmi = [(x - nmi_mean).^2 for x in stds_nmi]
    nmi_std = sqrt.(mean(stds_nmi))

    k_mean = mean(cluster_nums) 
    k_std = StatsBase.std(cluster_nums)
    
    ari_mean = mean(stds_ari) 
    ari_std = StatsBase.std(stds_ari)


    result = Dict(
        "timing" => cumsum(timing / times),
        "ll_mean" => ll_mean,
        "ll_std" => ll_std,
        "nmi_mean" => nmi_mean,
        "nmi_std" => nmi_std,
        "k_mean" => k_mean,
        "k_std" => k_std,
        "ari_mean" => ari_mean,
        "ari_std" => ari_std,
    )
    return result
end



# function run_and_save_experiment(
#     experiment::Symbol,
#     init_type,
#     alpha,
#     burnout_period = 5,
#     iterations = 1000,
#     save = false,
#     verbose = false,
# )
#     global datasets
#     global EXPERIMENTS_DIR

#     data, gt = load_data(datasets[experiment])
#     labels, clusters, weights, nmi, ll, timing, cluster_hist = DPMMSubClusters.fit(
#         Float32.(data),
#         Float32(alpha),
#         iters = iterations,
#         burnout = burnout_period,
#         init_type = init_type,
#         gt = gt,
#         verbose = verbose,
#     )
#     new_path = EXPERIMENTS_DIR * string(experiment) * "/" *
#                string(init_type)
#     Base.Filesystem.mkpath(new_path)
#     file_path = string(new_path * string(now()) * string(".h5"))
#     ispath(file_path) && rm(file_path)
#     if save
#         @save file_path nmi labels ll timing
#     end

#     return nmi, ll, timing, cluster_hist
# end


# function repeat_experiment(
#     experiment::Symbol,
#     times,
#     init_type,
#     alpha,
#     burnout_period = 5,
#     iterations = 200,
#     save = false,
#     verbose = false,
# )
#     timing, nmi, ll, cluster_num = run_and_save_experiment(
#         experiment,
#         init_type,
#         alpha,
#         burnout_period,
#         iterations,
#         save,
#         verbose,
#     )
#     stds = []
#     stds_nmi = []
#     cluster_nums = []
#     push!(stds, ll)
#     push!(stds_nmi, nmi)
#     push!(cluster_nums, cluster_num)
#     for j = 1:times-1
#         timing2, nmi2, ll2, cluster_num2 = run_and_save_experiment(
#             experiment,
#             init_type,
#             alpha,
#             burnout_period,
#             iterations,
#             save,
#             verbose
#             )
#             timing += timing2
#             ll += ll2
#             nmi += nmi2
#             push!(stds, ll2)
#             push!(stds_nmi, nmi2)
#             push!(cluster_nums, cluster_num2)
#         end

        
#     ll_mean = ll / (times)
#     stds = [(x - ll_mean).^2 for x in stds]
#     ll_std = sqrt.(mean(stds))

#     nmi_mean = nmi / (times)
#     stds_nmi = [(x - nmi_mean).^2 for x in stds_nmi]
#     nmi_std = sqrt.(mean(stds_nmi))

#     k_mean = mean(cluster_nums) 
#     k_std = StatsBase.std(cluster_nums)


#     result = Dict(
#         "timing" => cumsum(timing / times),
#         "ll_mean" => ll_mean,
#         "ll_std" => ll_std,
#         "nmi_mean" => nmi_mean,
#         "nmi_std" => nmi_std,
#         "k_mean" => k_mean,
#         "k_std" => k_std,
#     )
#     return result
# end

function plot_and_save_all(all_results, save_path)
    nmi_plot = plot_key_for_iters(all_results, "nmi", "NMI vs. Iterations")
    path = string(save_path) * "k" * string(K) * "_v" * string(Var) * "_nmi.png"
    Plots.savefig(nmi_plot, path)

    ari_plot = plot_key_for_iters(all_results, "ari", "ARI vs. Iterations")
    path = string(save_path) * "k" * string(K) * "_v" * string(Var) * "_ari.png"
    Plots.savefig(ari_plot, path)

    ll_plot = plot_key_for_iters(all_results, "ll", "Log-Likelihood vs. Iterations")
    path = string(save_path) * "k" * string(K) * "_v" * string(Var) * "_ll.png"
    Plots.savefig(ll_plot, path)

    k_plot = plot_key_for_iters(all_results, "k", "Inferred K vs. Iterations")
    path = string(save_path) * "k" * string(K) * "_v" * string(Var) * "_K.png"
    Plots.savefig(k_plot, path)
end

function plot_and_save_all_ds(all_results, save_path, ds_name)
    nmi_plot = plot_key_for_iters(all_results, "nmi", "NMI vs. Iterations")
    path = string(save_path) * "ds_" * string(ds_name) * "_nmi.png"
    Plots.savefig(nmi_plot, path)

    ari_plot = plot_key_for_iters(all_results, "ari", "ARI vs. Iterations")
    path = string(save_path) * "ds_" * string(ds_name) * "_ari.png"
    Plots.savefig(ari_plot, path)

    ll_plot = plot_key_for_iters(all_results, "ll", "Log-Likelihood vs. Iterations")
    path = string(save_path) * "ds_" * string(ds_name) * "_ll.png"
    Plots.savefig(ll_plot, path)

    k_plot = plot_key_for_iters(all_results, "k", "Inferred K vs. Iterations")
    path = string(save_path) * "ds_" * string(ds_name) * "_K.png"
    Plots.savefig(k_plot, path)
end

function plot_data(pts, labels)
    if size(x)[1] > 2
        M = MultivariateStats.fit(PCA, pts; maxoutdim=2) 
        pts = MultivariateStats.transform(M, pts)
    end
    plt=Plots.plot()
    Plots.plot!(pts[1,:],pts[2,:], seriestype=:scatter, color = Int64.(labels), markersize = 3, markerstrokewidth = 0.5, legend = false)
    return plt
end


function plot_key_for_iters(results::Dict, key="", title="", save_fig=false, for_time=false)
    fig = plot()
    for (name, res) in results

        if for_time
            x = res["timing"][1:end]
        else
            x = 1:length(results[name]["timing"])
        end
        plot!(
            fig,
            # 1:length(results[name]["timing"]),
            x,
            res[string(key)*"_mean"][1:end],
            ribbon = res[string(key)*"_std"][1:end],
            label = name,
            #markershape = :diamond,
            guidefontsize = 9,
            legend = :bottomright
        )
    end
    xlabel!("Iterations")
    ylabel!(string(key))
    if title != ""
        title!(string(title))
    else
        title!(string(key)*" vs. Iterations")
    end

    if save_fig
        fig_stamp = string(Dates.format(now(), "yyyy-mm-dd-HH:MM"))
        savefig(EXPERIMENTS_DIR*"/plots/"*fig_stamp*"-ll_plot.png")
    end
    return fig
end

# function plot_ll(results, partial = nothing)
#     gr()
#     final_index = (partial == nothing) ? length(results[1]["timing"]) : partial
#     plot(
#         eigen[1][1:final_index],
#         eigen[2][1:final_index],
#         ribbon = eigen[3][1:final_index],
#         label = "splitnet",
#         xaxis = :log,
#         markershape = :diamond,
#         guidefontsize = 9,
#         legend = :bottomright
#     )
#     plot!(random[1][1:final_index], random[2][1:final_index], ribbon = random[3][1:final_index], label = "vanilla", markershape = :star5)
#     fig_stamp = string(Dates.format(now(), "yyyy-mm-dd-HH:MM"))
#     savefig(EXPERIMENTS_DIR*"/plots/"*fig_stamp*"-ll_plot.png")
# end


# function plot_nmi_for_time(results::Dict, title="", save_fig=false)
#     fig = plot()
#     for (name, res) in results

#         plot!(
#             fig,
#             res["timing"][1:end],
#             res["nmi_mean"][1:end],
#             ribbon = res["nmi_std"][1:end],
#             label = name,
#             xaxis = :log,
#             #markershape = :diamond,
#             guidefontsize = 9,
#             legend = :bottomright
#         )
#     end
#     xlabel!("Time [ms?]")
#     ylabel!("NMI")
#     title!("NMI vs. Time | " *string(title))

#     if save_fig
#         fig_stamp = string(Dates.format(now(), "yyyy-mm-dd-HH:MM"))
#         savefig(EXPERIMENTS_DIR*"/plots/"*fig_stamp*"-nmi_plot.png")
#     end
#     return fig
# end


# function plot_nmi_for_iters(results::Dict, title="", save_fig=false)
#     fig = plot()
#     for (name, res) in results

#         plot!(
#             fig,
#             1:length(results[name]["timing"]),
#             res["nmi_mean"][1:end],
#             ribbon = res["nmi_std"][1:end],
#             label = name,
#             #markershape = :diamond,
#             guidefontsize = 9,
#             legend = :bottomright
#         )
#     end
#     xlabel!("Iterations")
#     ylabel!("NMI")
#     title!("NMI vs. Iterations | " *string(title))

#     if save_fig
#         fig_stamp = string(Dates.format(now(), "yyyy-mm-dd-HH:MM"))
#         savefig(EXPERIMENTS_DIR*"/plots/"*fig_stamp*"-nmi_plot.png")
#     end
#     return fig
# end


# function plot_ll_for_iters(results::Dict, title="", save_fig=false)
#     fig = plot()
#     for (name, res) in results

#         plot!(
#             fig,
#             1:length(results[name]["timing"]),
#             res["ll_mean"][1:end],
#             ribbon = res["ll_std"][1:end],
#             label = name,
#             #markershape = :diamond,
#             guidefontsize = 9,
#             legend = :bottomright
#         )
#     end
#     xlabel!("Iterations")
#     ylabel!("Log Likelihood")
#     title!("Log Likelihood vs. Iterations | " *string(title))

#     if save_fig
#         fig_stamp = string(Dates.format(now(), "yyyy-mm-dd-HH:MM"))
#         savefig(EXPERIMENTS_DIR*"/plots/"*fig_stamp*"-ll_plot.png")
#     end
#     return fig
# end



# function plot_K_for_iters(results::Dict, title="", save_fig=false)
#     fig = plot()
#     for (name, res) in results

#         plot!(
#             fig,
#             1:length(results[name]["timing"]),
#             res["k_mean"][1:end],
#             ribbon = res["k_std"][1:end],
#             label = name,
#             #markershape = :diamond,
#             guidefontsize = 9,
#             legend = :bottomright
#         )
#     end
#     xlabel!("Iterations")
#     ylabel!("Number of  Cluster")
#     title!("Numeber of Clusters vs. Iterations | " *string(title))

#     if save_fig
#         fig_stamp = string(Dates.format(now(), "yyyy-mm-dd-HH:MM"))
#         savefig(EXPERIMENTS_DIR*"/plots/"*fig_stamp*"-K_plot.png")
#     end
#     return fig
# end

# cifar_smart = repeat_experiment(:gauss2d,2,"splitnet",3, 30)
# cifar_random = repeat_experiment(:gauss2d,2,"none",3, 30)
# # cifar_random = repeat_experiment(:gauss3d,5,false,7, 300)
# plot_nmi(cifar_smart,cifar_random)
# plot_ll(cifar_smart,cifar_random)
# #
# #
#
# subc_nmi = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.07959427, 0.07970354, 0.07944535, 0.07849897, 0.0794486, 0.07982686, 0.1290593, 0.1368891, 0.139273, 0.1388824, 0.1227153, 0.11848, 0.1147086, 0.1127466, 0.1142877, 0.1159546, 0.1161205, 0.1170728, 0.1185645, 0.1204665, 0.1222847, 0.1242379, 0.1357441, 0.1394681, 0.1438078, 0.1436123, 0.1433871, 0.1437167, 0.1425898, 0.1559395, 0.1564514, 0.1560872, 0.1570587, 0.1566021, 0.1574379, 0.1595252, 0.1581669, 0.1597728, 0.1617535, 0.1628309, 0.1629995, 0.1634077, 0.1643861, 0.1635752, 0.1634904, 0.1632213, 0.1630481, 0.1622657, 0.1637352, 0.1619947, 0.1618271, 0.1619837, 0.1607221, 0.1613915, 0.1591221, 0.159582, 0.161543, 0.1605444, 0.1602328, 0.1605289, 0.160529, 0.1613139, 0.1598319, 0.1592019, 0.1595851, 0.1605065, 0.1599539, 0.1601576, 0.1608793, 0.161478, 0.1595351, 0.158717, 0.1594852, 0.1605048, 0.1605593, 0.1598349, 0.1592838, 0.1595072, 0.1597103, 0.1606406, 0.1601267, 0.1607425, 0.1602751, 0.1606777, 0.1582134, 0.1591221, 0.1594418, 0.159464, 0.1596887, 0.1592991, 0.1608755, 0.1594977, 0.159058, 0.1593717, 0.1580996]
#
# sub_timin

# using Clustering

# y_k = kmeans(x, 2).assignments

# dec=20
# plot_dp_2d(x[:, 1:dec:end],y_k[1:dec:end])


# y_s = DPMMSubClusters.torch_model_2d.infer(x)
# preds = argmax(y_s, dims=2)
# split_labels =  dropdims([ i[2] for i in preds ], dims=2)

# dec=20
# plot_dp_2d(x[:, 1:dec:end],split_labels[1:dec:end])

# 
# random_labels,clusters,weights = DPMMSubClusters.fit(x,100.0, iters = 1000, seed = 1234, gt=labels, init_type="kmeans", burnout=5)

