__author__ = "Simon Nilsson"

import warnings

warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", category=DeprecationWarning)
import functools
import multiprocessing
import os
import platform
from typing import Any, Dict, List, Optional, Tuple, Union

import cv2
import pandas as pd

try:
    from typing import Literal
except:
    from typing_extensions import Literal

import shutil

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas

from simba.mixins.config_reader import ConfigReader
from simba.mixins.plotting_mixin import PlottingMixin
from simba.mixins.train_model_mixin import TrainModelMixin
from simba.utils.checks import (check_file_exist_and_readable, check_float,
                                check_if_keys_exist_in_dict,
                                check_if_valid_rgb_tuple, check_int,
                                check_valid_boolean, check_valid_extension,
                                check_valid_tuple,
                                check_video_and_data_frm_count_align)
from simba.utils.data import plug_holes_shortest_bout
from simba.utils.enums import TextOptions
from simba.utils.printing import stdout_success
from simba.utils.read_write import (concatenate_videos_in_folder,
                                    create_directory, find_core_cnt,
                                    get_fn_ext, get_video_meta_data, read_df,
                                    read_pickle, write_df)


def _validation_video_mp(data: pd.DataFrame,
                        bp_dict: dict,
                        video_save_dir: str,
                        video_path: str,
                        text_thickness: int,
                        text_opacity: float,
                        font_size: int,
                        text_spacing: int,
                        circle_size: int,
                        show_pose: bool,
                        show_animal_names: bool,
                        gantt_setting: Union[int, None],
                        final_gantt: Optional[np.ndarray],
                        clf_data: np.ndarray,
                        clrs: List[List],
                        clf_name: str,
                        bouts_df: pd.DataFrame):

    def _put_text(img: np.ndarray,
                  text: str,
                  pos: Tuple[int, int],
                  font_size: int,
                  font_thickness: Optional[int] = 2,
                  font: Optional[int] = cv2.FONT_HERSHEY_DUPLEX,
                  text_color: Optional[Tuple[int, int, int]] = (255, 255, 255),
                  text_color_bg: Optional[Tuple[int, int, int]] = (0, 0, 0),
                  text_bg_alpha: float = 0.8):

        x, y = pos
        text_size, px_buffer = cv2.getTextSize(text, font, font_size, font_thickness)
        w, h = text_size
        overlay, output = img.copy(), img.copy()
        cv2.rectangle(overlay, (x, y-h), (x + w, y + px_buffer), text_color_bg, -1)
        cv2.addWeighted(overlay, text_bg_alpha, output, 1 - text_bg_alpha, 0, output)
        cv2.putText(output, text, (x, y), font, font_size, text_color, font_thickness)
        return output


    def _create_gantt(bouts_df: pd.DataFrame, clf_name: str, image_index: int, fps: int):

        fig, ax = plt.subplots(figsize=(final_gantt.shape[1] / dpi, final_gantt.shape[0] / dpi))
        matplotlib.font_manager._get_font.cache_clear()
        relRows = bouts_df.loc[bouts_df["End_frame"] <= image_index]
        for i, event in enumerate(relRows.groupby("Event")):
            data_event = event[1][["Start_time", "Bout_time"]]
            ax.broken_barh(data_event.values, (4, 4), facecolors="red")
        xLength = (round(image_index / fps)) + 1
        if xLength < 10:
            xLength = 10

        ax.set_xlim(0, xLength)
        ax.set_ylim([0, 12])
        ax.set_xlabel("Session (s)", fontsize=12)
        ax.set_ylabel(clf_name, fontsize=12)
        ax.set_title(f"{clf_name} GANTT CHART", fontsize=12)
        ax.set_yticks([])
        ax.yaxis.set_ticklabels([])
        ax.yaxis.grid(True)
        canvas = FigureCanvas(fig)
        canvas.draw()
        img = np.array(np.uint8(np.array(canvas.renderer._renderer)))[:, :, :3]
        plt.close(fig)
        return img

    dpi = plt.rcParams["figure.dpi"]
    fourcc, font = cv2.VideoWriter_fourcc(*"mp4v"), cv2.FONT_HERSHEY_DUPLEX
    cap = cv2.VideoCapture(video_path)
    video_meta_data = get_video_meta_data(video_path=video_path)
    batch_id, batch_data = data[0], data[1]
    start_frm, current_frm, end_frm = batch_data.index[0], batch_data.index[0], batch_data.index[-1]
    video_save_path = os.path.join(video_save_dir, f"{batch_id}.mp4")
    if gantt_setting is not None:
        video_size = (int(video_meta_data["width"] + final_gantt.shape[1]), int(video_meta_data["height"]))
        writer = cv2.VideoWriter(video_save_path, fourcc, video_meta_data["fps"], video_size)
    else:
        video_size = (int(video_meta_data["width"]), int(video_meta_data["height"]))
        writer = cv2.VideoWriter(video_save_path, fourcc, video_meta_data["fps"], video_size)
    cap.set(1, start_frm)
    while (current_frm <= end_frm) & (current_frm <= video_meta_data["frame_count"]):
        clf_frm_cnt = np.sum(clf_data[0:current_frm])
        ret, img = cap.read()
        if ret:
            if show_pose:
                for animal_cnt, (animal_name, animal_data) in enumerate(bp_dict.items()):
                    for bp_cnt, bp in enumerate(range(len(animal_data["X_bps"]))):
                        x_header, y_header = (animal_data["X_bps"][bp], animal_data["Y_bps"][bp])
                        animal_cords = tuple(batch_data.loc[current_frm, [x_header, y_header]])
                        cv2.circle(img, (int(animal_cords[0]), int(animal_cords[1])), circle_size, clrs[animal_cnt][bp_cnt], -1)
            if show_animal_names:
                for animal_cnt, (animal_name, animal_data) in enumerate(bp_dict.items()):
                    x_header, y_header = (animal_data["X_bps"][0], animal_data["Y_bps"][0],)
                    animal_cords = tuple(batch_data.loc[current_frm, [x_header, y_header]])
                    cv2.putText(img, animal_name, (int(animal_cords[0]), int(animal_cords[1])), font, font_size, clrs[animal_cnt][0], text_thickness)

            target_timer = round((1 / video_meta_data["fps"]) * clf_frm_cnt, 2)
            img = _put_text(img=img, text="BEHAVIOR TIMER:", pos=(TextOptions.BORDER_BUFFER_Y.value, text_spacing), font_size=font_size, font_thickness=TextOptions.TEXT_THICKNESS.value)
            addSpacer = 2
            img = _put_text(img=img, text=f"{clf_name} {target_timer}s", pos=(TextOptions.BORDER_BUFFER_Y.value, text_spacing * addSpacer), font_size=font_size, font_thickness=TextOptions.TEXT_THICKNESS.value, text_bg_alpha=text_opacity)
            addSpacer += 1
            img = _put_text(img=img, text="ENSEMBLE PREDICTION:", pos=(TextOptions.BORDER_BUFFER_Y.value, text_spacing * addSpacer), font_size=font_size, font_thickness=TextOptions.TEXT_THICKNESS.value, text_bg_alpha=text_opacity)
            addSpacer += 1
            if clf_data[current_frm] == 1:
                img = _put_text(img=img, text=clf_name, pos=(TextOptions.BORDER_BUFFER_Y.value, text_spacing * addSpacer), font_size=font_size, font_thickness=TextOptions.TEXT_THICKNESS.value, text_color=TextOptions.COLOR.value, text_bg_alpha=text_opacity)
            addSpacer += 1
            if gantt_setting == 1:
                img = np.concatenate((img, final_gantt), axis=1)
            elif gantt_setting == 2:
                gantt_img = _create_gantt(bouts_df, clf_name, current_frm, video_meta_data["fps"])
                img = np.concatenate((img, gantt_img), axis=1)
            img = cv2.resize(img, video_size, interpolation=cv2.INTER_LINEAR)
            writer.write(np.uint8(img))
            current_frm += 1
            print(f"Multi-processing video frame {current_frm} on core {batch_id}...")
    cap.release()
    writer.release()
    return batch_id


