-----------------------------------------------------------------------------
-- |
-- Module      :  DSP.Estimation.Frequency.FCI
-- Copyright   :  (c) Matthew Donadio 2003
-- License     :  GPL
--
-- Maintainer  :  m.p.donadio@ieee.org
-- Stability   :  experimental
-- Portability :  portable
--
-- This module contains a few simple algorithms for interpolating the
-- peak location of a DFT\/FFT.
--
-----------------------------------------------------------------------------

-- TODO: confirm that quinn2 needs log10 and not ln

module DSP.Estimation.Frequency.FCI (quinn1, quinn2, quinn3, jacobsen, macleod3, macleod5, rv) where

import DSP.Basic((^!))
import Data.Array
import Data.Complex

log10 :: Floating a => a -> a
log10 = logBase 10

-- | Quinn's First Estimator (FCI1)

quinn1 :: (Ix a, Integral a, RealFloat b) => Array a (Complex b) -- ^ X[k]
       -> a -- ^ k
       -> b -- ^ w

quinn1 x k = 2 * pi * ((fromIntegral k) + d) / (fromIntegral n)
    where d | dp > 0 && dm > 0 = dp
	    | otherwise        = dm
	  dp = -ap / (1 - ap)
 	  dm =  am / (1 - am)
 	  ap = magnitude (x!(k+1)) / magnitude (x!k)
 	  am = magnitude (x!(k-1)) / magnitude (x!k)
 	  n = snd (bounds x) + 1

-- | Quinn's Second Estimator (FCI2)

quinn2 :: (Ix a, Integral a, RealFloat b) => Array a (Complex b) -- ^ X[k]
       -> a -- ^ k
       -> b -- ^ w

quinn2 x k = 2 * pi * ((fromIntegral k) + d) / (fromIntegral n)
    where d = (dp + dm) / 2 + tau(dp^!2) - tau(dm^!2)
          dp = -ap / (1 - ap)
 	  dm =  am / (1 - am)
 	  ap = magnitude (x!(k+1)) / magnitude (x!k)
 	  am = magnitude (x!(k-1)) / magnitude (x!k)
 	  tau y = 0.25 * log10(3*y^!2 + 6 * y + 1) - (sqrt 6) / 24 * log10 ((y + 1 - sqrt (2/3)) / (y + 1 + sqrt (2/3)))
 	  n = snd (bounds x) + 1

-- | Quinn's Third Estimator (FCI3)

quinn3 :: (Ix a, Integral a, RealFloat b) => Array a (Complex b) -- ^ X[k]
       -> a -- ^ k
       -> b -- ^ w

quinn3 x k = 2 * pi * ((fromIntegral k) + d) / (fromIntegral n)
    where d = (dm + dp) / 2 + (dp - dm) * (3*dt^!3 + 2*dt) / (3*dt^!4+6*dt^!2+1)
	  dt | dm > 0 && dp > 0 = dp
	     | otherwise        = dm
	  dp = -ap / (1 - ap)
 	  dm =  am / (1 - am)
 	  ap = magnitude (x!(k+1)) / magnitude (x!k)
 	  am = magnitude (x!(k-1)) / magnitude (x!k)
 	  n = snd (bounds x) + 1

-- | Eric Jacobsen's Estimator

jacobsen :: (Ix a, Integral a, RealFloat b) => Array a (Complex b) -- ^ X[k]
       -> a -- ^ k
       -> b -- ^ w

jacobsen x k = 2 * pi * ((fromIntegral k) + d) / (fromIntegral n)
    where d = realPart ((x!(k-1) - x!(k+1)) / (2 * x!k - x!(k-1) - x!(k+1)))
 	  n = snd (bounds x) + 1

-- | MacLeod's Three Point Estimator

macleod3 :: (Ix a, Integral a, RealFloat b) => Array a (Complex b) -- ^ X[k]
       -> a -- ^ k
       -> b -- ^ w

macleod3 x k = 2 * pi * ((fromIntegral k) + d) / (fromIntegral n)
    where rm1 = realPart (x!(k-1) * conjugate (x!k))
 	  r   = realPart (x!k     * conjugate (x!k))
 	  rp1 = realPart (x!(k+1) * conjugate (x!k))
	  d = (sqrt (1 + 8 * g^!2) - 1) / 4 / g
 	  g = (rm1 - rp1) / (2 * r + rm1 + rp1)
 	  n = snd (bounds x) + 1

-- | MacLeod's Three Point Estimator

macleod5 :: (Ix a, Integral a, RealFloat b) => Array a (Complex b) -- ^ X[k]
       -> a -- ^ k
       -> b -- ^ w

macleod5 x k = 2 * pi * ((fromIntegral k) + d) / (fromIntegral n)
     where rm2 = realPart (x!(k-2) * conjugate (x!k))
	   rm1 = realPart (x!(k-1) * conjugate (x!k))
 	   r   = realPart (x!k     * conjugate (x!k))
 	   rp1 = realPart (x!(k+1) * conjugate (x!k))
 	   rp2 = realPart (x!(k+2) * conjugate (x!k))
	   d = 0.4041 * atan (2.93 * g)
 	   g = (4 * (rm1 - rp1) + 2 * (rm2 - rp2)) / (12 * r + 8 * (rm1 + rp1) + rm2 + rp2)
 	   n = snd (bounds x) + 1

-- | Rife and Vincent's Estimator

rv :: (Ix a, Integral a, RealFloat b) => Array a (Complex b) -- ^ X[k]
       -> a -- ^ k
       -> b -- ^ w

rv x k = 2 * pi * ((fromIntegral k) + d) / (fromIntegral n)
    where d = fromIntegral at * magnitude (x!(k+at) / x!k) / (1 + magnitude (x!(k+at) / x!k))
	  at | (magnitude (x!(k+1)))^!2 > (magnitude (x!(k-1)))^!2 =  1
	     | otherwise                                         = -1
 	  n = snd (bounds x) + 1