__author__ = "Simon Nilsson; sronilsson@gmail.com"

import functools
import multiprocessing
import os
import platform
from copy import deepcopy
from typing import Dict, List, Optional, Tuple, Union

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

import cv2
import numpy as np
import pandas as pd

from simba.mixins.config_reader import ConfigReader
from simba.mixins.geometry_mixin import GeometryMixin
from simba.mixins.plotting_mixin import PlottingMixin
from simba.mixins.train_model_mixin import TrainModelMixin
from simba.utils.checks import (check_float, check_if_valid_rgb_tuple,
                                check_int, check_nvidea_gpu_available,
                                check_str, check_that_column_exist,
                                check_that_hhmmss_start_is_before_end,
                                check_valid_boolean, check_valid_dict,
                                check_video_and_data_frm_count_align)
from simba.utils.data import (check_if_string_value_is_valid_video_timestamp,
                              create_color_palette, detect_bouts,
                              find_frame_numbers_from_time_stamp, get_cpu_pool,
                              terminate_cpu_pool)
from simba.utils.enums import ConfigKey, Dtypes, Options, TagNames, TextOptions
from simba.utils.errors import (InvalidInputError, NoDataError,
                                NoSpecifiedOutputError)
from simba.utils.printing import (SimbaTimer, log_event, stdout_information,
                                  stdout_success)
from simba.utils.read_write import (concatenate_videos_in_folder,
                                    create_directory,
                                    find_all_videos_in_project, find_core_cnt,
                                    get_fn_ext, get_video_meta_data,
                                    read_config_entry, read_df,
                                    seconds_to_timestamp)
from simba.utils.warnings import FrameRangeWarning

START_TIME, END_TIME = 'start_time', 'end_time'


