import numpy as np
import joblib

from tqdm import tqdm
from sklearn.tree import DecisionTreeClassifier

from sklearn.ensemble import AdaBoostClassifier

import ethologger.utils.auxiliary as aux

from ethologger.experiments.profile import ActiveFrames
from ethologger.experiments.profile import DormantIntervals
from ethologger.utils.visualization import VisualizationWrapper as VizWrap
from ethologger.project.helper import ProjectIO, Logger


class FrameExtraction(ProjectIO):
    def __init__(
        self,
        fps,
        project_path,
        expt_path_dict,
        annotation_path_dict,
        moving_threshold_args,
        dormant_intervals_args,
        activity_datums_list,
        activity_threshold_args,
        decision_tree_args=None,
    ):
        self.fps = fps
        self.project_path = project_path
        self.expt_path_dict = expt_path_dict
        self.annotation_path_dict = annotation_path_dict
        self.moving_threshold_args = moving_threshold_args
        self.dormant_intervals_args = dormant_intervals_args
        self.activity_datums_list = activity_datums_list
        self.activity_threshold_args = activity_threshold_args
        self.decision_tree_args = decision_tree_args
        if self.decision_tree_args is None:
            self.decision_tree_args = {}
        self.logger = Logger()
        ProjectIO.__init__(self)

    @aux.timeit
    def extract_dormant_intervals(self):
        ExtDormant = DormantIntervals(fps=self.fps)

        pbar = tqdm(self.expt_path_dict.items())
        for name, expt_path in pbar:
            pbar.set_description(f"Extracting dormant intervals of {name}...")

            expt_item = self._load_expt_item(expt_path)
            dfDelta = self._load_delta(expt_path)
            delta_names_dict = self._load_delta_names_dict(expt_path)

            ExtDormant.delta_names_dict = delta_names_dict
            frame_val = ExtDormant.get_movement_values(
                dfDelta, datums=list(delta_names_dict.values()), value_type="delta"
            )
            threshold = ExtDormant.get_threshold(frame_val, self.moving_threshold_args)

            _ = ExtDormant.set_bouts(frame_val, threshold)
            dormant_moving_labels = ExtDormant.get_labels(**self.dormant_intervals_args)

            expt_item.mask_dormant = dormant_moving_labels == 0
            self.logger.info_saved("Mask of dormant frames")

            viz = VizWrap(tick_stepsize=1.5)
            fig, _ = viz.ethogram(
                dormant_moving_labels,
                ylabel="state",
                title=f"Dormant-Moving Ethogram of {name}",
                label_name_dict=ExtDormant.id_to_state,
            )
            dormant_moving_figpath = (
                expt_path / "figures" / "dormant_moving_ethogram.png"
            )
            fig.savefig(dormant_moving_figpath, bbox_inches="tight")
            self.logger.file_saved(
                "Dormant-Moving ethogram", dormant_moving_figpath, sd=3
            )

            dormant_intervals = []
            intervals = aux.cont_intvls(dormant_moving_labels)
            for i in range(1, intervals.shape[0]):
                if dormant_moving_labels[intervals[i - 1]] == 0:
                    dormant_intervals.append((intervals[i - 1], intervals[i]))

            expt_item.dormant_intervals = dormant_intervals
            self.logger.info_saved("Dormant intervals")

            self._save_expt_item(expt_item, expt_path)

    @staticmethod
    def _get_threshold_labels(mask_list, shape):
        labels = np.zeros(shape, dtype=int)
        for i, mask in enumerate(mask_list):
            labels[mask] = i + 1
        return labels

    def _get_thresholds(self, ActiveFrames_instance, dfCoef, mask_dormant, datums_list):
        thresholds = []
        for datums in datums_list:
            frame_val = ActiveFrames_instance.get_activity_values(
                dfCoef, datums=datums, value_type="wavelet-coef"
            )
            th = ActiveFrames_instance.get_threshold(
                frame_val[mask_dormant], self.activity_threshold_args
            )
            thresholds.append(th)
        return thresholds

    @aux.timeit
    def extract_active_frames(self):
        ExtActive = ActiveFrames(fps=self.fps)

        pbar = tqdm(self.expt_path_dict.items())
        for name, expt_path in pbar:
            pbar.set_description(f"Extracting active frames of {name}...")

            expt_item = self._load_expt_item(expt_path)
            mask_dormant = expt_item.mask_dormant

            ftws = self._load_ws(expt_path)
            snap_names_dict = self._load_snap_names_dict(expt_path)

            ExtActive.snap_names_dict = snap_names_dict
            dfCoef = ExtActive.get_dfCoef(ftws, log=True)
            datums_list = list(self.activity_datums_list)

            thresholds = self._get_thresholds(
                ExtActive, dfCoef, mask_dormant, datums_list
            )
            _, mask_active_p = ExtActive.get_mask_active(
                dfCoef, datums_list, thresholds, value_type="wavelet-coef"
            )

            threshold_labels = self.__class__._get_threshold_labels(
                mask_active_p, ftws.shape[0]
            )
            mask_active = threshold_labels != 0

            activity_labels = np.where(mask_active, 1, 0).astype(int)
            self.logger.number_of_dormant_active(
                name, sum(activity_labels[mask_dormant])
            )

            expt_item.threshold_labels = threshold_labels
            self.logger.info_saved("Threshold labels")

            expt_item.mask_active = mask_active
            self.logger.info_saved("Mask of active frames")

            viz = VizWrap()
            fig, _ = viz.ethogram(
                threshold_labels,
                title=f"Threshold Labels of {name}",
            )
            threshold_lbl_figpath = expt_path / "figures" / "threshold_labels.png"
            fig.savefig(threshold_lbl_figpath, bbox_inches="tight")
            self.logger.file_saved(
                "Broken-hbar plot of threshold labels", threshold_lbl_figpath, sd=3
            )

            fig, _ = viz.ethogram(
                activity_labels,
                title=f"Activity Labels of {name}",
            )
            activity_lbl_figpath = expt_path / "figures" / "activity_labels.png"
            fig.savefig(activity_lbl_figpath, bbox_inches="tight")
            self.logger.file_saved(
                "Broken-hbar plot of activity labels",
                activity_lbl_figpath,
                sd=3,
            )

            self._save_expt_item(expt_item, expt_path)

    @aux.timeit
    def construct_activity_decision_tree(self):
        ExtActive = ActiveFrames(fps=self.fps)

        X_activity_list = []
        y_activity_list = []

        annotation_path_items = self.annotation_path_dict.items()
        if not annotation_path_items:
            raise RuntimeError(
                "At least one annotated experiment must be provided.\n"
                "Activity decision tree can not be constructed.\n"
                "Please use 'extract_active_frames' directly to extract active frames."
            )

        pbar = tqdm(annotation_path_items)
        for name, _ in pbar:
            pbar.set_description(f"Processing {name}...")

            expt_path = self.expt_path_dict[name]

            annotations = self._load_annotations(expt_path)
            expt_item = self._load_expt_item(expt_path)
            ftws = self._load_ws(expt_path)
            snap_names_dict = self._load_snap_names_dict(expt_path)

            mask_dormant = expt_item.mask_dormant

            ExtActive.snap_names_dict = snap_names_dict
            dfCoef = ExtActive.get_dfCoef(ftws, log=True)
            datums_list = list(self.activity_datums_list)

            thresholds = self._get_thresholds(
                ExtActive, dfCoef, mask_dormant, datums_list
            )

            _, mask_active_p = ExtActive.get_mask_active(
                dfCoef, datums_list, thresholds, value_type="wavelet-coef"
            )

            X_activity_list.append(dfCoef.to_numpy()[mask_dormant, :])

            threshold_labels = self.__class__._get_threshold_labels(
                mask_active_p, ftws.shape[0]
            )

            expt_item.threshold_labels = threshold_labels
            self.logger.info_saved("Threshold labels")

            # Alternative way.
            mask_active = threshold_labels != 0
            mask_active[annotations == 7] = False
            # Sensible way of constructing ground-truth.
            # Backup-way to go.
            ## mask_active = annotations != 7
            ## mask_active[threshold_labels == 0] = False
            ## mask_active[annotations == 2] = True
            # Alternative, "more" unsupervised
            ## mask_active = threshold_labels != 0

            y_temp = np.where(mask_active, 1, 0).astype(int)
            y_activity_list.append(y_temp[mask_dormant])

            # Might redundant, can be removed.
            expt_item.gt_activity_labels = y_temp
            self.logger.info_saved("Ground-truth activity labels")

            viz = VizWrap()
            fig, _ = viz.ethogram(
                threshold_labels,
                title=f"Threshold Labels of {name}",
            )
            threshold_lbl_figpath = expt_path / "figures" / "threshold_labels.png"
            fig.savefig(threshold_lbl_figpath, bbox_inches="tight")
            self.logger.file_saved(
                "Broken-hbar plot of threshold labels", threshold_lbl_figpath, sd=3
            )

            fig, _ = viz.ethogram(
                y_temp,
                title=f"Ground-truth Activity Labels of {name}",
            )
            gt_activity_lbl_figpath = expt_path / "figures" / "gt_activity_labels.png"
            fig.savefig(gt_activity_lbl_figpath, bbox_inches="tight")
            self.logger.file_saved(
                "Broken-hbar plot of ground-truth activity labels",
                gt_activity_lbl_figpath,
                sd=3,
            )

            self._save_expt_item(expt_item, expt_path)

        X = np.concatenate(X_activity_list)
        y = np.concatenate(y_activity_list)
        decision_tree = DecisionTreeClassifier(**self.decision_tree_args)
        clf = AdaBoostClassifier(base_estimator=decision_tree, n_estimators=20)
        ## clf = decision_tree
        clf.fit(X=X, y=y)
        decision_tree_path = self.project_path / "activity_decision_tree.z"
        joblib.dump(clf, decision_tree_path)
        self.logger.file_saved("Activity decision tree", decision_tree_path, sd=3)

    @aux.timeit
    def classify_active_frames(self):
        ExtActive = ActiveFrames(fps=self.fps)

        clf = self._load_activity_decision_tree(self.project_path)

        pbar = tqdm(self.expt_path_dict.items())
        for name, expt_path in pbar:
            pbar.set_description(f"Classifying active frames of {name}...")

            expt_item = self._load_expt_item(expt_path)
            mask_dormant = expt_item.mask_dormant

            ftws = self._load_ws(expt_path)
            snap_names_dict = self._load_snap_names_dict(expt_path)

            ExtActive.snap_names_dict = snap_names_dict

            dfCoef = ExtActive.get_dfCoef(ftws, log=True)
            X_expt = dfCoef.to_numpy()

            y_pred = clf.predict(X_expt)
            self.logger.number_of_dormant_active(name, sum(y_pred[mask_dormant]))

            expt_item.mask_active = y_pred != 0
            self.logger.info_saved("Mask of active frames")
            self._save_expt_item(expt_item, expt_path)

            viz = VizWrap()
            fig, _ = viz.ethogram(
                y_pred,
                title=f"Activiy Labels of {name}",
            )
            activity_lbl_figpath = expt_path / "figures" / "activity_labels.png"
            fig.savefig(activity_lbl_figpath, bbox_inches="tight")
            self.logger.file_saved(
                "Broken-hbar plot of activity labels", activity_lbl_figpath, sd=3
            )

            if expt_item.hasAnn:
                y_true = expt_item.gt_activity_labels
                self.logger.dormant_active_training_fscore(
                    name, y_true[mask_dormant], y_pred[mask_dormant]
                )
