{-# LANGUAGE BangPatterns #-}

module Statistics.LinearRegression (linearRegressionRSqr, linearRegression, correl, covar) where

import qualified Data.Vector.Unboxed as U
import qualified Statistics.Sample as S

--- * Simple linear regression

-- | Covariance of two samples
covar :: S.Sample -> S.Sample -> Double
covar xs ys = U.sum (U.zipWith (*) (U.map f1 xs) (U.map f2 ys)) / (n-1)
    where
          !n = fromIntegral $ U.length xs
          !m1 = S.mean xs
          !m2 = S.mean ys
          f1 = \x -> (x - m1)
          f2 = \x -> (x - m2)
{-# INLINE covar #-}


-- | Pearson's product-moment correlation coefficient
correl :: S.Sample -> S.Sample -> Double
correl xs ys = let !c = covar xs ys
                   !sx = S.stdDev xs
                   !sy = S.stdDev ys
               in c / (sx * sy)
{-# INLINE correl #-}

-- | Simple linear regression between 2 samples.
--   Takes two vectors Y={yi} and X={xi} and returns
--   (alpha, beta, r*r) such that Y = alpha + beta*X
--   and where r is the Pearson product-moment correlation
--   coefficient
linearRegressionRSqr :: S.Sample -> S.Sample -> (Double, Double, Double)
linearRegressionRSqr xs ys = (alpha, beta, r*r)
    where 
          !c                   = U.sum (U.zipWith (*) (U.map (subtract m1) xs) (U.map (subtract m2) ys)) / (n-1)
          !r                   = c / (sx * sy)
          !m1                  = S.mean xs 
          !m2                  = S.mean ys
          !sx                  = S.stdDev xs
          !sy                  = S.stdDev ys
          !n                   = fromIntegral $ U.length xs
          !beta                = r * sy / sx
          !alpha               = m2 - beta * m1
{-# INLINE linearRegressionRSqr #-}
          
-- | Simple linear regression between 2 samples.
--   Takes two vectors Y={yi} and X={xi} and returns
--   (alpha, beta, r*r) such that Y = alpha + beta*X          
linearRegression :: S.Sample -> S.Sample -> (Double, Double)
linearRegression xs ys = (alpha, beta)
    where 
          !c                   = U.sum (U.zipWith (*) (U.map (subtract m1) xs) (U.map (subtract m2) ys)) / (n-1)
          !r                   = c / (sx * sy)
          !m1                  = S.mean xs 
          !m2                  = S.mean ys
          !sx                  = S.stdDev xs
          !sy                  = S.stdDev ys
          !n                   = fromIntegral $ U.length xs
          !beta                = r * sy / sx
          !alpha               = m2 - beta * m1
{-# INLINE linearRegression #-}