import ClientConstants as CC
import cStringIO
import numpy.core.multiarray # important this comes before cv!
import cv
import cv2
import HydrusConstants as HC
import HydrusExceptions
import HydrusThreading
import lz4
import numpy
import os
from PIL import _imaging
from PIL import Image as PILImage
import shutil
import struct
import threading
import time
import traceback
import wx

#LINEAR_SCALE_PALETTE = [ 0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4, 5, 5, 5, 6, 6, 6, 7, 7, 7, 8, 8, 8, 9, 9, 9, 10, 10, 10, 11, 11, 11, 12, 12, 12, 13, 13, 13, 14, 14, 14, 15, 15, 15, 16, 16, 16, 17, 17, 17, 18, 18, 18, 19, 19, 19, 20, 20, 20, 21, 21, 21, 22, 22, 22, 23, 23, 23, 24, 24, 24, 25, 25, 25, 26, 26, 26, 27, 27, 27, 28, 28, 28, 29, 29, 29, 30, 30, 30, 31, 31, 31, 32, 32, 32, 33, 33, 33, 34, 34, 34, 35, 35, 35, 36, 36, 36, 37, 37, 37, 38, 38, 38, 39, 39, 39, 40, 40, 40, 41, 41, 41, 42, 42, 42, 43, 43, 43, 44, 44, 44, 45, 45, 45, 46, 46, 46, 47, 47, 47, 48, 48, 48, 49, 49, 49, 50, 50, 50, 51, 51, 51, 52, 52, 52, 53, 53, 53, 54, 54, 54, 55, 55, 55, 56, 56, 56, 57, 57, 57, 58, 58, 58, 59, 59, 59, 60, 60, 60, 61, 61, 61, 62, 62, 62, 63, 63, 63, 64, 64, 64, 65, 65, 65, 66, 66, 66, 67, 67, 67, 68, 68, 68, 69, 69, 69, 70, 70, 70, 71, 71, 71, 72, 72, 72, 73, 73, 73, 74, 74, 74, 75, 75, 75, 76, 76, 76, 77, 77, 77, 78, 78, 78, 79, 79, 79, 80, 80, 80, 81, 81, 81, 82, 82, 82, 83, 83, 83, 84, 84, 84, 85, 85, 85, 86, 86, 86, 87, 87, 87, 88, 88, 88, 89, 89, 89, 90, 90, 90, 91, 91, 91, 92, 92, 92, 93, 93, 93, 94, 94, 94, 95, 95, 95, 96, 96, 96, 97, 97, 97, 98, 98, 98, 99, 99, 99, 100, 100, 100, 101, 101, 101, 102, 102, 102, 103, 103, 103, 104, 104, 104, 105, 105, 105, 106, 106, 106, 107, 107, 107, 108, 108, 108, 109, 109, 109, 110, 110, 110, 111, 111, 111, 112, 112, 112, 113, 113, 113, 114, 114, 114, 115, 115, 115, 116, 116, 116, 117, 117, 117, 118, 118, 118, 119, 119, 119, 120, 120, 120, 121, 121, 121, 122, 122, 122, 123, 123, 123, 124, 124, 124, 125, 125, 125, 126, 126, 126, 127, 127, 127, 128, 128, 128, 129, 129, 129, 130, 130, 130, 131, 131, 131, 132, 132, 132, 133, 133, 133, 134, 134, 134, 135, 135, 135, 136, 136, 136, 137, 137, 137, 138, 138, 138, 139, 139, 139, 140, 140, 140, 141, 141, 141, 142, 142, 142, 143, 143, 143, 144, 144, 144, 145, 145, 145, 146, 146, 146, 147, 147, 147, 148, 148, 148, 149, 149, 149, 150, 150, 150, 151, 151, 151, 152, 152, 152, 153, 153, 153, 154, 154, 154, 155, 155, 155, 156, 156, 156, 157, 157, 157, 158, 158, 158, 159, 159, 159, 160, 160, 160, 161, 161, 161, 162, 162, 162, 163, 163, 163, 164, 164, 164, 165, 165, 165, 166, 166, 166, 167, 167, 167, 168, 168, 168, 169, 169, 169, 170, 170, 170, 171, 171, 171, 172, 172, 172, 173, 173, 173, 174, 174, 174, 175, 175, 175, 176, 176, 176, 177, 177, 177, 178, 178, 178, 179, 179, 179, 180, 180, 180, 181, 181, 181, 182, 182, 182, 183, 183, 183, 184, 184, 184, 185, 185, 185, 186, 186, 186, 187, 187, 187, 188, 188, 188, 189, 189, 189, 190, 190, 190, 191, 191, 191, 192, 192, 192, 193, 193, 193, 194, 194, 194, 195, 195, 195, 196, 196, 196, 197, 197, 197, 198, 198, 198, 199, 199, 199, 200, 200, 200, 201, 201, 201, 202, 202, 202, 203, 203, 203, 204, 204, 204, 205, 205, 205, 206, 206, 206, 207, 207, 207, 208, 208, 208, 209, 209, 209, 210, 210, 210, 211, 211, 211, 212, 212, 212, 213, 213, 213, 214, 214, 214, 215, 215, 215, 216, 216, 216, 217, 217, 217, 218, 218, 218, 219, 219, 219, 220, 220, 220, 221, 221, 221, 222, 222, 222, 223, 223, 223, 224, 224, 224, 225, 225, 225, 226, 226, 226, 227, 227, 227, 228, 228, 228, 229, 229, 229, 230, 230, 230, 231, 231, 231, 232, 232, 232, 233, 233, 233, 234, 234, 234, 235, 235, 235, 236, 236, 236, 237, 237, 237, 238, 238, 238, 239, 239, 239, 240, 240, 240, 241, 241, 241, 242, 242, 242, 243, 243, 243, 244, 244, 244, 245, 245, 245, 246, 246, 246, 247, 247, 247, 248, 248, 248, 249, 249, 249, 250, 250, 250, 251, 251, 251, 252, 252, 252, 253, 253, 253, 254, 254, 254, 255, 255, 255 ]

