module GraphUtil (
        Format(..),
        plot, plotf,
        renderToFile,
        myAutoAxis,
        theAxis,
        theOtherAxis
    ) where

import StatUtil

import Bio.Util.Numeric             ( wilson )
import Data.List                    ( group, sort )
import Diagrams.Backend.SVG         ( renderSVG )
import Diagrams.Backend.Rasterific  ( renderRasterific )
import Diagrams.Prelude
import Diagrams_Rasterific          ( Rasterific(..), Options(RasterificOptions) )
import Graphics.Rasterific          ( renderDrawingAtDpiToPDF )
import Numeric                      ( showFFloat )

import qualified Data.ByteString.Lazy                       as B
import qualified Graphics.Rendering.Chart                   as C
import qualified Graphics.Rendering.Chart.Backend.Diagrams  as C

data Format = Txt | Svg | Pdf | Png

plotf :: FilePath -> Format -> String -> [Curve Ratio] -> IO ()
plotf fn format t cs = let Just p = plot myAutoAxis (-10,39) t cs
                       in renderToFile fn format defaultWidth defaultHeight p

toCairoStyle :: LineStyle -> Maybe C.LineStyle
toCairoStyle (Solid r g b) = Just $ C.solidLine 1 (opaque $ sRGB r g b)
toCairoStyle (Dashed p r g b) = Just $ C.dashedLine 1 p (opaque $ sRGB r g b)
toCairoStyle Hidden = Nothing

toBackgroundStyle :: LineStyle -> Maybe C.FillStyle
toBackgroundStyle (Solid r g b) = Just $ C.solidFillStyle $ sRGB r g b `withOpacity` 0.1
toBackgroundStyle (Dashed _ r g b) = Just $ C.solidFillStyle $ sRGB r g b `withOpacity` 0.05
toBackgroundStyle Hidden = Nothing

myAutoAxis :: C.AxisFn Double
myAutoAxis = mod_labels . C.autoScaledAxis def

theAxis :: Double -> C.AxisFn Double
theAxis y = mod_labels . const (C.autoScaledAxis def [0,y])

theOtherAxis :: Double -> C.AxisFn Double
theOtherAxis y = mod_labels . const (C.autoScaledAxis def [l,h])
  where l  = max 0 $ 0.25 - y
        h  = min 1 $ 0.25 + y

mod_labels :: C.AxisData Double -> C.AxisData Double
mod_labels = C.axis_labels %~ map (format' . map fst)
 where
    format' vs | ulength vs == length vs = format 1 vs
               | otherwise               = zip vs (tryFmt 2 vs)

    format n vs | ulength lbls == length lbls = zip vs lbls
                | otherwise                   = format (n+1) vs
      where lbls = tryFmt n vs

    ulength xs = length . group . sort $ xs
    tryFmt   n = map $ \v -> showFFloat (Just n) v []

plot :: C.AxisFn Double -> (Int,Int) -> String -> [Curve Ratio] -> Maybe (C.Renderable ())
plot yaxis (begs,ends) title curves
    | null (graph ^. C.layout_plots) = Nothing
    | otherwise                      = Just $ C.toRenderable graph
  where
    graph =
        C.layout_title .~ title $
        C.layout_y_axis . C.laxis_style . C.axis_label_style . C.font_size %~ (*2) $
        C.layout_x_axis . C.laxis_style . C.axis_label_style . C.font_size %~ (*2) $
        C.layout_y_axis . C.laxis_generate .~ yaxis $
        C.layout_plots .~ plots $ def

    plots = [ C.toPlot $ C.plot_fillbetween_title .~ fnname
                       $ C.plot_fillbetween_style .~ style
                       $ C.plot_fillbetween_values .~ values $ def
            | (Curve fnname style0 row, style1) <- zip curves line_styles
            , style <- maybe [] return . toBackgroundStyle $ maybe style1 id style0
            , let values = [ (fromIntegral x :: Double, case wilson 0.05 n d of (u,_,v) -> (u,v))
                           | x <- [begs..ends], (n,d) <- [row `atS` x], d > 32 ]
            , (_:_:_) <- [ values ] ] ++
            [ C.toPlot $ C.plot_lines_title .~ fnname
                       $ C.plot_lines_style .~ style
                       $ C.plot_lines_values .~ [ values ] $ def
            | (Curve fnname style0 row, style1) <- zip curves line_styles
            , style <- maybe [] return . toCairoStyle $ maybe style1 id style0
            , let values = [ (fromIntegral x :: Double, fromIntegral n / fromIntegral d :: Double)
                            | x <- [begs..ends], (n,d) <- [row `atS` x], d > 32 ]
            , (_:_:_) <- [ values ] ]


line_styles :: [ LineStyle ]
line_styles = [ solid_red, solid_green, solid_blue
              , Solid 0.7 0.7 0, Solid 1 0 1, Solid 0 1 1
              , Dashed [3,3] 1 0 0, Dashed [3,3] 0 1 0
              , Dashed [3,3] 0 0 1, Dashed [3,3] 0.7 0.7 0
              , Dashed [3,3] 1 0 1, Dashed [3,3] 0 1 1 ]


renderToFile :: FilePath -> Format -> Int -> Int -> C.Renderable () -> IO ()
renderToFile fp fmt x y r = do
    env <- C.defaultEnv C.vectorAlignmentFns (fromIntegral x) (fromIntegral y) :: IO (C.DEnv Double)
    let (w,h) = C.envOutputSize env
        cb    = C.render r (w,h)
    case fmt of
        Txt -> fail $ "renderToFile should not be called for Txt format"

        Pdf -> B.writeFile (fp ++ ".pdf") . renderDrawingAtDpiToPDF x y 600 .
               renderDia Rasterific (RasterificOptions (dims (V2 w h))) . fst $ C.runBackend env cb

        Png -> renderRasterific (fp ++ ".png") (dims (V2 w h)) . fst $ C.runBackend env cb

        Svg -> renderSVG (fp ++ ".svg") (dims (V2 w h)) . fst $ C.runBackend env cb