# (c) 2019 Molecular Systems Lab
# Wyss Institute for Biologically-Inspired Engineering
# Harvard University
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
# 
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

from PIL import Image
import glob
import numpy as np
import colorsys
import os

MAX_16BIT = 2 ** 16 - 1 # 65535
MAX_8BIT = 2 ** 8 - 1 # 255

rgb_to_hls = np.vectorize(colorsys.rgb_to_hls)
hls_to_rgb = np.vectorize(colorsys.hls_to_rgb)

def overlayImages(files, prepend, cropCoords, setsToSave, contrastMap, 
                  colorMap, overlayMode='screen', backgroundColor=None,
                  in16Bit=True):
  '''Retrieves the specificied set of image files, contrasts, and overlays them
     according to which sets should be saved, and writes them to disk.'''

  (_, _, channels) = readFiles(files, contrastMap)
  
  overlayChannels(channels, files, prepend, cropCoords, setsToSave, colorMap, \
                  overlayMode, backgroundColor, in16Bit=in16Bit)


def alignAndOverlay(files, prepend, cropCoords, setsToSave, contrastMap,
                    alignmentMap, colorMap, dapiThreshold, maxShift, shiftInc,
                    maxShift2, shiftInc2, overlayMode='alpha',
                    backgroundColor=(0, 0, 0), in16Bit=True):
  '''Retrieves the specificied set of image files, aligns them according to the
     alignment map and parameters, contrasts the images, overlays them
     according to which sets should be saved, and writes them to disk.'''

  (aligners, alignDiffs, channels) = readFiles(files, contrastMap, 
                                               alignMap=alignmentMap)

  # Compute alignment values based off first aligner.
  baseID = aligners.keys()[0]
  base = aligners[baseID]
  alignDiffs[baseID] = (0, 0)
  for align in aligners.keys()[1:]:
    print 'Calculating alignment %d' % align
    alignDiffs[align] = crossCorrelate(base, aligners[align], \
                                       (dapiThreshold, maxShift, shiftInc, \
                                        maxShift2, shiftInc2))

  overlayChannels(channels, files, prepend, cropCoords, setsToSave, colorMap, \
                  overlayMode, backgroundColor, alignDiffs, in16Bit=in16Bit)

class Channel:
  def __init__(self, chanID, alignID, con0, con1, name, im):
    self.cID = chanID
    self.aID = alignID
    self.c0 = con0
    self.c1 = con1
    self.n = name
    self.img = im

def readFiles(files, contrastMap = None, alignMap = None):
  '''Reads a given set of files and creates Channel objects that store the
     image, contrast, and alignment information.'''

  # File format:
  # id_name.tif
  # id = (unique) ID of this image
  # name = remainder of string ignored
  #
  # ALIGN_alignID_name.tif
  # align = align (dapi) used for alignment
  # name = remainder of string ignored

  print files

  aligners = {} # aligners keyed on unique ID
  alignDiffs = {} # alignment diffs (x, y) keyed on unique ID
  channels = {} # channels

  for f in files:
    print f
    img = Image.open(f)

    data = os.path.basename(f).split('_')

    if data[0] == 'ALIGN':
      aligners[int(data[1])] = img

    else:
      cID = int(data[0])
      name = data[1:]

      if alignMap is None:
        aID = None
      else:
        aID = alignMap[cID]

      if contrastMap is None:
        (c0, c1) = (None, None)
      else:
        (c0, c1) = contrastMap[cID]

      channels[cID] = Channel(cID, aID, c0, c1, name, img)
  
  return (aligners, alignDiffs, channels)

