import matplotlib.pyplot as plt
import pandas as pd 
import numpy as np
import numbers as nb
import os 

x_0 = 116
y_0 = 156 
x_factor = 1.24/1648  # Einheit m/px 
y_factor = 0.81/748  # Einheit m/px 


def coordinate_transformation(position):
    x = (position[0] - x_0) * x_factor
    y = (position[1] - y_0) * y_factor
    return (x, y) #in m 


class TrackingResult():
    
    def __init__(self, results_file) -> None:
        super().__init__()
        if not os.path.exists(results_file):
            raise ValueError("File %s does not exist!" % results_file)
        self._file_name = results_file
        
        self._data_frame = pd.read_hdf(results_file)
        self._level_shape = self._data_frame.columns.levshape
        self._scorer = self._data_frame.columns.levels[0].values
        self._bodyparts = self._data_frame.columns.levels[1].values if self._level_shape[1] > 0 else []
        self._positions = self._data_frame.columns.levels[2].values if self._level_shape[2] > 0 else []

        
    @property
    def filename(self):
        return self._file_name
    
    @property
    def dataframe(self):
        return self._data_frame
    
    @property
    def scorer(self):
        return self._scorer
    
    @property
    def bodyparts(self):
        return self._bodyparts
    
    @property
    def positions(self):
        return self._positions
    
    def position_values(self, scorer=0, bodypart=0, framerate=30):
        """returns the x and y positions in m and the likelihood of the positions.

        Args:
            scorer (int, optional): [description]. Defaults to 0.
            bodypart (int, optional): [description]. Defaults to 0.
            framerate (int, optional): [description]. Defaults to 30.

        Raises:
            ValueError: [description]
            ValueError: [description]

        Returns:
            [type]: [description]
        """

        if  isinstance(scorer, nb.Number):
            sc = self._scorer[scorer]
        elif isinstance(scorer, str) and scorer in self._scorer:
            sc = scorer
        else:
            raise ValueError("Scorer %s is not in dataframe!" % scorer)
        if  isinstance(bodypart, nb.Number):
            bp = self._bodyparts[bodypart]
        elif isinstance(bodypart, str) and bodypart in self._bodyparts:
            bp = bodypart
        else:
            raise ValueError("Bodypart %s is not in dataframe!" % bodypart)
    
        x = self._data_frame[sc][bp]["x"] if "x" in self._positions else []
        x = (np.asarray(x) - x_0) * x_factor
        y = self._data_frame[sc][bp]["y"] if "y" in self._positions else []
        y = (np.asarray(y) - y_0) * y_factor
        l = self._data_frame[sc][bp]["likelihood"] if "likelihood" in self._positions else []
        time = np.arange(len(self._data_frame))/framerate
        return time, x, y, l, bp
    
    def plot(self, scorer=0, bodypart=0, threshold=0.9, framerate=30):
        t, x, y, l, name  = self.position_values(scorer=scorer, bodypart=bodypart, framerate=framerate)
        plt.scatter(x[l > threshold], y[l > threshold], c=t[l > threshold], label=name)
        plt.plot(x[l > threshold], y[l > threshold])
        plt.xlabel("x position")
        plt.ylabel("y position")
        plt.gca().invert_yaxis()
        bar = plt.colorbar()
        bar.set_label("time [s]")
        plt.legend()
        plt.show()