if __name__ == "__main__":
    import os
    import datetime
    import time
    from importlib.metadata import version
    import carbatpy as cb
    import pandas as pd

    from carbatpy.optimizations.helpers_optimization import (
        extract_boundary_with_path,
        opti_Problem,
        CombinedTermination,
        get_data_from_file,
        save_data,
    )

    from pymoo.algorithms.moo.nsga2 import NSGA2

    from pymoo.optimize import minimize
    from pymoo.operators.survival.rank_and_crowding import RankAndCrowding
    from pymoo.operators.sampling.lhs import LHS

    from pymoo.visualization.scatter import Scatter

    # Multiprocessing
    import multiprocessing
    from multiprocessing import freeze_support
    from pymoo.parallelization.starmap import StarmapParallelization

    freeze_support()

    # %%
    # Variables
    # Flags General
    save_res = True
    save_history = True
    verbose = True
    plotting = False

    # Optimization Flags
    same_fluid = True
    spp_machines = True
    hs_water = False
    wf_butane = False

    # power and dt_subcool since they deviate from standard io-data
    power = 1e6
    dt_subcool = 0.5

    p_high_bnds = [10e5, 40e5]
    p_low_orc_bnds = [1e5, 10e5]
    dt_min_bnds = [2.0, 20.0]

    # Optimization
    n_processes = 44

    fun_to_optimize = "cb"
    to_optimize = ["RTE", "CAPEX"]
    n_ieq_constr = 3

    pop_size = 200
    n_offspring = None
    sampling = LHS(iterations=50)
    survival = RankAndCrowding(crowding_func="pcd")  # Encouraged by pymoo for n_obj=2

    n_gen = 3000
    period = 50
    ftol = 1e-4

    # %%
    # Directories
    if wf_butane:
        fluidscreening_file = os.path.join("Fluidscreening", "butane_results.yaml")
    else:
        fluidscreening_file = os.path.join("Fluidscreening", "hexane_results.yaml")

    dir_data_opti = cb.CB_DEFAULTS["General"]["CB_DATA"]
    dir_hp_config = os.path.join(dir_data_opti, "io-hp-data.yaml")
    dir_orc_config = os.path.join(dir_data_opti, "io-orc-data.yaml")
    dir_fluidscreening = os.path.join(
        os.environ["CARBATPY_RES_DIR"], fluidscreening_file
    )
    dir_res_root = os.path.join(os.environ["CARBATPY_RES_DIR"], "opti_ma_dreis")

    if save_res:
        date_str = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
        dir_results = os.path.join(dir_res_root, f"cb_opti_nsga2_{date_str}")
        print(f"Create results directory: {dir_results}")
        os.makedirs(dir_results, exist_ok=True)

    # %%
    # Fluidscreening
    fluidscreening = get_data_from_file(dir_fluidscreening)

    if fluidscreening is not None:
        print("Using values from Fluidscreening")
        workingfluid = fluidscreening["parameters"]["working_fluid"]
        frac_bounds = fluidscreening["results"]["max_values"]
        dt_superheat_low = fluidscreening["parameters"]["dt_superheat"]
        print(f"Workingfluid from Fluidscreening: {workingfluid}")
    else:
        print("No fluidscreening found, using standard manually defined values")
        workingfluid = "pentane * butane * propane * ethane"
        frac_bounds = [[0, 1], [0, 1], [0, 1], [0, 1]]
        dt_superheat_low = 1.0

    frac_bounds = frac_bounds[:-1]  # last fraction gets calculated automatically

    # Calc type of machines
    if spp_machines:
        calc_type_compressor = "turbo * RadialCompressor"
        calc_type_expander = "turbo * RadialTurbine"
    else:
        calc_type_compressor = "const_eta"
        calc_type_expander = "const_eta"

    # Hot Storage Data
    if hs_water:
        fluid_hs = "water"
        temp_hs_high = 368.15
        print("Using water in Hot Storage")
    else:
        fluid_hs = "MD4M"
        temp_hs_high = 403.15
        print("Using MD4M in Hot Storage")

    # Config Data
    hp_config = get_data_from_file(dir_hp_config)
    orc_config = get_data_from_file(dir_orc_config)

    hp_config["process"]["fixed"]["compressor"]["power"] = power
    hp_config["working_fluid"]["species"] = workingfluid
    hp_config["compressor"]["calc_type"] = calc_type_compressor
    hp_config["hot_storage"]["species"] = fluid_hs
    hp_config["hot_storage"]["temp_high"] = temp_hs_high

    orc_config["start"]["dt_subcool"] = dt_subcool
    orc_config["working_fluid"]["species"] = workingfluid
    orc_config["expander"]["calc_type"] = calc_type_expander
    orc_config["hot_storage"]["species"] = fluid_hs
    orc_config["hot_storage"]["temp_high"] = temp_hs_high

    dir_combined = {"hp": hp_config, "orc": orc_config}

    # Define Boundaries
    bounds_hp = {
        "start": {"dt_superheat": [dt_superheat_low, 5.0]},
        "condenser": {"dt_min": dt_min_bnds},
        "evaporator": {"dt_min": dt_min_bnds},
        "working_fluid": {"p_high": p_high_bnds, "fractions": frac_bounds},
    }

    bounds_orc = {
        "condenser": {"dt_min": dt_min_bnds},
        "evaporator": {"dt_min": dt_min_bnds},
        "working_fluid": {
            "p_high": p_high_bnds,
            "p_low": p_low_orc_bnds,
            #'fractions': frac_bounds # When not same fluid
        },
        "environment_storage": {"temp_high": [288.15, 320]},
    }

    bounds = {"hp": bounds_hp, "orc": bounds_orc}
    paths, lower_bnds, upper_bnds = extract_boundary_with_path(bounds)

    # %%
    pool = multiprocessing.Pool(n_processes, maxtasksperchild=20)
    runner = StarmapParallelization(pool.starmap)

    start_time = time.time()

    manager = multiprocessing.Manager()
    gen_shared = manager.Value("i", 1)

    n_var = len(paths)
    n_obj = len(to_optimize)
    problem = opti_Problem(
        n_var=n_var,
        n_obj=n_obj,
        n_ieq_constr=n_ieq_constr,
        opti_fun=fun_to_optimize,
        dir_config=dir_combined,
        paths=paths,
        xl=lower_bnds,
        xu=upper_bnds,
        elementwise_runner=runner,
        same_fluid=same_fluid,
    )

    algorithm = NSGA2(
        pop_size=pop_size,
        n_offsprings=n_offspring,
        sampling=sampling,
        survival=survival,
        eliminate_duplicates=True,
    )

    termination = CombinedTermination(max_gen=n_gen, period=period, ftol=ftol)

    result = minimize(
        problem,
        algorithm,
        save_history=save_history,
        termination=termination,
        return_least_infeasible=False,
        verbose=verbose,
    )

    pool.close()
    pool.join()

    end_time = time.time()

    # %%
    # Save results to the results directory
    if save_res and "dir_results" in locals():

        def path_to_colname(p):
            if isinstance(p, (list, tuple)):
                return ".".join(map(str, p))
            return str(p)

        # -----------------------------
        # 1) Pareto-Front (Objectives)
        # -----------------------------
        F = result.F
        obj_cols = to_optimize
        df_F = pd.DataFrame(F, columns=obj_cols)

        # -----------------------------
        # 2) Pareto-Front (Variables)
        # -----------------------------
        X = result.X
        var_cols = [path_to_colname(p) for p in paths]
        df_X = pd.DataFrame(X, columns=var_cols)

        # Combined F and X as CSV
        df_all = pd.concat([df_X, df_F], axis=1)
        all_csv_file = os.path.join(dir_results, "solutions_with_objectives.csv")
        save_data(all_csv_file, df_all)

        # -----------------------------
        # 3. Setup file: paths, bounds, optimization metadata, and timestamps
        # -----------------------------
        setup_file = os.path.join(dir_results, "setup.yaml")
        setup_data = {
            "paths": paths,
            "bounds": {
                "lower_bounds": list(lower_bnds),
                "upper_bounds": list(upper_bnds),
            },
            "objective_names": to_optimize,
            "same_fluid": same_fluid,
            "optimization": {
                "n_gen": n_gen,
                "period": period,
                "pop_size": pop_size,
                "n_offspring": n_offspring,
                "algorithm": result.algorithm.__class__.__name__,
                "problem": {
                    "n_var": getattr(result.problem, "n_var", None),
                    "n_obj": getattr(result.problem, "n_obj", None),
                    "n_ieq_constr": getattr(result.problem, "n_ieq_constr", None),
                },
            },
            "timestamps": {
                "start": datetime.datetime.fromtimestamp(start_time).isoformat(),
                "end": datetime.datetime.fromtimestamp(end_time).isoformat(),
                "exec_time": result.exec_time,
            },
            "software": {"carbatpy version": version("carbatpy")},
        }
        save_data(setup_file, setup_data)

        # -----------------------------
        # 4. Optimization history (res.history) as pickle
        # -----------------------------
        if save_history:
            try:
                history_data = []
                for i, algo in enumerate(result.history):
                    eps = algo.output.eps.value
                    if eps is not None:
                        eps = float(eps)

                    try:
                        # Get the number of function evaluations
                        n_evals = algo.evaluator.n_eval
                        # Get the optimum from the algorithm
                        pop = algo.pop

                        # Store the data for this generation
                        gen_data = {
                            "n_gen": getattr(algo, "n_gen", None),
                            "n_evals": n_evals,
                            "F": pop.get("F").tolist(),
                            "X": pop.get("X").tolist(),
                            "CV": pop.get("CV").tolist(),
                            "G": pop.get("G").tolist(),
                            "eps": eps,
                            "indicator": algo.output.indicator.value,
                        }

                        history_data.append(gen_data)
                    except Exception as e:
                        print(f"Error processing generation {i}: {e}")
                        continue
                history_file_all_data = os.path.join(
                    dir_results, "history_all_data.pkl"
                )
                save_data(history_file_all_data, history_data)

            except Exception as e:
                print(f"Error in history: {e}")
                pass

        # -----------------------------
        # 5. Save config files (hp and orc)
        # -----------------------------
        hp_config_file = os.path.join(dir_results, "hp_process.yaml")
        orc_config_file = os.path.join(dir_results, "orc_process.yaml")
        save_data(hp_config_file, hp_config)
        save_data(orc_config_file, orc_config)

        print(f"All results have been saved to {dir_results}.")

    if plotting:
        plot = Scatter()
        plot.add(result.F, alpha=0.8)
        plot.show(block=True)