using Random
using DPMMSubClusters
using Glob
import YAML
using Pandas
include("run_experiments.jl")



# load config:
cfg = YAML.load_file("/home/vlad/projects/DPMMSubClusters.jl/benchmarks/configs/basic.yaml")

results_dir = ARGS[1]

# Extract params:
Ns = cfg["data"]["Ns"]
Ds = cfg["data"]["Ds"]
Ks = cfg["data"]["Ks"]
SEED = cfg["data"]["seed"]
DS_PATH = cfg["data"]["path"]
data_path = DS_PATH * "seed_" * string(SEED)
real_datasets = cfg["data"]["real_datasets"]

REPEATS = cfg["models"]["repeats"]
BURNOUT = cfg["models"]["burnout"]
ITERS = cfg["models"]["iterations"]
ALPHA = cfg["models"]["alpha"]

Random.seed!(SEED)


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




# # Synthetic Data:

# iterate over number of points:
for l in 1:length(Ns)
    N = Ns[l]

    # iterate over dimensions:
    for i in 1:length(Ds)
        D = Ds[i]

        # iterate over number of clusters
        for j in 1:length(Ks)
            K = Ks[j]

            filename = "*D=" * string(D) * "*K=" * string(K) *"*.h5"

            # get path:
            datapath = glob(filename, data_path)[1]
            print("loading from: ")
            println(datapath)
            
            # load data:
            x, y = load_data(datapath)
            
            # run methods:
            

            # Random:
            # check if results exist:
            csv_results_random = "DPMM_D=" * string(D) * "_K=" * string(K) * "_results_random.csv"
            if length(glob(csv_results_random, data_path * "/results/")) > 0
                println("Found existing results for: " * csv_results_random)
            else
                println("No existing results for: " * csv_results_random * ", running inference:")
                
                res_random = repeat_experiment_xy(x, y, ALPHA, REPEATS,"none", BURNOUT, ITERS)
                # save results
                df_random = DataFrame(res_random)
                df_random["iter"] = 1:length(df_random)
                df_random["method"] = "random"

                save_path =  data_path * "/results/" * csv_results_random
                println("Saving results at: " * save_path)
                to_csv(df_random, save_path)
            end

            # kmeans
            csv_results_kmeans = "DPMM_D=" * string(D) * "_K=" * string(K) * "_results_kmeans.csv"
            if length(glob(csv_results_kmeans, data_path * "/results/")) > 0
                println("Found existing results for: " * csv_results_kmeans)
            else
                println("No existing results for: " * csv_results_kmeans * ", running inference:")
                
                res_kmeans = repeat_experiment_xy(x, y, ALPHA, REPEATS,"kmeans", BURNOUT, ITERS)
                # save results
                
                df_kmeans = DataFrame(res_kmeans)
                df_kmeans["iter"] = 1:length(df_kmeans)
                df_kmeans["method"] = "kmeans"

                save_path =  data_path * "/results/" * csv_results_kmeans
                println("Saving results at: " * save_path)
                to_csv(df_kmeans, save_path)
            end
            
            # Splitnet 
            if D == 2
                res_splitnet = repeat_experiment_xy(x, y, ALPHA, REPEATS, "splitnet_2d", BURNOUT, ITERS)
            elseif D == 10
                res_splitnet = repeat_experiment_xy(x, y, ALPHA, REPEATS, "splitnet_10d", BURNOUT, ITERS)
            elseif D == 128
                res_splitnet = repeat_experiment_xy(x, y, ALPHA, REPEATS, "splitnet_128d", BURNOUT, ITERS)
            end
            
            
            
            df_splitnet = DataFrame(res_splitnet)
            df_splitnet["iter"] = 1:length(df_splitnet)
            df_splitnet["method"] = "splitnet"
            
            
            csv_save_path = results_dir * "/DPMM_D=" * string(D) * "_K=" * string(K) * "_results_splitnet.csv"
            to_csv(df_splitnet, csv_save_path)
            print("Julia: Saved results to: ")
            println(csv_save_path)
            
        end
    end
end