class ValidateModelOneVideoMultiprocess(ConfigReader, PlottingMixin, TrainModelMixin):
    """
    Create classifier validation video for a single input video. Results are stored in the
    ``project_folder/frames/output/validation`` directory.

    :param str config_path: path to SimBA project config file in Configparser format
    :param str feature_file_path: path to SimBA file (parquet or CSV) containing pose-estimation and feature fields.
    :param str model_path: path to pickled classifier object
    :param float discrimination_threshold: classification threshold.
    :param int shortest_bout: Allowed classified bout length expressed in milliseconds. E.g., `1000` will shift frames classified
        as containing the behavior, but occuring in a bout shorter than `1000`, from `target present to `target absent`.
    :param str create_gantt:
        If SimBA should create gantt charts alongside the validation video. OPTIONS: 'None', 'Gantt chart: final frame only (slightly faster)',
        'Gantt chart: video'.
    :param dict settings: User style settings for video. E.g., {'pose': True, 'animal_names': True, 'styles': None}
    :param int cores: Number of cores to use.

    :example:
    >>> test = ValidateModelOneVideoMultiprocess(config_path=r'/Users/simon/Desktop/envs/simba/troubleshooting/two_black_animals_14bp/project_folder/project_config.ini',
    >>>                              feature_file_path='/Users/simon/Desktop/envs/simba/troubleshooting/two_black_animals_14bp/project_folder/csv/features_extracted/Together_1.csv',
    >>>                              model_path='/Users/simon/Desktop/envs/simba/troubleshooting/two_black_animals_14bp/models/generated_models/Attack.sav',
    >>>                              discrimination_threshold=0.6,
    >>>                              shortest_bout=50,
    >>>                              cores=6,
    >>>                              settings={'pose': True, 'animal_names': True, 'styles': None},
    >>>                              create_gantt=None)
    >>> test.run()
    """

    def __init__(self,
                 config_path: Union[str, os.PathLike],
                 feature_path: Union[str, os.PathLike],
                 model_path: Union[str, os.PathLike],
                 show_pose: bool = True,
                 show_animal_names: bool = False,
                 font_size: Optional[bool] = None,
                 circle_size: Optional[int] = None,
                 text_spacing: Optional[int] = None,
                 text_thickness: Optional[int] = None,
                 text_opacity: Optional[float] = None,
                 discrimination_threshold: float = 0.0,
                 shortest_bout: int = 0.0,
                 core_cnt: int = -1,
                 create_gantt: Optional[Union[None, int]] = None):


        ConfigReader.__init__(self, config_path=config_path)
        PlottingMixin.__init__(self)
        TrainModelMixin.__init__(self)
        check_file_exist_and_readable(file_path=config_path)
        check_file_exist_and_readable(file_path=feature_path)
        check_file_exist_and_readable(file_path=model_path)
        check_valid_boolean(value=[show_pose], source=f'{self.__class__.__name__} show_pose', raise_error=True)
        check_valid_boolean(value=[show_animal_names], source=f'{self.__class__.__name__} show_animal_names', raise_error=True)
        check_int(name=f"{self.__class__.__name__} core_cnt", value=core_cnt, min_value=-1, unaccepted_vals=[0])
        if font_size is not None: check_int(name=f'{self.__class__.__name__} font_size', value=font_size)
        if circle_size is not None: check_int(name=f'{self.__class__.__name__} circle_size', value=circle_size)
        if text_spacing is not None: check_int(name=f'{self.__class__.__name__} text_spacing', value=text_spacing)
        if text_opacity is not None: check_float(name=f'{self.__class__.__name__} text_opacity', value=text_opacity, min_value=0.1)
        if text_thickness is not None: check_float(name=f'{self.__class__.__name__} text_thickness', value=text_thickness, min_value=0.1)
        check_float(name=f"{self.__class__.__name__} discrimination_threshold", value=discrimination_threshold, min_value=0, max_value=1.0)
        check_int(name=f"{self.__class__.__name__} shortest_bout", value=shortest_bout, min_value=0)
        if create_gantt is not None:
            check_int(name=f"{self.__class__.__name__} create gantt", value=create_gantt, max_value=2, min_value=1)
        if not os.path.exists(self.single_validation_video_save_dir):
            os.makedirs(self.single_validation_video_save_dir)
        _, self.feature_filename, ext = get_fn_ext(feature_path)
        self.video_path = self.find_video_of_file(self.video_dir, self.feature_filename)
        self.video_meta_data = get_video_meta_data(video_path=self.video_path)
        self.clf_name, self.feature_file_path = (os.path.basename(model_path).replace(".sav", ""), feature_path)
        self.vid_output_path = os.path.join(self.single_validation_video_save_dir, f"{self.feature_filename} {self.clf_name}.mp4")
        self.clf_data_save_path = os.path.join(self.clf_data_validation_dir, f"{self.feature_filename }.csv")
        self.show_pose, self.show_animal_names = show_pose, show_animal_names
        self.font_size, self.circle_size, self.text_spacing = font_size, circle_size, text_spacing
        self.text_opacity, self.text_thickness = text_opacity, text_thickness
        self.clf = read_pickle(data_path=model_path, verbose=True)
        self.data_df = read_df(feature_path, self.file_type)
        self.x_df = self.drop_bp_cords(df=self.data_df)
        self.discrimination_threshold, self.shortest_bout, self.create_gantt = float(discrimination_threshold), shortest_bout, create_gantt
        check_video_and_data_frm_count_align(video=self.video_path, data=self.data_df, name=self.feature_filename, raise_error=False)
        self.core_cnt = find_core_cnt()[0] if core_cnt == -1 or core_cnt > find_core_cnt()[0] else core_cnt
        self.temp_dir = os.path.join(self.single_validation_video_save_dir, "temp")
        self.video_save_path = os.path.join(self.single_validation_video_save_dir, f"{self.feature_filename}.mp4")
        create_directory(paths=self.temp_dir, overwrite=True)
        if platform.system() == "Darwin":
            multiprocessing.set_start_method("spawn", force=True)

    def _get_styles(self):
        self.video_text_thickness = TextOptions.TEXT_THICKNESS.value if self.text_thickness is None else int(max(self.text_thickness, 1))
        longest_str = str(max(['TIMERS:', 'ENSEMBLE PREDICTION:'] + self.clf_names, key=len))
        optimal_font_size, _, optimal_spacing_scale = self.get_optimal_font_scales(text=longest_str, accepted_px_width=int(self.video_meta_data["width"] / 3), accepted_px_height=int(self.video_meta_data["height"] / 10), text_thickness=self.video_text_thickness)
        optimal_circle_size = self.get_optimal_circle_size(frame_size=(self.video_meta_data["width"], self.video_meta_data["height"]), circle_frame_ratio=100)
        self.video_circle_size = optimal_circle_size if self.circle_size is None else int(self.circle_size)
        self.video_font_size = optimal_font_size if self.font_size is None else self.font_size
        self.video_space_size = optimal_spacing_scale if self.text_spacing is None else int(max(self.text_spacing, 1))
        self.video_text_opacity = 0.8 if self.text_opacity is None else float(self.text_opacity)

    def run(self):
        self.prob_col_name = f"Probability_{self.clf_name}"
        self.data_df[self.prob_col_name] = self.clf_predict_proba(clf=self.clf, x_df=self.x_df, model_name=self.clf_name, data_path=self.feature_file_path)
        self.data_df[self.clf_name] = np.where(self.data_df[self.prob_col_name] > self.discrimination_threshold, 1, 0)
        if self.shortest_bout > 1:
            self.data_df = plug_holes_shortest_bout(data_df=self.data_df, clf_name=self.clf_name, fps=self.video_meta_data['fps'], shortest_bout=self.shortest_bout)
        _ = write_df(df=self.data_df, file_type=self.file_type, save_path=self.clf_data_save_path)
        print(f"Predictions created for video {self.feature_filename}...")
        if self.create_gantt is not None:
            self.bouts_df = self.get_bouts_for_gantt(data_df=self.data_df, clf_name=self.clf_name, fps=self.video_meta_data['fps'])
            self.final_gantt_img = self.create_gantt_img(self.bouts_df,self.clf_name,len(self.data_df), self.video_meta_data['fps'],"Behavior gantt chart (entire session)")
            self.final_gantt_img = self.resize_gantt(self.final_gantt_img, self.video_meta_data["height"])
        else:
            self.bouts_df, self.final_gantt_img = None, None
        self._get_styles()


        self.data_df = self.data_df.head(min(len(self.data_df), self.video_meta_data["frame_count"]))
        data = np.array_split(self.data_df, self.core_cnt)
        data = [(i, j) for i, j in enumerate(data)]

        with multiprocessing.Pool(self.core_cnt, maxtasksperchild=self.maxtasksperchild) as pool:
            constants = functools.partial(_validation_video_mp,
                                          bp_dict=self.animal_bp_dict,
                                          video_save_dir=self.temp_dir,
                                          text_thickness=self.video_text_thickness,
                                          text_opacity=self.video_text_opacity,
                                          font_size=self.video_font_size,
                                          text_spacing=self.video_space_size,
                                          circle_size=self.video_circle_size,
                                          video_path=self.video_path,
                                          show_pose=self.show_pose,
                                          show_animal_names=self.show_animal_names,
                                          gantt_setting=self.create_gantt,
                                          final_gantt=self.final_gantt_img,
                                          clf_data=self.data_df[self.clf_name].values,
                                          clrs=self.clr_lst,
                                          clf_name=self.clf_name,
                                          bouts_df=self.bouts_df)
            for cnt, result in enumerate(pool.imap(constants, data, chunksize=self.multiprocess_chunksize)):
                print(f"Image batch {result} complete, Video {self.feature_filename}...")
        pool.terminate()
        pool.join()
        concatenate_videos_in_folder(in_folder=self.temp_dir, save_path=self.video_save_path)
        self.timer.stop_timer()
        stdout_success(msg=f"Video complete, saved at {self.video_save_path}", elapsed_time=self.timer.elapsed_time_str)


