__author__ = "Simon Nilsson"

import os
from typing import List, Union

import cv2
import numpy as np

from simba.mixins.config_reader import ConfigReader
from simba.mixins.geometry_mixin import GeometryMixin
from simba.mixins.plotting_mixin import PlottingMixin
from simba.utils.checks import (
    check_all_file_names_are_represented_in_video_log,
    check_filepaths_in_iterable_exist, check_int, check_str,
    check_valid_dataframe, check_valid_dict)
from simba.utils.enums import Formats, Options
from simba.utils.errors import InvalidInputError, NoSpecifiedOutputError
from simba.utils.printing import SimbaTimer, stdout_success
from simba.utils.read_write import get_fn_ext, read_df


class HeatMapperClfSingleCore(ConfigReader, PlottingMixin):
    """
    Create heatmaps representing the locations of the classified behavior.

    .. note::
       `GitHub visualizations tutorial <https://github.com/sgoldenlab/simba/blob/master/docs/tutorial.md#step-11-visualizations>`__.
       For improved run-time, see :meth:`simba.heat_mapper_clf_mp.HeatMapperClfMultiprocess` for multiprocess class.

    .. image:: _static/img/heatmap.png
       :width: 500
       :align: center

    :param str config_path: path to SimBA project config file in Configparser format
    :param bool final_img_setting: If True, then create a single image representing the last frame of the input video
    :param bool video_setting: If True, then create a video of heatmaps.
    :param bool frame_setting: If True, then create individual heatmap frames.
    :param str clf_name: The name of the classified behavior.
    :param str bodypart: The name of the body-part used to infer the location of the classified behavior
    :param Dict style_attr: Dict containing settings for colormap, bin-size, max scale, and smooothing operations. For example: {'palette': 'jet', 'shading': 'gouraud', 'bin_size': 50, 'max_scale': 'auto'}.

    :example:
    >>> test = HeatMapperClfSingleCore(config_path=r"C:\troubleshooting\RAT_NOR\project_folder\project_config.ini",
    >>>                  style_attr = {'palette': 'jet', 'shading': 'gouraud', 'bin_size': 50, 'max_scale': 'auto'},
    >>>                  final_img_setting=True,
    >>>                  video_setting=True,
    >>>                  frame_setting=False,
    >>>                  bodypart='Ear_left',
    >>>                  clf_name='straub_tail',
    >>>                  data_paths=[r"C:\troubleshooting\RAT_NOR\project_folder\csv\test\2022-06-20_NOB_DOT_4.csv"])
    >>> test.run()
    """

    def __init__(self,
                 config_path: Union[str, os.PathLike],
                 bodypart: str,
                 clf_name: str,
                 data_paths: List[str],
                 style_attr: dict,
                 final_img_setting: bool = True,
                 video_setting: bool = False,
                 frame_setting: bool = False):



        ConfigReader.__init__(self, config_path=config_path)
        PlottingMixin.__init__(self)
        if (not frame_setting) and (not video_setting) and (not final_img_setting):
            raise NoSpecifiedOutputError(msg="Please select either heatmap videos, frames, and/or final image.")
        check_filepaths_in_iterable_exist(file_paths=data_paths, name=f'{self.__class__.__name__} data_paths')
        check_str(name=f'{self.__class__.__name__} clf_name', value=clf_name)
        check_str(name=f'{self.__class__.__name__} bodypart', value=bodypart)
        check_valid_dict(x=style_attr, required_keys=('max_scale', 'bin_size', 'shading', 'palette'))
        self.frame_setting, self.video_setting, self.final_img_setting = frame_setting, video_setting, final_img_setting
        self.bin_size, self.max_scale, self.palette, self.shading = (style_attr["bin_size"], style_attr["max_scale"], style_attr["palette"], style_attr["shading"])
        check_str(name=f'{self.__class__.__name__} shading', value=style_attr["shading"], options=Options.HEATMAP_SHADING_OPTIONS.value)
        check_int(name=f'{self.__class__.__name__} bin_size', value=style_attr["bin_size"], min_value=1)
        self.clf_name, self.data_paths, self.bp = clf_name, data_paths, bodypart
        if not os.path.exists(self.heatmap_clf_location_dir): os.makedirs(self.heatmap_clf_location_dir)
        self.bp_lst = [f"{self.bp}_x", f"{self.bp}_y"]
        self.timer = SimbaTimer(start=True)

    def __calculate_max_scale(self, clf_array: np.array):
        return np.round(np.max(np.max(clf_array[-1], axis=0)), 3)

    def run(self):
        print(f"Processing heatmaps for {len(self.data_paths)} video(s)...")
        check_all_file_names_are_represented_in_video_log(video_info_df=self.video_info_df, data_paths=self.data_paths)
        for file_cnt, file_path in enumerate(self.data_paths):
            video_timer = SimbaTimer(start=True)
            _, self.video_name, _ = get_fn_ext(file_path)
            print(f'Plotting heatmap classification map for video {self.video_name}...')
            self.video_info, self.px_per_mm, self.fps = self.read_video_info(video_name=self.video_name)
            self.width, self.height = int(self.video_info["Resolution_width"].values[0]), int(self.video_info["Resolution_height"].values[0])
            if self.video_setting:
                self.fourcc = cv2.VideoWriter_fourcc(*Formats.MP4_CODEC.value)
                self.video_save_path = os.path.join(self.heatmap_clf_location_dir, f"{self.video_name}.mp4")
                self.writer = cv2.VideoWriter(self.video_save_path, self.fourcc, self.fps, (self.width, self.height))
            if self.frame_setting:
                self.save_video_folder = os.path.join(self.heatmap_clf_location_dir, self.video_name)
                if not os.path.exists(self.save_video_folder):
                    os.makedirs(self.save_video_folder)
            self.data_df = read_df(file_path=file_path, file_type=self.file_type)
            check_valid_dataframe(df=self.data_df, required_fields=[self.clf_name] + self.bp_lst, valid_dtypes=Formats.NUMERIC_DTYPES.value)
            bp_data = self.data_df[self.bp_lst].values.astype(np.int32)
            clf_data = self.data_df[self.clf_name].values.astype(np.int32)
            if len(np.unique(clf_data)) == 1:
                raise InvalidInputError(msg=f'Cannot plot heatmap for behavior {self.clf_name} in video {self.video_name}. The behavior is classified as {np.unique(clf_data)} in every single frame.')
            grid, aspect_ratio = GeometryMixin.bucket_img_into_grid_square(img_size=(self.width, self.height), bucket_grid_size_mm=self.bin_size, px_per_mm=self.px_per_mm, add_correction=False)
            clf_data = GeometryMixin().cumsum_bool_geometries(data=bp_data, geometries=grid, bool_data=clf_data, fps=self.fps, verbose=False)
            if self.max_scale == "auto":
                self.max_scale = max(1, self.__calculate_max_scale(clf_array=clf_data))
            if self.final_img_setting:
                file_name = os.path.join(self.heatmap_clf_location_dir, f"{self.video_name}_final_frm.png")
                self.make_location_heatmap_plot(frm_data=clf_data[-1:, :, :][0],
                                                max_scale=self.max_scale,
                                                palette=self.palette,
                                                aspect_ratio=aspect_ratio,
                                                file_name=file_name,
                                                shading=self.shading,
                                                img_size=(self.width, self.height))
                print(f"Final heatmap image saved at {file_name}.")

            if self.video_setting or self.frame_setting:
                for frm_cnt, cumulative_frm_idx in enumerate(range(clf_data.shape[0])):
                    frm_data = clf_data[cumulative_frm_idx, :, :]
                    img = self.make_location_heatmap_plot(frm_data=frm_data,
                                                          max_scale=self.max_scale,
                                                          palette=self.palette,
                                                          aspect_ratio=aspect_ratio,
                                                          shading=self.shading,
                                                          img_size=(self.width, self.height))[:,:,:3]

                    if self.video_setting:
                        self.writer.write(img)
                    if self.frame_setting:
                        frame_save_path = os.path.join(self.save_video_folder, f"{frm_cnt}.png")
                        cv2.imwrite(frame_save_path, img)
                    print(f"Created heatmap frame: {frm_cnt+1} / {len(self.data_df)}. Video: {self.video_name} ({file_cnt + 1}/{len(self.data_paths)})")
            if self.video_setting:
                self.writer.release()
            video_timer.stop_timer()
            print(f"Heatmap plot for video {self.video_name} saved at {self.heatmap_clf_location_dir} (elapsed time: {video_timer.elapsed_time_str}s)...")

        self.timer.stop_timer()
        stdout_success(msg=f"All heatmap visualizations created in {self.heatmap_clf_location_dir} directory", elapsed_time=self.timer.elapsed_time_str)


# test = HeatMapperClfSingleCore(config_path=r"C:\troubleshooting\RAT_NOR\project_folder\project_config.ini",
#                      style_attr = {'palette': 'jet', 'shading': 'gouraud', 'bin_size': 50, 'max_scale': 'auto'},
#                      final_img_setting=True,
#                      video_setting=True,
#                      frame_setting=False,
#                      bodypart='Ear_left',
#                      clf_name='straub_tail',
#                      data_paths=[r"C:\troubleshooting\RAT_NOR\project_folder\csv\test\2022-06-20_NOB_DOT_4.csv"])
# test.run()





# test = HeatMapperClfSingleCore(config_path='/Users/simon/Desktop/envs/troubleshooting/Two_animals_16bps/project_folder/project_config.ini',
#                      style_attr = {'palette': 'jet', 'shading': 'gouraud', 'bin_size': 75, 'max_scale': 'auto'},
#                      final_img_setting=False,
#                      video_setting=True,
#                      frame_setting=False,
#                      bodypart='Nose_1',
#                      clf_name='Attack',
#                      files_found=['/Users/simon/Desktop/envs/troubleshooting/Two_animals_16bps/project_folder/csv/machine_results/Together_3.csv'])
# test.run()
