# %% imports
from autogaita.resources.utils import write_issues_to_textfile
from autogaita.common2D.common2D_1_preparation import (
    check_and_expand_cfg,
    flip_mouse_body,
)
from autogaita.dlc.dlc_utils import prepare_DLC_df
from autogaita.common2D.common2D_constants import FILE_ID_STRING_ADDITIONS
import os
import shutil
import json
import pandas as pd
import numpy as np

# %% constants
from autogaita.resources.constants import (
    ISSUES_TXT_FILENAME,
    TIME_COL,
    CONFIG_JSON_FILENAME,
)
from autogaita.common2D.common2D_constants import (
    DIRECTION_DLC_THRESHOLD,
)

# %% workflow step #1 - preparation


def some_prep(info, folderinfo, cfg):
    """Preparation of the data & cfg file for dlc analyses"""

    # ............................  unpack stuff  ......................................
    # => DON'T unpack (joint) cfg-keys that are tested later by check_and_expand_cfg
    name = info["name"]
    results_dir = info["results_dir"]
    data_string = folderinfo["data_string"]
    beam_string = folderinfo["beam_string"]
    sampling_rate = cfg["sampling_rate"]
    subtract_beam = cfg["subtract_beam"]
    convert_to_mm = cfg["convert_to_mm"]
    pixel_to_mm_ratio = cfg["pixel_to_mm_ratio"]
    standardise_y_at_SC_level = cfg["standardise_y_at_SC_level"]
    invert_y_axis = cfg["invert_y_axis"]
    flip_gait_direction = cfg["flip_gait_direction"]
    analyse_average_x = cfg["analyse_average_x"]
    standardise_x_coordinates = cfg["standardise_x_coordinates"]
    standardise_y_to_a_joint = cfg["standardise_y_to_a_joint"]

    # .............................  move data  ........................................
    # => see if we can delete a previous runs results folder if existant. if not, it's a
    #    bit ugly since we only update results if filenames match...
    # => for example if angle acceleration not wanted in current run, but was stored in
    #    previous run, the previous run's figure is in the folder
    # => inform the user and leave this as is
    if os.path.exists(results_dir):
        try:
            shutil.rmtree(results_dir)
            move_data_to_folders(info, folderinfo)
        except OSError:
            move_data_to_folders(info, folderinfo)
            unable_to_rm_resdir_error = (
                "\n***********\n! WARNING !\n***********\n"
                + "Unable to remove previous Results subfolder of ID: "
                + name
                + "!\n Results will only be updated if filenames match!"
            )
            print(unable_to_rm_resdir_error)
            write_issues_to_textfile(unable_to_rm_resdir_error, info)
    else:
        move_data_to_folders(info, folderinfo)

    # .......  initialise Issues.txt & quick check for file existence  .................
    # Issues.txt - delete if saved in a previous run
    issues_txt_path = os.path.join(results_dir, ISSUES_TXT_FILENAME)
    if os.path.exists(issues_txt_path):
        os.remove(issues_txt_path)
    # read data & beam
    if not os.listdir(results_dir):
        no_files_error = (
            "\n******************\n! CRITICAL ERROR !\n******************\n"
            + "Unable to identify ANY RELEVANT FILES for "
            + name
            + "!\nThis is likely due to issues with unique file name identifiers.. "
            + "check capitalisation!"
        )
        write_issues_to_textfile(no_files_error, info)
        print(no_files_error)
        return

    # ............................  import data  .......................................
    datadf = pd.DataFrame(data=None)  # prep stuff for error handling
    datadf_duplicate_error = ""
    if subtract_beam:
        if data_string == beam_string:
            beam_and_data_string_error_message = (
                "\n******************\n! CRITICAL ERROR !\n******************\n"
                + "Your data & baseline (beam) identifiers ([G] in our "
                + "file  naming convention) are identical. "
                + "\nNote that they must be different! \nTry again"
            )
            write_issues_to_textfile(beam_and_data_string_error_message, info)
            return
        beamdf = pd.DataFrame(data=None)
        beamdf_duplicate_error = ""
    for filename in os.listdir(results_dir):  # import
        if filename.endswith(".csv"):
            if data_string in filename:
                if datadf.empty:
                    datadf = pd.read_csv(os.path.join(results_dir, filename))
                else:
                    datadf_duplicate_error = (
                        "\n******************\n! CRITICAL ERROR !\n******************\n"
                        + "Two DATA csv-files found for "
                        + name
                        + "!\nPlease ensure your root directory only has one datafile "
                        + "per video!"
                    )
            if subtract_beam:
                if beam_string in filename:
                    if beamdf.empty:
                        beamdf = pd.read_csv(os.path.join(results_dir, filename))
                    else:
                        beamdf_duplicate_error = (
                            "\n******************\n! CRITICAL ERROR !\n***************"
                            + "***\nTwo BEAM csv-files found for "
                            + name
                            + "!\nPlease ensure your root directory only has one "
                            + "beamfile per video!"
                        )
    # handle import errors
    # => append to empty strings to handle multiple issues at once seemlessly
    import_error_message = ""
    if datadf_duplicate_error:
        import_error_message += datadf_duplicate_error
    if datadf.empty:
        import_error_message += (  # if pd didn't raise errors but dfs still empty
            "\n******************\n! CRITICAL ERROR !\n******************\n"
            + "Unable to load a DATA csv file for "
            + name
            + "!\nTry again!"
        )
    if subtract_beam:
        if beamdf_duplicate_error:
            import_error_message += beamdf_duplicate_error
        if beamdf.empty:
            import_error_message += (
                "\n******************\n! CRITICAL ERROR !\n******************\n"
                + "Unable to load a BEAM csv file for "
                + name
                + "!\nTry again!"
            )
    if import_error_message:  # see if there was any issues with import, if so: stop
        print(import_error_message)
        write_issues_to_textfile(import_error_message, info)
        return

    # ....  finalise import: rename cols, get rid of unnecessary elements, floatit  ....
    datadf = prepare_DLC_df(datadf)
    if subtract_beam:  # beam df
        beamdf = prepare_DLC_df(beamdf)
        data = pd.concat([datadf, beamdf], axis=1)
    else:
        data = datadf.copy(deep=True)

    # ................  final data checks, conversions & additions  ....................
    # IMPORTANT - MAIN TESTS OF USER-INPUT VALIDITY OCCUR HERE!
    # => UNPACK VARS FROM CFG THAT ARE TESTED BY check_and_expand HERE, NOT EARLIER!
    cfg = check_and_expand_cfg(data, cfg, info)
    if cfg is None:  # some critical error occured
        return
    hind_joints = cfg["hind_joints"]
    fore_joints = cfg["fore_joints"]
    angles = cfg["angles"]
    beam_hind_jointadd = cfg["beam_hind_jointadd"]
    beam_fore_jointadd = cfg["beam_fore_jointadd"]
    direction_joint = cfg["direction_joint"]
    # important to unpack to vars and not to cfg since cfg is overwritten in multiruns!
    x_standardisation_joint = cfg["x_standardisation_joint"][0]
    y_standardisation_joint = cfg["y_standardisation_joint"][0]
    # store config json file @ group path
    # !!! NU - do this @ mouse path!
    group_path = results_dir.split(name)[0]
    config_json_path = os.path.join(group_path, CONFIG_JSON_FILENAME)
    config_vars_to_json = {
        "sampling_rate": sampling_rate,
        "convert_to_mm": convert_to_mm,
        "standardise_y_at_SC_level": standardise_y_at_SC_level,
        "analyse_average_x": analyse_average_x,
        "standardise_x_coordinates": standardise_x_coordinates,
        "x_standardisation_joint": x_standardisation_joint,
        "standardise_y_to_a_joint": standardise_y_to_a_joint,
        "y_standardisation_joint": y_standardisation_joint,
        "hind_joints": hind_joints,
        "fore_joints": fore_joints,
        "angles": angles,
        "tracking_software": "DLC",
    }
    # note - using "w" will overwrite/truncate file, thus no need to remove it if exists
    with open(config_json_path, "w") as config_json_file:
        json.dump(config_vars_to_json, config_json_file, indent=4)
    # a little test to see if columns make sense, i.e., same number of x/y/likelihood
    x_col_count = len([c for c in data.columns if c.endswith(" x")])
    y_col_count = len([c for c in data.columns if c.endswith(" y")])
    likelihood_col_count = len([c for c in data.columns if c.endswith(" likelihood")])
    if x_col_count == y_col_count == likelihood_col_count:
        pass
    else:
        cols_are_weird_message = (
            "\n***********\n! WARNING !\n***********\n"
            + "We detected an unequal number of columns ending with x, y or "
            + "likelihood!\nCounts were:\n"
            + "x: "
            + str(x_col_count)
            + ", y: "
            + str(y_col_count)
            + ", likelihood: "
            + str(likelihood_col_count)
            + "!\n\n"
            + "We continue with the analysis but we strongly suggest you have another "
            + "look at your dataset, this should not happen.\n"
        )
        print(cols_are_weird_message)
        write_issues_to_textfile(cols_are_weird_message, info)
    # if wanted: fix that deeplabcut inverses y
    if invert_y_axis:
        for col in data.columns:
            if col.endswith(" y"):
                data[col] = data[col] * -1
    # if we don't have a beam to subtract, standardise y to a joint's or global ymin = 0
    if not subtract_beam:
        y_min = float("inf")
        y_cols = [col for col in data.columns if col.endswith("y")]
        if standardise_y_to_a_joint:
            y_min = data[y_standardisation_joint + "y"].min()
        else:
            y_min = data[y_cols].min().min()
        data[y_cols] -= y_min
    # convert pixels to millimeters
    if convert_to_mm:
        for column in data.columns:
            if not column.endswith("likelihood"):
                data[column] = data[column] / pixel_to_mm_ratio
    # quick warning if cfg is set to not flip gait direction but to standardise x
    if not flip_gait_direction and standardise_x_coordinates:
        message = (
            "\n***********\n! WARNING !\n***********\n"
            + "You are standardising x-coordinates without standardising the direction "
            + "of gait (e.g. all walking from right to left)."
            + "\nThis can be correct if you are doing things like treadmill walking "
            + "but can lead to unexpected behaviour otherwise!"
            + "\nMake sure you know what you are doing!"
        )
        print(message)
        write_issues_to_textfile(message, info)
    # check gait direction & DLC file validity
    data = check_gait_direction(data, direction_joint, flip_gait_direction, info)
    if data is None:  # this means DLC file is broken
        return
    # subtract the beam from the joints to standardise y
    # => bc. we simulate that all mice run from left to right, we can write:
    #     (note that we also flip beam x columns, but never y-columns!)
    # => & bc. we multiply y values by *-1 earlier, it's a neg_num - - neg_num
    #    pushing it towards zero.
    # => using list(set()) to ensure that we don't have duplicate values (if users
    #    should have provided them in both cfg vars by misstake)
    # => beam_col_left and right is provided by users
    if subtract_beam:
        # note beam_col_left/right are always lists in cfg!
        beam_col_left = cfg["beam_col_left"][0]
        beam_col_right = cfg["beam_col_right"][0]
        for joint in list(set(hind_joints + beam_hind_jointadd)):
            data[joint + "y"] = data[joint + "y"] - data[beam_col_left + "y"]
        for joint in list(set(fore_joints + beam_fore_jointadd)):
            data[joint + "y"] = data[joint + "y"] - data[beam_col_right + "y"]
        data.drop(columns=list(beamdf.columns), inplace=True)  # beam not needed anymore
    # add Time
    data[TIME_COL] = data.index * (1 / sampling_rate)
    # reorder the columns we added
    cols = [TIME_COL, "Flipped"]
    data = data[cols + [c for c in data.columns if c not in cols]]
    return data