# if __name__ == "__main__":
#     test = ValidateModelOneVideoMultiprocess(config_path=r"D:\troubleshooting\mitra\project_folder\project_config.ini",
#                                              feature_path=r"D:\troubleshooting\mitra\project_folder\csv\features_extracted\592_MA147_CNO1_0515.csv",
#                                              model_path=r"C:\troubleshooting\mitra\models\generated_models\grroming_undersample_2_1000\grooming.sav",
#                                              create_gantt=2,
#                                              show_pose=True,
#                                              show_animal_names=True,
#                                              core_cnt=5)
#     test.run()


#
# if __name__ == "__main__":
#     test = ValidateModelOneVideoMultiprocess(config_path=r"C:\troubleshooting\mitra\project_folder\project_config.ini",
#                                     feature_file_path=r"C:\troubleshooting\mitra\project_folder\csv\features_extracted\844_MA131_gq_CNO_0624.csv",
#                                     model_path=r"C:\troubleshooting\mitra\models\validations\lay-on-belly_1\lay-on-belly.sav",
#                                     discrimination_threshold=0.35,
#                                     shortest_bout=200,
#                                     cores=-1,
#                                     settings={'pose': True, 'animal_names': False, 'styles': None},
#                                     create_gantt=2)
#     test.run()





# test = ValidateModelOneVideoMultiprocess(config_path=r'/Users/simon/Desktop/envs/simba/troubleshooting/mouse_open_field/project_folder/project_config.ini',
#                              feature_file_path='/Users/simon/Desktop/envs/simba/troubleshooting/mouse_open_field/project_folder/csv/features_extracted/SI_DAY3_308_CD1_PRESENT.csv',
#                              model_path='/Users/simon/Desktop/envs/simba/troubleshooting/mouse_open_field/models/generated_models/Running.sav',
#                              discrimination_threshold=0.6,
#                              shortest_bout=50,
#                              cores=6,
#                              settings={'pose': True, 'animal_names': True, 'styles': None},
#                              create_gantt=None)
# test.run()

# test = ValidateModelOneVideoMultiprocess(config_path=r'/Users/simon/Desktop/envs/simba/troubleshooting/two_black_animals_14bp/project_folder/project_config.ini',
#                              feature_file_path='/Users/simon/Desktop/envs/simba/troubleshooting/two_black_animals_14bp/project_folder/csv/features_extracted/Together_1.csv',
#                              model_path='/Users/simon/Desktop/envs/simba/troubleshooting/two_black_animals_14bp/models/generated_models/Attack.sav',
#                              discrimination_threshold=0.6,
#                              shortest_bout=50,
#                              cores=6,
#                              settings={'pose': True, 'animal_names': True, 'styles': None},
#                              create_gantt=None)
# test.run()