# "Real" Datasets:
for i in 1:length(real_datasets)
    ds_type = real_datasets[i]
    ds_name = ds_type["ds_name"]
    ds_path_r = DS_PATH * "/real_datasets/" * ds_type["path"]
    
    # load data:
    x, y = load_data(ds_path_r)
    
    # Random:
    # check if results exist:
    csv_results_random = "DPMM_" * string(ds_name) * "_results_random.csv"
    if length(glob(csv_results_random, DS_PATH * "/real_datasets/results/")) > 0
        println("Found existing results for: " * csv_results_random)
    else
        println("No existing results for: " * csv_results_random * ", running inference:")
        
        res_random = repeat_experiment_xy(x, y, ALPHA, 1,"none", 4, 10)
        # res_random = repeat_experiment_xy(x, y, ALPHA, REPEATS,"none", BURNOUT, ITERS)
        # save results
        df_random = DataFrame(res_random)
        df_random["iter"] = 1:length(df_random)
        df_random["method"] = "random"

        save_path =  DS_PATH * "/real_datasets/results/" * csv_results_random
        println("Saving results at: " * save_path)
        to_csv(df_random, save_path)
    end

    # kmeans
    csv_results_kmeans = "DPMM_" * string(ds_name) * "_results_kmeans.csv"
    if length(glob(csv_results_kmeans, DS_PATH * "/real_datasets/results/")) > 0
        println("Found existing results for: " * csv_results_kmeans)
    else
        println("No existing results for: " * csv_results_kmeans * ", running inference:")
        
        res_kmeans = repeat_experiment_xy(x, y, ALPHA, 1,"none", 4, 10)
        # res_kmeans = repeat_experiment_xy(x, y, ALPHA, REPEATS,"kmeans", BURNOUT, ITERS)
        
        # save results        
        df_kmeans = DataFrame(res_kmeans)
        df_kmeans["iter"] = 1:length(df_kmeans)
        df_kmeans["method"] = "kmeans"

        save_path =  DS_PATH * "/real_datasets/results/" * csv_results_kmeans
        println("Saving results at: " * save_path)
        to_csv(df_kmeans, save_path)
    end

    # it's always D=128
    res_splitnet = repeat_experiment_xy(x, y, ALPHA, 1,"splitnet_128d", 4, 10)
    
    # save results        
    df_splitnet = DataFrame(res_splitnet)
    df_splitnet["iter"] = 1:length(df_splitnet)
    df_splitnet["method"] = "splitnet"

    save_path =  results_dir * "/DPMM_" * string(ds_name) * "_results_splitnet.csv"
    println("Saving results at: " * save_path)
    to_csv(df_splitnet, save_path)


end



# REPEATS=1
# BURNOUT=5


# ITERS=40
# # ITERS=5
# Ks = 5
# x,y,clusters = generate_gaussian_data(10^5,2,Ks,80.0)


# run(`rm -rf $dir`)


# run(`mkdir $dir`)
# # labels,clusters,weights = fit(x,50.0, iters = ITERS, seed = SEED, gt=y, init_type="splitnet-gpu", burnout=BURNOUT)
# gauss2d_splitnet = repeat_experiment_xy(x,y,REPEATS,"splitnet_2d",BURNOUT, ITERS)
# # run(`convert -delay 50 -loop 0 ./test_dir/clusters_*.png ./compare_methods_gifs/clusters_K-"$Ks"_B-"$BURNOUT"_seed_"$SEED"_splitnet.gif`)
# run(`rm -rf $dir`)

# run(`mkdir $dir`)
# gauss2d_random = repeat_experiment_xy(x,y,REPEATS,"none",BURNOUT, ITERS)
# # labels,clusters,weights = DPMMSubClusters.fit(x,100.0, iters = ITERS, seed = SEED, gt=y, init_type="none", burnout=BURNOUT)
# # run(`convert -delay 50 -loop 0 ./test_dir/clusters_*.png ./compare_methods_gifs/clusters_K-"$Ks"_B-"$BURNOUT"_seed_"$SEED"_random.gif`)
# run(`rm -rf $dir`)

# run(`mkdir $dir`)
# # # labels,clusters,weights = fit(x,50.0, iters = ITERS, seed = SEED, gt=y, init_type="kmeans", burnout=BURNOUT)
# gauss2d_kmeans = repeat_experiment_xy(x,y,REPEATS,"kmeans",BURNOUT, ITERS)
# # run(`convert -delay 50 -loop 0 ./test_dir/clusters_*.png ./compare_methods_gifs/clusters_K-"$Ks"_B-"$BURNOUT"_seed_"$SEED"_kmeans.gif`)
# run(`rm -rf $dir`)




# all_results = Dict(
#     "rabdom" => gauss2d_random,
#     "kmeans" => gauss2d_kmeans,
#     "splitnet_gpu" => gauss2d_splitnet,
# )

# title = "K-" * string(Ks) * "_Burnout-" * string(BURNOUT) * "seed-" * string(SEED) * "_model_"*string(model)


# fig = plot_nmi_for_iters(all_results, title)
# savefig("./compare_methods_gifs/result_"* string(title) *".png")



# ITERS=40
# Ks = 10
# x,y,clusters = generate_gaussian_data(10^5,2,Ks,80.0)


