import sys
import glob
import os
import subprocess
import questionary

from esm_parser import user_error

class oasis:
    def __init__(
        self,
        nb_of_couplings=1,
        coupled_execs=["echam", "fesom"],
        runtime=1,
        debug_level=1,
        nnorest="F",
        mct_version="4.0",
        lucia=False,
    ):
        if isinstance(mct_version, tuple):
            pass
        elif isinstance(mct_version, int):
            mct_version = (mct_version, 0)
        elif isinstance(mct_version, float):
            mct_version = tuple(int(x) for x in str(mct_version).split("."))
        # PG: Fixup version to be tuples:
        elif isinstance(mct_version, str):
            mct_version = tuple(int(x) for x in mct_version.split("."))
        else:
            print(
                "Init of Oasis needs the argument mct_version to be either a tuple or a string!",
                flush=True,
            )
            sys.exit(1)
        self.namcouple = [
            "# This namcouple was automatically generated by the esm-tools (Python)"
        ]
        self.namcouple += [" $NFIELDS", "            " + str(nb_of_couplings), " $END"]
        exec_entry = ""
        for exe in coupled_execs:
            exec_entry = exec_entry + " " + exe
        exec_entry = str(len(coupled_execs)) + exec_entry
        self.namcouple += [" $NBMODEL", "            " + str(exec_entry), " $END"]
        self.namcouple += [" $RUNTIME", "           " + str(runtime), " $END"]
        # seb-wahl: add lucia support
        if lucia:
            self.namcouple += [" $NLOGPRT", "           " + "1 -1", " $END"]
        else:
            self.namcouple += [" $NLOGPRT", "           " + str(debug_level), " $END"]
        if mct_version >= (4, 0):
            # If true, OASIS can start without restart files
            self.namcouple += [" $NNOREST", "           " + str(nnorest), " $END "]
        self.namcouple += [" $STRINGS"]
        self.namcouple += [
            "###############################################################################"
        ]
        self.namcouple += [
            "###############################################################################"
        ]
        self.next_coupling = 1
        self.name = "oasis3mct"

    def add_input_coupling(self, field_name, freq, field_filepath):
        self.namcouple += ["#"]
        nb = self.next_coupling
        self.namcouple += [
            f"{field_name} {field_name} {nb} {freq} 0 {field_filepath} INPUT"
        ]
        self.namcouple += "#"
        self.next_coupling += 1

    def add_coupling(
        self,
        lefts,
        lgrid,
        rights,
        rgrid,
        direction,
        transformation,
        restart_file,
        time_step,
        lresume,
        export_mode="DEFAULT",
    ):

        self.namcouple += ["#"]

        nb = self.next_coupling

        left = sep = ""
        for lefty in lefts:

            restart_out_file = lefty + "_"

            left += sep + lefty
            self.next_coupling += 1
            sep = ":"

        right = sep = ""
        for righty in rights:
            right += sep + righty
            sep = ":"

        if export_mode == "DEFAULT":
            if bool(lresume) is False:
                export_mode = "EXPOUT"
            else:
                export_mode = "EXPORTED"

        if bool(lresume) is False:
            lag = str(0)
        else:
            lag = direction.get("lag", "0")

        # if a transformation method for CONSERV (e.g. GLOBAL) is set below,
        # increase seq (=number of lines describing the transformation) by 1
        # this part was broke, confused seq with nb of trafo lines

        seq = int(direction.get("seq", "2"))
        # if transformation.get("postprocessing", {}).get("conserv", {}).get("method"):
        #    seq += 1

        p_rgrid = p_lgrid = "0"
        if "number_of_overlapping_points" in rgrid:
            p_rgrid = str(rgrid["number_of_overlapping_points"])
        if "number_of_overlapping_points" in lgrid:
            p_lgrid = str(lgrid["number_of_overlapping_points"])

        trafo_line = ""
        trafo_details = []

        alltimes = transformation.get("time_transformation", "bla")
        if not type(alltimes) == list:
            alltimes = [alltimes]
        for time in alltimes:
            detail_line = ""
            if time.lower() in ["instant", "accumul", "average", "t_min", "t_max"]:
                trafo_line = "LOCTRANS"
                detail_line = time.upper()
                trafo_details.append(detail_line.strip())

        allpres = transformation.get("preprocessing", "bla")
        if not type(allpres) == list:
            allpres = [allpres]
        for pre in allpres:
            detail_line = ""
            if type(pre) == dict:
                pre = list(pre.keys())[0]
            if pre.lower() == "checkin":
                trafo_line += " CHECKIN"
                detail_line = "INT = 1"
                trafo_details.append(detail_line.strip())
            elif pre.lower() == "blasold":
                trafo_line += " BLASOLD"
                coefficient = transformation["preprocessing"][pre].get("xmult", None)
                if not coefficient:
                    print(
                        "xmult needs to be defined for preprocessing BLASOLD",
                        flush=True,
                    )
                    sys.exit(2)
                add_scalar = transformation["preprocessing"][pre].get(
                    "add_scalar", None
                )
                if not add_scalar:
                    print(
                        "add_scalar needs to be defined (0 or 1) for preprocessing BLASOLD",
                        flush=True,
                    )
                    sys.exit(2)
                detail_line = str(coefficient) + " " + str(add_scalar)
                trafo_details.append(detail_line.strip())
                if str(add_scalar) == "1":
                    scalar_to_add = transformation["preprocessing"][pre].get(
                        "scalar_to_add", None
                    )
                    if not add_scalar:
                        print(
                            "scalar_to_add needs to be defined if add_scalar is set to  1 for preprocessing BLASOLD",
                            flush=True,
                        )
                        sys.exit(2)
                    detail_line = " CONSTANT" + str(add_scalar)
                    trafo_details.append(detail_line.strip())

        alltrans = transformation.get("remapping", {"bla": "blub"})
        if not type(alltrans) == list:
            alltrans = [alltrans]
        for thistrans in alltrans:
            (trans, transform) = list(thistrans.items())[0]
            detail_line = ""
            if "mapping" == trans.lower():
                trafo_line += " MAPPING"
                mapname = transform.get("mapname", None)
                if not mapname:
                    print(
                        "mapname needs to be defined for transformation MAPPING",
                        flush=True,
                    )
                    sys.exit(2)
                maploc = transform.get("map_regrid_on", "")
                mapstrat = transform.get("mapstrategy", "")
                detail_line = mapname + " " + maploc + " " + mapstrat
                trafo_details.append(detail_line.strip())

            elif trans.lower() in [
                "distwgt",
                "bicubic",
                "bilinear",
                "gauswgt",
                "conserv",
                "loccunif",
            ]:
                trafo_line += " SCRIPR"
                srcgridtype = str(rgrid["oasis_grid_type"]).upper()
                search_bin = transform.get("search_bin", None)
                if not search_bin:
                    print(
                        "search_bin (LATITUDE or LATLON) needs to be defined for transformations DISTWGT, GAUSWGT, BILINEAR, BICUBIC, LOCCUNIF",
                        flush=True,
                    )
                    sys.exit(2)
                bins = transform.get("nb_of_search_bins", "1")
                detail_line = (
                    trans.upper()
                    + " "
                    + srcgridtype.upper()
                    + " SCALAR "
                    + search_bin.upper()
                    + " "
                    + str(bins)
                )
                if trans.lower() in ["distwgt", "gauswgt", "loccunif"]:
                    nb_of_neighbours = transform.get("nb_of_neighbours", None)
                    if not nb_of_neighbours:
                        print(
                            "nb_of_neighbours needs to be defined for transformations DISTWGT, GAUSWGT and LOCCUNIF",
                            flush=True,
                        )
                        sys.exit(2)
                    detail_line += " " + str(nb_of_neighbours)
                if trans.lower() == "gauswgt":
                    weight = transform.get("weight", None)
                    if not weight:
                        print(
                            "weight needs to be defined for transformation GAUSWGT",
                            flush=True,
                        )
                        sys.exit(2)
                    detail_line += " " + str(weight)
                if trans.lower() == "conserv":
                    normalization = transform.get("normalization", None)
                    if not normalization:
                        print(
                            "normalization (FRACAREA, DESTAREA or FRACNNEI) needs to be defined for transformations CONSERV",
                            flush=True,
                        )
                        sys.exit(2)
                    order = transform.get("order", None)
                    if not order:
                        print(
                            "order (FIRST or SECOND) needs to be defined for transformation CONSERV",
                            flush=True,
                        )
                        sys.exit(2)
                    detail_line += " " + normalization.upper() + " " + order.upper()
                trafo_details += [detail_line.strip()]

        allpost = transformation.get("postprocessing", "bla")
        if not type(allpost) == list:
            allpost = list(allpost)
        for post in allpost:
            detail_line = ""
            if post.lower() == "conserv":
                trafo_line += " CONSERV"
                method = transformation["postprocessing"][post].get("method", None)
                if not method:
                    print(
                        " a method (GLOBAL, GLBPOS, BASBAL or BASPOS) needs to be defined for postprocessing CONSERV",
                        flush=True,
                    )
                    sys.exit(2)
                algorithm = transformation["postprocessing"][post].get("algorithm", "")
                detail_line = method.upper() + " " + algorithm.lower()
                trafo_details.append(detail_line.strip())
            elif post.lower() == "checkout":
                trafo_line += " CHECKOUT"
                detail_line = "INT = 1"
                trafo_details.append(detail_line.strip())
            elif post.lower() == "blasnew":
                trafo_line += " BLASNEW"
                coefficient = transformation["postprocessing"][post].get("xmult", None)
                if not coefficient:
                    print(
                        "xmult needs to be defined for postprocessing BLASNEW",
                        flush=True,
                    )
                    sys.exit(2)
                add_scalar = transformation["postprocessing"][post].get(
                    "add_scalar", None
                )
                if not add_scalar:
                    print(
                        "add_scalar needs to be defined (0 or 1) for preprocessing BLASOLD",
                        flush=True,
                    )
                    sys.exit(2)
                detail_line = str(coefficient) + " " + str(add_scalar)
                trafo_detail.append(detail_line.strip())
                if str(add_scalar) == "1":
                    scalar_to_add = transformation["postprocessing"][post].get(
                        "scalar_to_add", None
                    )
                    if not add_scalar:
                        print(
                            "scalar_to_add needs to be defined if add_scalar is set to  1 for postprocessing BLASNEW",
                            flush=True,
                        )
                        sys.exit(2)
                    detail_line = " CONSTANT" + str(add_scalar)
                    trafo_details.append(detail_line.strip())

        # Remove the first space in cases such as trafo_line = ' SCRIPR CONSERV'
        if trafo_line[0] == " ":
            trafo_line = trafo_line[1:]

        nb_of_trafo_lines = len(trafo_details)

        self.namcouple += [
            right
            + " "
            + left
            + " "
            + str(nb)
            + " "
            + str(time_step)
            + " "
            + str(nb_of_trafo_lines)
            + " "
            + str(restart_file)
            + " "
            + export_mode
        ]
        if lgrid and rgrid:
            self.namcouple += [
                str(rgrid["nx"])
                + " "
                + str(rgrid["ny"])
                + " "
                + str(lgrid["nx"])
                + " "
                + str(lgrid["ny"])
                + " "
                + rgrid["name"]
                + " "
                + lgrid["name"]
                + " LAG="
                + str(lag)
            ]

        self.namcouple += ["P " + p_rgrid + " P " + p_lgrid]
        self.namcouple += [trafo_line]
        for line in trafo_details:
            self.namcouple += [line]

        self.namcouple += ["#"]
        self.namcouple += ["#"]
        self.namcouple += ["#"]
        self.namcouple += [
            "###############################################################################"
        ]

    def print_config_files(self):
        for line in self.namcouple:
            print(line, flush=True)

    def add_output_file(self, lefts, rights, leftmodel, rightmodel, config):
        out_file = []

        coupling = self.next_coupling

        if self.next_coupling < 10:
            this_coupling = "0" + str(coupling)
        else:
            this_coupling = str(coupling)

        for lefty in lefts:
            out_file.append(lefty + "_" + leftmodel + "_" + this_coupling + ".nc")
        for righty in rights:
            out_file.append(righty + "_" + rightmodel + "_" + this_coupling + ".nc")

        self.next_coupling += 1

        if "outdata_files" not in config:
            config["outdata_files"] = {}
        if "outdata_in_work" not in config:
            config["outdata_in_work"] = {}
        if "outdata_sources" not in config:
            config["outdata_sources"] = {}

        for thisfile in out_file:

            config["outdata_files"][thisfile] = thisfile
            config["outdata_in_work"][thisfile] = thisfile
            config["outdata_sources"][thisfile] = thisfile

    def add_restart_files(self, restart_file_label, fconfig):
        """
        Handles the special restart case of the coupling fields.

        Cases
        -----
        1. If this run is a restart but not a branch-off experiment, set the source to
            be the same as defined by the user in ``restart_in_sources`` or the same
            as the name coming from ``coupling_<target/input>_fields``, if the first is
            missing (the normal case).
        2. Same as case 1 but with the time stamp added to the name of the restart file
            to make sure the correct file (and not a link to the last restart file made)
            is loaded for the branch-off experiment. This option uses the
            non-timestamped version of the file when only one file is found (e.g. the
            parent simulation only has one run, or the files are taken for the first
            run from the pool, as in AWICM3).

        Parameters
        ----------
        restart_file_label : str
            The file's label (not the file name itself!). Used to retrieve the
            file's source and target path. As defined in the keys of
            ``coupling_target_fields`` or ``coupling_input_fields`` in the yamls
        fconfig : ConfigSetup
            The complete simulation configuration.
        """

        config = fconfig[self.name]
        gconfig = fconfig["general"]
        is_runtime = gconfig["run_or_compile"] == "runtime"
        enddate = "_" + gconfig["end_date"].format(
            form=9, givenph=False, givenpm=False, givenps=False
        )
        parentdate = "_" + config["parent_date"].format(
            form=9, givenph=False, givenpm=False, givenps=False
        )

        if "restart_out_files" not in config:
            config["restart_out_files"] = {}
        if "restart_out_in_work" not in config:
            config["restart_out_in_work"] = {}
        if "restart_out_sources" not in config:
            config["restart_out_sources"] = {}

        if "restart_in_files" not in config:
            config["restart_in_files"] = {}
        if "restart_in_in_work" not in config:
            config["restart_in_in_work"] = {}
        if "restart_in_sources" not in config:
            config["restart_in_sources"] = {}

        # Find the actual path of the restart
        restart_file_path = config["restart_in_sources"].get(restart_file_label, None)
        # Find the actual name of the restart: if a path is given in restart_in_sources
        # get the basename of that path, otherwise assign the file label also as name
        # of the file (coming from ``coupling_<target/input>_fields``)
        if restart_file_path:
            restart_file = os.path.basename(restart_file_label)
        else:
            restart_file = restart_file_label

        config["restart_out_files"][restart_file_label] = restart_file
        config["restart_out_files"][restart_file_label + "_recv"] = restart_file + "_recv"

        config["restart_out_in_work"][restart_file_label] = restart_file  # + enddate
        config["restart_out_in_work"][restart_file_label + "_recv"] = (
            restart_file + "_recv"
        )  # + enddate

        config["restart_out_sources"][restart_file_label] = restart_file
        config["restart_out_sources"][restart_file_label + "_recv"] = restart_file + "_recv"

        config["restart_in_files"][restart_file_label] = restart_file
        config["restart_in_in_work"][restart_file_label] = restart_file

        # In case of a branch-off experiment -> use the correct oasis restart files:
        # Not the soft link to the last, but the actual one for the branch-off date
        if gconfig["run_number"] == 1 and config["lresume"] and gconfig["jobtype"] == "prepcompute":
            # If they do not exist, define ``ini_restart_date`` and ``ini_restart_dir``
            # based on ``ini_parent_date`` and ``ini_parent_dir``
            if "ini_parent_date" in config and "ini_restart_date" not in config:
                config["ini_restart_date"] = config["ini_parent_date"]
            if "ini_parent_dir" in config and "ini_restart_dir" not in config:
                config["ini_restart_dir"] = config["ini_parent_dir"]
            # If the restart file path is not defined, or it's not an absolute path to
            # the file, set it to be the same as the ini_restart_dir
            if not restart_file_path or restart_file_path == restart_file:
                restart_file_path = f"{config['ini_restart_dir']}/{restart_file}"
            # If set in config (oasis):
            if "ini_restart_dir" in config and "ini_restart_date" in config:
                # check if restart file with ini_restart_date in filename is in the restart
                # folder of the parent experiment to be branched off from:
                glob_search_file = (
                    f"{restart_file_path}*"
                    f"{config['ini_restart_date'].year}"
                    f"{config['ini_restart_date'].month:02}"
                    f"{config['ini_restart_date'].day:02}"
                )
            else:
                glob_search_file = restart_file_path

            glob_restart_file = glob.glob(glob_search_file)
            glob_restart_file.sort()
            if restart_file and is_runtime:
                # If there are more than one file found let the user decide which one to take
                if len(glob_restart_file) == 1:
                    restart_file = os.path.basename(glob_restart_file[0])
                elif len(glob_restart_file) == 0:
                    restart_file = restart_file_path
                    if not os.path.isfile(restart_file):
                        user_error(
                            "Restart file missing",
                            f"No OASIS restart file for ``{restart_file_label}`` found "
                            f"matching the pattern ``{glob_search_file}`` nor "
                            f"``{restart_file}``"
                        )
                else:
                    if not gconfig["isinteractive"]:
                        # If more than one restart file found that matches ini_restart_date,
                        # ask the user to select from the result list:
                        message = (
                            "More than one OASIS restart file was found for "
                            "your branchoff experiment that matches the "
                            "ini_restart_date you selected. Please select "
                            "one of the following OASIS restart files:"
                        )
                        answers = questionary.form(
                            restarts = questionary.select(message, choices=glob_restart_file)
                        ).ask()
                        restart_file = answers["restarts"]

            config["restart_in_sources"][restart_file_label] = restart_file

        if restart_file_label not in config["restart_in_sources"]:
            config["restart_in_sources"][restart_file_label] = restart_file


    def prepare_restarts(self, restart_file, all_fields, models, config):
        enddate = "_" + config["general"]["end_date"].format(
            form=9, givenph=False, givenpm=False, givenps=False
        )
        # enddate = "_" + str(config["general"]["end_date"].year) + str(config["general"]["end_date"].month) + str(config["general"]["end_date"].day)

        print("Preparing oasis restart files from initial run...", flush=True)
        # Assign an exe per model
        exes = [config[model]["executable"] for model in models]
        print(restart_file, all_fields, models, exes, flush=True)
        cwd = os.getcwd()
        os.chdir(config["general"]["thisrun_work_dir"])
        filelist = ""
        # Loop through the fields and their corresponding models and exes
        for field, model, exe in zip(all_fields, models, exes):
            print(field + "-" + model, flush=True)
            thesefiles = glob.glob(field + "_" + exe + "_*.nc")
            print(thesefiles, flush=True)
            for thisfile in thesefiles:
                print("cdo showtime " + thisfile + " 2>/dev/null | head -n 1 | wc -w", flush=True)
                lasttimestep = (
                    subprocess.check_output(
                        "cdo showtime " + thisfile + " 2>/dev/null | head -n 1 | wc -w", shell=True
                    )
                    .decode("utf-8")
                    .rstrip()
                )
                # print (lasttimestep)

                print(
                    "cdo -O seltimestep,"
                    + str(lasttimestep)
                    + " "
                    + thisfile
                    + " onlyonetimestep.nc",
                    flush=True,
                )
                os.system(
                    "cdo -O seltimestep,"
                    + str(lasttimestep)
                    + " "
                    + thisfile
                    + " onlyonetimestep.nc"
                )
                print(
                    "ncwa -O -a time onlyonetimestep.nc notimestep_" + field + ".nc",
                    flush=True,
                )
                os.system(
                    "ncwa -O -a time onlyonetimestep.nc notimestep_" + field + ".nc"
                )
                filelist += "notimestep_" + field + ".nc "
                print(filelist)
        # MA: -O flag added to overwrite oasis restart files in case oasis creats them
        # before (i.e. when using LOCTRANS)
        if os.path.isfile(restart_file) and config["general"]["verbose"]:
            print(f"{restart_file} already exits, overwriting", flush=True)
        print("cdo -O merge " + filelist + " " + restart_file, flush=True)  # + enddate)
        os.system("cdo -O merge " + filelist + " " + restart_file)  # + enddate)
        rmlist = glob.glob("notimestep*")
        rmlist.append("onlyonetimestep.nc")
        for rmfile in rmlist:
            print("rm " + rmfile, flush=True)
            os.system("rm " + rmfile)
        os.chdir(cwd)

    def finalize(self, destination_dir):
        self.namcouple += [" $END"]
        endline = ""
        with open(destination_dir + "/namcouple", "w") as namcouple:
            for line in self.namcouple:
                namcouple.write(endline)
                namcouple.write(line)
                endline = "\n"
