import os
import numpy as np
from PIL import Image

import acl
import acllite_utils as utils
import acllite_logger as acl_log
import constants as const

class AclLiteImage(object):
 """Image data and operation class
 Wrap image data and operation method, support jpeg, png, yuv file and
 memory data

 Attributes:
 _run_mode: device run mode
 _data: image binary data or numpy array
 _memory_type: the data in which memory, include dvpp,
 device and np array
 width: image width
 height: image height
 align_width: align image width
 align_height: align image height
 _encode_format: image format
 _load_ok: load image success or not

 """
 _run_mode, _ = acl.rt.get_run_mode()

 def __init__(self, image, width=0, height=0, align_width=0, align_height=0,
 size=0, memory_type=const.MEMORY_NORMAL):
 """Create AclLiteImage instance
 Args:
 image: image data, binary, numpy array or file path
 width: image width. if image is jpeg or png file,
 this arg is not nesscessary
 height: image height. if image is jpeg or png file, this arg is
 not nesscessary
 size: image data size. if image is file path, this arg is not
 nesscessary
 memory_type: memory type of image data. if image is file path, this
 arg is not nesscessary
 """
 self._data = None
 self._bytes_data = None
 self._memory_type = memory_type
 self.width = 0
 self.height = 0
 self.align_width = 0
 self.align_height = 0
 self.size = 0
 self._encode_format = const.ENCODE_FORMAT_UNKNOW
 self._load_ok = True

 if isinstance(image, str):
 self._instance_by_image_file(image, width, height)
 elif isinstance(image, int):
 self._instance_by_buffer(image, width, height, align_width, align_height, size)
 elif isinstance(image, np.ndarray):
 self._instance_by_nparray(image, width, height, align_width, align_height)
 else:
 acl_log.log_error("Create instance failed for "
 "unknow image data type")

 def __del__(self):
 self.destroy()

 def is_loaded(self):
 """Image file load result
 When create image instance by file, call this method to check
 file load success or not

 Returns:
 True: load success
 False: load failed
 """
 return self._load_ok

 def byte_data_to_np_array(self):
 """Trans image data to np array"""
 if self._type == const.IMAGE_DATA_NUMPY:
 return self._data.copy()

 return utils.copy_data_as_numpy(self._data, self.size,
 self._memory_type, AclLiteImage._run_mode)

 def data(self):
 """Get image binary data"""
 if self._type == const.IMAGE_DATA_NUMPY:
 if "bytes_to_ptr" in dir(acl.util):
 self._bytes_data = self._data.tobytes()
 factor_ptr = acl.util.bytes_to_ptr(self._bytes_data)
 else:
 factor_ptr = acl.util.numpy_to_ptr(self._data)
 return factor_ptr
 else:
 return self._data

 def copy_to_dvpp(self):
 """Copy image data to dvpp"""
 device_ptr = utils.copy_data_to_dvpp(self.data(), self.size,
 self._run_mode)
 if device_ptr is None:
 acl_log.log_error("Copy image to dvpp failed")
 return None
 return AclLiteImage(device_ptr, self.width, self.height, 0, 0,
 self.size, const.MEMORY_DVPP)

 def copy_to_host(self):
 """"Copy data to host"""
 if self._type == const.IMAGE_DATA_NUMPY:
 data_np = self._data.copy()
 return AclLiteImage(data_np, self.width, self.height, 0, 0)

 data = None
 mem_type = const.MEMORY_HOST
 if AclLiteImage._run_mode == const.ACL_HOST:
 if self.is_local():
 data = utils.copy_data_host_to_host(self._data, self.size)
 else:
 data = utils.copy_data_device_to_host(self._data, self.size)
 else:
 data = utils.copy_data_device_to_device(self._data, self.size)
 mem_type = const.MEMORY_DEVICE
 if data is None:
 acl_log.log_error("Copy image to host failed")
 return None

 return AclLiteImage(data, self.width, self.height, 0, 0, self.size, mem_type)

 def is_local(self):
 """Image data is in host server memory and access directly or not"""
 # in atlas200dk, all kind memory can access directly
 if AclLiteImage._run_mode == const.ACL_DEVICE:
 return True
 # in atlas300, only acl host memory or numpy array can access directly
 elif ((AclLiteImage._run_mode == const.ACL_HOST) and
 ((self._memory_type == const.MEMORY_HOST) or
 (self._memory_type == const.MEMORY_NORMAL))):
 return True
 else:
 return False

 def save(self, filename):
 """Save image as file"""
 image_np = self.byte_data_to_np_array()
 image_np.tofile(filename)

 def destroy(self):
 """Release image memory"""
 if (self._data is None) or (self.size == 0):
 acl_log.log_error("Release image abnormaly, data is None")
 return

 if self._memory_type == const.MEMORY_DEVICE:
 acl.rt.free(self._data)
 elif self._memory_type == const.MEMORY_HOST:
 acl.rt.free_host(self._data)
 elif self._memory_type == const.MEMORY_DVPP:
 acl.media.dvpp_free(self._data)
 # numpy no need release
 self._data = None
 self.size = 0

 def _instance_by_image_file(self, image_path, width, height):
 # Get image format by filename suffix
 self._encode_format = self._get_image_format_by_suffix(image_path)
 if self._encode_format == const.ENCODE_FORMAT_UNKNOW:
 acl_log.log_error("Load image %s failed" % (image_path))
 self._load_ok = False
 return

 # Read image data from file to memory
 self._data = np.fromfile(image_path, dtype=np.byte)
 self._type = const.IMAGE_DATA_NUMPY
 self.size = self._data.itemsize * self._data.size
 self._memory_type = const.MEMORY_NORMAL

 # Get image parameters of jpeg or png file by pillow
 if ((self._encode_format == const.ENCODE_FORMAT_JPEG) or
 (self._encode_format == const.ENCODE_FORMAT_PNG)):
 image = Image.open(image_path)
 self.width, self.height = image.size
 else:
 # pillow can not decode yuv, so need input widht and height args
 self.width = width
 self.height = height

 def _get_image_format_by_suffix(self, filename):
 suffix = os.path.splitext(filename)[-1].strip().lower()
 if (suffix == ".jpg") or (suffix == ".jpeg"):
 image_format = const.ENCODE_FORMAT_JPEG
 elif suffix == ".png":
 image_format = const.ENCODE_FORMAT_PNG
 elif suffix == ".yuv":
 image_format = const.ENCODE_FORMAT_YUV420_SP
 else:
 acl_log.log_error("Unsupport image format: ", suffix)
 image_format = const.ENCODE_FORMAT_UNKNOW

 return image_format

 def _instance_by_buffer(self, image_buffer, width, height, align_width, align_height, size):
 self.width = width
 self.height = height
 self.align_height = align_height
 self.align_width = align_width
 self.size = size
 self._data = image_buffer
 self._type = const.IMAGE_DATA_BUFFER

 def _instance_by_nparray(self, data, width, height, align_width, align_height):
 self.width = width
 self.height = height
 self.align_height = align_height
 self.align_width = align_width
 self.size = data.itemsize * data.size
 self._data = data
 self._type = const.IMAGE_DATA_NUMPY
 self._memory_type = const.MEMORY_NORMAL