import numpy as np
import pandas as pd
import pickle as pkl
import os.path
import cv2
from PyQt5 import QtCore, QtGui


class DataContainer():

    def __init__(self, movie_fileName, data_fileName):
        self.movie_fileName = movie_fileName
        self.data_fileName = data_fileName

        self.frameCount = 0
        self.fps = 0
        self.frameSize = [0, 0]

        self.dataGaze = None
        self.dataMovement = None
        self.number_ids = None
        self.dataRTGene = None
        self.image = None

    def readData(self, verbose=False, rtgene=False):
        if (verbose):
            print("## Start Reading Data")

        # Read Video Data
        f = self.movie_fileName
        print(f)
        if os.path.isfile(f):
            cap = cv2.VideoCapture(f)
            ret, frame = cap.read()
            self.image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

            h, w, ch = self.image.shape
            bytesPerLine = ch * w
            convertToQtFormat = QtGui.QImage(self.image.data, w, h, bytesPerLine, QtGui.QImage.Format_RGB888)

            self.fps = cap.get(cv2.CAP_PROP_FPS)
            self.frameCount = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)  # float
            height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)  # float
            self.frameSize = [width, height]
            print(self.fps, self.frameCount)

            if (verbose):
                print("Video frameCount %i" % self.frameCount)
                duration = self.frameCount / self.fps
                minutes = int(duration / 60)
                seconds = duration % 60
                print('Video duration (M:S) = ' + str(minutes) + ':' + str(seconds))
        else:
            print("WARNING: no video avaibale.")

        # read data file
        with open(self.data_fileName, 'rb') as f:
            data = pkl.load(f)

        # Read RT-Gene Gaze
        if rtgene:
            f = "%s_Gaze.pkl" % self.dataPath[0]
            df_gt = pd.read_pickle('exampledata/G2_VID1_GroundTruth.pkl')
            df_gt = df_gt[['Frame', 'ID0_target_spher', 'ID1_target_spher']]
            if os.path.isfile(f):
                df = pd.read_pickle(f)
                self.number_ids = len(df.PId.unique())

                self.dataRTGene = df.pivot(index='Frame', columns="PId", values=["GazeTheta", "GazePhi", "HeadCenter",
                                                                                 "HeadPoseYaw", "HeadPoseTheta",
                                                                                 "Phi", "Theta"])
                lst = []
                for label in ["GazeTheta", "GazePhi", "Head", "HeadPoseYaw", "HeadPoseTheta", "Phi", "Theta"]:
                    for head_id in range(self.number_ids):
                        lst.append("ID%i_%s" % (head_id, label))
                self.dataRTGene.columns = lst
                self.dataRTGene = self.dataRTGene.reset_index()
                self.dataRTGene = pd.merge(self.dataRTGene, df_gt, on=['Frame'])
                self.dataRTGene = self.dataRTGene.rename(
                    columns={'ID0_target_spher': 'ID1_target', 'ID1_target_spher': 'ID0_target'},
                    errors='raise')

                print('Detected %i IDs in video' % self.number_ids)
                if (verbose):
                    print("Gaze sample count %i" % len(self.dataRTGene))
            else:
                print("WARNING: no RT-Gene data avaibale.")

        # Read Gaze Data
        if "HeadPose" in data:
            self.dataGaze = data["HeadPose"]
            self.number_ids = len([col for col in self.dataGaze.columns if 'head' in col])

            print('Detected %i IDs in video' % self.number_ids)
            if (verbose):
                print("Gaze sample count %i" % len(self.dataGaze))
        else:
            print("WARNING: no gaze data avaibale.")

        """
        # Read OpenPose Data
        f = self.__get_filename_with_substring("OpenPose")
        if os.path.isfile(f):
            self.dataPose = pd.read_pickle(f)
            if (verbose):
                print("Pose sample count %i" % len(self.dataPose))
        else:
            print("WARNING: no pose data avaibale.")
        """

        # Read Movement Data
        if "BodyMovement" in data:
            self.dataMovement = data["BodyMovement"]
            if verbose:
                print('Body movement sample count %i' % len(self.dataMovement))
        else:
            print('WARNING: no body movement data available.')

    def getFrameCount(self):
        return self.frameCount

    def getFrameSize(self):
        return self.frameSize

    def getFPS(self):
        return self.fps

    def getVideo(self):
        return self.movie_fileName

    def getGazeData(self):
        return self.dataGaze

    def getFrame(self, frameIdx):
        return frameIdx

    def getFrameCurrent(self):
        return 1

    def getNumberIDs(self):
        return self.number_ids

    def getMovementData(self):
        return self.dataMovement

    def getRTGeneData(self):
        return self.dataRTGene

    def getImage(self):
        return self.image