#!/usr/bin/env python
import csv
from itertools import combinations
from multiprocessing import Pool

import cv2
import matplotlib
import numpy as np
import pandas as pd
import sleap
import yaml

from naps.sleap_utils import load_tracks_from_slp
from naps.utils.interactor import Interactor


class TrackingArray:
    def __init__(
        self,
        frames_coords_array="",
        tracks=[],
        nodes=[],
        interactor_model_dict={},
        chunk_size=5000,
        threads=1,
        **kwargs,
    ):

        # Assign the frame coordsinates array
        self._frames_coords_array = frames_coords_array

        # Assign the tracks and nodes
        self._tracks = tracks
        self._nodes = nodes

        # Store the Interactor model
        self._interactor_model_dict = interactor_model_dict

        # Assign the frame chunk size
        self._chunk_size = chunk_size

        # Assign empty variables for processing
        self._process_frame_start = 0
        self._process_frame_end = 0
        self._process_coords_array = None

        # Assign general variables
        self._threads = threads

    def assign(self, out_file, frame_start=0, frame_end=-1):

        # Update end frame
        frame_end = frame_end if frame_end != -1 else self._frames_coords_array.shape[0]

        # Assign the pair-wise combinations of each track
        def _yield_combinations():
            for _tA, _tB in combinations(range(len(self._tracks)), 2):
                yield _tA, _tB

        with open(out_file, "w") as outfile:
            csv_writer = csv.DictWriter(
                outfile,
                fieldnames=[
                    "Origin interactor",
                    "Destination interactor",
                    "Interaction Name",
                    "Interaction Frame",
                ],
                delimiter="\t",
            )
            csv_writer.writeheader()

            for chunk_frame_start in range(frame_start, frame_end, self._chunk_size):

                # Store the frame start for the current processing chunk
                self._process_frame_start = chunk_frame_start

                # Store the frame end for the current processing chunk
                chunk_end = chunk_frame_start + self._chunk_size
                self._process_frame_end = (
                    frame_end if frame_end < chunk_end else chunk_end
                )

                # Store the processing chunked coordsinates array
                self._process_coords_array = self._frames_coords_array[
                    self._process_frame_start : self._process_frame_end, :, :, :
                ]

                if self._threads > 1:
                    with Pool(processes=self._threads) as pool:
                        for report_list in pool.starmap(
                            self._checkForInteractions, _yield_combinations()
                        ):
                            for report in report_list:
                                csv_writer.writerow(report)
                else:
                    for tagA, tagB in _yield_combinations():
                        for report in self._checkForInteractions(tagA, tagB):
                            csv_writer.writerow(report)

    def plot(
        self,
        interactions_filename,
        models_to_plot,
        video_filename,
        out_file,
        frame_start=0,
        frame_end=-1,
        limit_interactions=[],
        plot_color=None,
        line_width=2,
    ):

        # Convert to int32 for cv2 coords
        def int_coords(coords):
            return np.array(coords).round().astype(np.int32)

        # Update end frame
        frame_end = frame_end if frame_end != -1 else self._frames_coords_array.shape[0]
        frame_end = (
            frame_end - 1
            if frame_end >= self._frames_coords_array.shape[0]
            else frame_end
        )

        # Create dataframe of instances to plot per frame
        frame_instances_dataframe = self.processInteractionsFile(
            interactions_filename,
            frame_start,
            frame_end,
            limit_interactions=limit_interactions,
        )

        # Assign color for plotting the model:
        # 1) By default, plot red
        # 2) If given a list or array, check and assign the color
        # 3) If given a str, assign the color from matplotlib
        if plot_color == None:
            plot_color = (255, 0, 0)
        elif isinstance(plot_color, (list, tuple)):
            if len(plot_color) != 3:
                raise Exception(
                    f"Unable to assign plot color: {plot_color}. Iterables must be of length 3"
                )
            plot_color = tuple(plot_color)
        elif isinstance(plot_color, str):
            plot_color = matplotlib.colors.to_rgb(plot_color)
            plot_color = tuple(_v * 255 for _v in plot_color)

        # Color the RGB to BGR used by CV2
        plot_color = plot_color[::-1]

        # Set the current frame of the job
        current_frame = frame_start

        # Initialize OpenCV, then set the starting frame (0-based)
        video = cv2.VideoCapture(video_filename)
        video.set(cv2.CAP_PROP_POS_FRAMES, current_frame)

        # Store the width, height, and fps of the input video
        video_width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
        video_height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
        video_fps = int(video.get(cv2.CAP_PROP_FPS))

        # Open the output video using the input video values
        out_video = cv2.VideoWriter(
            out_file,
            cv2.VideoWriter_fourcc(*"mp4v"),
            video_fps,
            (video_width, video_height),
        )

        # Read in the frame, confirm it was successful
        frame_read, frame_img = video.read()
        while frame_read and current_frame <= frame_end:

            # Assign the instances to plot
            instances_to_plot = frame_instances_dataframe[
                frame_instances_dataframe["Interaction Frame"] == current_frame
            ]["Instances"].values
            if len(instances_to_plot) > 1:
                raise Exception("Error parsing interactions for instances")
            if len(instances_to_plot) > 0:
                instances_to_plot = instances_to_plot[0]

            # Loop each instance for this frame
            for instance_to_plot in instances_to_plot:

                # Create an interactor using the instance
                instance_interactor = Interactor.withModelDict(
                    instance_to_plot,
                    current_frame,
                    self._nodes,
                    self._frames_coords_array[
                        current_frame, :, :, self._tracks.index(instance_to_plot)
                    ],
                    self._interactor_model_dict,
                )

                # Loop and plot each model for the instance
                for model_to_plot in models_to_plot:
                    instance_model_coords = int_coords(
                        instance_interactor._models["Head_w_antennae"].exterior.coords
                    )
                    cv2.polylines(
                        frame_img, [instance_model_coords], True, plot_color, line_width
                    )

            # Write the frame
            out_video.write(frame_img)

            # Read the next frame and advance the current frame
            frame_read, frame_img = video.read()
            current_frame += 1

        # Release the video
        out_video.release()

    def _checkForInteractions(self, trackA, trackB):

        # Create list to store the interactions
        interaction_list = []

        # Subset the array by tracks
        filtered_coords_array = np.take(self._process_coords_array, [trackA, trackB], 3)

        # Check each row for an interaction
        for frame, row in enumerate(filtered_coords_array):

            # Skip if missing data
            if np.isnan(row).any():
                continue

            # Adjust the frame to account for the start frame
            adj_frame = self._process_frame_start + frame

            # Create the interactors
            interactor1 = Interactor.withModelDict(
                self._tracks[trackA],
                adj_frame,
                self._nodes,
                row[:, :, 0],
                self._interactor_model_dict,
            )
            interactor2 = Interactor.withModelDict(
                self._tracks[trackB],
                adj_frame,
                self._nodes,
                row[:, :, 1],
                self._interactor_model_dict,
            )

            # Check if there is an interaction
            interaction_list.extend(interactor1.interacts(interactor2))

        return interaction_list

    @classmethod
    def fromSLEAP(cls, sleap_file, model_dict={}, model_yaml="", **kwargs):

        # Assign the tracks from the input file
        labels = sleap.Labels.load_file(sleap_file)
        tracks = [_k.name for _k in labels.tracks]
        labels = None

        # Create an array of [frames, nodes, (x,y), tracks] from the input file
        frames_coords_array, nodes = load_tracks_from_slp(sleap_file)

        # Confirm one a single model was given
        if model_dict and model_yaml:
            raise Exception("Only a single Interactor model may be specified")

        # Assign the model dict, if specified
        if model_yaml:
            interactor_model_dict = TrackingArray.loadYAML(model_yaml)
        elif model_dict:
            interactor_model_dict = model_dict

        return cls(
            frames_coords_array=frames_coords_array,
            tracks=tracks,
            nodes=nodes,
            interactor_model_dict=interactor_model_dict,
            **kwargs,
        )

    @staticmethod
    def loadYAML(yaml_filename):

        # Read in the YAML file and return as dict
        with open(yaml_filename, "r") as yaml_file:
            try:
                yaml_dict = yaml.safe_load(yaml_file)
            except yaml.YAMLError as err:
                raise Exception(err)
        return yaml_dict

    @staticmethod
    def processInteractionsFile(
        interactions_file,
        frame_start,
        frame_end,
        limit_interactions=[],
        expected_cols=[
            "Origin interactor",
            "Destination interactor",
            "Interaction Name",
            "Interaction Frame",
        ],
    ):

        # Process the interactions file
        interactions_origin_dataframe = pd.read_csv(interactions_file, sep="\t")[
            expected_cols
        ]
        if limit_interactions:
            interactions_origin_dataframe = interactions_origin_dataframe[
                interactions_origin_dataframe["Interaction Name"].isin(
                    limit_interactions
                )
            ]
        interactions_dest_dataframe = interactions_origin_dataframe.copy()

        # Assign the origin instances for each frame
        interactions_origin_dataframe = interactions_origin_dataframe.drop(
            columns=["Destination interactor"]
        )
        interactions_origin_dataframe = interactions_origin_dataframe.rename(
            columns={"Origin interactor": "Instance"}
        )

        # Assign the destination instances for each frame
        interactions_dest_dataframe = interactions_dest_dataframe.drop(
            columns=["Origin interactor"]
        )
        interactions_dest_dataframe = interactions_dest_dataframe.rename(
            columns={"Destination interactor": "Instance"}
        )

        # Create the instance dataframe
        instance_dataframe = interactions_origin_dataframe.append(
            interactions_dest_dataframe, ignore_index=True
        )
        instance_dataframe = instance_dataframe[
            instance_dataframe["Interaction Frame"].between(frame_start, frame_end)
        ]

        # Check the presence of interactions within the given frames
        if instance_dataframe.shape[0] == 0:
            raise Exception(
                f"No interactions found between frames {frame_start} to {frame_end}"
            )

        instance_dataframe = (
            instance_dataframe.groupby("Interaction Frame")["Instance"]
            .apply(list)
            .reset_index(name="Instances")
        )

        return instance_dataframe

    @staticmethod
    def interpolate(coords_array):
        import pandas as pd

        for i in range(coords_array.shape[3]):

            coords_df = pd.DataFrame(coords_array[:, :, 0, i])
            coords_df = coords_df.interpolate(method="linear", axis=0)
            diff_df = coords_df.diff().abs().ge(5)
            coords_df[diff_df] = np.nan
            coords_array[:, :, 0, i] = coords_df.to_numpy()

            coords_df = pd.DataFrame(coords_array[:, :, 1, i])
            coords_df = coords_df.interpolate(method="linear", axis=0)
            diff_df = coords_df.diff().abs().ge(5)
            coords_df[diff_df] = np.nan
            coords_array[:, :, 1, i] = coords_df.to_numpy()

        return coords_array