def ConvertToPngIfBmp( path ):
    
    with open( path, 'rb' ) as f: header = f.read( 2 )
    
    if header == 'BM':
        
        temp_path = HC.GetTempPath()
        
        shutil.move( path, temp_path )
        
        pil_image = GeneratePILImage( temp_path )
        
        pil_image = pil_image.convert( 'P' )
        
        pil_image.save( path, 'PNG' )
        
        os.remove( temp_path )
        
    
def EfficientlyResizeNumpyImage( numpy_image, ( target_x, target_y ) ):
    
    ( im_y, im_x, depth ) = numpy_image.shape
    
    if target_x >= im_x and target_y >= im_y: return numpy_image
    
    result = numpy_image
    
    # this seems to slow things down a lot, at least for cv!
    #if im_x > 2 * target_x and im_y > 2 * target_y: result = cv2.resize( numpy_image, ( 2 * target_x, 2 * target_y ), interpolation = cv2.INTER_NEAREST )
    
    return cv2.resize( result, ( target_x, target_y ), interpolation = cv2.INTER_LINEAR )
    
def EfficientlyResizePILImage( pil_image, ( target_x, target_y ) ):
    
    ( im_x, im_y ) = pil_image.size
    
    if target_x >= im_x and target_y >= im_y: return pil_image
    
    #if pil_image.mode == 'RGB': # low quality resize screws up alpha channel!
    #    
    #    if im_x > 2 * target_x and im_y > 2 * target_y: pil_image.thumbnail( ( 2 * target_x, 2 * target_y ), PILImage.NEAREST )
    #    
    
    return pil_image.resize( ( target_x, target_y ), PILImage.ANTIALIAS )
    
def EfficientlyThumbnailNumpyImage( numpy_image, ( target_x, target_y ) ):
    
    ( im_y, im_x, depth ) = numpy_image.shape
    
    if target_x >= im_x and target_y >= im_y: return numpy_image
    
    ( target_x, target_y ) = GetThumbnailResolution( ( im_x, im_y ), ( target_x, target_y ) )
    
    return cv2.resize( numpy_image, ( target_x, target_y ), interpolation = cv2.INTER_AREA )
    
