{-# LANGUAGE ScopedTypeVariables #-}

module LevMar.Chart
    ( FileType(..)
    , levmarChart
    , levmarChartFile
    , plotResults
    , plotToFile
    , Graphics.Rendering.Chart.PlotValue
    ) where

import Control.Monad                ( fmap )
import Data.Accessor                ( (^=), (.>) )
import Data.Colour                  ( withOpacity, opaque )
import Data.Colour.Names            ( red, blue )
import Graphics.Rendering.Chart
import Graphics.Rendering.Chart.Gtk ( renderableToWindow )
import LevMar.Fitting
import NFunction                    ( ($*) )
import Text.Printf                  ( printf )
import System.IO                    ( FilePath )

data FileType = PDF  -- ^Portable Document Format
              | PNG  -- ^Portable Network Graphics
              | PS   -- ^Postscript
              | SVG  -- ^Scalable Vector Graphics
                deriving (Show)

-- |Apply the Levenbarg-Marquardt algorithm and plot the results in a
-- window.
levmarChart :: forall n k r a.
               ( Nat n
               , Nat k
               , PlotValue r, LevMarable r
               , PlotValue a, Fractional a
               )
            => (Model n r a)                   -- ^ Model
            -> Maybe (Jacobian n r a)          -- ^ Optional jacobian
            -> SizedList n r                   -- ^ Initial parameters
            -> [(a, r)]                        -- ^ Samples
            -> Integer                         -- ^ Maximum number of iterations
            -> Options r                       -- ^ Minimalization options
            -> Maybe (SizedList n r)           -- ^ Optional lower bounds
            -> Maybe (SizedList n r)           -- ^ Optional upper bounds
            -> Maybe (LinearConstraints k n r) -- ^ Optional linear
            -> Maybe (SizedList n r)           -- ^ Optional weights
            -> IO (Either LevMarError (SizedList n r, Info r, CovarMatrix n r))
levmarChart model mJac params ys itMax opts mLowBs mUpBs mLinC mWghts =
    let results = levmar model
                         mJac
                         params
                         ys
                         itMax
                         opts
                         mLowBs
                         mUpBs
                         mLinC
                         mWghts
    in withRightM (plotResults model ys) results

-- |Apply the Levenbarg-Marquardt algorithm and plot the results in a
-- file.
levmarChartFile :: forall n k r a.
               ( Nat n
               , Nat k
               , PlotValue r, LevMarable r
               , PlotValue a, Fractional a
               )
               => (Model n r a)                   -- ^ Model
               -> Maybe (Jacobian n r a)          -- ^ Optional jacobian
               -> SizedList n r                   -- ^ Initial parameters
               -> [(a, r)]                        -- ^ Samples
               -> Integer                         -- ^ Maximum number of iterations
               -> Options r                       -- ^ Minimalization options
               -> Maybe (SizedList n r)           -- ^ Optional lower bounds
               -> Maybe (SizedList n r)           -- ^ Optional upper bounds
               -> Maybe (LinearConstraints k n r) -- ^ Optional linear
               -> Maybe (SizedList n r)           -- ^ Optional weights
               -> FileType                        -- ^ Type of image file to produce
               -> FilePath                        -- ^ Destination path (must include extension)
               -> Int                             -- ^ Width of the image
               -> Int                             -- ^ Height of the image
               -> IO (Either LevMarError (SizedList n r, Info r, CovarMatrix n r))
levmarChartFile model mJac params ys itMax opts mLowBs mUpBs mLinC mWghts ft path width height =
    let results = levmar model
                         mJac
                         params
                         ys
                         itMax
                         opts
                         mLowBs
                         mUpBs
                         mLinC
                         mWghts
    in withRightM (\r -> plotToFile model ys r ft path width height) results

-------------------------------------------------------------------------------

-- |Plots the results of the Levenberg-Marquardt algorithm in a window.
plotResults :: forall n r a. (PlotValue r, Fractional a, PlotValue a)
            => (Model n r a)
            -> [(a, r)]
            -> (SizedList n r, Info r, CovarMatrix n r)
            -> IO ()
plotResults model samples result = renderableToWindow (renderResults model samples result) 800 600

-- |Plots the results of the Levenberg-Marquardt algorithm in a file.
plotToFile :: forall n r a. (PlotValue r, Fractional a, PlotValue a)
           => (Model n r a)
           -> [(a, r)]
           -> (SizedList n r, Info r, CovarMatrix n r)
           -> FileType
           -> FilePath
           -> Int -- width
           -> Int -- height
           -> IO ()
plotToFile model samples result ft path width height = renderToFile ft renderable
                                                                    width height path
    where
      renderable = renderResults model samples result

      renderToFile PDF = renderableToPDFFile
      renderToFile PNG = renderableToPNGFile
      renderToFile PS  = renderableToPSFile
      renderToFile SVG = renderableToSVGFile

-------------------------------------------------------------------------------

renderResults :: forall n r a. (PlotValue r, Fractional a, PlotValue a)
              => (Model n r a)
              -> [(a, r)]
              -> (SizedList n r, Info r, CovarMatrix n r)
              -> Renderable ()
renderResults model samples (params, info, _) = toRenderable r
    where xs = fmap fst samples

          r :: Layout1 a r
          r = layout1_title ^= title
            $ layout1_plots ^= fmap Left [samplePts, fitted]
            $ layout1_left_axis   .> laxis_title ^= "y-axis"
            $ layout1_bottom_axis .> laxis_title ^= "x-axis"
            $ defaultLayout1

          title :: String
          title = printf (  "LevMar Fit - "
                         ++ "%d iters - %d func evals - "
                         ++ "%d jacob evals - %d lin systems solved - "
                         ++ "stop reason: %s"
                         )
                         (infNumIter           info)
                         (infNumFuncEvals      info)
                         (infNumJacobEvals     info)
                         (infNumLinSysSolved   info)
                         (show $ infStopReason info)

          samplePts = toPlot
                    $ plot_points_values ^= samples
                    $ plot_points_style  ^= filledCircles 2 (opaque blue)
                    $ defaultPlotPoints

          fitted = toPlot
                 $ plot_lines_values ^= [zip fittedXs $ fmap (model $* params) fittedXs]
                 $ plot_lines_style .> line_color ^= (red `withOpacity` 0.5)
                 $ defaultPlotLines

          fittedXs :: [a]
          fittedXs | null xs   = []
                   | otherwise = [ minX + (fromIntegral i / fromIntegral lenXs) * abs (maxX - minX)
                                 | i <- [0 .. lenXs]
                                 ]
              where minX   = minimum xs
                    maxX   = maximum xs
                    lenXs  = length xs

-------------------------------------------------------------------------------

withRightM :: Monad m => (r -> m ()) -> Either l r -> m (Either l r)
withRightM _ e@(Left  _) = return e
withRightM f e@(Right y) = f y >> return e