{-# LANGUAGE QuasiQuotes #-}

-- | example comparing the convergence of <http://en.wikipedia.org/wiki/Romberg's_method romberg integration>
-- in terms of the built-in error estimate as well as an \"exact\" solution ("Data.Number.CReal")
--
-- originally submitted to uwaterloo's CHE 322. Adapted plot to be done
-- with R.
module Main where

import Data.Number.CReal
import Data.List
import Data.Typeable
import RlangQQ

import Data.HList.CommonMain
import Control.Monad.State
import Data.Int


main = do
      mkPlt 22 (1 :: Double)
      mkPlt 22 (1 :: Float)


-- number of traces on the plots
-- ie. 1 = trapezoid only, 2 = trapezoid + simpsons
-- 3 = trapezoid + simpsons + simpsons 3/8
-- 4 = ...
nTraces = [n| 5 |]

mkPlt nMax proxy =
    let exact = exp 10 - 1
        calc = rombM exp 0 10


        xs = listToRecN nTraces $ map (take nMax . map (\x -> toDouble $
                            abs (tfr x - exact) / exact))
           (calc `asTypeOf` [[proxy]])
        es = listToRecN nTraces $ (map (take nMax . map tfr) $ getRombErr calc :: [[Double]])

        outputfile = "romberg_" ++ show (typeOf proxy) ++ ".pdf"

        ls = [1 .. hNat2Integral nTraces :: Int32]
        ns = [1 .. fromIntegral nMax     :: Double]

    in
 [r|

        library(lattice); library(reshape); library(plyr);

        conv <- function(hsvar) { # should be a better way to get a data.frame
            x <- data.frame(hsvar)
            names(x) <- hs_ls
            x <- ddply(melt(x), .(variable), function(x)
                    data.frame(x, refinements = hs_ns   ))
            within(x, refinements <- as.integer(variable) + refinements)
        }

        es <- conv(hs_es)
        xs <- conv(hs_xs)

        pdf(hs_outputfile)

        plot(xyplot( value ~ refinements, xs, groups = variable,
                scales= list(y=list(log=T),x=list(log=T)), type='l',
                ylab='error', auto.key=list(title='order', type='l',
                                            columns= hs_ls[length(hs_ls)] ),
                main = hs_outputfile))

        plot(xyplot( xs$value ~ es$value, groups=xs$variable, type='b',
                scales=list(x=list(log=T), y=list(log=T)),
                ylab='actual error',
                xlab='estimated error',
                panel= function(...){
                    panel.xyplot(...)
                    panel.abline(a=0, b=1, col='black')
                }))

        dev.off()

    |]



-- | trapezoid rule. doesn't reuse previous values
trap f a b n =
    let h = (b - a) / fromIntegral n
        xs = take (n-1) $ iterate (h+) (a+h)
            -- easier to see as [a+h, a+2*h .. b-h])
            -- but some interval arithmetic doesn't support [a..b]
    in h * (foldl' (\s x -> s + f x) ((f a + f b) / 2) xs)

rombM f a b =
    let is = map (\n -> trap f a b (2^n)) [0..]
    in map snd
        $ iterate
            (\(n,s) -> (n+1, zipWith (improve n) (tail s) s))
            (2,is)

romb f a b = map head $ rombM f a b

improve n x y = let c = 4^(n-1) in (c*x - y) / (c - 1)

getRombErr a = unStripe . map (\x -> zipWith (\a b -> abs (a - b)) x (tail x))
                $ stripe a

-- shorthand conversions
tfr x = fromRational (toRational x)
toDouble x = read (showCReal 100 x) :: Double


-- succesful with infinite stripes only
unStripe ss = unfoldr (\xs -> case unzip (map (splitAt 1) xs) of
                        (a,b) -> Just (concat a, b))
                    ss

-- stolen from Control.Monad.Omega
stripe [] = []; stripe ([]:xss) = stripe xss
stripe ((x:xs):xss) = [x] : zipCons xs (stripe xss)

zipCons [] ys = ys; zipCons xs [] = map (:[]) xs
zipCons (x:xs) (y:ys) = (x:y) : zipCons xs ys