#!/usr/bin/env python
import rospy
import numpy as np
import yaml
from geometry_msgs.msg import Pose


# Keys for dimensions return mask
CAMERA_PX = 0b0001
POSE_SHIFT = 0b0010
REAL_MM = 0b0100
RESIZE_TO_PX = 0b1000

CURRENT_DIR = rospy.get_param("robocheckers/work_dir")

def save_camera_params(mtx, dist):
    """Saves the camera calibration data into the yaml file, old data save as backup"""
    # Choose file path for storing the old data
    bckp_data_file = open(CURRENT_DIR + "/calibration/calib_data_backup.yml", "w")
    
    try:
        # Load the old chessboard params
        mtx_bckp, dist_bckp = load_camera_params()

        # Creating calibration data variable ready to save into yaml file
        camera_calib_bckp = {'camera-matrix' : np.asarray(mtx_bckp).tolist(),
                            'distorsion-coefficients' : np.asarray(dist_bckp).tolist()}

        # Saving actual calibration data backup
        yaml.dump(camera_calib_bckp, bckp_data_file, default_flow_style=None)
    except Exception:
        print "ERROR: Could not backup, old data corrupted!"

    # Closing file with saved calibration data backup
    bckp_data_file.close()

    # Choose file path
    file_data = open(CURRENT_DIR + "/calibration/calib_data.yml", "w")

    # Creating calibration data variable ready to save into yaml file
    calib_output = {'camera-matrix' : np.asarray(mtx).tolist(),
                    'distorsion-coefficients' : np.asarray(dist).tolist()}

    # Saving actual calibration data
    yaml.dump(calib_output, file_data, default_flow_style=None)

    # Closing file with saved calibration data
    file_data.close()


def load_camera_params():
    """Returns the camera calibration data from the yaml file"""
    # Open the file and read the data
    file_data = open(CURRENT_DIR + "/calibration/calib_data.yml", "r")
    loaded = yaml.load(file_data)

    # Find matrix and distorsion values in loaded data
    tmp = loaded.get('camera-matrix')
    if tmp:
        mtx = np.array(tmp)

    tmp = loaded.get('distorsion-coefficients')
    if tmp:
        dist = np.array(tmp)

    # Closing file with saved calibration data
    file_data.close()

    # If any of the loaded data is None, return both as None
    if mtx is None or dist is None:
        return None, None
    return mtx, dist


def save_chessboard_params(center, angle, apexes):
    """Saves the chessboard parameters into the yaml file, old data save as backup"""
    # Choose file path for storing the old data
    bckp_data = open(CURRENT_DIR + "/calibration/chessboard_data_backup.yml", "w")

    # Load the old chessboard params
    center_bckp, angle_bckp, apexes_bckp = load_chessboard_params()

    # Creating calibration data variable ready to save into yaml file
    chessboard_params_bckp = {'center' : np.asarray(center_bckp).tolist(),
                              'angle' : np.asarray(angle_bckp).tolist(),
                              'apexes' : np.asarray(apexes_bckp).tolist()}

    # Saving actual calibration data backup
    yaml.dump(chessboard_params_bckp, bckp_data, default_flow_style=None)

    # Closing file with saved calibration data backup
    bckp_data.close()

    # Choose file path
    file_data = open(CURRENT_DIR + "/calibration/chessboard_data.yml", "w")

    # Creating calibration data variable ready to save into yaml file
    chessboard_params = {'center' : np.asarray(center).tolist(),
                         'angle' : np.asarray(angle).tolist(),
                         'apexes' : np.asarray(apexes).tolist()}

    # Saving actual calibration data
    yaml.dump(chessboard_params, file_data, default_flow_style=None)

    # Closing file with saved calibration data
    file_data.close()


def load_chessboard_params():
    """Returns the chessboard parameters from the yaml file"""
    # Open the file and read the data
    file_data = open(CURRENT_DIR + "/calibration/chessboard_data.yml", "r")
    loaded = yaml.load(file_data)

    # Find center, angle and apexes values in loaded data
    tmp = loaded.get('center')
    if tmp:
        center = tuple(tmp)

    tmp = loaded.get('angle')
    if tmp:
        angle = tmp

    tmp = loaded.get('apexes')
    if tmp:
        apexes = list(tmp)

    # Closing file with saved calibration data
    file_data.close()

    # If any of the loaded data is None, return both as None
    if center is None or angle is None or apexes is None:
        return None, None, None
    return center, angle, apexes


