# =============================================================================
# Authors: PAR Government
# Organization: DARPA
#
# Copyright (c) 2016 PAR Government
# All rights reserved.
# ==============================================================================

from PIL import Image
import cv2
import numpy as np
import os
import inspect
import subprocess
from pkg_resources import iter_entry_points
from cachetools import LRUCache
from threading import RLock
import logging
import os
import png
import exif
from numpngw import write_png
from maskgen import config

Image.MAX_IMAGE_PIXELS = 10000000000


image_lock = config.getAndSet('image_lock', RLock())
image_cache = config.getAndSet('image_cache', LRUCache(maxsize=24))

try:
    from skimage.external.tifffile import TiffFile,imsave
except ImportError:
    try:
        from tifffile import TiffFile, imsave
    except ImportError:
        def imsave(filename, img, **newargs):
            import scipy.misc
            scipy.misc.imsave(filename, img)


        class TiffFile:
            def __init__(self, filename):
                pass

            def __array__(self):
                return []

            def __exit__(self, x, y, z):
                return []

            def __enter__(self):
                return []


def _processRaw(filename, raw, isMask=False, args=None):
    import rawpy
    def _open_from_rawpy(raw,args=None):
        if args is not None and 'Bits per Channel' in args:
            bits = int(args['Bits per Channel'])
        else:
            bits = 8
        use_camera_wb = args is None or \
                        'White Balance' not in args or \
                        args['White Balance'] == 'camera'
        use_auto_wb = args is not None and \
                      'White Balance' in args and \
                      args['White Balance'] == 'auto'
        colorspace = rawpy.ColorSpace.raw
        if args is not None and 'Color Space' in args:
            v = args['Color Space']
            cs_mapping = {'Adobe': rawpy.ColorSpace.Adobe,
                          'sRGB': rawpy.ColorSpace.sRGB,
                          'XYZ': rawpy.ColorSpace.XYZ,
                          'Wide': rawpy.ColorSpace.Wide,
                          'ProPhoto': rawpy.ColorSpace.ProPhoto,
                          'default': rawpy.ColorSpace.raw
                          }
            colorspace = cs_mapping[v]
        if args is not None and 'Demosaic Algorithm' in args and args['Demosaic Algorithm'] != 'default':
            v = args['Demosaic Algorithm']
            mapping = {'AAHD': rawpy.DemosaicAlgorithm.AAHD,
                       'AFD': rawpy.DemosaicAlgorithm.AFD,
                       'AMAZE': rawpy.DemosaicAlgorithm.AMAZE,
                       'DCB': rawpy.DemosaicAlgorithm.DCB,
                       'DHT': rawpy.DemosaicAlgorithm.DHT,
                       'LMMSE': rawpy.DemosaicAlgorithm.LMMSE,
                       'MODIFIED_AHD': rawpy.DemosaicAlgorithm.MODIFIED_AHD,
                       'PPG': rawpy.DemosaicAlgorithm.PPG,
                       'VCD': rawpy.DemosaicAlgorithm.VCD,
                       'LINEAR': rawpy.DemosaicAlgorithm.LINEAR,
                       'VCD_MODIFIED_AHD': rawpy.DemosaicAlgorithm.VCD_MODIFIED_AHD,
                       'VNG': rawpy.DemosaicAlgorithm.VNG}
            return raw.postprocess(demosaic_algorithm=mapping[v],
                                                output_bps=bits,
                                                use_camera_wb=use_camera_wb,
                                                use_auto_wb=use_auto_wb,
                                                output_color=colorspace)

        return raw.postprocess(output_bps=bits,
                                            use_camera_wb=use_camera_wb,
                                            use_auto_wb=use_auto_wb,
                                            output_color=colorspace)
    try:
        rawdata = _open_from_rawpy(raw,args=args)
        return ImageWrapper(rawdata,to_mask=isMask,isRaw=True)
    except Exception as e:
        logging.getLogger('maskgen').error('Raw Open: ' + str(e))
        return None


