--    The University of New Mexico's Haskell Image Processing Library
--    Copyright (C) 2013 Joseph Collard
--
--    This program is free software: you can redistribute it and/or modify
--    it under the terms of the GNU General Public License as published by
--    the Free Software Foundation, either version 3 of the License, or
--    (at your option) any later version.
--
--    This program is distributed in the hope that it will be useful,
--    but WITHOUT ANY WARRANTY; without even the implied warranty of
--    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
--    GNU General Public License for more details.
--
--    You should have received a copy of the GNU General Public License
--    along with this program.  If not, see <http://www.gnu.org/licenses/>.

{-# LANGUAGE ViewPatterns, FlexibleInstances, TypeFamilies, FlexibleContexts #-}
{-# OPTIONS -O2 #-}
module Data.Image.Complex(-- * Complex Images
                          ComplexPixel(..),
                          -- * Discrete Fourier Transformations
                          makeFilter,
                          fft, ifft,
                          -- * Converting Complex Images
                          complex,
                          realPart,
                          imagPart,
                          magnitude,
                          angle,
                          shrink,
                          complexImageToRectangular,
                          complexImageToPolar,
                          module Data.Complex) where

import Data.Image.Internal(Image(..), PixelOp, imageMap,dimensions)

--base>=4
import qualified Data.Complex as C
import Data.Complex(Complex((:+)))
import Data.List(transpose)
import Data.Bits

--vector>=0.10.0.1
import qualified Data.Vector as V

class RealFloat (Value px) => ComplexPixel px where
  type Value px :: *
  toComplex ::  px -> C.Complex (Value px)
  fromComplex :: C.Complex (Value px) -> px
  
{-| Given a positive integer m, a positive integer n, and a function 
    returning a pixel value, makeFilter returns an image with m rows 
    and n columns. Let x equal i if i is less than m/2 and i - m otherwise 
    and let y equal j if j is less than n/2 and j - n otherwise. To match 
    the periodicity of the 2D discrete Fourier spectrum, the value of the 
    result image at location (i, j) is computed by applying the function to x 
    and y, e.g., the value at location (0, 0) is the result of applying the 
    function to 0 and 0, the value at (m-1, n-1) is the result of applying 
    function to -1 and -1.

    >>> makeFilter 128 128 (\ r c -> fromIntegral (r + c)) :: GrayImage
    < Image 128x128 >

    <https://raw.github.com/jcollard/unm-hip/master/examples/makefilter.jpg>

    @
    laplacianOfGaussian :: Double -> Int -> Int -> Double
    laplacianOfGaussian stddev i j = ((-pi) / (stddev*stddev)) * (1 - x) * (exp (-x)) where
      r = fromIntegral $ i*i + j*j
      x = r / (2.0*stddev)
    @

    >>>let d2g = makeFilter 128 128 (laplacianOfGaussian 8) :: GrayImage
 
    <https://raw.github.com/jcollard/unm-hip/master/examples/d2g.jpg>

    @
    gaussian :: Double -> Int -> Int -> Double 
    gaussian variance i j = exp (-x) where
      r = fromIntegral $ i*i + j*j
      x = r / (2.0*pi*variance)
    @

    >>>let g = makeFilter 128 128 (gaussian 8) :: GrayImage
    
    <https://raw.github.com/jcollard/unm-hip/master/examples/g.jpg>
-}
makeFilter :: (Image img) => Int -> Int -> (PixelOp (Pixel img)) -> img
makeFilter rows cols func = makeImage rows cols func' where
  func' r c = let x = if r < (rows `div` 2) then r else r-rows 
                  y = if c < (cols `div` 2) then c else c-cols
              in func x y

{-| Given an image whose pixels can be converted to a complex value, 
    fft returns an image with complex pixels representing its 2D discrete 
    Fourier transform (DFT). Because the DFT is computed using the Fast Fourier 
    Transform (FFT) algorithm, the number of rows and columns of the image 
    must both be powers of two, i.e., 2K where K is an integer.

    >>>frog <- readImage "images/frog.pgm"
    >>>let frogpart = crop 64 64 128 128 frog
    
    <https://raw.github.com/jcollard/unm-hip/master/examples/frog.jpg>
    
    <https://raw.github.com/jcollard/unm-hip/master/examples/frogpart.jpg>

    >>>magnitude . imageMap log . fft $ frogpart

    <https://raw.github.com/jcollard/unm-hip/master/examples/fft.jpg>  
    
    >>>fft d2g
 
    <https://raw.github.com/jcollard/unm-hip/master/examples/fftd2g.jpg>  
    
    >>>fft g
    
    <https://raw.github.com/jcollard/unm-hip/master/examples/fftg.jpg>      
-}
fft :: (Image img,
        ComplexPixel (Pixel img),
        Image img',
        Pixel img' ~ C.Complex (Value (Pixel img))) => img -> img'
fft img@(dimensions -> (rows, cols)) = check where
  check = if (isPowerOfTwo rows && isPowerOfTwo cols) then makeImage rows cols fftimg else error "Image is not a power of 2 in rows and cols"
  fftimg r c = fft' V.! (r*cols + c)
  vector = V.map toComplex . V.fromList . pixelList $ img
  fft' = fftv rows cols vector

{-| Given an image, ifft returns a complex image representing its 2D 
    inverse discrete Fourier transform (DFT). Because the inverse DFT is 
    computed using the Fast Fourier Transform (FFT) algorithm, the number 
    of rows and columns of <image> must both be powers of two, i.e., 2K 
    where K is an integer. 

    >>>ifft ((fft frogpart) * (fft d2g))

    <https://raw.github.com/jcollard/unm-hip/master/examples/ifft.jpg>
    
    >>>ifft ((fft frogpart) * (fft g))

    <https://raw.github.com/jcollard/unm-hip/master/examples/ifft2.jpg>
 -}
ifft :: (Image img,
        ComplexPixel (Pixel img),
        Image img',
        Pixel img' ~ C.Complex (Value (Pixel img))) => img -> img'
ifft img@(dimensions -> (rows, cols)) = check where
  check = if (isPowerOfTwo rows && isPowerOfTwo cols) then makeImage rows cols fftimg else error "Image is not a power of 2 in rows and cols"
  fftimg r c = fft' V.! (r*cols + c)
  vector = V.map toComplex . V.fromList . pixelList $ img
  fft' = ifftv rows cols vector

{-| Given a complex image, returns a real image representing
    the real part of the image.

    @
    harmonicSignal :: Double -> Double -> Int -> Int -> C.Complex Double
    harmonicSignal u v m n = exp (-pii*2.0 * var) where 
      pii = 0.0 C.:+ pi
      var = (u*m' + v*n') C.:+ 0.0
      [m',n'] = map fromIntegral [m, n]
    @

    >>> let signal = makeImage 128 128 (harmonicSignal (3/128) (2/128)) :: ComplexImage

    <https://raw.github.com/jcollard/unm-hip/master/examples/signal.jpg>

    >>>let cosine = realPart signal

    <https://raw.github.com/jcollard/unm-hip/master/examples/cosine.jpg>

    >>>realPart . ifft $ (fft frogpart) * (fft d2g)
 
    <https://raw.github.com/jcollard/unm-hip/master/examples/realpart.jpg>
    
    >>>realPart . ifft $ (fft frogpart) * (fft g)
 
    <https://raw.github.com/jcollard/unm-hip/master/examples/realpart2.jpg>

-}
realPart :: (Image img,
             ComplexPixel (Pixel img),
             Image img',
             Pixel img' ~ Value (Pixel img)) => img -> img'
realPart = imageMap (C.realPart . toComplex)

{-| Given a complex image, returns a real image representing
   the imaginary part of the image

   >>>let sine = imagPart signal
 
    <https://raw.github.com/jcollard/unm-hip/master/examples/sine.jpg>

 -}
imagPart :: (Image img,
             ComplexPixel (Pixel img),
             Image img',
             Pixel img' ~ Value (Pixel img)) => img -> img'
imagPart = imageMap (C.imagPart . toComplex)

{-| Given a complex image, returns a real image representing
    the magnitude of the image.

    >>>magnitude signal

 -}
magnitude :: (Image img,
              ComplexPixel (Pixel img),
              Image img',
              Pixel img' ~ Value (Pixel img)) => img -> img'
magnitude = imageMap (C.magnitude . toComplex)

{-| Given a complex image, returns a real image representing
    the angle of the image 
   
    >>>angle signal

    <https://raw.github.com/jcollard/unm-hip/master/examples/angle.jpg>
-}
angle :: (Image img,
          ComplexPixel (Pixel img),
          Image img',
          Pixel img' ~ Value (Pixel img)) => img -> img'
angle = imageMap (C.phase . toComplex)

{-| Given a complex image, returns a pair of real images each
    representing the component (magnitude, phase) of the image

    >>>leftToRight' . complexImageToPolar $ signal
 
    <https://raw.github.com/jcollard/unm-hip/master/examples/compleximagetopolar.jpg>
-}
complexImageToPolar :: (Image img,
                        ComplexPixel (Pixel img),
                        Image img',
                        Pixel img' ~ Value (Pixel img)) => img -> (img', img')
complexImageToPolar img@(dimensions -> (rows, cols)) = (mkImg mag, mkImg phs) where
  mkImg = makeImage rows cols
  ref' r c = C.polar . toComplex . ref img r $ c
  mag r c = fst (ref' r c)
  phs r c = snd (ref' r c)

{-| Given an image representing the real part of a complex image, and
    an image representing the imaginary part of a complex image, returns
    a complex image.

    >>>complex cosine sine

    <https://raw.github.com/jcollard/unm-hip/master/examples/signal.jpg>
 -}
complex :: (Image img,
            Image img',
            Pixel img' ~ C.Complex (Pixel img)) => img -> img -> img'
complex real imag@(dimensions -> (rows, cols)) = makeImage rows cols ri where
  ri r c = (ref real r c) C.:+ (ref imag r c)

{-| Given a complex image, return a pair of real images each representing
    a component of the complex image (real, imaginary).

    >>>leftToRight' . complexImageToRectangular $ signal

    <https://raw.github.com/jcollard/unm-hip/master/examples/complexsignaltorectangular.jpg>
 -}
complexImageToRectangular :: (Image img,
                              ComplexPixel (Pixel img),
                              Image img',
                              Pixel img' ~ Value (Pixel img)) => img -> (img', img')
complexImageToRectangular img = (realPart img, imagPart img)

{-| Given a complex image and a real positive number x, shrink returns 
    a complex image with the same dimensions. Let z be the value of the 
    image at location (i, j). The value of the complex result image at 
    location (i, j) is zero if |z| < x, otherwise the result has the 
    same phase as z but the amplitude is decreased by x.
   
 -}
shrink :: (Image img,
           ComplexPixel (Pixel img)) => (Value (Pixel img)) -> img -> img
shrink x img@(dimensions -> (rows, cols)) = 
  makeImage rows cols shrink' where
    shrink' r c = fromComplex . helper $ px where
      px = toComplex . ref img r $ c
      helper px
        | (C.magnitude px) < x = 0.0 C.:+ 0.0
        | otherwise = real C.:+ imag where
          mag = C.magnitude px
          x' = (mag - x)/mag
          real = x'*(C.realPart px)
          imag = x'*(C.imagPart px)


type Vector a = V.Vector (Complex a)
type FFT a = [Int] -> Vector a -> Int -> Int -> [Complex a] 

-- FFT support code

fftv :: (RealFloat a) => Int -> Int -> Vector a -> Vector a
fftv = fft' fftRange

ifftv :: (RealFloat a) => Int -> Int -> Vector a -> Vector a
ifftv rows cols vec = V.map (/fromIntegral (rows*cols)) . fft' ifftRange rows cols $ vec

isPowerOfTwo :: Int -> Bool
isPowerOfTwo n = n /= 0 && (n .&. (n-1)) == 0

fft' :: FFT a -> Int -> Int -> Vector a -> Vector a
fft' range rows cols orig = if check then fromRows rows' else err where 
  check = and . map isPowerOfTwo $ [rows, cols]
  err = error "FFT can only be applied to images with dimensions 2^k x 2^j where k and j are integers."
  (fromColumns -> cols') = map (fftc range rows cols 0 (rows-1) orig) [0..cols-1] -- FFT on each col
  rows' = map (fftr range cols 0 (cols-1) cols') [0..rows-1] -- FFT on each row

fromColumns :: [[Complex a]] -> V.Vector (Complex a)
fromColumns = fromRows . transpose

fromRows :: [[Complex a]] -> V.Vector (Complex a)
fromRows = V.fromList . concat

fftc :: FFT a -> Int -> Int -> Int -> Int -> Vector a -> Int -> [Complex a]
fftc fftfunc rows cols sIx eIx orig row = fftfunc indices orig rows 1 where
  indices = map ((+row) . (*cols)) $ [sIx..eIx]

fftr :: FFT a -> Int -> Int -> Int -> Vector a ->  Int -> [Complex a]
fftr fftfunc cols sIx eIx orig row = fftfunc indices orig cols 1 where
  indices = map (+ (row*cols)) $ [sIx..eIx]

fftRange :: (RealFloat a) => FFT a
fftRange = range (-2*pii)

ifftRange :: (RealFloat a) => FFT a
ifftRange = range (2*pii)

range :: (RealFloat a) => Complex a -> FFT a
range e ix vec n s 
  | n == 1 = [vec V.! (head ix)]
  | otherwise = fft' where
    fft' = seperate data'
    fi = fromIntegral
    ix0 = range e ix vec (n `div` 2) (2*s)
    ix1 = range e (drop s ix) vec (n `div` 2) (2*s)
    data' = (flip map) (zip3 ix0 ix1 [0..]) (\ (ix0, ix1, k) -> do
      let e' = exp (e * ((fi k) / (fi n)))
          ix0' = ((ix0 + e' * ix1))
          ix1' = ((ix0 - e' * ix1))
        in (ix0', ix1'))
    
seperate :: [(a, a)] -> [a]
seperate = seperate' [] [] where
  seperate' acc0 acc1 [] = (reverse acc0) ++ (reverse acc1)
  seperate' acc0 acc1 ((a, b):xs) = seperate' (a:acc0) (b:acc1) xs

pii :: (Floating a) => Complex a
pii = 0 :+ pi

-- End FFT support code