def overlayChannels(channels, files, prepend, cropCoords, setsToSave,
                    colorMaps, overlayMode, backgroundColor, alignDiffs=None,
                    in16Bit=True):
  '''Overlays channels according to specificied mode.
       'screen': should be same as Photoshop screen method.
       'alpha': alpha composite, transparencies scaled by pixel intensity.
       'max': takes max value of each color (R/G/B) for each pixel.
       'sum': directly adds values from each channel.'''

  for cID in channels:
    print 'Looking at channel %d' % cID
    c = channels[cID]
    pixels = np.array(c.img) # Get array of pixel values
    if in16Bit:
      pixels = changeContrast16Bit(pixels, c)
      pixels = convert16to8bit(pixels)
    else:
      pixels = changeContrast8Bit(pixels, c)
    pixels = np.array(createRGBAfromGreyscaleArray(pixels))
    if alignDiffs is not None:
      pixels = shiftAndLeaveZerosRGBA(pixels, alignDiffs[c.aID][0], \
                                      alignDiffs[c.aID][1])
    if overlayMode == 'alpha':
      pixels = changeAlphaPixelValues(pixels)
      pixels = changeAllPixelsToHue(pixels, colorMaps[cID])
    elif overlayMode in ['screen', 'sum', 'max']:
      pixels = scalePixelsToHue(pixels, colorMaps[cID])
    c.img = createRGBAfromRGBAArray(pixels)

  # Alpha composite method
  if overlayMode == 'alpha':
    for imgSet in setsToSave:
      toSave = Image.new('RGBA', (2048, 2048), backgroundColor)
      for cID in imgSet:
        toSave = Image.alpha_composite(toSave, channels[cID].img)
      toSave.save(prepend + '_'.join([str(i) for i in imgSet]) + '.tif')
      for i, coords in enumerate(cropCoords):
        toSave.crop(coords).save('%sCrop%d_%s.tif' % \
                                 (prepend, i, \
                                  '_'.join([str(i) for i in imgSet])))
  # Screen/sum/max blend mode
  elif overlayMode in ['screen', 'sum', 'max']:
    for setInd, imgSet in enumerate(setsToSave):
      pixels = np.array(channels[imgSet[0]].img)
      for cID in imgSet[1:]:
        if overlayMode == 'screen':
          pixels = screenBlend(pixels, np.array(channels[cID].img))
        elif overlayMode == 'sum':
          pixels = addBlend(pixels, np.array(channels[cID].img))
        elif overlayMode == 'max':
          pixels = maxBlend(pixels, np.array(channels[cID].img))

      toSave = createRGBAfromRGBAArray(np.rint(pixels))
      toSave.save(prepend + '_'.join([str(i) for i in imgSet]) + '.tif')
      for i, coords in enumerate(cropCoords[setInd]):
        toSave.crop(coords).save('%sCrop%d_%s.tif' % \
                                 (prepend, i, \
                                  '_'.join([str(i) for i in imgSet])))

def hlsTo8BitRGB(h, l, s):
  '''Converts HLS color space to RGB. H, L, and S should all be floating point
     numbers between 0 and 1.'''
  return np.rint(np.array(hls_to_rgb(h, l, s)) * 255.0)

def getIndividualSets(channels):
  return [[c] for c in channels]

def getIndividualSetsPlusOne(channels, baseline):
  return [[baseline, c] for c in channels]

def scalePixelsToHue(pixVals, colorInfo):
  '''Scales Red, Green, and Blue channel values for each pixel by the registered
     intensity. Colors may be in HLS or RGB space. NOTE: assumes images have
     just been converted from greyscale without modification to R/G/B values.'''
  r, g, b, a = pixVals.T

  if colorInfo[0] == 'RGB':
    (newR, newG, newB) = (float(colorInfo[1]), float(colorInfo[2]), 
                          float(colorInfo[3]))

  elif colorInfo[0] == 'HLS':
    newR, newG, newB = hlsTo8BitRGB(colorInfo[1], colorInfo[2], colorInfo[3])

  print (colorInfo, newR, newG, newB)

  r = r * newR / 255.0
  g = g * newG / 255.0
  b = b * newB / 255.0

  return np.dstack((r.T, g.T, b.T, a.T))

def changeAlphaPixelValues(pixVals):
  '''Scales transparency (alpha) values according to pixel intensities. NOTE:
     assumes hue has not been changed yet after conversion to RGBA from
     greyscale.'''
  r, g, b, a = pixVals.T
  a = r.copy() # Set alpha to value
  return np.dstack((r.T, g.T, b.T, a.T))

def changeAllPixelsToHue(pixVals, colorInfo):
  '''Sets all pixels to the same hue.'''
  r, g, b, a = pixVals.T

  if colorInfo[0] == 'RGB':
    r.fill(colorInfo[1])
    g.fill(colorInfo[2])
    b.fill(colorInfo[3])
  elif colorInfo[0] == 'HLS':
    newR, newG, newB = hlsTo8BitRGB(colorInfo[1], colorInfo[2], colorInfo[3])
    r.fill(int(newR))
    g.fill(int(newG))
    b.fill(int(newB))

  return np.dstack((r.T, g.T, b.T, a.T))

def changeContrast16Bit(pixVals, channel):
  '''Adjusts pixel values according to contrast limits.'''

  def contrast(c):
    if channel.c0 is None or channel.c1 is None:
      return c
    if c < channel.c0:
      return 0
    if c > channel.c1:
      return MAX_16BIT
    return MAX_16BIT * (c - channel.c0) / (channel.c1 - channel.c0)

  vectContrast = np.vectorize(contrast)
  pixVals = vectContrast(pixVals)

  return pixVals

def changeContrast8Bit(pixVals, channel):

  def contrast(c):
    if channel.c0 is None or channel.c1 is None:
      return c
    if c < channel.c0:
      return 0
    if c > channel.c1:
      return MAX_8BIT
    return MAX_8BIT * (c - channel.c0) / (channel.c1 - channel.c0)

  vectContrast = np.vectorize(contrast)
  pixVals = vectContrast(pixVals)

  return pixVals

def convert16to8bit(pixVals):
  '''Converts 16 bit pixel values into 8 bit (does not change dtype) by integer
     dividing to create 256 equal bins.'''
  return pixVals / 256

def createRGBAfromGreyscaleArray(pixVals):
  '''Given an 8 bit pixel array, returns a new RGBA image.'''
  return Image.fromarray(pixVals.astype('uint8')).convert('RGBA')