def EfficientlyThumbnailPILImage( pil_image, ( target_x, target_y ) ):
    
    ( im_x, im_y ) = pil_image.size
    
    #if pil_image.mode == 'RGB': # low quality resize screws up alpha channel!
    #    
    #    if im_x > 2 * target_x or im_y > 2 * target_y: pil_image.thumbnail( ( 2 * target_x, 2 * target_y ), PILImage.NEAREST )
    #    
    
    pil_image.thumbnail( ( target_x, target_y ), PILImage.ANTIALIAS )
    
def GenerateNumpyImage( path ):
    
    numpy_image = cv2.imread( path, flags = -1 ) # flags = -1 loads alpha channel, if present
    
    ( y, x, depth ) = numpy_image.shape
    
    if depth == 4: raise Exception( 'CV is bad at alpha!' )
    else: numpy_image = cv2.cvtColor( numpy_image, cv2.COLOR_BGR2RGB )
    
    return numpy_image
    
def GenerateHydrusBitmap( path ):
    
    try:
        
        numpy_image = GenerateNumpyImage( path )
        
        return GenerateHydrusBitmapFromNumPyImage( numpy_image )
        
    except:
        
        pil_image = GeneratePILImage( path )
        
        return GenerateHydrusBitmapFromPILImage( pil_image )
        
    
def GenerateHydrusBitmapFromNumPyImage( numpy_image ):
    
    ( y, x, depth ) = numpy_image.shape
    
    if depth == 4: return HydrusBitmap( numpy_image.data, wx.BitmapBufferFormat_RGBA, ( x, y ) )
    else: return HydrusBitmap( numpy_image.data, wx.BitmapBufferFormat_RGB, ( x, y ) )
    