# ..............................  helper functions  ....................................


def move_data_to_folders(info, folderinfo):
    """Find files, copy data, video, beamdata & beamvideo to new results_dir"""
    # unpack
    results_dir = info["results_dir"]
    postmouse_string = folderinfo["postmouse_string"]
    postrun_string = folderinfo["postrun_string"]
    os.makedirs(results_dir)  # important to do this outside of loop!
    # check if user forgot some underscores or dashes in their filenames
    # => two levels of string additions for two post FILE-ID strings
    # => in theory if the user has some strange cases in which this double forloop
    # would be true twice (because one file is called -6DLC and another is called _6DLC
    # for some reason) it will break after the first time and always ignore the second
    # one - keep this in mind if it should come up but it should be very unlikely
    for mouse_string_addition in FILE_ID_STRING_ADDITIONS:
        candidate_postmouse_string = mouse_string_addition + postmouse_string
        for run_string_addition in FILE_ID_STRING_ADDITIONS:
            candidate_postrun_string = run_string_addition + postrun_string
            found_it = check_this_filename_configuration(
                info,
                folderinfo,
                candidate_postmouse_string,
                candidate_postrun_string,
                results_dir,
            )
            if found_it:  # if our search was successful, stop searching and continue
                break


def check_this_filename_configuration(
    info, folderinfo, postmouse_string, postrun_string, results_dir
):
    # unpack
    name = info["name"]
    mouse_num = info["mouse_num"]
    run_num = info["run_num"]
    root_dir = folderinfo["root_dir"]
    data_string = folderinfo["data_string"]
    beam_string = folderinfo["beam_string"]
    premouse_string = folderinfo["premouse_string"]
    prerun_string = folderinfo["prerun_string"]
    whichvideo = ""  # initialise
    found_it = False
    for filename in os.listdir(root_dir):
        # the following condition is True for data & beam csv
        if (
            (premouse_string + str(mouse_num) + postmouse_string in filename)
            and (prerun_string + str(run_num) + postrun_string in filename)
            and (filename.endswith(".csv"))
        ):
            found_it = True
            # Copy the Excel file to the new subfolder
            shutil.copy2(
                os.path.join(root_dir, filename), os.path.join(results_dir, filename)
            )
            # Check if there is a video and if so copy it too
            vidname = filename[:-4] + "_labeled.mp4"
            vidpath = os.path.join(root_dir, vidname)
            if os.path.exists(vidpath):
                shutil.copy2(vidpath, os.path.join(results_dir, vidname))
            else:
                if data_string in vidname:
                    whichvideo = "Data"
                elif beam_string in vidname:
                    whichvideo = "Beam"
                this_message = (
                    "\n***********\n! WARNING !\n***********\n"
                    + "No "
                    + whichvideo
                    + "video for "
                    + name
                    + "!"
                )
                print(this_message)
                write_issues_to_textfile(this_message, info)
    return found_it