def save_dimensions_data(dx, dy, chessboard_px):
    """Saves the dimensions data into the yaml file, old data save as backup"""
    # Opening file to read data and file to store them
    file_data = open(CURRENT_DIR + "/config/dimension_data.yml", "r")
    bckp_data = open(CURRENT_DIR + "/config/dimension_data_backup.yml", "w")

    # Load the old data
    loaded = yaml.load(file_data)

    # Backup the old data
    yaml.dump(loaded, bckp_data, default_flow_style=False)

    # Closing both files after backup
    file_data.close()
    bckp_data.close()

    # Load real data
    tmp = loaded.get('real_mm')
    if tmp:
        chessboard_mm = tmp.get('chessboard')
        pos_x = tmp.get('pos_x_max')
        pos_y = tmp.get('pos_y_max')

    # Load real data
    tmp = loaded.get('pose_shift')
    if tmp:
        off_x = tmp.get('offset_x')
        off_y = tmp.get('offset_y')

    # Load resize to data
    tmp = loaded.get('resize_to_px')
    if tmp:
        resized_board_px = tmp.get('chessboard')

    # Calculating new data
    coef = -float(chessboard_mm)/float(resized_board_px) # minus, because both asxises are reversed

    # Choose file path and open for write
    file_data = open(CURRENT_DIR + "/config/dimension_data.yml", "w")

    # Creating data variable ready to save into yaml file
    dim_data = {'camera_px': {'dx': float(dx), 'dy': float(dy), 'chessboard': float(chessboard_px)},
                'pose_shift': {'coef': float(coef), 'offset_x': off_x, 'offset_y': off_y},
                'real_mm': {'chessboard': chessboard_mm, 'pos_x_max': pos_x, 'pos_y_max': pos_y},
                'resize_to_px': {'chessboard': resized_board_px}}

    # Saving existing variable with all dimension data
    yaml.dump(dim_data, file_data, default_flow_style=False)

    # Closing file
    file_data.close()


def load_dimensions_data(return_mask=0b1111):
    """Returns the dimensions data from the yaml file"""
    # Open the file and read the data
    file_data = open(CURRENT_DIR + "/config/dimension_data.yml", "r")
    loaded = yaml.load(file_data)

    # Create empty return_list variable and then fill it in dependecy on return_mask
    return_list = []

    if return_mask & POSE_SHIFT:
        tmp = loaded.get('pose_shift')
        if tmp:
            return_list.append(tmp)

    if return_mask & CAMERA_PX:
        tmp = loaded.get('camera_px')
        if tmp:
            return_list.append(tmp)

    if return_mask & REAL_MM:
        tmp = loaded.get('real_mm')
        if tmp:
            return_list.append(tmp)

    if return_mask & RESIZE_TO_PX:
        tmp = loaded.get('resize_to_px')
        if tmp:
            return_list.append(tmp)

    # Closing file with saved calibration data
    file_data.close()

    if len(return_list) == 1:
        return return_list[0]
    return return_list


def load_static_positions():
    """Returns all static positions"""
    # Open the file and read the data
    file_data = open(CURRENT_DIR + "/config/static_positions.yml", "r")
    loaded = yaml.load(file_data)

    formatted = {}
    for item in loaded:
        position = loaded.get(item).get("position")
        orientation = loaded.get(item).get("orientation")
    
        pose = Pose()
        # Get position
        pose.position.x = position.get("x")
        pose.position.y = position.get("y")
        pose.position.z = position.get("z")
        # Get orientation
        pose.orientation.x = orientation.get("x")
        pose.orientation.y = orientation.get("y")
        pose.orientation.z = orientation.get("z")
        pose.orientation.w = orientation.get("w")

        # Append pose to pose_list
        formatted[item] = pose

    # Closing file with saved calibration data
    file_data.close()

    return formatted
