import os
import shutil
import sys
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../'))) # Add the path to the fiqus package to the system path
import fiqus.data.DataFiQuSPancake3D as Pancake3D
from fiqus.data.DataFiQuS import FDM
from fiqus.utils.Utils import FilesAndFolders as Util
from fiqus import MainFiQuS as mf


# Generate reference files for the models below:
model_names = [
    "TEST_Pancake3D_REF",
    "TEST_Pancake3D_REFStructured",
    "TEST_Pancake3D_TSA",
    "TEST_Pancake3D_TSAStructured",
    "TEST_Pancake3D_TSAInsulating"
]

for model_name in model_names:
    input_file = os.path.join(
        os.path.dirname(os.path.dirname(__file__)),
        "_inputs",
        model_name,
        f"{model_name}.yaml",
    )
    # select _references folder as the output folder:
    output_folder = os.path.join(
        os.path.dirname(os.path.dirname(__file__)), "_references", model_name
    )

    # Cast input yaml file to FDM
    data_model: FDM = Util.read_data_from_yaml(input_file, FDM)
    data_model.magnet.postproc = Pancake3D.Pancake3DPostprocess()
    data_model.run.overwrite = True
    data_model.run.launch_gui = False

    # if the output folder exists, remove it and create a new one:
    if os.path.exists(output_folder):
        shutil.rmtree(output_folder)
        # Create the output folder:
        os.makedirs(output_folder)

    # Solve the same model three times with different solve types:
    solve_types = [
        "weaklyCoupled",
        "stronglyCoupled",
        "electromagnetic",
    ]
    for i, solve_type in enumerate(solve_types):
        if i == 0:
            # Make the run type start_from_yaml:
            data_model.run.type = "start_from_yaml"
        else:
            # Make the run type solve_only:
            data_model.run.type = "solve_only"

        # data_model.run.type = "post_process"
        # data_model.run.type = "solve_only"

        data_model.magnet.solve.type = solve_type
        data_model.run.solution = solve_type
        if solve_type in ["weaklyCoupled", "stronglyCoupled"]:
            data_model.magnet.solve.save = [
                Pancake3D.Pancake3DSolveSaveQuantity(
                    quantity="magneticField",
                ),
                Pancake3D.Pancake3DSolveSaveQuantity(
                    quantity="currentDensity",
                ),
                Pancake3D.Pancake3DSolveSaveQuantity(
                    quantity="temperature",
                ),
            ]
        elif solve_type == "electromagnetic":
            data_model.magnet.solve.save = [
                Pancake3D.Pancake3DSolveSaveQuantity(
                    quantity="magneticField",
                ),
                Pancake3D.Pancake3DSolveSaveQuantity(
                    quantity="currentDensity",
                ),
            ]

        fiqus_instance = mf.MainFiQuS(fdm=data_model, model_folder=output_folder)

        # remove fiqus_instance to avoid memory issues:
        del fiqus_instance

        # # remove logs directory and run_log.csv inside the output folder:
        # shutil.rmtree(os.path.join(output_folder, "logs"))
        # os.remove(os.path.join(output_folder, "run_log.csv"))