def check_gait_direction(data, direction_joint, flip_gait_direction, info):
    """Check direction of gait - reverse it if needed

    Note
    ----
    Also using this check to check for DLC files being broken
    flip_gait_direction is only used after the check for DLC files being broken
    """

    data["Flipped"] = False
    enterframe = 0
    idx = 0
    flip_error_message = ""
    while enterframe == 0:  # first find out when mouse was in the video frame.
        if (
            np.mean(data[direction_joint + "likelihood"][idx : idx + 5])
            > DIRECTION_DLC_THRESHOLD
        ):  # +5 to increase conf.
            enterframe = idx + 5
        idx += 1
        if (idx > len(data)) | (enterframe > len(data)):
            flip_error_message += (
                "\n******************\n! CRITICAL ERROR !\n******************\n"
                + "Unable to determine gait direction!"
                + "\nThis hints a critical issue with DLC tracking, e.g., likelihood "
                + "\ncolumns being low everywhere or tables being suspiciously short!"
                + "\nTo be sure, we cancel everything here."
                + "\nPlease check your input DLC csv files for correctness & try again!"
            )
            break
    leaveframe = 0
    idx = 1
    while leaveframe == 0:  # see where mouse left frame (same logic from back)
        if (
            np.mean(data[direction_joint + "likelihood"][-idx - 5 : -idx])
            > DIRECTION_DLC_THRESHOLD
        ):
            leaveframe = len(data) - idx - 5
        idx += 1
        if idx > len(data):
            if not flip_error_message:
                flip_error_message += (
                    "\n******************\n! CRITICAL ERROR !\n******************\n"
                    + "Unable to determine gait direction!"
                    + "\nThis hints a critical issue with DLC tracking, e.g., "
                    + "likelihood \ncolumns being low everywhere or tables being "
                    + "suspiciously short!\nTo be sure, we cancel everything here."
                    + "\nPlease check your input DLC csv files for correctness & try "
                    + "again!"
                )
            break
    if flip_error_message:
        write_issues_to_textfile(flip_error_message, info)
        print(flip_error_message)
        return
    if (
        data[direction_joint + "x"][enterframe]
        > data[direction_joint + "x"][leaveframe]
    ):  # i.e.: right to left
        # simulate that mouse ran from left to right (only if user wants it)
        if flip_gait_direction:
            data = flip_mouse_body(data, info)
            data["Flipped"] = True
    return data