def createRGBAfromRGBAArray(pixVals):
  '''Given an 8 bit RGBA array, returns a new RGBA image.'''
  return Image.fromarray(pixVals.astype('uint8'), 'RGBA')

def shiftAndLeaveZeros(imgarr, xDiff, yDiff):
  '''Shifts an image array by a given number of pixels in the x and y
     directions, leaving zeros behind.'''

  if xDiff < 0:
    xDiffTuple = (0, -xDiff)
    xDiffKeep = (-xDiff, len(imgarr[0]) - xDiff)
  elif xDiff == 0:
    xDiffTuple = (0, 0)
    xDiffKeep = (0, len(imgarr[0]))
  elif xDiff > 0:
    xDiffTuple = (xDiff, 0)
    xDiffKeep = (0, -xDiff)

  
  if yDiff < 0:
    yDiffTuple = (0, -yDiff)
    yDiffKeep = (-yDiff, len(imgarr) - yDiff)
  elif yDiff == 0:
    yDiffTuple = (0, 0)
    yDiffKeep = (0, len(imgarr))
  elif yDiff > 0:
    yDiffTuple = (yDiff, 0)
    yDiffKeep = (0, -yDiff)

  return np.pad(imgarr,(yDiffTuple,xDiffTuple),
                mode='constant')[yDiffKeep[0]:yDiffKeep[1],
                                 xDiffKeep[0]:xDiffKeep[1]]

def shiftAndLeaveZerosRGBA(pixVals, xDiff, yDiff):
  '''Shifts an RGBA array by a given number of pixels in the x and y
     directions, leaving zeros behind.'''
  r, g, b, a = pixVals.T

  # NOTE: inverted due to transform    
  r = shiftAndLeaveZeros(r, yDiff, xDiff)
  g = shiftAndLeaveZeros(g, yDiff, xDiff)
  b = shiftAndLeaveZeros(b, yDiff, xDiff)
  a = shiftAndLeaveZeros(a, yDiff, xDiff)

  return np.dstack((r.T, g.T, b.T, a.T))

def screenBlend(pixVals1, pixVals2):
  '''Given two pixel value arrays, computes their screen blending result.'''
  # Typically, 1 - (1 - pix1)(1 - pix2), for 0 <= pix <= 1
  # We scale to 255 here
  return (255.0 * (1.0 - \
                   (np.multiply((1.0 - pixVals1 / 255.0), \
                                (1.0 - pixVals2 / 255.0)))))

def addBlend(pixVals1, pixVals2):
  '''Given two pixel value arrays, compute their add blend result. NOTE: values
     are converted to 32 bit to avoid wraparound.'''
  return np.clip(pixVals1.astype('uint32') + pixVals2.astype('uint32'), 
                 0, 255).astype('uint8')

def maxBlend(pixVals1, pixVals2):
  '''Given two pixel value arrays, computes their max blend result.'''
  return np.maximum(pixVals1, pixVals2)

def crossCorrelate(img1, img2, params):
  '''Determines how far an image (img2) must be shifted to maximize the overlap
     with a given base image (img1). Typically, these images correspond to DAPI
     or some other marker that is present in all sample images.'''
  (threshold, maxShift, shiftInc, maxShift2, shiftInc2) = params

  # Return amount that img2 should be shifted (xdiff, ydiff) to align with img1
  img1arr = np.array(img1)
  img2arr = np.array(img2)

  print '  Step 1 (of 2)...'
  maxOverlapTupleIntermediate = maximizeCrossCorrelation(img1arr, img2arr,
                                                         threshold, maxShift,
                                                         shiftInc)

  print '  Step 2 (of 2)...'
  return maximizeCrossCorrelation(img1arr, img2arr, threshold, maxShift2,
                                  shiftInc2,
                                  initXDiff=maxOverlapTupleIntermediate[0],
                                  initYDiff=maxOverlapTupleIntermediate[1])

def maximizeCrossCorrelation(img1arr, img2arr, threshold, maxShift, shiftInc,
                             initXDiff=0, initYDiff=0):
  '''Finds the maximal cross-correlation of two image channel arrays in a
     grid search area.'''
  maxOverlap = 0
  maxOverlapTuple = None
  for xDiff in range(-maxShift + initXDiff, maxShift + 1 + initXDiff, shiftInc):
    for yDiff in range(-maxShift + initYDiff,
                       maxShift + 1 + initYDiff, shiftInc):
      img2arrShifted = shiftAndLeaveZeros(img2arr, xDiff, yDiff)

      pixelDiffs = np.absolute(img2arrShifted - img1arr)

      # Calculate the overlap between the signals, by first thresholding and
      # then counting the number of pixels in both masks.
      overlap = np.sum((img2arrShifted > threshold) & \
                       (img1arr > threshold))

      if overlap > maxOverlap:
        maxOverlap = overlap
        maxOverlapTuple = (xDiff, yDiff)

  return maxOverlapTuple