def _multiprocess_sklearn_video(data: pd.DataFrame,
                                bp_dict: dict,
                                video_save_dir: str,
                                frame_save_dir: str,
                                clf_cumsum: dict,
                                rotate: bool,
                                video_path: str,
                                print_timers: bool,
                                video_setting: bool,
                                frame_setting: bool,
                                pose_threshold: float,
                                clf_confidence: Union[dict, None],
                                show_pose: bool,
                                show_animal_names: bool,
                                bbox: Optional[Literal['axis-aligned', 'animal-aligned']],
                                circle_size: int,
                                font_size: int,
                                space_size: int,
                                text_thickness: int,
                                text_opacity: float,
                                text_bg_clr: Tuple[int, int, int],
                                text_color: Tuple[int, int, int],
                                pose_clr_lst: List[Tuple[int, int, int]],
                                show_gantt: Optional[int],
                                bouts_df: Optional[pd.DataFrame],
                                final_gantt: Optional[np.ndarray],
                                gantt_clrs: List[Tuple[float, float, float]],
                                clf_names: List[str],
                                verbose:bool):

    fourcc, font = cv2.VideoWriter_fourcc(*"mp4v"), cv2.FONT_HERSHEY_DUPLEX
    video_meta_data = get_video_meta_data(video_path=video_path)
    if rotate:
        video_meta_data["height"], video_meta_data["width"] = (video_meta_data['width'], video_meta_data['height'])
    cap = cv2.VideoCapture(video_path)
    batch, data = data
    start_frm, current_frm, end_frm = (data["index"].iloc[0], data["index"].iloc[0], data["index"].iloc[-1])
    if video_setting:
        video_save_path = os.path.join(video_save_dir, f"{batch}.mp4")
        if show_gantt is None:
            video_writer = cv2.VideoWriter(video_save_path, fourcc, video_meta_data["fps"], (video_meta_data["width"], video_meta_data["height"]))
        else:
            video_writer = cv2.VideoWriter(video_save_path, fourcc, video_meta_data["fps"], (int(video_meta_data["width"] + final_gantt.shape[1]), video_meta_data["height"]))
    cap.set(1, start_frm)
    while current_frm < end_frm:
        ret, img = cap.read()
        if ret:
            clr_cnt = 0
            for animal_cnt, (animal_name, animal_data) in enumerate(bp_dict.items()):
                if show_pose:
                    for bp_no in range(len(animal_data["X_bps"])):
                        x_bp, y_bp, p_bp = (animal_data["X_bps"][bp_no], animal_data["Y_bps"][bp_no], animal_data["P_bps"][bp_no])
                        bp_cords = data.loc[current_frm, [x_bp, y_bp, p_bp]]
                        if bp_cords[p_bp] >= pose_threshold:
                            img = cv2.circle(img, (int(bp_cords[x_bp]), int(bp_cords[y_bp])), circle_size, pose_clr_lst[clr_cnt], -1)
                        clr_cnt += 1
                if show_animal_names:
                    x_bp, y_bp, p_bp = (animal_data["X_bps"][0], animal_data["Y_bps"][0], animal_data["P_bps"][0])
                    bp_cords = data.loc[current_frm, [x_bp, y_bp, p_bp]]
                    img = cv2.putText(img, animal_name, (int(bp_cords[x_bp]), int(bp_cords[y_bp])), font, font_size, pose_clr_lst[0],  text_thickness)
                if bbox is not None:
                    animal_headers = [val for pair in zip(animal_data["X_bps"], animal_data["Y_bps"]) for val in pair]
                    animal_cords = data.loc[current_frm, animal_headers].values.reshape(-1, 2).astype(np.int32)
                    try:
                        if bbox == Options.AXIS_ALIGNED.value:
                            animal_bbox = GeometryMixin().keypoints_to_axis_aligned_bounding_box(keypoints=animal_cords.reshape(-1, len(animal_cords), 2).astype(np.int32))
                        else:
                            animal_bbox = GeometryMixin().minimum_rotated_rectangle(shape=animal_cords, buffer=None)
                            animal_bbox = np.round(np.array(animal_bbox.exterior.coords)).astype(np.int32)
                        img = cv2.polylines(img, [animal_bbox], True, pose_clr_lst[animal_cnt], thickness=circle_size, lineType=cv2.LINE_AA)
                    except Exception as e:
                        pass
            if rotate:
                img = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)
            if show_gantt == 1:
                img = np.concatenate((img, final_gantt), axis=1)
            elif show_gantt == 2:
                bout_rows = bouts_df.loc[bouts_df["End_frame"] <= current_frm]
                gantt_plot = PlottingMixin().make_gantt_plot(x_length=current_frm + 1,
                                                       bouts_df=bout_rows,
                                                       clf_names=clf_names,
                                                       fps=video_meta_data['fps'],
                                                       width=video_meta_data['width'],
                                                       height=video_meta_data['height'],
                                                       font_size=12,
                                                       font_rotation=45,
                                                       x_tick_lbl_rotation=45,
                                                       hhmmss=True,
                                                       video_name=video_meta_data['video_name'],
                                                       save_path=None,
                                                       palette=gantt_clrs)
                img = np.concatenate((img, gantt_plot), axis=1)
            if print_timers:
                img = PlottingMixin().put_text(img=img, text="TIMERS:", pos=(TextOptions.BORDER_BUFFER_Y.value, ((video_meta_data["height"] - video_meta_data["height"]) + space_size)), font_size=font_size, font_thickness=text_thickness, font=font, text_bg_alpha=text_opacity, text_color_bg=text_bg_clr, text_color=text_color)
            add_spacer = 2
            for clf_name, clf_time_df in clf_cumsum.items():
                frame_results = clf_time_df.loc[current_frm]
                clf_time = round(frame_results / video_meta_data['fps'], 2)
                if print_timers:
                    img = PlottingMixin().put_text(img=img, text=f"{clf_name} {clf_time}",pos=(TextOptions.BORDER_BUFFER_Y.value, ((video_meta_data["height"] - video_meta_data["height"]) + space_size * add_spacer)), font_size=font_size,  font_thickness=text_thickness, font=font, text_bg_alpha=text_opacity, text_color_bg=text_bg_clr, text_color=text_color)
                    add_spacer += 1
                if clf_confidence is not None:
                    conf = round(clf_confidence[clf_name][current_frm], 4)
                    frm_clf_conf_txt = f'{clf_name} CONFIDENCE: {conf:.4f}'
                    img = PlottingMixin().put_text(img=img, text=frm_clf_conf_txt,pos=(TextOptions.BORDER_BUFFER_Y.value, ((video_meta_data["height"] - video_meta_data["height"]) + space_size * add_spacer)), font_size=font_size,  font_thickness=text_thickness, font=font, text_bg_alpha=text_opacity, text_color_bg=text_bg_clr, text_color=text_color)
                    add_spacer += 1

            img = PlottingMixin().put_text(img=img, text="ENSEMBLE PREDICTION:", pos=(TextOptions.BORDER_BUFFER_Y.value, ((video_meta_data["height"] - video_meta_data["height"]) + space_size * add_spacer)), font_size=font_size, font_thickness=text_thickness, font=font, text_bg_alpha=text_opacity, text_color_bg=text_bg_clr, text_color=text_color)
            add_spacer += 1
            for clf_name in clf_cumsum.keys():
                if data.loc[current_frm, clf_name] == 1:
                    img = PlottingMixin().put_text(img=img, text=clf_name, pos=(TextOptions.BORDER_BUFFER_Y.value, (video_meta_data["height"] - video_meta_data["height"]) + space_size * add_spacer), font_size=font_size, font_thickness=text_thickness, font=font, text_color=TextOptions.FLAMINGO.value, text_bg_alpha=text_opacity)
                    add_spacer += 1
            if video_setting:
                video_writer.write(img.astype(np.uint8))
            if frame_setting:
                frame_save_name = os.path.join(frame_save_dir, f"{current_frm}.png")
                cv2.imwrite(frame_save_name, img)
            current_frm += 1
            time = seconds_to_timestamp(seconds= current_frm/ video_meta_data['fps'])
            if verbose: stdout_information(msg=f"Multi-processing video frame {current_frm} (time-stamp: {time}, core batch: {batch}, video name: {video_meta_data['video_name']})...")
        else:
            FrameRangeWarning(msg=f'Could not read frame {current_frm} in video {video_path}. Stopping video creation.')
            break

    cap.release()
    if video_setting:
        video_writer.release()
    return batch


