import argparse
from utils._plots import violinplot_compare_metrics,\
    violinplot_scenario_vs_vanilla, violinplot_scenario,\
    params_vanilla_transfer, params_default_transfer,\
    execution_time

def main(scenarios, noise, metrics, graph_type, sample_size, plot_type, density):
    # plot_vanilla_fixed_noise(noise=noise, metric=metric)
    if plot_type == "scenario":
        violinplot_scenario(
            scenarios=scenarios, noise=noise, metrics=metrics, 
            base_fig_dir=f"/home/ec2-user/causal-benchmark/tmp/metrics_plots/scenario/{graph_type}/{noise}",
            graph_type=graph_type
        )
    elif plot_type == "vs-vanilla":
        density = "full" if graph_type == "FC" else density
        violinplot_scenario_vs_vanilla(
            scenarios=scenarios, noise=noise, metrics=metrics,
            base_fig_dir=f"/home/ec2-user/causal-benchmark/tmp/metrics_plots/scenario_vs_vanilla/{graph_type}/{noise}",
            graph_type=graph_type,
            density=density
        )
    elif plot_type == "compare-metrics":
        density = "full" if graph_type == "FC" else density
        violinplot_compare_metrics(
            scenarios=scenarios, noise=noise,
            metrics=metrics, base_fig_dir=f"/home/ec2-user/causal-benchmark/tmp/metrics_plots/compare_metrics/{graph_type}/{noise}",
            graph_type=graph_type,
            density=density
        )
    elif plot_type == "stability":
        density = "full" if graph_type == "FC" else density
        params_vanilla_transfer(
            scenarios=scenarios,
            base_fig_dir=f"/home/ec2-user/causal-benchmark/tmp/metrics_plots/stability/{graph_type}/{sample_size}_samples",
            graph_type=graph_type,
            density=density,
            sample_size=sample_size
        )
    elif plot_type == "default-param":
        density = "full" if graph_type == "FC" else density
        params_default_transfer(
            scenarios=scenarios,
            base_fig_dir=f"/home/ec2-user/causal-benchmark/tmp/metrics_plots/stability/{graph_type}/{sample_size}_samples",
            metric=metrics[0],
            graph_type=graph_type,
            density=density,
            sample_size=sample_size
        )
    elif plot_type == "time":
        density = "full" if graph_type == "FC" else density
        execution_time(
            scenarios=scenarios,
            base_fig_dir=f"/home/ec2-user/causal-benchmark/tmp/metrics_plots/time/{graph_type}/{sample_size}_samples",
            graph_type=graph_type,
            density=density,
            sample_size=sample_size
        )
    else:
        raise ValueError("Unkown graph type")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Dataset generation")
    parser.add_argument(
        "--noise", 
        default="gauss", 
        type=str, 
        help="Noise distribution."
    )
    parser.add_argument(
        "--graph_type", 
        default="ER", 
        type=str, 
        help="Algorithm for generation of synthetic graphs."
    )
    parser.add_argument(
        "--sample_size", 
        default=1000, 
        type=int, 
        help="Algorithm for generation of synthetic graphs."
    )
    parser.add_argument(
        "--density", 
        default="dense",
        type=str, 
        help="Specify density of the graphs used in the plots",
    )
    parser.add_argument(
        "--plot_type", 
        type=str, 
        help="Admitted values: ['scenario', 'vs-vanilla', 'dtop-vs-fnr']",
        required=True
    )
    parser.add_argument(
        "--metrics", 
        nargs='+',
        type=str, 
        help="Metrics of interest.",
    )
    parser.add_argument(
        "--scenarios", 
        nargs='+',
        type=str, 
        help="Experimental scenarios."
    )

    args = parser.parse_args()
    main(args.scenarios, args.noise, args.metrics, args.graph_type, args.sample_size,args.plot_type, args.density)