# run(`mkdir $dir`)
# # labels,clusters,weights = fit(x,50.0, iters = ITERS, seed = SEED, gt=y, init_type="splitnet-gpu", burnout=BURNOUT)
# gauss2d_splitnet = repeat_experiment_xy(x,y,REPEATS,"splitnet-gpu",BURNOUT, ITERS)
# run(`convert -delay 50 -loop 0 ./test_dir/clusters_*.png ./compare_methods_gifs/clusters_K-"$Ks"_B-"$BURNOUT"_seed_"$SEED"_splitnet.gif`)
# run(`rm -rf $dir`)

# run(`mkdir $dir`)
# gauss2d_random = repeat_experiment_xy(x,y,REPEATS,"none",BURNOUT, ITERS)
# # labels,clusters,weights = DPMMSubClusters.fit(x,100.0, iters = ITERS, seed = SEED, gt=y, init_type="none", burnout=BURNOUT)
# run(`convert -delay 50 -loop 0 ./test_dir/clusters_*.png ./compare_methods_gifs/clusters_K-"$Ks"_B-"$BURNOUT"_seed_"$SEED"_random.gif`)
# run(`rm -rf $dir`)


# run(`mkdir $dir`)
# # # labels,clusters,weights = fit(x,50.0, iters = ITERS, seed = SEED, gt=y, init_type="kmeans", burnout=BURNOUT)
# gauss2d_kmeans = repeat_experiment_xy(x,y,REPEATS,"kmeans",BURNOUT, ITERS)
# run(`convert -delay 50 -loop 0 ./test_dir/clusters_*.png ./compare_methods_gifs/clusters_K-"$Ks"_B-"$BURNOUT"_seed_"$SEED"_kmeans.gif`)
# run(`rm -rf $dir`)




# all_results = Dict(
#     "rabdom" => gauss2d_random,
#     "kmeans" => gauss2d_kmeans,
#     "splitnet_gpu" => gauss2d_splitnet,
# )

# title = "K-" * string(Ks) * "_Burnout-" * string(BURNOUT) * "seed-" * string(SEED) * "_model_"*string(model)


# fig = plot_nmi_for_iters(all_results, title)
# savefig("./compare_methods_gifs/result_"* string(title) *".png")



# ITERS=100
# Ks = 40
# x,y,clusters = generate_gaussian_data(2*10^5,2,Ks,160.0)


# run(`mkdir $dir`)
# # labels,clusters,weights = fit(x,50.0, iters = ITERS, seed = SEED, gt=y, init_type="splitnet-gpu", burnout=BURNOUT)
# gauss2d_splitnet = repeat_experiment_xy(x,y,REPEATS,"splitnet-gpu",BURNOUT, ITERS)
# run(`convert -delay 50 -loop 0 ./test_dir/clusters_*.png ./compare_methods_gifs/clusters_K-"$Ks"_B-"$BURNOUT"_seed_"$SEED"_splitnet.gif`)
# run(`rm -rf $dir`)

# run(`mkdir $dir`)
# gauss2d_random = repeat_experiment_xy(x,y,REPEATS,"none",BURNOUT, ITERS)
# # labels,clusters,weights = DPMMSubClusters.fit(x,100.0, iters = ITERS, seed = SEED, gt=y, init_type="none", burnout=BURNOUT)
# run(`convert -delay 50 -loop 0 ./test_dir/clusters_*.png ./compare_methods_gifs/clusters_K-"$Ks"_B-"$BURNOUT"_seed_"$SEED"_random.gif`)
# run(`rm -rf $dir`)


# run(`mkdir $dir`)
# # # labels,clusters,weights = fit(x,50.0, iters = ITERS, seed = SEED, gt=y, init_type="kmeans", burnout=BURNOUT)
# gauss2d_kmeans = repeat_experiment_xy(x,y,REPEATS,"kmeans",BURNOUT, ITERS)
# run(`convert -delay 50 -loop 0 ./test_dir/clusters_*.png ./compare_methods_gifs/clusters_K-"$Ks"_B-"$BURNOUT"_seed_"$SEED"_kmeans.gif`)
# run(`rm -rf $dir`)




# all_results = Dict(
#     "rabdom" => gauss2d_random,
#     "kmeans" => gauss2d_kmeans,
#     "splitnet_gpu" => gauss2d_splitnet,
# )

# title = "K-" * string(Ks) * "_Burnout-" * string(BURNOUT) * "seed-" * string(SEED) * "_model_"*string(model)


# fig = plot_nmi_for_iters(all_results, title)
# savefig("./compare_methods_gifs/result_"* string(title) *".png")