def openImageMaskFile(directory, filename):
    return openImageFile(os.path.join(directory, filename), isMask=True).to_array()


def openRaw(filename, isMask=False, args=None):
    logger = logging.getLogger('maskgen')
    try:
        import rawpy
        with rawpy.imread(filename) as raw:
            if type(args) == list:
                result = {}
                for argitem in args:
                    rawim = _processRaw(filename, raw, isMask=isMask, args=argitem)
                    if rawim is not None:
                        if 'outputname' in argitem:
                            rawim.save(argitem['outputname'], format='PNG')
                            result[argitem['outputname']] = rawim
                        else:
                            result[str(argitem)] = rawim
            else:
                result = _processRaw(filename,raw,args=args)
            if logger.isEnabledFor(logging.DEBUG):
                logging.debug('Opened {} as raw'.format(filename))
            return result
    except Exception as e:
        logger.error('Raw Open: ' + str(e))
        return None


def _openCV2(filename):
    img = cv2.imread(filename, flags=cv2.IMREAD_UNCHANGED)
    if len(img.shape) > 2:
        if img.shape[2] > 3:
            result = cv2.cvtColor(img, cv2.COLOR_BGRA2RGBA)
        else:
            result = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    else:
        result = img
    return result

def openTiff(filename, isMask=False, args=None):
    raw = openRaw(filename, isMask=isMask, args=args)
    info = {}
    if filename.lower().find('tif') >= 0:
        try:
            with TiffFile(filename) as tiffdata:
                for page in tiffdata:
                    for tag in page.tags.values():
                        t, v = tag.name, tag.value
                        info[t] = v
        except:
            pass
        try:
            nonRaw = ImageWrapper(_openCV2(filename),
                                  info=info,
                                  to_mask=isMask,
                                  filename=filename)
            if raw is not None and raw.size[0] > nonRaw.size[0] and raw.size[1] > nonRaw.size[1]:
                return raw
            return nonRaw
        except:
            return raw
    return raw


def wand_image_extractor(filename, isMask=False,args=None):
    import pgmagick
    im = pgmagick.Image(filename)
    myPilImage = Image.new('RGB', (im.GetWidth(), im.GetHeight()))
    myPilImage.fromstring(im.GetData())
    return ImageWrapper(np.asarray(myPilImage), mode=myPilImage.mode, info=myPilImage.info,
                        to_mask=isMask,
                        filename=filename)

def pdf2_image_extractor(filename, isMask=False,args=None):
    import PyPDF2
    import io
    from PyPDF2 import generic
    def find_image(page0):
        xObject = page0['/Resources']['/XObject'].getObject()
        rotate = page0['/Rotate'] if '/Rotate' in page0 else 0
        for obj in xObject:
            if xObject[obj]['/Subtype'] == '/Image':
                size = (xObject[obj]['/Width'], xObject[obj]['/Height']
)
                mode = 'RGB'
                if '/ColorSpace' in xObject[obj]:
                    mode = "P" if xObject[obj]['/ColorSpace'] == '/DeviceGray' else 
mode
                    mode = "CMYK" if xObject[obj]['/ColorSpace'] == '/DeviceCMYK' else mode
                hasJPEGOld = len([x for x in xObject[obj]['/Filter'] if x == '/DCTDecode']) > 0
                hasJPEG = xObject[obj]['/Filter'] in ['/DCTDecode', '/JBIG2Decode']
                #isBig =  xObject[obj]['/Filter'] == '/JBIG2Decode'
                if hasJPEG or hasJPEGOld:
                    if type(xObject[obj]['/Filter']) == generic.ArrayObject:
                        xObject[obj].update({generic.NameObject('/Filter'):
                                                 generic.ArrayObject([x for x in xObject[obj]['/Filter']
                                                                      i