def GenerateNumPyImageFromPILImage( pil_image ):
    
    if pil_image.mode == 'RGBA' or ( pil_image.mode == 'P' and pil_image.info.has_key( 'transparency' ) ):
        
        if pil_image.mode == 'P': pil_image = pil_image.convert( 'RGBA' )
        
    else:
        
        if pil_image.mode != 'RGB': pil_image = pil_image.convert( 'RGB' )
        
    
    ( w, h ) = pil_image.size
    
    s = pil_image.tostring()
    
    return numpy.fromstring( s, dtype = 'uint8' ).reshape( ( h, w, len( s ) // ( w * h ) ) )
    
def GenerateHydrusBitmapFromPILImage( pil_image ):
    
    if pil_image.mode == 'RGBA' or ( pil_image.mode == 'P' and pil_image.info.has_key( 'transparency' ) ):
        
        if pil_image.mode == 'P': pil_image = pil_image.convert( 'RGBA' )
        
        format = wx.BitmapBufferFormat_RGBA
        
    else:
        
        if pil_image.mode != 'RGB': pil_image = pil_image.convert( 'RGB' )
        
        format = wx.BitmapBufferFormat_RGB
        
    
    return HydrusBitmap( pil_image.tostring(), format, pil_image.size )
    
def GeneratePerceptualHash( path ):
    
    numpy_image = cv2.imread( path, cv2.CV_LOAD_IMAGE_UNCHANGED )
    
    ( y, x, depth ) = numpy_image.shape
    
    if depth == 4:
        
        # create a white greyscale canvas
        
        white = numpy.ones( ( x, y ) ) * 255
        
        # create weight and transform numpy_image to greyscale
        
        numpy_alpha = numpy_image[ :, :, 3 ]
        
        numpy_image_bgr = numpy_image[ :, :, :3 ]
        
        numpy_image_gray = cv2.cvtColor( numpy_image_bgr, cv2.COLOR_BGR2GRAY )
        
        numpy_image_result = numpy.empty( ( y, x ), numpy.float32 )
        
        # paste greyscale onto the white
        
        # can't think of a better way to do this!
        # cv2.addWeighted only takes a scalar for weight!
        for i in range( y ):
            
            for j in range( x ):
                
                opacity = float( numpy_alpha[ i, j ] ) / 255.0
                
                grey_part = numpy_image_gray[ i, j ] * opacity
                white_part = 255 * ( 1 - opacity )
                
                pixel = grey_part + white_part
                
                numpy_image_result[ i, j ] = pixel
                
            
        
        numpy_image_gray = numpy_image_result
        
        # use 255 for white weight, alpha for image weight
        
    else:
        
        numpy_image_gray = cv2.cvtColor( numpy_image, cv2.COLOR_BGR2GRAY )
        
    
    numpy_image_tiny = cv2.resize( numpy_image_gray, ( 32, 32 ), interpolation = cv2.INTER_AREA )
    
    # convert to float and calc dct
    
    numpy_image_tiny_float = numpy.float32( numpy_image_tiny )
    
    dct = cv2.dct( numpy_image_tiny_float )
    
    # take top left 8x8 of dct
    
    dct_88 = dct[:8,:8]
    
    # get mean of dct, excluding [0,0]
    
    mask = numpy.ones( ( 8, 8 ) )
    
    mask[0,0] = 0
    
    average = numpy.average( dct_88, weights = mask )
    
    # make a monochromatic, 64-bit hash of whether the entry is above or below the mean
    
    bytes = []
    
    for i in range( 8 ):
        
        byte = 0
        
        for j in range( 8 ):
            
            byte <<= 1 # shift byte one left
            
            value = dct_88[i,j]
            
            if value > average: byte |= 1
            
        
        bytes.append( byte )
        
    
    answer = str( bytearray( bytes ) )
    
    # we good
    
    return answer
    
def old_GeneratePerceptualHash( path ):
    
    # I think what I should be doing here is going cv2.imread( path, flags = cv2.CV_LOAD_IMAGE_GRAYSCALE )
    # then efficiently resize
    
    thumbnail = GeneratePILImage( path )
    
    # convert to 32 x 32 greyscale
    
    if thumbnail.mode == 'P':
        
        thumbnail = thumbnail.convert( 'RGBA' ) # problem with some P images converting to L without RGBA step in between
        
    
    if thumbnail.mode == 'RGBA':
        
        # this is some code i picked up somewhere
        # another great example of PIL failing; it turns all alpha to pure black on a RGBA->RGB
        
        thumbnail.load()
        
        canvas = PILImage.new( 'RGB', thumbnail.size, ( 255, 255, 255 ) )
        
        canvas.paste( thumbnail, mask = thumbnail.split()[3] )
        
        thumbnail = canvas
        
    
    thumbnail = thumbnail.convert( 'L' )
    
    thumbnail = thumbnail.resize( ( 32, 32 ), PILImage.ANTIALIAS )
    
    # convert to mat
    
    numpy_thumbnail_8 = cv.CreateMatHeader( 32, 32, cv.CV_8UC1 )
    
    cv.SetData( numpy_thumbnail_8, thumbnail.tostring() )
    
    numpy_thumbnail_32 = cv.CreateMat( 32, 32, cv.CV_32FC1 )
    
    cv.Convert( numpy_thumbnail_8, numpy_thumbnail_32 )
    
    # compute dct
    
    dct = cv.CreateMat( 32, 32, cv.CV_32FC1 )
    
    cv.DCT( numpy_thumbnail_32, dct, cv.CV_DXT_FORWARD )
    
    # take top left 8x8 of dct
    
    dct = cv.GetSubRect( dct, ( 0, 0, 8, 8 ) )
    
    # get mean of dct, excluding [0,0]
    
    mask = cv.CreateMat( 8, 8, cv.CV_8U )
    
    cv.Set( mask, 1 )
    
    mask[0,0] = 0
    
    channel_averages = cv.Avg( dct, mask )
    
    average = channel_averages[0]
    
    # make a monochromatic, 64-bit hash of whether the entry is above or below the mean
    
    bytes = []
    
    for i in range( 8 ):
        
        byte = 0
        
        for j in range( 8 ):
            
            byte <<= 1 # shift byte one left
            
            value = dct[i,j]
            
            if value > average: byte |= 1
            
        
        bytes.append( byte )
        
    
    answer = str( bytearray( bytes ) )
    
    # we good
    
    return answer
    
def GeneratePILImage( path ): return PILImage.open( path )

def GeneratePILImageFromNumpyImage( numpy_image ):
    
    ( h, w, depth ) = numpy_image.shape
    
    if depth == 3: format = 'RGB'
    elif depth == 4: format = 'RGBA'
    
    pil_image = PILImage.fromstring( format, ( w, h ), numpy_image.data )
    
    return pil_image
    
def GetGIFFrameDurations( path ):
    
    pil_image_for_duration = GeneratePILImage( path )
    
    frame_durations = []
    
    i = 0
    
    while True:
        
        try: pil_image_for_duration.seek( i )
        except: break
        
        if 'duration' not in pil_image_for_duration.info: duration = 40 # 25 fps default when duration is missing or too funky to extract. most stuff looks ok at this.
        else:
            
            duration = pil_image_for_duration.info[ 'duration' ]
            
            if duration == 0: duration = 40
            
        
        frame_durations.append( duration )
        
        i += 1
        
    
    return frame_durations
    
def GetHammingDistance( phash1, phash2 ):
    
    distance = 0
    
    phash1 = bytearray( phash1 )
    phash2 = bytearray( phash2 )
    
    for i in range( len( phash1 ) ):
        
        xor = phash1[i] ^ phash2[i]
        
        while xor > 0:
            
            distance += 1
            xor &= xor - 1
            
        
    
    return distance
    
def GetImageProperties( path ):
    
    ( ( width, height ), num_frames ) = GetResolutionAndNumFrames( path )
    
    if num_frames > 1:
        
        durations = GetGIFFrameDurations( path )
        
        duration = sum( durations )
        
    else:
        
        duration = None
        num_frames = None
        
    
    return ( ( width, height ), duration, num_frames )
    
def GetResolutionAndNumFrames( path ):
    
    pil_image = GeneratePILImage( path )
    
    ( x, y ) = pil_image.size
    
    try:
        
        pil_image.seek( 1 )
        pil_image.seek( 0 )
        
        num_frames = 1
        
        while True:
            
            try:
                
                pil_image.seek( pil_image.tell() + 1 )
                num_frames += 1
                
            except: break
            
        
    except: num_frames = 1
    
    return ( ( x, y ), num_frames )
    
def GetThumbnailResolution( ( im_x, im_y ), ( target_x, target_y ) ):
    
    im_x = float( im_x )
    im_y = float( im_y )
    
    target_x = float( target_x )
    target_y = float( target_y )
    
    x_ratio = im_x / target_x
    y_ratio = im_y / target_y
    
    ratio_to_use = max( x_ratio, y_ratio )
    
    target_x = int( im_x / ratio_to_use )
    target_y = int( im_y / ratio_to_use )
    
    return ( target_x, target_y )
    
''' # old pil code

def _GetCurrentFramePIL( pil_image, target_resolution, canvas ):
    
    current_frame = EfficientlyResizePILImage( pil_image, target_resolution )
    
    if pil_image.mode == 'P' and 'transparency' in pil_image.info:
        
        # I think gif problems are around here somewhere; the transparency info is not converted to RGBA properly, so it starts drawing colours when it should draw nothing
        
        current_frame = current_frame.convert( 'RGBA' )
        
        if canvas is None: canvas = current_frame
        else: canvas.paste( current_frame, None, current_frame ) # yeah, use the rgba image as its own mask, wut.
        
    else: canvas = current_frame
    
    return canvas
    

def _GetFramePIL( self, index ):
    
    pil_image = self._image_object
    
    pil_image.seek( index )
    
    canvas = self._GetCurrentFramePIL( pil_image, self._target_resolution, canvas )
    
    return GenerateHydrusBitmapFromPILImage( canvas )
    

def _GetFramesPIL( self ):
    
    pil_image = self._image_object
    
    canvas = None
    
    global_palette = pil_image.palette
    
    dirty = pil_image.palette.dirty
    mode = pil_image.palette.mode
    rawmode = pil_image.palette.rawmode
    
    # believe it or not, doing this actually fixed a couple of gifs!
    pil_image.seek( 1 )
    pil_image.seek( 0 )
    
    while True:
        
        canvas = self._GetCurrentFramePIL( pil_image, self._target_resolution, canvas )
        
        yield GenerateHydrusBitmapFromPILImage( canvas )
        
        try:
            
            pil_image.seek( pil_image.tell() + 1 )
            
            if pil_image.palette == global_palette: # for some reason, when we fall back to global palette (no local-frame palette), we reset bunch of important variables!
                
                pil_image.palette.dirty = dirty
                pil_image.palette.mode = mode
                pil_image.palette.rawmode = rawmode
                
            
        except: break
        
    
'''

# the cv code was initially written by @fluffy_cub
class HydrusBitmap( object ):
    
    def __init__( self, data, format, size ):
        
        self._data = lz4.dumps( data )
        self._format = format
        self._size = size
        
    
    def GetWxBitmap( self ):
        
        ( width, height ) = self._size
        
        if self._format == wx.BitmapBufferFormat_RGB: return wx.BitmapFromBuffer( width, height, lz4.loads( self._data ) )
        else: return wx.BitmapFromBufferRGBA( width, height, lz4.loads( self._data ) )
        
    
    def GetWxImage( self ):
        
        ( width, height ) = self._size
        
        if self._format == wx.BitmapBufferFormat_RGB: return wx.ImageFromBuffer( width, height, lz4.loads( self._data ) )
        else:
            
            bitmap = wx.BitmapFromBufferRGBA( width, height, lz4.loads( self._data ) )
            
            image = wx.ImageFromBitmap( bitmap )
            
            wx.CallAfter( bitmap.Destroy )
            
            return image
            
        
    
    def GetEstimatedMemoryFootprint( self ): return len( self._data )
    
    def GetSize( self ): return self._size
    
class RasterContainer( object ):
    
    def __init__( self, media, target_resolution = None ):
        
        if target_resolution is None: target_resolution = media.GetResolution()
        
        self._media = media
        self._target_resolution = target_resolution
        
        hash = self._media.GetHash()
        mime = self._media.GetMime()
        
        self._path = CC.GetFilePath( hash, mime )
        
        ( original_width, original_height ) = self._media.GetResolution()
        
        ( my_width, my_height ) = target_resolution
        
        width_zoom = my_width / float( original_width )
        height_zoom = my_height / float( original_height )
        
        self._zoom = min( ( width_zoom, height_zoom ) )
        
        if self._zoom > 1.0: self._zoom = 1.0
        
    
class ImageContainer( RasterContainer ):
    
    def __init__( self, media, target_resolution = None ):
        
        RasterContainer.__init__( self, media, target_resolution )
        
        self._hydrus_bitmap = None
        
        HydrusThreading.CallToThread( self.THREADRender )
        
    
    def _GetHydrusBitmap( self ):
        
        try:
            
            numpy_image = GenerateNumpyImage( self._path )
            
            resized_numpy_image = EfficientlyResizeNumpyImage( numpy_image, self._target_resolution )
            
            return GenerateHydrusBitmapFromNumPyImage( resized_numpy_image )
            
        except:
            
            pil_image = GeneratePILImage( self._path )
            
            resized_pil_image = EfficientlyResizePILImage( pil_image, self._target_resolution )
            
            return GenerateHydrusBitmapFromPILImage( resized_pil_image )
            
        
    
    def THREADRender( self ):
        
        time.sleep( 0.00001 ) # thread yield
        
        wx.CallAfter( self.SetHydrusBitmap, self._GetHydrusBitmap() )
        
        HC.pubsub.pub( 'finished_rendering', self.GetKey() )
        
    
    def GetEstimatedMemoryFootprint( self ): return self._hydrus_bitmap.GetEstimatedMemoryFootprint()
    
    def GetHash( self ): return self._media.GetHash()
    
    def GetHydrusBitmap( self ): return self._hydrus_bitmap
    
    def GetKey( self ): return ( self._media.GetHash(), self._target_resolution )
    
    def GetNumFrames( self ): return self._media.GetNumFrames()
    
    def GetResolution( self ): return self._media.GetResolution()
    
    def GetSize( self ): return self._target_resolution
    
    def GetZoom( self ): return self._zoom
    
    def IsRendered( self ): return self._hydrus_bitmap is not None
    
    def IsScaled( self ): return self._zoom != 1.0
    
    def SetHydrusBitmap( self, hydrus_bitmap ): self._hydrus_bitmap = hydrus_bitmap