class PlotSklearnResultsMultiProcess(ConfigReader, TrainModelMixin, PlottingMixin):
    """
    Plot classification results on videos using multiprocessing. Results are stored in the
    `project_folder/frames/output/sklearn_results` directory of the SimBA project.

    This class creates annotated videos/frames showing classifier predictions overlaid on pose-estimation data,
    with optional Gantt charts, timers, and bounding boxes. Processing is parallelized across multiple CPU cores
    for faster rendering of large video datasets.

    .. seealso::
       `Tutorial <https://github.com/sgoldenlab/simba/blob/master/docs/tutorial.md#step-10-sklearn-visualization>`__.
       For single-core processing, see :meth:`simba.plotting.plot_clf_results.PlotSklearnResultsSingleCore`.

    .. image:: _static/img/sklearn_visualization.gif
       :width: 600
       :align: center

    .. video:: _static/img/T1.webm
       :width: 1000
       :autoplay:
       :loop:
       :muted:
       :align: center

    ..  youtube:: Frq6mMcaHBc
       :width: 640
       :height: 480
       :align: center

    :param Union[str, os.PathLike] config_path: Path to SimBA project config file in Configparser format.
    :param bool video_setting: If True, creates compressed MP4 videos. Default True.
    :param bool frame_setting: If True, saves individual annotated frames as PNG images. Default False.
    :param Optional[Union[List[Union[str, os.PathLike]], Union[str, os.PathLike]]] video_paths: Path(s) to video file(s) to process. If None, processes all videos found in the project's video directory. Default None.
    :param bool rotate: If True, rotates output videos 90 degrees clockwise. Default False.
    :param bool animal_names: If True, displays animal names on the video frames. Default False.
    :param bool show_pose: If True, overlays pose-estimation keypoints on the video. Default True.
    :param Optional[Union[int, float]] font_size: Font size for text overlays. If None, auto-computed based on video resolution. Default None.
    :param Optional[Union[int, float]] space_size: Vertical spacing between text lines. If None, auto-computed. Default None.
    :param Optional[Union[int, float]] text_thickness: Thickness of text characters. If None, uses default. Default None.
    :param Optional[Union[int, float]] text_opacity: Opacity of text background (0.0-1.0). If None, defaults to 0.8. Default None.
    :param Optional[Union[int, float]] circle_size: Radius of pose keypoint circles. If None, auto-computed based on video resolution. Default None.
    :param Optional[str] pose_palette: Name of color palette for pose keypoints. Must be from :class:`simba.utils.enums.Options.PALETTE_OPTIONS_CATEGORICAL` or :class:`simba.utils.enums.Options.PALETTE_OPTIONS`. Default 'Set1'.
    :param bool print_timers: If True, displays cumulative time for each classifier behavior on each frame. Default True.
    :param Optional[Literal['axis-aligned', 'animal-aligned']] bbox: If not None, draw bounding boxes around each animal. ``'axis-aligned'`` = rectangle aligned with video axes; ``'animal-aligned'`` = minimum rotated rectangle aligned with the animal's orientation. Default None (no bounding boxes).
    :param Optional[int] show_gantt: If 1, appends static Gantt chart to video. If 2, appends dynamic Gantt chart that updates per frame. If None, no Gantt chart. Default None.
    :param Tuple[int, int, int] text_clr: RGB color tuple for text foreground. Default (255, 255, 255) (white).
    :param Tuple[int, int, int] text_bg_clr: RGB color tuple for text background. Default (0, 0, 0) (black).
    :param bool gpu: If True, uses GPU acceleration for video concatenation (requires CUDA-capable GPU). Default False.
    :param int core_cnt: Number of CPU cores to use for parallel processing. Pass -1 to use all available cores. Default -1.

    :example:
    >>> clf_plotter = PlotSklearnResultsMultiProcess(
    ...     config_path='/Users/simon/Desktop/envs/simba/troubleshooting/beepboop174/project_folder/project_config.ini',
    ...     video_setting=True,
    ...     frame_setting=False,
    ...     video_paths='Trial_10.mp4',
    ...     rotate=False,
    ...     show_pose=True,
    ...     bbox='axis-aligned',
    ...     print_timers=True,
    ...     show_gantt=1,
    ...     core_cnt=5
    ... )
    >>> clf_plotter.run()
    """

    def __init__(self,
                 config_path: Union[str, os.PathLike],
                 video_setting: bool = True,
                 frame_setting: bool = False,
                 video_paths: Optional[Union[List[Union[str, os.PathLike]], Union[str, os.PathLike]]] = None,
                 rotate: bool = False,
                 animal_names: bool = False,
                 show_pose: bool = True,
                 show_confidence: bool = False,
                 font_size: Optional[Union[int, float]] = None,
                 space_size: Optional[Union[int, float]] = None,
                 text_thickness: Optional[Union[int, float]] = None,
                 text_opacity: Optional[Union[int, float]] = None,
                 circle_size: Optional[Union[int, float]] = None,
                 pose_palette: Optional[str] = 'Set1',
                 print_timers: bool = True,
                 bbox: Optional[Literal['axis-aligned', 'animal-aligned']] = None,
                 time_slice: Optional[Dict[str, str]] = None,
                 show_gantt: Optional[int] = None,
                 text_clr: Tuple[int, int, int] = (255, 255, 255),
                 text_bg_clr: Tuple[int, int, int] = (0, 0, 0),
                 gpu: bool = False,
                 verbose: bool = True,
                 core_cnt: int = -1):


        ConfigReader.__init__(self, config_path=config_path)
        TrainModelMixin.__init__(self)
        PlottingMixin.__init__(self)
        log_event(logger_name=str(__class__.__name__), log_type=TagNames.CLASS_INIT.value, msg=self.create_log_msg_from_init_args(locals=locals()))
        for i in [video_setting, frame_setting, rotate, print_timers, animal_names, show_pose, gpu, show_confidence]:
            check_valid_boolean(value=i, source=self.__class__.__name__, raise_error=True)
        if (not video_setting) and (not frame_setting):
            raise NoSpecifiedOutputError(msg="Please choose to create a video and/or frames. SimBA found that you ticked neither video and/or frames", source=self.__class__.__name__)
        if font_size is not None: check_float(name=f'{self.__class__.__name__} font_size', value=font_size, min_value=0.1)
        if space_size is not None: check_float(name=f'{self.__class__.__name__} space_size', value=space_size, 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)
        if circle_size is not None: check_float(name=f'{self.__class__.__name__} text_thickness', value=circle_size, min_value=0.1)
        if circle_size is not None: check_float(name=f'{self.__class__.__name__} text_thickness', value=circle_size, min_value=0.1)
        if text_opacity is not None: check_float(name=f'{self.__class__.__name__} text_opacity', value=text_opacity, min_value=0.1)
        pose_palettes = Options.PALETTE_OPTIONS_CATEGORICAL.value + Options.PALETTE_OPTIONS.value
        check_str(name=f'{self.__class__.__name__} pose_palette', value=pose_palette, options=pose_palettes)
        self.clr_lst = create_color_palette(pallete_name=pose_palette, increments=len(self.body_parts_lst)+1)
        check_if_valid_rgb_tuple(data=text_clr, source=f'{self.__class__.__name__} text_clr')
        check_if_valid_rgb_tuple(data=text_bg_clr, source=f'{self.__class__.__name__} text_bg_clr')
        check_valid_boolean(value=verbose, source=f'{self.__class__.__name__} verbose', raise_error=True)
        if show_gantt is not None: check_int(name=f"{self.__class__.__name__} show_gantt", value=show_gantt, max_value=2, min_value=1)
        if time_slice is not None:
            check_valid_dict(x=time_slice, valid_key_dtypes=(str,), valid_values_dtypes=(str,), valid_keys=(START_TIME, END_TIME), required_keys=(START_TIME, END_TIME),)
            check_if_string_value_is_valid_video_timestamp(value=time_slice[START_TIME], name='START TIME', raise_error=True)
            check_if_string_value_is_valid_video_timestamp(value=time_slice[END_TIME], name='END TIME', raise_error=True)
            check_that_hhmmss_start_is_before_end(start_time=time_slice[START_TIME], end_time=time_slice[END_TIME], name=f'TIME SLICE', raise_error=True)
        if bbox is not None:
            check_str(name=f'{self.__class__.__name__} bbox', value=bbox, options=Options.BBOX_OPTIONS.value, allow_blank=False, raise_error=True)
        self.video_setting, self.frame_setting, self.rotate, self.print_timers = video_setting, frame_setting, rotate, print_timers
        self.circle_size, self.font_size, self.animal_names, self.text_opacity = circle_size, font_size, animal_names, text_opacity
        self.text_thickness, self.space_size, self.show_pose, self.pose_palette, self.verbose = text_thickness, space_size, show_pose, pose_palette, verbose
        self.text_color, self.text_bg_color, self.bbox, self.show_gantt, self.show_confidence = text_clr, text_bg_clr, bbox, show_gantt, show_confidence
        self.gpu = True if check_nvidea_gpu_available() and gpu else False
        self.pose_threshold = read_config_entry(self.config, ConfigKey.THRESHOLD_SETTINGS.value, ConfigKey.SKLEARN_BP_PROB_THRESH.value, Dtypes.FLOAT.value, 0.00)
        self.time_slice = time_slice
        if not os.path.exists(self.sklearn_plot_dir):
            os.makedirs(self.sklearn_plot_dir)
        if isinstance(video_paths, str): self.video_paths = [video_paths]
        elif isinstance(video_paths, list): self.video_paths = video_paths
        elif video_paths is None:
            self.video_paths = find_all_videos_in_project(videos_dir=self.video_dir)
            if len(self.video_paths) == 0:
                raise NoDataError(msg=f'Cannot create classification videos. No videos exist in {self.video_dir} directory', source=self.__class__.__name__)
        else:
            raise InvalidInputError(msg=f'video_paths has to be a path of a list of paths. Got {type(video_paths)}', source=self.__class__.__name__)

        for video_path in self.video_paths:
            video_name = get_fn_ext(filepath=video_path)[1]
            data_path = os.path.join(self.machine_results_dir, f'{video_name}.{self.file_type}')
            if not os.path.isfile(data_path): raise NoDataError(msg=f'Cannot create classification videos for {video_name}. Expected classification data at location {data_path} but file does not exist', source=self.__class__.__name__)
        check_int(name=f'{self.__class__.__name__} core_cnt', value=core_cnt, min_value=-1, unaccepted_vals=[0])
        self.core_cnt = find_core_cnt()[0] if int(core_cnt) == -1 or int(core_cnt) > find_core_cnt()[0] else int(core_cnt)
        self.conf_cols = [f'Probability_{x}' for x in self.clf_names]
        if platform.system() == "Darwin":
            multiprocessing.set_start_method("spawn", force=True)

    def __get_print_settings(self):
        optimal_circle_size = self.get_optimal_circle_size(frame_size=(self.video_meta_data["width"], self.video_meta_data["height"]), circle_frame_ratio=100)
        longest_str = str(max(['TIMERS:', 'ENSEMBLE PREDICTION:'] + self.clf_names, key=len))
        self.video_text_thickness = TextOptions.TEXT_THICKNESS.value if self.text_thickness is None else int(max(self.text_thickness, 1))
        optimal_font_size, _, optimal_spacing_scale = self.get_optimal_font_scales(text=longest_str, accepted_px_width=int(self.video_meta_data["width"] / 2), accepted_px_height=int(self.video_meta_data["height"] / 5), text_thickness=self.video_text_thickness)
        self.video_circle_size = optimal_circle_size if self.circle_size is None else int(max(1, 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.space_size is None else int(max(self.space_size, 1))
        self.video_text_opacity = 0.8 if self.text_opacity is None else float(self.text_opacity)

    def run(self):
        if self.verbose: stdout_information(msg=f'Creating {len(self.video_paths)} classification visualization(s) using {self.core_cnt} cores...')
        self.pool = get_cpu_pool(core_cnt=self.core_cnt, source=self.__class__.__name__, )
        for video_cnt, video_path in enumerate(self.video_paths):
            video_timer = SimbaTimer(start=True)
            _, self.video_name, _ = get_fn_ext(video_path)
            if self.verbose: stdout_information(msg=f"Creating classification visualization for video {self.video_name}...")
            self.data_path = os.path.join(self.machine_results_dir, f'{self.video_name}.{self.file_type}')
            self.data_df = read_df(self.data_path, self.file_type).reset_index(drop=True).fillna(0)
            if self.show_pose: check_that_column_exist(df=self.data_df, column_name=self.bp_col_names, file_name=self.data_path)
            if self.show_confidence: check_that_column_exist(df=self.data_df, column_name=self.conf_cols, file_name=self.data_path)
            self.video_meta_data = get_video_meta_data(video_path=video_path)
            if self.time_slice is not None:
                check_if_string_value_is_valid_video_timestamp(value=self.time_slice[START_TIME], name=f'START TIME {START_TIME}')
                check_if_string_value_is_valid_video_timestamp(value=self.time_slice[END_TIME], name=f'END TIME {END_TIME}')
                frm_ids = find_frame_numbers_from_time_stamp(start_time=self.time_slice[START_TIME], end_time=self.time_slice[END_TIME], fps=int(self.video_meta_data['fps']))
                self.data_df = self.data_df.loc[frm_ids]
            height, width = deepcopy(self.video_meta_data["height"]), deepcopy(self.video_meta_data["width"])
            self.save_path = os.path.join(self.sklearn_plot_dir, f"{self.video_name}.mp4")
            self.video_frame_dir, self.video_temp_dir = None, None
            if self.video_setting:
                self.video_save_path = os.path.join(self.sklearn_plot_dir, f"{self.video_name}.mp4")
                self.video_temp_dir = os.path.join(self.sklearn_plot_dir, self.video_name, "temp")
                create_directory(paths=self.video_temp_dir, overwrite=True)
            if self.frame_setting:
                self.video_frame_dir = os.path.join(self.sklearn_plot_dir, self.video_name)
                create_directory(paths=self.video_temp_dir, overwrite=True)
            if self.rotate:
                self.video_meta_data["height"], self.video_meta_data["width"] = (width, height)
            if self.time_slice is None: check_video_and_data_frm_count_align(video=video_path, data=self.data_df, name=self.video_name, raise_error=False)
            check_that_column_exist(df=self.data_df, column_name=self.clf_names, file_name=self.data_path)
            self.__get_print_settings()
            if self.show_gantt is not None:
                self.gantt_clrs = create_color_palette(pallete_name=self.pose_palette, increments=len(self.clf_names) + 1, as_int=True, as_rgb_ratio=True)
                self.bouts_df = detect_bouts(data_df=self.data_df, target_lst=list(self.clf_names), fps=int(self.video_meta_data["fps"]))
                self.final_gantt_img = PlottingMixin().make_gantt_plot(x_length=len(self.data_df) + 1,
                                                                       bouts_df=self.bouts_df,
                                                                       clf_names=self.clf_names,
                                                                       fps=self.video_meta_data["fps"],
                                                                       width=self.video_meta_data["width"],
                                                                       height=self.video_meta_data["height"],
                                                                       font_size=12,
                                                                       font_rotation=45,
                                                                       x_tick_lbl_rotation=45,
                                                                       hhmmss=True,
                                                                       video_name=self.video_meta_data["video_name"],
                                                                       save_path=None,
                                                                       palette=self.gantt_clrs)
                self.final_gantt_img = self.resize_gantt(self.final_gantt_img, self.video_meta_data["height"])
            else:
                self.bouts_df, self.final_gantt_img, self.gantt_clrs = None, None, None


            self.clf_cumsums, self.clf_p = {}, {} if self.show_confidence else None
            for clf_name in self.clf_names:
                self.clf_cumsums[clf_name] = self.data_df[clf_name].cumsum()
                if self.show_confidence: self.clf_p[clf_name] = np.round(self.data_df[f'Probability_{clf_name}'].values.reshape(-1), 4)

            self.data_df["index"] = self.data_df.index
            data = np.array_split(self.data_df, self.core_cnt)
            data = [(cnt, x) for (cnt, x) in enumerate(data)]

            constants = functools.partial(_multiprocess_sklearn_video,
                                          bp_dict=self.animal_bp_dict,
                                          video_save_dir=self.video_temp_dir,
                                          frame_save_dir=self.video_frame_dir,
                                          clf_cumsum=self.clf_cumsums,
                                          rotate=self.rotate,
                                          video_path=video_path,
                                          clf_confidence=self.clf_p,
                                          print_timers=self.print_timers,
                                          video_setting=self.video_setting,
                                          frame_setting=self.frame_setting,
                                          pose_threshold=self.pose_threshold,
                                          show_pose=self.show_pose,
                                          show_animal_names=self.animal_names,
                                          circle_size=self.video_circle_size,
                                          font_size=self.video_font_size,
                                          space_size=self.video_space_size,
                                          text_thickness=self.video_text_thickness,
                                          text_opacity=self.video_text_opacity,
                                          text_bg_clr=self.text_bg_color,
                                          text_color=self.text_color,
                                          pose_clr_lst=self.clr_lst,
                                          bbox=self.bbox,
                                          show_gantt=self.show_gantt,
                                          bouts_df=self.bouts_df,
                                          final_gantt=self.final_gantt_img,
                                          gantt_clrs=self.gantt_clrs,
                                          clf_names=self.clf_names,
                                          verbose=self.verbose)

            for cnt, result in enumerate(self.pool.imap(constants, data, chunksize=self.multiprocess_chunksize)):
                if self.verbose: stdout_information(f"Image batch {result} complete, Video {(video_cnt + 1)}/{len(self.video_paths)}...")

            if self.video_setting:
                if self.verbose: stdout_information(msg=f"Joining {self.video_name} multiprocessed video...")
                concatenate_videos_in_folder(in_folder=self.video_temp_dir, save_path=self.video_save_path, gpu=self.gpu, verbose=self.verbose)
                video_timer.stop_timer()
                stdout_information(msg=f"Video {self.video_name} complete (elapsed time: {video_timer.elapsed_time_str}s)...")

        terminate_cpu_pool(pool=self.pool, force=False, source=self.__class__.__name__)
        self.timer.stop_timer()
        if self.video_setting:
            stdout_success(msg=f"{len(self.video_paths)} video(s) saved in {self.sklearn_plot_dir} directory", elapsed_time=self.timer.elapsed_time_str, source=self.__class__.__name__)
        if self.frame_setting:
            stdout_success(f"Frames for {len(self.video_paths)} videos saved in sub-folders within {self.sklearn_plot_dir} directory", elapsed_time=self.timer.elapsed_time_str, source=self.__class__.__name__)



if __name__ == "__main__":
    clf_plotter = PlotSklearnResultsMultiProcess(config_path=r"E:\troubleshooting\mitra_pbn\mitra_pbn\project_folder\project_config.ini",
                                                 video_paths=None, #r"E:\troubleshooting\mitra_pbn\mitra_pbn\project_folder\videos\2026-01-05 14-17-54 box1_1143_0_Gq_sal.mp4",
                                                 video_setting=True,
                                                 frame_setting=False,
                                                 print_timers=True,
                                                 rotate=False,
                                                 show_confidence=False,
                                                 core_cnt=16,
                                                 animal_names=False,
                                                 #time_slice={START_TIME: '00:00:00', END_TIME: '00:01:00'},
                                                 bbox=None,
                                                 text_opacity=0.6,
                                                 show_gantt=2)
    clf_plotter.run()




# if __name__ == "__main__":
#     clf_plotter = PlotSklearnResultsMultiProcess(config_path=r"/Users/simon/Desktop/envs/simba/troubleshooting/two_black_animals_14bp/project_folder/project_config.ini",
#                                                 video_setting=True,
#                                                 frame_setting=False,
#                                                 video_paths=r"/Users/simon/Desktop/envs/simba/troubleshooting/two_black_animals_14bp/project_folder/videos/Together_1.mp4",
#                                                 print_timers=True,
#                                                 rotate=False,
#                                                 animal_names=False,
#                                                 show_bbox=True,
#                                                 show_gantt=None)
#     clf_plotter.run()




#text_settings = {'circle_scale': 5, 'font_size': 0.528, 'spacing_scale': 28, 'text_thickness': 2}
# clf_plotter = PlotSklearnResultsMultiProcess(config_path='/Users/simon/Desktop/envs/simba/troubleshooting/beepboop174/project_folder/project_config.ini',
#                                              video_setting=True,
#                                              frame_setting=False,
#                                              rotate=False,
#                                              video_file_path='592_MA147_Gq_CNO_0515.mp4',
#                                              cores=-1,
#                                              text_settings=False)
# clf_plotter.run()
#

# clf_plotter = PlotSklearnResultsMultiProcess(config_path='/Users/simon/Desktop/envs/troubleshooting/DLC_2_Black_animals/project_folder/project_config.ini', video_setting=True, frame_setting=False, rotate=False, video_file_path='Together_1.avi', cores=5)
# clf_plotter.run()

# if __name__ == "__main__":
#     clf_plotter = PlotSklearnResultsMultiProcess(config_path=r"C:\troubleshooting\mitra\project_folder\project_config.ini",
#                                                  video_setting = True,
#                                                  frame_setting = False,
#                                                  rotate = False,
#                                                  core_cnt = 6,
#                                                  show_confidence=True,
#                                                  video_paths=r"C:\troubleshooting\mitra\project_folder\videos\501_MA142_Gi_CNO_0521.mp4")
#     clf_plotter.run()