module Statistics.LinearRegression (linearRegression) where

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

--- * Simple linear regression

-- | Covariance
covar :: Sample -> Sample -> (Double, Double, Double)
covar xs ys = (U.sum (U.zipWith (*) (U.map f1 xs) (U.map f2 ys)) / (n-1), m1, m2)
    where
          n = fromIntegral $ U.length xs
          m1 = mean xs
          m2 = mean ys
          f1 = \x -> (x - m1)
          f2 = \x -> (x - m2)

-- | Pearson's product-moment correlation coefficient
correl :: Sample -> Sample -> (Double, Double, Double, Double, Double)
correl xs ys = let (c, m1, m2) = covar xs ys
                   sx = stdDev xs
                   sy = stdDev ys
               in (c / (stdDev xs * stdDev ys), m1, m2, sx, sy)

-- | Simple linear regression between 2 samples.
--   Takes two vectors Y={yi} and X={xi} and returns
--   (alpha, beta) such that Y = alpha + betaX
linearRegression :: Sample -> Sample -> (Double, Double)
linearRegression xs ys = (alpha, beta)
    where 
          (r, m1, m2, sx, sy) = correl xs ys
          beta                = r * sy / sx
          alpha               = m2 - beta * m1