Use separate drawing + placing renderer for track

module Genetics.Browser.Track.Backend where

import Prelude

import Color (Color, black, white)
import Color.Scheme.Clrs (aqua, blue, fuchsia, green, lime, maroon, navy, olive, orange, purple, red, teal, yellow)
import Control.Monad.Aff (Aff)
import Control.Monad.Eff.Exception (error)
import Control.Monad.Error.Class (throwError)
import Control.Monad.Reader (class MonadReader, ask)
import Data.Argonaut (Json, _Array, _Number, _Object, _String)
import Data.Array ((..))
import Data.Array as Array
import Data.Bifunctor (bimap)
import Data.Filterable (class Filterable, filterMap)
import Data.Foldable (class Foldable, fold, foldMap, foldl, length, maximum)
import Data.FunctorWithIndex (mapWithIndex)
import Data.Int as Int
import Data.Lens (Fold', Traversal', foldMapOf, to, traversed, (^.), (^?))
import Data.Lens.Index (ix)
import Data.List (List)
import Data.List as List
import Data.Map (Map)
import Data.Map as Map
import Data.Maybe (Maybe(Just, Nothing), fromMaybe)
import Data.Monoid (class Monoid, mempty)
import Data.Newtype (class Newtype, unwrap, wrap)
import Data.Pair (Pair(..))
import Data.Profunctor.Strong (fanout)
import Data.Record as Record
import Data.Symbol (SProxy(..))
import Data.Traversable (traverse)
import Data.Tuple (Tuple(Tuple), snd, uncurry)
import Data.Unfoldable (class Unfoldable, none)
import Genetics.Browser.Types (Bp(Bp), ChrId(..), Point)
import Genetics.Browser.Types.Coordinates (BrowserPoint, CoordInterval, CoordSys(CoordSys), Interval, Normalized(Normalized), _BrowserIntervals, _CoordSys, _Index, _Interval, intervalToScreen, intervalsToMap, nPointToFrame, normPoint, viewIntervals)
import Genetics.Browser.View (Pixels)
import Graphics.Drawing (Drawing, circle, fillColor, filled, lineWidth, outlineColor, outlined, rectangle, scale, translate)
import Graphics.Drawing as Drawing
import Graphics.Drawing.Font (font, sansSerif)
import Network.HTTP.Affjax as Affjax
import Type.Prelude (class RowLacks)
import Unsafe.Coerce (unsafeCoerce)

type BrowserView = Interval BrowserPoint

intersection :: forall k a b.
                Ord k
             => Map k a
             -> Map k b
             -> Map k a
intersection a b = Map.filterKeys (flip Map.member b) a

zipMaps :: forall k a b.
           Ord k
        => Map k a
        -> Map k b
        -> Map k (Tuple a b)
zipMaps a b =
  let kas = Array.unzip $ Map.toAscUnfoldable $ a `intersection` b
      kbs = Array.unzip $ Map.toAscUnfoldable $ b `intersection` a
      kabs = uncurry Array.zip $ map (Array.zip (snd kas)) kbs
  in Map.fromFoldable kabs

zipMapsWith :: forall k a b c.
               Ord k
            => (a -> b -> c)
            -> Map k a
            -> Map k b
            -> Map k c
zipMapsWith f a b = uncurry f <$> zipMaps a b

type ChrCtx r = Map ChrId (Record r)

-- type Dimensions x = x -> BrowserView -> { x :: Pixels, width :: Pixels }

type ToPixels feature = feature -> BrowserView -> Interval Pixels

type Draw feature = feature -> Interval Pixels -> Drawing

placeWide :: forall a r.
             { chrSize :: a -> Maybe Bp
             , min     :: a -> Bp
             , max     :: a -> Bp
             | r }
          -> a
          -> Maybe (Interval (Normalized Number))
placeWide get a = do
  size <- get.chrSize a
  let l = unwrap $ (get.min a) / size
      r = unwrap $ (get.max a) / size
  pure $ Normalized <$> Pair l r

type DrawFeature a  = a -> Drawing
type PointFeature a = a -> Maybe (Normalized Point)

type PureRenderer a = a -> Maybe { drawing :: Drawing
                                 , point   :: Normalized Point
                                 }

type PureRendererNeue a = { drawing :: DrawFeature a
                          , point   :: PointFeature a }

type GWASFeature r = { score :: Number
                     , pos :: Bp
                     , chrId :: ChrId | r }

gwasDraw :: Color -> DrawFeature (GWASFeature _)
gwasDraw color =
  let r = 2.2
      c = circle 0.0 0.0 r
      out = outlined (outlineColor color) c
      fill = filled (fillColor color) c
  in const $ out <> fill

gwasPoint :: forall m rf rctx.
             MonadReader (ChrCtx (size :: Bp | rctx)) m
          => {min :: Number, max :: Number, sig :: Number}
          -> m (PointFeature (GWASFeature rf))
gwasPoint {min, max} = do
  ctx <- ask
  -- let place f = do
  pure \f -> do
        size <- _.size <$> Map.lookup f.chrId ctx
        let x = unwrap $ f.pos / size
            y = (f.score - min) / (max - min)
        normPoint {x, y}
  -- pure place

mkGwasRenderer :: forall m rf rctx.
                  MonadReader (ChrCtx (size :: Bp, color :: Color | rctx)) m
               => {min :: Number, max :: Number, sig :: Number}
               -> m (PureRenderer (GWASFeature rf))
mkGwasRenderer {min, max} = do
  ctx <- ask
  let renderer f = do
        size <- _.size <$> Map.lookup f.chrId ctx
        color <- _.color <$> Map.lookup f.chrId ctx
        let r = 2.2
            x = unwrap $ f.pos / size
            y = (f.score - min) / (max - min)
            c = circle x y r
            out = outlined (outlineColor color) c
            fill = filled (fillColor color) c
            drawing = out <> fill

        point <- normPoint {x, y}
        pure { drawing, point }
  pure renderer

-- annotDraw
annotDraw :: forall rf.
             DrawFeature (Annot rf)
annotDraw =
  -- ctx <- ask
  let font' = font sansSerif 12 mempty
      renderer f =
        let rad = 5.0
            c = circle zero zero rad
            out  = outlined (outlineColor maroon <> lineWidth 3.0) c
            fill = filled   (fillColor red) c
            text' = Drawing.text font' (7.5) (2.5) (fillColor black) f.name
        in out <> fill <> text'
  in renderer

annotPoint :: forall m rf rctx.
              MonadReader (ChrCtx (size :: Bp | rctx)) m
           => m (PointFeature (Annot rf))
annotPoint = do
  ctx <- ask
  let place f = do
        size <- _.size <$> Map.lookup f.chrId ctx

        let x = unwrap $ f.pos / size
            y = 0.0

        normPoint {x, y}

  pure place

mkAnnotRenderer :: forall m rf rctx.
                   MonadReader (ChrCtx (size :: Bp | rctx)) m
                => m (PureRenderer (Annot rf))
mkAnnotRenderer = do
  ctx <- ask
  let font' = font sansSerif 12 mempty
  let renderer f = do
        size <- _.size <$> Map.lookup f.chrId ctx

        let rad = 5.0
            x = unwrap $ f.pos / size
            y = 0.0
            c = circle x y rad
            out  = outlined (outlineColor maroon <> lineWidth 3.0) c
            fill = filled   (fillColor red) c
            -- the canvas is flipped so y-axis increases upward, need to flip the text too
            text' = scale 1.0 (-1.0)
                    $ Drawing.text font' (x+7.5) (y+2.5) (fillColor black) f.name
            drawing = out <> fill <> text'

        point <- normPoint {x, y}
        pure { drawing, point }

  pure renderer

pureRender :: forall f a.
              Filterable f
           => PureRenderer a
           -> f a
           -> f { drawing :: Drawing, point :: Normalized Point }
pureRender r as = filterMap r as

drawPureInterval :: forall f r .
                    Foldable f
                 => { height :: Pixels | r }
                 -> { width :: Pixels, offset :: Pixels }
                 -> f { drawing :: Drawing, point :: Normalized Point }
                 -> Drawing
drawPureInterval {height} {width, offset} = foldMap render'
  where render' {drawing, point} =
          let {x,y} = nPointToFrame width height $ point
          in translate (x + offset) y drawing

drawNeue :: forall f r a.
            Foldable f
         => { height :: Pixels | r }
         -> PureRendererNeue a
         -> { width :: Pixels, offset :: Pixels }
         -> f a
         -> Drawing
drawNeue {height} {drawing, point} {width, offset} =
  -- TODO this can be cleaned up quite a bit more
  fromMaybe mempty <<< foldMap \a -> do
    (Normalized {x,y}) <- point a
                                     -- Flipping y-axis so it increases upward
    let p = nPointToFrame width height $ Normalized {x, y: 1.0 - y}
    pure $ translate (p.x + offset) p.y (drawing a)

viewportIntervals :: forall i a r.
                     Ord i
                  => CoordSys i BrowserPoint
                  -> { width :: Pixels | r }
                  -> Interval BrowserPoint
                  -> Map i { width :: Pixels, offset :: Pixels }
viewportIntervals cs@(CoordSys s) canvasBox bView =
  let f :: Interval BrowserPoint
        -> CoordInterval i BrowserPoint
        -> { width :: Pixels, offset :: Pixels }
      f iv c = (intervalToScreen canvasBox iv) (c^._Interval)

  in f bView <$> intervalsToMap cs

renderNeue :: forall f i a r.
              Ord i
           => Foldable f
           => { height :: Pixels | r }
           -> PureRendererNeue a
           -> Map i {width :: Pixels, offset :: Pixels}
           -> Map i (f a)
           -> Map i Drawing
renderNeue h r = zipMapsWith (drawNeue h r)

drawDemoNeue :: forall f r.
                Foldable f
             => Filterable f
             => CoordSys ChrId BrowserPoint
             -> { min :: Number, max :: Number, sig :: Number }
             -> { width :: Pixels, height :: Pixels }
             -> { gwas   :: Map ChrId (f (GWASFeature ()))
                , annots :: Map ChrId (f (Annot (minY :: Number))) }
             -> Interval BrowserPoint
             -> Drawing
drawDemoNeue cs s canvasBox {gwas, annots} =
  let renderers = renderersDemo' s
      ivals v = viewportIntervals cs canvasBox v
      -- ruler' = ruler s red

      drawGwas v   = fold $ renderNeue canvasBox renderers.gwas   (ivals v) gwas
      drawAnnots v = fold $ renderNeue canvasBox renderers.annots (ivals v) annots
      -- drawGwas   = drawData cs canvasBox renderers.gwas gwas
      -- drawAnnots = drawData cs canvasBox renderers.annots annots
  -- in unsafeCoerce unit
  in \view -> (drawAnnots view) <> (drawGwas view)
  -- in foldMap (drawAnnots <> drawGwas)
  -- in \view -> fold $ (drawGwas view) <> (drawAnnots view)

drawFrames :: forall i a.
              Ord i
           => CoordSys i BrowserPoint
           -> { width :: Pixels, height :: Pixels }
           -> Map i (Array { drawing :: _, point :: _ })
           -> Interval BrowserPoint
           -> Drawing
drawFrames cs@(CoordSys s) canvasBox drawings =
  let f :: Interval _  -> CoordInterval i _ -> { width :: _, offset :: _ }
      f iv c = (intervalToScreen canvasBox iv) (c^._Interval)

      g :: CoordInterval i _ -> Array { drawing :: _, point :: _ }
      g c = fromMaybe [] $ Map.lookup (c^._Index) drawings

      toDraw :: Interval BrowserPoint -> Array (Tuple _ (Array _))
      toDraw iv = (\x -> Tuple (f iv x) (g x)) <$> (viewIntervals cs iv)

      drawIvs :: Array (Tuple _ (Array _)) -> Drawing
      drawIvs = foldMap (uncurry $ drawPureInterval canvasBox)

      drawIvs' :: Array (Tuple _ (Array _)) -> Array Drawing
      drawIvs' = map (uncurry $ drawPureInterval canvasBox)

  in \bView -> translate 0.0 canvasBox.height
                $ scale 1.0 (-1.0)
                $ drawIvs (toDraw bView)

drawData :: forall i c f a.
            Ord i
         => Foldable f
         => Filterable f
         => CoordSys i BrowserPoint
         -> { width :: Pixels, height :: Pixels }
         -> PureRenderer a
         -> Map i (f a)
         -> Interval BrowserPoint
         -> Drawing
drawData cs@(CoordSys s) canvasBox renderer dat =
  let drawings :: Map i (Array { drawing :: _, point :: _ })
      drawings = map (Array.fromFoldable <<< pureRender renderer) dat

      f :: Interval _  -> CoordInterval i _ -> { width :: _, offset :: _ }
      f iv c = (intervalToScreen canvasBox iv) (c^._Interval)

      g :: CoordInterval i _ -> Array { drawing :: _, point :: _ }
      g c = fromMaybe [] $ Map.lookup (c^._Index) drawings

      toDraw :: Interval BrowserPoint -> Array (Tuple _ (Array _))
      toDraw iv = (\x -> Tuple (f iv x) (g x)) <$> (viewIntervals cs iv)

      drawIvs :: Array (Tuple _ (Array _)) -> Drawing
      drawIvs = foldMap (uncurry $ drawPureInterval canvasBox)

      drawIvs' :: Array (Tuple _ (Array _)) -> Array Drawing
      drawIvs' = map (uncurry $ drawPureInterval canvasBox)

  in \bView -> translate 0.0 canvasBox.height
                $ scale 1.0 (-1.0)
                $ drawIvs (toDraw bView)

shiftPlaceFeatureMinY :: forall rf.
                         RowLacks "minY" rf
                      => Number
                      -> { min :: Number, max :: Number, sig :: Number }
                      -> PointFeature (Record rf)
                      -> PointFeature (Record (minY :: Number | rf))
shiftPlaceFeatureMinY dist s pointF f = do
  point <- pointF (Record.delete (SProxy :: SProxy "minY") f)
  let {x,y} = unwrap point
      normY = (f.minY - s.min) / (s.max - s.min)
      y' = min (normY + dist) 0.95

  normPoint {x, y: y'}

shiftRendererMinY :: forall rf.
                      RowLacks "minY" rf
                   => Number
                   -> { min :: Number, max :: Number, sig :: Number }
                   -> PureRenderer (Record rf)
                   -> PureRenderer (Record (minY :: Number | rf))
shiftRendererMinY dist s render f = do
  {drawing, point} <- render (Record.delete (SProxy :: SProxy "minY") f)
  let {x,y} = unwrap point
      normY = (f.minY - s.min) / (s.max - s.min)
      y' = min (normY + dist) 0.95

  point' <- normPoint {x, y: y'}
  pure {drawing, point: point'}

groupToChrs :: forall a f rData.
               Monoid (f {chrId :: ChrId | rData})
            => Foldable f
            => Applicative f
            => f { chrId :: ChrId | rData }
            -> Map ChrId (f { chrId :: ChrId | rData })
groupToChrs = foldl (\chrs r@{chrId} -> Map.alter (add r) chrId chrs ) mempty
  where add x Nothing   = Just $ pure x
        add x (Just xs) = Just $ pure x <> xs

fetchJSON :: forall a.
             (Json -> Maybe a)
          -> String
          -> Aff _ (Array a)
fetchJSON p url = do
  json <- _.response <$> Affjax.get url

  case json ^? _Array of
    Nothing -> throwError $ error "Parse error: JSON is not an array"
    Just as -> case traverse p as of
      Nothing -> throwError $ error "Failed to parse JSON features"
      Just fs -> pure $ fs

gemmaJSONParse :: Json -> Maybe (GWASFeature ())
gemmaJSONParse j = do
  obj <- j ^? _Object
  chrId <- ChrId <$> obj ^? ix "chr" <<< _String
  pos   <- Bp    <$> obj ^? ix "ps"  <<< _Number
  score <-           obj ^? ix "af"  <<< _Number
  pure {pos, score, chrId}

fetchGemmaJSON :: String -> Aff _ (Array (GWASFeature ()))
fetchGemmaJSON = fetchJSON gemmaJSONParse

type GeneRow r = ( geneID :: String
                 , desc :: String
                 , start :: Bp
                 , end :: Bp
                 , name :: String
                 , chrId :: ChrId | r )

type Gene r = Record (GeneRow r)

geneJSONParse :: Json -> Maybe (Gene ())
geneJSONParse j = do
  obj    <- j ^? _Object
  geneID <- obj ^? ix "Gene stable ID" <<< _String
  desc   <- obj ^? ix "Gene description" <<< _String
  name   <- obj ^? ix "Gene name" <<< _String
  start  <- Bp    <$> obj ^? ix "Gene start (bp)"  <<< _Number
  end    <- Bp    <$> obj ^? ix "Gene end (bp)"  <<< _Number
  chrId  <- ChrId <$> obj ^? ix "ChrId" <<< _String
  pure {geneID, desc, name, start, end, chrId}

fetchGeneJSON :: String -> Aff _ (Array (Gene ()))
fetchGeneJSON = fetchJSON geneJSONParse

type AnnotRow r = ( geneID :: String
                  , desc :: String
                  , pos :: Bp
                  , name :: String
                  , chrId :: ChrId | r )

type Annot r = Record (AnnotRow r)

fetchAnnotJSON :: String -> Aff _ (Array (Annot ()))
fetchAnnotJSON str = map toAnnot <$> fetchJSON geneJSONParse str
  where toAnnot :: Gene () -> Annot ()
        toAnnot gene = { geneID: gene.geneID
                       , desc: gene.desc
                       , pos: gene.start
                       , name: gene.name
                       , chrId: gene.chrId }

bumpAnnots :: forall rAnnot rGWAS.
             RowLacks "minY" (AnnotRow rAnnot)
          => Bp
          -> Map ChrId (List (GWASFeature rGWAS))
          -> Map ChrId (List (Annot rAnnot))
          -> Map ChrId (List (Annot (minY :: Number | rAnnot)))
bumpAnnots radius = zipMapsWith (map <<< bumpAnnot)
  where maxScoreIn :: Bp -> Bp -> List (GWASFeature rGWAS) -> Number
        maxScoreIn l r gwas = fromMaybe 0.5 $ maximum
                              $ map (\g -> if g.pos >= l && g.pos <= r
                                           then g.score else 0.0) gwas
        bumpAnnot :: List (GWASFeature rGWAS) -> Annot rAnnot -> Annot (minY :: Number | rAnnot)
        bumpAnnot l g = let minY = maxScoreIn (g.pos - radius) (g.pos + radius) l
                       in Record.insert (SProxy :: SProxy "minY") minY g

getDataDemo :: { gwas :: String
               , annots :: String }
            -> Aff _ { gwas  :: Map ChrId (List (GWASFeature ()       ))
                     , annots :: Map ChrId (List (Annot (minY :: Number)))
                     }
getDataDemo urls = do
  gwas <- fetchGemmaJSON urls.gwas
  annots' <- fetchAnnotJSON urls.annots
  -- Divide data by chromosomes
  let gwasChr :: Map ChrId (List (GWASFeature ()))
      gwasChr = List.fromFoldable <$> groupToChrs gwas
  let annotsChr :: Map ChrId (List (Annot ()))
      annotsChr = List.fromFoldable <$> groupToChrs annots'
  let bumpedAnnots :: Map ChrId (List (Annot (minY :: Number)))
      bumpedAnnots = bumpAnnots (Bp 1000000.0) gwasChr annotsChr

  pure { gwas: gwasChr, annots: bumpedAnnots }

renderersDemo' :: forall r.
                 { min :: Number, max :: Number, sig :: Number }
              -> { gwas   :: PureRendererNeue _
                 , annots :: PureRendererNeue _
                 }
renderersDemo' s = { gwas, annots }
  where gwas :: PureRendererNeue _
        gwas = { drawing: gwasDraw navy
                , point:   gwasPoint s mouseChrCtx }

        annots :: PureRendererNeue (Annot (minY :: Number))
        annots = { drawing: annotDraw
                 , point:   shiftPlaceFeatureMinY 0.06 s $ annotPoint mouseChrCtx }

renderersDemo :: forall r.
                 { min :: Number, max :: Number, sig :: Number }
              -> { gwas   :: PureRenderer (GWASFeature ()       )
                 , annots :: PureRenderer (Annot (minY :: Number))
                 }
renderersDemo s = { gwas, annots }
  where colorsCtx = zipMapsWith (\r {color} -> Record.insert (SProxy :: SProxy "color") color r)
                      mouseChrCtx mouseColors
        gwas :: PureRenderer (GWASFeature ())
        gwas = mkGwasRenderer s colorsCtx

        annots :: PureRenderer (Annot (minY :: Number))
        annots = shiftRendererMinY 0.06 s $ mkAnnotRenderer mouseChrCtx

ruler :: forall r r1.
         { min :: Number, max :: Number, sig :: Number }
      -> Color
      -> { width :: Pixels, height :: Pixels }
      -> Drawing
ruler {min, max, sig} color f = outlined outline rulerDrawing
  where normY = (sig - min) / (max - min)
        y = f.height - (normY * f.height)
        outline = outlineColor color
        rulerDrawing = Drawing.path [{x: 0.0, y}, {x: f.width, y}]

chrLabels :: forall c.
             CoordSys ChrId c
          -> Map ChrId (Array { drawing :: _, point :: _ })
chrLabels cs = Map.fromFoldable results
  where mkLabel :: ChrId -> _
        mkLabel chr = { drawing: chrText chr, point: Normalized {x: 1.0, y: -0.1} }

        font' = font sansSerif 12 mempty

        chrText :: ChrId -> Drawing
        chrText chr = scale 1.0 (-1.0)
                    $ Drawing.text font' zero zero (fillColor black) (unwrap chr)

        results :: Array _
        results = map (\x -> Tuple (x^._Index) [mkLabel (x^._Index)]) $ cs^._BrowserIntervals

        results' :: Traversal' (CoordSys ChrId c) ChrId
        results' = _BrowserIntervals <<< traversed <<< _Index

        test :: forall r. Monoid r => Fold' r _ _
        test = results' <<< (to (fanout id mkLabel))

        output :: _
        output = Map.fromFoldable $ (foldMapOf test pure cs :: Array _)

        -- xxx :: Traversal' (CoordSys ChrId _) ChrId
        -- xxx = results' <<< _Index

        -- results2 :: _ -> Array _
        -- results2 = traverseOf results' pure

        -- results3 :: _ -> Array _
        -- results3 = foldMapOf results' pure

        -- results4 :: _
        -- results4 = foldMapOf results'
                     -- (fanout id chrText)

drawDemo :: forall f r.
            Foldable f
         => Filterable f
         => CoordSys _ _
         -> { min :: Number, max :: Number, sig :: Number }
         -> { width :: Pixels, height :: Pixels }
         -> { gwas   :: Map ChrId (f (GWASFeature ()))
            , annots :: Map ChrId (f (Annot (minY :: Number))) }
         -> Interval BrowserPoint
         -> Drawing
drawDemo cs s canvasBox {gwas, annots} =
  let renderers = renderersDemo s
      ruler' = ruler s red
      drawGwas   = drawData cs canvasBox renderers.gwas gwas
      drawAnnots = drawData cs canvasBox renderers.annots annots
  in \view -> (drawGwas view) <> (drawAnnots view) <> ruler' canvasBox

{-
TODO
there is currently no connection between the vertical scale's
position and the track. Need to take stuff like vertical
offset and padding into account; calculate once and use
in both places. However don't just use canvasheight - yoffset
or whatever; actually derive the top and bottom of the track.
-}

drawVScale :: forall r.
              Number
           -> Number
           -> { min :: Number, max :: Number, sig :: Number }
           -> Color
           -> Drawing
drawVScale width height vs col =
  -- should have some padding here too; hardcode for now
  let hPad = width  / 5.0
      vPad = height / 10.0
      x = width / 2.0

      y1 = 0.0
      y2 = height

      n = (_ * 0.1) <<< Int.toNumber <$> (0 .. 10)

      p = Drawing.path [{x, y:y1}, {x, y:y2}]

      bar w y = Drawing.path [{x:x-w, y}, {x, y}]

      ps = foldMap (\i -> bar
                          (if i == 0.0 || i == 1.0 then 8.0 else 3.0)
                          (y1 + i*(y2-y1))) n

      ft = font sansSerif 10 mempty
      mkT y = Drawing.text ft (x-12.0) y (fillColor black)

      topLabel = mkT (y1-(0.5*vPad)) $ show vs.max
      btmLabel = mkT (y2+vPad) $ show vs.min

  in outlined (outlineColor col <> lineWidth 2.0) (p <> ps)
  <> topLabel <> btmLabel

drawLegendItem :: Number
               -> { text :: String, icon :: Drawing }
               -> Drawing
drawLegendItem w {text, icon} =
  let ft = font sansSerif 12 mempty
      t = Drawing.text ft 12.0 0.0 (fillColor black) text
  in icon <> t

drawLegend :: Number
           -> Number
           -> Array { text :: String, icon :: Drawing }
           -> Drawing
drawLegend width height icons =
  let hPad = width  / 5.0
      vPad = height / 5.0
      n :: Int
      n = length icons
      x = hPad
      f :: Number -> { text :: String, icon :: Drawing } -> Drawing
      f y ic = translate x y $ drawLegendItem (width - 2.0*hPad) ic
      d = (height - 2.0*vPad) / (length icons)
      ds = mapWithIndex (\i ic -> f (vPad+(vPad*(Int.toNumber i))) ic) icons
  in fold ds

mkIcon :: Color -> String -> { text :: String, icon :: Drawing }
mkIcon c text =
  let sh = circle (-2.5) (-2.5) 5.0
      icon = outlined (outlineColor c <> lineWidth 2.0) sh <>
             filled (fillColor c) sh
  in {text, icon}

demoLegend :: Array {text :: String, icon :: Drawing}
demoLegend =
  [ mkIcon red "Red thing"
  , mkIcon blue "blue thing"
  , mkIcon purple "boop"
  ]

              -- the first 2 are used by all parts of the browser
demoBrowser :: forall f.
               Foldable f
            => Filterable f
            => CoordSys ChrId BrowserPoint
            -> { width :: Pixels, height :: Pixels }
            -> Pixels
            -- this by the vertical scale and the tracks
            -> { min :: Number, max :: Number, sig :: Number }
            -- these define the width of the non-track parts of the browser;
            -- the browser takes up the remaining space.
            -> { vScaleWidth :: Pixels
               , legendWidth :: Pixels }
            -- vScale only
            -> Color
            -- to be drawn in the legend on the RHS
            -> Array { text :: String, icon :: Drawing }
            -- finally, used by the main track
            -> { gwas :: Map ChrId (f (GWASFeature ()))
               , annots :: Map ChrId (f (Annot (minY :: Number))) }
            -> BrowserView
            -- the drawing produced contains all 3 parts.
            -> { track :: Drawing, overlay :: Drawing }
demoBrowser cs canvas vpadding vscale sizes vscaleColor legend {gwas, annots} =
  let height = canvas.height - 2.0 * vpadding
      width  = canvas.width

      trackCanvas = { width: canvas.width - sizes.vScaleWidth - sizes.legendWidth
                    , height }

      bg w' h' = filled (fillColor white) $ rectangle 0.0 0.0 w' h'

      -- TODO unify vertical offset by padding further; do it as late as possible

      vScale :: _ -> Drawing
      vScale _ = let w = sizes.vScaleWidth
                 in  translate zero vpadding
                   $ bg w height
                  <> drawVScale w height vscale vscaleColor

      track v = translate sizes.vScaleWidth vpadding
                  $ drawDemoNeue cs vscale trackCanvas {gwas, annots} v

      chrs :: _ -> Drawing
      chrs v = translate zero vpadding
                 $ drawFrames cs trackCanvas (chrLabels cs) v

      legendD :: _ -> Drawing
      legendD _ = let w = sizes.legendWidth
                  in translate (canvas.width - w) vpadding
                    $ bg w height
                   <> drawLegend sizes.legendWidth canvas.height legend

  in \view -> { track: track view <> chrs view
              , overlay: vScale view <> legendD view }

mouseChrIds :: Array ChrId
mouseChrIds =
            [ (ChrId "1")
            , (ChrId "2")
            , (ChrId "3")
            , (ChrId "4")
            , (ChrId "5")
            , (ChrId "6")
            , (ChrId "7")
            , (ChrId "8")
            , (ChrId "9")
            , (ChrId "10")
            , (ChrId "11")
            , (ChrId "12")
            , (ChrId "13")
            , (ChrId "14")
            , (ChrId "15")
            , (ChrId "16")
            , (ChrId "17")
            , (ChrId "18")
            , (ChrId "19")
            , (ChrId "X")
            , (ChrId "Y")
            ]

mouseChrs :: Map ChrId Bp
mouseChrs = Map.fromFoldable $ Array.zip mouseChrIds $
            [ (Bp 195471971.0)
            , (Bp 182113224.0)
            , (Bp 160039680.0)
            , (Bp 156508116.0)
            , (Bp 151834684.0)
            , (Bp 149736546.0)
            , (Bp 145441459.0)
            , (Bp 129401213.0)
            , (Bp 124595110.0)
            , (Bp 130694993.0)
            , (Bp 122082543.0)
            , (Bp 120129022.0)
            , (Bp 120421639.0)
            , (Bp 124902244.0)
            , (Bp 104043685.0)
            , (Bp 98207768.0)
            , (Bp 94987271.0)
            , (Bp 90702639.0)
            , (Bp 61431566.0)
            , (Bp 17103129.0)
            , (Bp 9174469.0)
            ]

mouseColors :: ChrCtx (color :: Color)
mouseColors = Map.fromFoldable $ map (\x -> Tuple x {color: navy}) mouseChrIds

mouseColors' :: ChrCtx (color :: Color)
mouseColors' = Map.fromFoldable $ Array.zip mouseChrIds $ map (\x -> {color: x})
               [ navy, blue, aqua, teal, olive
               , green, lime, yellow, orange, red
               , maroon, fuchsia, purple, navy, blue
               , aqua, teal, olive, green, lime
               , yellow ]

mouseChrCtx :: ChrCtx (size :: Bp)
mouseChrCtx = map (\size -> { size }) mouseChrs

module Genetics.Browser.Track.Backend where

import Prelude

import Color (Color, black, white)
import Color.Scheme.Clrs (aqua, blue, fuchsia, green, lime, maroon, navy, olive, orange, purple, red, teal, yellow)
import Control.Monad.Aff (Aff)
import Control.Monad.Eff.Exception (error)
import Control.Monad.Error.Class (throwError)
import Control.Monad.Reader (class MonadReader, ask)
import Data.Argonaut (Json, _Array, _Number, _Object, _String)
import Data.Array ((..))
import Data.Array as Array
import Data.Bifunctor (bimap)
import Data.Filterable (class Filterable, filterMap)
import Data.Foldable (class Foldable, fold, foldMap, foldl, length, maximum)
import Data.FunctorWithIndex (mapWithIndex)
import Data.Int as Int
import Data.Lens (Fold', Traversal', foldMapOf, to, traversed, (^.), (^?))
import Data.Lens.Index (ix)
import Data.List (List)
import Data.List as List
import Data.Map (Map)
import Data.Map as Map
import Data.Maybe (Maybe(Just, Nothing), fromMaybe)
import Data.Monoid (class Monoid, mempty)
import Data.Newtype (class Newtype, unwrap, wrap)
import Data.Pair (Pair(..))
import Data.Profunctor.Strong (fanout)
import Data.Record as Record
import Data.Symbol (SProxy(..))
import Data.Traversable (traverse)
import Data.Tuple (Tuple(Tuple), snd, uncurry)
import Data.Unfoldable (class Unfoldable, none)
import Genetics.Browser.Types (Bp(Bp), ChrId(..), Point)
import Genetics.Browser.Types.Coordinates (BrowserPoint, CoordInterval, CoordSys(CoordSys), Interval, Normalized(Normalized), _BrowserIntervals, _CoordSys, _Index, _Interval, intervalToScreen, intervalsToMap, nPointToFrame, normPoint, viewIntervals)
import Genetics.Browser.View (Pixels)
import Graphics.Drawing (Drawing, circle, fillColor, filled, lineWidth, outlineColor, outlined, rectangle, scale, translate)
import Graphics.Drawing as Drawing
import Graphics.Drawing.Font (font, sansSerif)
import Network.HTTP.Affjax as Affjax
import Type.Prelude (class RowLacks)
import Unsafe.Coerce (unsafeCoerce)

type BrowserView = Interval BrowserPoint

intersection :: forall k a b.
                Ord k
             => Map k a
             -> Map k b
             -> Map k a
intersection a b = Map.filterKeys (flip Map.member b) a

zipMaps :: forall k a b.
           Ord k
        => Map k a
        -> Map k b
        -> Map k (Tuple a b)
zipMaps a b =
  let kas = Array.unzip $ Map.toAscUnfoldable $ a `intersection` b
      kbs = Array.unzip $ Map.toAscUnfoldable $ b `intersection` a
      kabs = uncurry Array.zip $ map (Array.zip (snd kas)) kbs
  in Map.fromFoldable kabs

zipMapsWith :: forall k a b c.
               Ord k
            => (a -> b -> c)
            -> Map k a
            -> Map k b
            -> Map k c
zipMapsWith f a b = uncurry f <$> zipMaps a b

type ChrCtx r = Map ChrId (Record r)

-- type Dimensions x = x -> BrowserView -> { x :: Pixels, width :: Pixels }

type ToPixels feature = feature -> BrowserView -> Interval Pixels

type Draw feature = feature -> Interval Pixels -> Drawing

placeWide :: forall a r.
             { chrSize :: a -> Maybe Bp
             , min     :: a -> Bp
             , max     :: a -> Bp
             | r }
          -> a
          -> Maybe (Interval (Normalized Number))
placeWide get a = do
  size <- get.chrSize a
  let l = unwrap $ (get.min a) / size
      r = unwrap $ (get.max a) / size
  pure $ Normalized <$> Pair l r

type DrawFeature a  = a -> Drawing
type PointFeature a = a -> Maybe (Normalized Point)

type PureRenderer a = a -> Maybe { drawing :: Drawing
                                 , point   :: Normalized Point
                                 }

type PureRendererNeue a = { drawing :: DrawFeature a
                          , point   :: PointFeature a }

type GWASFeature r = { score :: Number
                     , pos :: Bp
                     , chrId :: ChrId | r }

gwasDraw :: Color -> DrawFeature (GWASFeature _)
gwasDraw color =
  let r = 2.2
      c = circle 0.0 0.0 r
      out = outlined (outlineColor color) c
      fill = filled (fillColor color) c
  in const $ out <> fill

gwasPoint :: forall m rf rctx.
             MonadReader (ChrCtx (size :: Bp | rctx)) m
          => {min :: Number, max :: Number, sig :: Number}
          -> m (PointFeature (GWASFeature rf))
gwasPoint {min, max} = do
  ctx <- ask
  -- let place f = do
  pure \f -> do
        size <- _.size <$> Map.lookup f.chrId ctx
        let x = unwrap $ f.pos / size
            y = (f.score - min) / (max - min)
        normPoint {x, y}
  -- pure place

mkGwasRenderer :: forall m rf rctx.
                  MonadReader (ChrCtx (size :: Bp, color :: Color | rctx)) m
               => {min :: Number, max :: Number, sig :: Number}
               -> m (PureRenderer (GWASFeature rf))
mkGwasRenderer {min, max} = do
  ctx <- ask
  let renderer f = do
        size <- _.size <$> Map.lookup f.chrId ctx
        color <- _.color <$> Map.lookup f.chrId ctx
        let r = 2.2
            x = unwrap $ f.pos / size
            y = (f.score - min) / (max - min)
            c = circle x y r
            out = outlined (outlineColor color) c
            fill = filled (fillColor color) c
            drawing = out <> fill

        point <- normPoint {x, y}
        pure { drawing, point }
  pure renderer

-- annotDraw
annotDraw :: forall rf.
             DrawFeature (Annot rf)
annotDraw =
  -- ctx <- ask
  let font' = font sansSerif 12 mempty
      renderer f =
        let rad = 5.0
            c = circle zero zero rad
            out  = outlined (outlineColor maroon <> lineWidth 3.0) c
            fill = filled   (fillColor red) c
            -- the canvas is flipped so y-axis increases upward, need to flip the text too
            text' = scale 1.0 (-1.0)
                    $ Drawing.text font' (7.5) (2.5) (fillColor black) f.name
        in out <> fill <> text'
  in renderer

annotPoint :: forall m rf rctx.
              MonadReader (ChrCtx (size :: Bp | rctx)) m
           => m (PointFeature (Annot rf))
annotPoint = do
  ctx <- ask
  let place f = do
        size <- _.size <$> Map.lookup f.chrId ctx

        let x = unwrap $ f.pos / size
            y = 0.0
            -- the canvas is flipped so y-axis increases upward, need to flip the text too

        normPoint {x, y}

  pure place

mkAnnotRenderer :: forall m rf rctx.
                   MonadReader (ChrCtx (size :: Bp | rctx)) m
                => m (PureRenderer (Annot rf))
mkAnnotRenderer = do
  ctx <- ask
  let font' = font sansSerif 12 mempty
  let renderer f = do
        size <- _.size <$> Map.lookup f.chrId ctx

        let rad = 5.0
            x = unwrap $ f.pos / size
            y = 0.0
            c = circle x y rad
            out  = outlined (outlineColor maroon <> lineWidth 3.0) c
            fill = filled   (fillColor red) c
            -- the canvas is flipped so y-axis increases upward, need to flip the text too
            text' = scale 1.0 (-1.0)
                    $ Drawing.text font' (x+7.5) (y+2.5) (fillColor black) f.name
            drawing = out <> fill <> text'

        point <- normPoint {x, y}
        pure { drawing, point }

  pure renderer

pureRender :: forall f a.
              Filterable f
           => PureRenderer a
           -> f a
           -> f { drawing :: Drawing, point :: Normalized Point }
pureRender r as = filterMap r as

drawPureInterval :: forall f r .
                    Foldable f
                 => { height :: Pixels | r }
                 -> { width :: Pixels, offset :: Pixels }
                 -> f { drawing :: Drawing, point :: Normalized Point }
                 -> Drawing
drawPureInterval {height} {width, offset} = foldMap render'
  where render' {drawing, point} =
          let {x,y} = nPointToFrame width height $ point
          in translate (x + offset) y drawing

drawNeue :: forall f r a.
            Foldable f
         => { height :: Pixels | r }
         -> PureRendererNeue a
         -> { width :: Pixels, offset :: Pixels }
         -> f a
         -> Drawing
drawNeue {height} {drawing, point} {width, offset} =
  -- TODO this can be cleaned up quite a bit more
  fromMaybe mempty <<< foldMap \a -> do
    p <- point a
    let {x,y} = nPointToFrame width height p
    pure $ translate (x + offset) y (drawing a)

viewportIntervals :: forall i a r.
                     Ord i
                  => CoordSys i BrowserPoint
                  -> { width :: Pixels | r }
                  -> Interval BrowserPoint
                  -> Map i { width :: Pixels, offset :: Pixels }
viewportIntervals cs@(CoordSys s) canvasBox bView =
  let f :: Interval BrowserPoint
        -> CoordInterval i BrowserPoint
        -> { width :: Pixels, offset :: Pixels }
      f iv c = (intervalToScreen canvasBox iv) (c^._Interval)

  in f bView <$> intervalsToMap cs

renderNeue :: forall f i a r.
              Ord i
           => Foldable f
           => { height :: Pixels | r }
           -> PureRendererNeue a
           -> Map i {width :: Pixels, offset :: Pixels}
           -> Map i (f a)
           -> Map i Drawing
renderNeue h r = zipMapsWith (drawNeue h r)

drawFrames :: forall i a.
              Ord i
           => CoordSys i BrowserPoint
           -> { width :: Pixels, height :: Pixels }
           -> Map i (Array { drawing :: _, point :: _ })
           -> Interval BrowserPoint
           -> Drawing
drawFrames cs@(CoordSys s) canvasBox drawings =
  let f :: Interval _  -> CoordInterval i _ -> { width :: _, offset :: _ }
      f iv c = (intervalToScreen canvasBox iv) (c^._Interval)

      g :: CoordInterval i _ -> Array { drawing :: _, point :: _ }
      g c = fromMaybe [] $ Map.lookup (c^._Index) drawings

      toDraw :: Interval BrowserPoint -> Array (Tuple _ (Array _))
      toDraw iv = (\x -> Tuple (f iv x) (g x)) <$> (viewIntervals cs iv)

      drawIvs :: Array (Tuple _ (Array _)) -> Drawing
      drawIvs = foldMap (uncurry $ drawPureInterval canvasBox)

      drawIvs' :: Array (Tuple _ (Array _)) -> Array Drawing
      drawIvs' = map (uncurry $ drawPureInterval canvasBox)

  in \bView -> translate 0.0 canvasBox.height
                $ scale 1.0 (-1.0)
                $ drawIvs (toDraw bView)

drawData :: forall i c f a.
            Ord i
         => Foldable f
         => Filterable f
         => CoordSys i BrowserPoint
         -> { width :: Pixels, height :: Pixels }
         -> PureRenderer a
         -> Map i (f a)
         -> Interval BrowserPoint
         -> Drawing
drawData cs@(CoordSys s) canvasBox renderer dat =
  let drawings :: Map i (Array { drawing :: _, point :: _ })
      drawings = map (Array.fromFoldable <<< pureRender renderer) dat

      f :: Interval _  -> CoordInterval i _ -> { width :: _, offset :: _ }
      f iv c = (intervalToScreen canvasBox iv) (c^._Interval)

      g :: CoordInterval i _ -> Array { drawing :: _, point :: _ }
      g c = fromMaybe [] $ Map.lookup (c^._Index) drawings

      toDraw :: Interval BrowserPoint -> Array (Tuple _ (Array _))
      toDraw iv = (\x -> Tuple (f iv x) (g x)) <$> (viewIntervals cs iv)

      drawIvs :: Array (Tuple _ (Array _)) -> Drawing
      drawIvs = foldMap (uncurry $ drawPureInterval canvasBox)

      drawIvs' :: Array (Tuple _ (Array _)) -> Array Drawing
      drawIvs' = map (uncurry $ drawPureInterval canvasBox)

  in \bView -> translate 0.0 canvasBox.height
                $ scale 1.0 (-1.0)
                $ drawIvs (toDraw bView)

shiftPlaceFeatureMinY :: forall rf.
                         RowLacks "minY" rf
                      => Number
                      -> { min :: Number, max :: Number, sig :: Number }
                      -> PointFeature (Record rf)
                      -> PointFeature (Record (minY :: Number | rf))
shiftPlaceFeatureMinY dist s pointF f = do
  point <- pointF (Record.delete (SProxy :: SProxy "minY") f)
  let {x,y} = unwrap point
      normY = (f.minY - s.min) / (s.max - s.min)
      y' = min (normY + dist) 0.95

  normPoint {x, y: y'}

shiftRendererMinY :: forall rf.
                      RowLacks "minY" rf
                   => Number
                   -> { min :: Number, max :: Number, sig :: Number }
                   -> PureRenderer (Record rf)
                   -> PureRenderer (Record (minY :: Number | rf))
shiftRendererMinY dist s render f = do
  {drawing, point} <- render (Record.delete (SProxy :: SProxy "minY") f)
  let {x,y} = unwrap point
      normY = (f.minY - s.min) / (s.max - s.min)
      y' = min (normY + dist) 0.95

  point' <- normPoint {x, y: y'}
  pure {drawing, point: point'}

groupToChrs :: forall a f rData.
               Monoid (f {chrId :: ChrId | rData})
            => Foldable f
            => Applicative f
            => f { chrId :: ChrId | rData }
            -> Map ChrId (f { chrId :: ChrId | rData })
groupToChrs = foldl (\chrs r@{chrId} -> Map.alter (add r) chrId chrs ) mempty
  where add x Nothing   = Just $ pure x
        add x (Just xs) = Just $ pure x <> xs

fetchJSON :: forall a.
             (Json -> Maybe a)
          -> String
          -> Aff _ (Array a)
fetchJSON p url = do
  json <- _.response <$> Affjax.get url

  case json ^? _Array of
    Nothing -> throwError $ error "Parse error: JSON is not an array"
    Just as -> case traverse p as of
      Nothing -> throwError $ error "Failed to parse JSON features"
      Just fs -> pure $ fs

gemmaJSONParse :: Json -> Maybe (GWASFeature ())
gemmaJSONParse j = do
  obj <- j ^? _Object
  chrId <- ChrId <$> obj ^? ix "chr" <<< _String
  pos   <- Bp    <$> obj ^? ix "ps"  <<< _Number
  score <-           obj ^? ix "af"  <<< _Number
  pure {pos, score, chrId}

fetchGemmaJSON :: String -> Aff _ (Array (GWASFeature ()))
fetchGemmaJSON = fetchJSON gemmaJSONParse

type GeneRow r = ( geneID :: String
                 , desc :: String
                 , start :: Bp
                 , end :: Bp
                 , name :: String
                 , chrId :: ChrId | r )

type Gene r = Record (GeneRow r)

geneJSONParse :: Json -> Maybe (Gene ())
geneJSONParse j = do
  obj    <- j ^? _Object
  geneID <- obj ^? ix "Gene stable ID" <<< _String
  desc   <- obj ^? ix "Gene description" <<< _String
  name   <- obj ^? ix "Gene name" <<< _String
  start  <- Bp    <$> obj ^? ix "Gene start (bp)"  <<< _Number
  end    <- Bp    <$> obj ^? ix "Gene end (bp)"  <<< _Number
  chrId  <- ChrId <$> obj ^? ix "ChrId" <<< _String
  pure {geneID, desc, name, start, end, chrId}

fetchGeneJSON :: String -> Aff _ (Array (Gene ()))
fetchGeneJSON = fetchJSON geneJSONParse

type AnnotRow r = ( geneID :: String
                  , desc :: String
                  , pos :: Bp
                  , name :: String
                  , chrId :: ChrId | r )

type Annot r = Record (AnnotRow r)

fetchAnnotJSON :: String -> Aff _ (Array (Annot ()))
fetchAnnotJSON str = map toAnnot <$> fetchJSON geneJSONParse str
  where toAnnot :: Gene () -> Annot ()
        toAnnot gene = { geneID: gene.geneID
                       , desc: gene.desc
                       , pos: gene.start
                       , name: gene.name
                       , chrId: gene.chrId }

bumpAnnots :: forall rAnnot rGWAS.
             RowLacks "minY" (AnnotRow rAnnot)
          => Bp
          -> Map ChrId (List (GWASFeature rGWAS))
          -> Map ChrId (List (Annot rAnnot))
          -> Map ChrId (List (Annot (minY :: Number | rAnnot)))
bumpAnnots radius = zipMapsWith (map <<< bumpAnnot)
  where maxScoreIn :: Bp -> Bp -> List (GWASFeature rGWAS) -> Number
        maxScoreIn l r gwas = fromMaybe 0.5 $ maximum
                              $ map (\g -> if g.pos >= l && g.pos <= r
                                           then g.score else 0.0) gwas
        bumpAnnot :: List (GWASFeature rGWAS) -> Annot rAnnot -> Annot (minY :: Number | rAnnot)
        bumpAnnot l g = let minY = maxScoreIn (g.pos - radius) (g.pos + radius) l
                       in Record.insert (SProxy :: SProxy "minY") minY g

getDataDemo :: { gwas :: String
               , annots :: String }
            -> Aff _ { gwas  :: Map ChrId (List (GWASFeature ()       ))
                     , annots :: Map ChrId (List (Annot (minY :: Number)))
                     }
getDataDemo urls = do
  gwas <- fetchGemmaJSON urls.gwas
  annots' <- fetchAnnotJSON urls.annots
  -- Divide data by chromosomes
  let gwasChr :: Map ChrId (List (GWASFeature ()))
      gwasChr = List.fromFoldable <$> groupToChrs gwas
  let annotsChr :: Map ChrId (List (Annot ()))
      annotsChr = List.fromFoldable <$> groupToChrs annots'
  let bumpedAnnots :: Map ChrId (List (Annot (minY :: Number)))
      bumpedAnnots = bumpAnnots (Bp 1000000.0) gwasChr annotsChr

  pure { gwas: gwasChr, annots: bumpedAnnots }

renderersDemo' :: forall r.
                 { min :: Number, max :: Number, sig :: Number }
              -> { gwas   :: PureRendererNeue _
                 , annots :: PureRendererNeue _
                 }
renderersDemo' s = { gwas, annots }
  where gwas :: PureRendererNeue _
        gwas = { drawing: gwasDraw navy
                , point:   gwasPoint s mouseChrCtx }

        annots :: PureRendererNeue (Annot (minY :: Number))
        annots = { drawing: annotDraw
                  , point:   shiftPlaceFeatureMinY 0.06 s $ annotPoint mouseChrCtx }

renderersDemo :: forall r.
                 { min :: Number, max :: Number, sig :: Number }
              -> { gwas   :: PureRenderer (GWASFeature ()       )
                 , annots :: PureRenderer (Annot (minY :: Number))
                 }
renderersDemo s = { gwas, annots }
  where colorsCtx = zipMapsWith (\r {color} -> Record.insert (SProxy :: SProxy "color") color r)
                      mouseChrCtx mouseColors
        gwas :: PureRenderer (GWASFeature ())
        gwas = mkGwasRenderer s colorsCtx

        annots :: PureRenderer (Annot (minY :: Number))
        annots = shiftRendererMinY 0.06 s $ mkAnnotRenderer mouseChrCtx

ruler :: forall r r1.
         { min :: Number, max :: Number, sig :: Number }
      -> Color
      -> { width :: Pixels, height :: Pixels }
      -> Drawing
ruler {min, max, sig} color f = outlined outline rulerDrawing
  where normY = (sig - min) / (max - min)
        y = f.height - (normY * f.height)
        outline = outlineColor color
        rulerDrawing = Drawing.path [{x: 0.0, y}, {x: f.width, y}]

chrLabels :: forall c.
             CoordSys ChrId c
          -> Map ChrId (Array { drawing :: _, point :: _ })
chrLabels cs = Map.fromFoldable results
  where mkLabel :: ChrId -> _
        mkLabel chr = { drawing: chrText chr, point: Normalized {x: 1.0, y: -0.1} }

        font' = font sansSerif 12 mempty

        chrText :: ChrId -> Drawing
        chrText chr = scale 1.0 (-1.0)
                    $ Drawing.text font' zero zero (fillColor black) (unwrap chr)

        results :: Array _
        results = map (\x -> Tuple (x^._Index) [mkLabel (x^._Index)]) $ cs^._BrowserIntervals

        results' :: Traversal' (CoordSys ChrId c) ChrId
        results' = _BrowserIntervals <<< traversed <<< _Index

        test :: forall r. Monoid r => Fold' r _ _
        test = results' <<< (to (fanout id mkLabel))

        output :: _
        output = Map.fromFoldable $ (foldMapOf test pure cs :: Array _)

        -- xxx :: Traversal' (CoordSys ChrId _) ChrId
        -- xxx = results' <<< _Index

        -- results2 :: _ -> Array _
        -- results2 = traverseOf results' pure

        -- results3 :: _ -> Array _
        -- results3 = foldMapOf results' pure

        -- results4 :: _
        -- results4 = foldMapOf results'
                     -- (fanout id chrText)

drawDemo :: forall f r.
            Foldable f
         => Filterable f
         => CoordSys _ _
         -> { min :: Number, max :: Number, sig :: Number }
         -> { width :: Pixels, height :: Pixels }
         -> { gwas   :: Map ChrId (f (GWASFeature ()))
            , annots :: Map ChrId (f (Annot (minY :: Number))) }
         -> Interval BrowserPoint
         -> Drawing
drawDemo cs s canvasBox {gwas, annots} =
  let renderers = renderersDemo s
      ruler' = ruler s red
      drawGwas   = drawData cs canvasBox renderers.gwas gwas
      drawAnnots = drawData cs canvasBox renderers.annots annots
  in \view -> (drawGwas view) <> (drawAnnots view) <> ruler' canvasBox

{-
TODO
there is currently no connection between the vertical scale's
position and the track. Need to take stuff like vertical
offset and padding into account; calculate once and use
in both places. However don't just use canvasheight - yoffset
or whatever; actually derive the top and bottom of the track.
-}

drawVScale :: forall r.
              Number
           -> Number
           -> { min :: Number, max :: Number, sig :: Number }
           -> Color
           -> Drawing
drawVScale width height vs col =
  -- should have some padding here too; hardcode for now
  let hPad = width  / 5.0
      vPad = height / 10.0
      x = width / 2.0

      y1 = 0.0
      y2 = height

      n = (_ * 0.1) <<< Int.toNumber <$> (0 .. 10)

      p = Drawing.path [{x, y:y1}, {x, y:y2}]

      bar w y = Drawing.path [{x:x-w, y}, {x, y}]

      ps = foldMap (\i -> bar
                          (if i == 0.0 || i == 1.0 then 8.0 else 3.0)
                          (y1 + i*(y2-y1))) n

      ft = font sansSerif 10 mempty
      mkT y = Drawing.text ft (x-12.0) y (fillColor black)

      topLabel = mkT (y1-(0.5*vPad)) $ show vs.max
      btmLabel = mkT (y2+vPad) $ show vs.min

  in outlined (outlineColor col <> lineWidth 2.0) (p <> ps)
  <> topLabel <> btmLabel

drawLegendItem :: Number
               -> { text :: String, icon :: Drawing }
               -> Drawing
drawLegendItem w {text, icon} =
  let ft = font sansSerif 12 mempty
      t = Drawing.text ft 12.0 0.0 (fillColor black) text
  in icon <> t

drawLegend :: Number
           -> Number
           -> Array { text :: String, icon :: Drawing }
           -> Drawing
drawLegend width height icons =
  let hPad = width  / 5.0
      vPad = height / 5.0
      n :: Int
      n = length icons
      x = hPad
      f :: Number -> { text :: String, icon :: Drawing } -> Drawing
      f y ic = translate x y $ drawLegendItem (width - 2.0*hPad) ic
      d = (height - 2.0*vPad) / (length icons)
      ds = mapWithIndex (\i ic -> f (vPad+(vPad*(Int.toNumber i))) ic) icons
  in fold ds

mkIcon :: Color -> String -> { text :: String, icon :: Drawing }
mkIcon c text =
  let sh = circle (-2.5) (-2.5) 5.0
      icon = outlined (outlineColor c <> lineWidth 2.0) sh <>
             filled (fillColor c) sh
  in {text, icon}

demoLegend :: Array {text :: String, icon :: Drawing}
demoLegend =
  [ mkIcon red "Red thing"
  , mkIcon blue "blue thing"
  , mkIcon purple "boop"
  ]

              -- the first 2 are used by all parts of the browser
demoBrowser :: forall f.
               Foldable f
            => Filterable f
            => CoordSys ChrId BrowserPoint
            -> { width :: Pixels, height :: Pixels }
            -> Pixels
            -- this by the vertical scale and the tracks
            -> { min :: Number, max :: Number, sig :: Number }
            -- these define the width of the non-track parts of the browser;
            -- the browser takes up the remaining space.
            -> { vScaleWidth :: Pixels
               , legendWidth :: Pixels }
            -- vScale only
            -> Color
            -- to be drawn in the legend on the RHS
            -> Array { text :: String, icon :: Drawing }
            -- finally, used by the main track
            -> { gwas :: Map ChrId (f (GWASFeature ()))
               , annots :: Map ChrId (f (Annot (minY :: Number))) }
            -> BrowserView
            -- the drawing produced contains all 3 parts.
            -> { track :: Drawing, overlay :: Drawing }
demoBrowser cs canvas vpadding vscale sizes vscaleColor legend {gwas, annots} =
  let height = canvas.height - 2.0 * vpadding
      width  = canvas.width

      trackCanvas = { width: canvas.width - sizes.vScaleWidth - sizes.legendWidth
                    , height }

      bg w' h' = filled (fillColor white) $ rectangle 0.0 0.0 w' h'

      -- TODO unify vertical offset by padding further; do it as late as possible

      vScale :: _ -> Drawing
      vScale _ = let w = sizes.vScaleWidth
                 in  translate zero vpadding
                   $ bg w height
                  <> drawVScale w height vscale vscaleColor

      track v = translate sizes.vScaleWidth vpadding
                  $ drawDemo cs vscale trackCanvas {gwas, annots} v

      chrs :: _ -> Drawing
      chrs v = translate zero vpadding
                 $ drawFrames cs trackCanvas (chrLabels cs) v

      legendD :: _ -> Drawing
      legendD _ = let w = sizes.legendWidth
                  in translate (canvas.width - w) vpadding
                    $ bg w height
                   <> drawLegend sizes.legendWidth canvas.height legend

  in \view -> { track: track view <> chrs view
              , overlay: vScale view <> legendD view }

mouseChrIds :: Array ChrId
mouseChrIds =
            [ (ChrId "1")
            , (ChrId "2")
            , (ChrId "3")
            , (ChrId "4")
            , (ChrId "5")
            , (ChrId "6")
            , (ChrId "7")
            , (ChrId "8")
            , (ChrId "9")
            , (ChrId "10")
            , (ChrId "11")
            , (ChrId "12")
            , (ChrId "13")
            , (ChrId "14")
            , (ChrId "15")
            , (ChrId "16")
            , (ChrId "17")
            , (ChrId "18")
            , (ChrId "19")
            , (ChrId "X")
            , (ChrId "Y")
            ]

mouseChrs :: Map ChrId Bp
mouseChrs = Map.fromFoldable $ Array.zip mouseChrIds $
            [ (Bp 195471971.0)
            , (Bp 182113224.0)
            , (Bp 160039680.0)
            , (Bp 156508116.0)
            , (Bp 151834684.0)
            , (Bp 149736546.0)
            , (Bp 145441459.0)
            , (Bp 129401213.0)
            , (Bp 124595110.0)
            , (Bp 130694993.0)
            , (Bp 122082543.0)
            , (Bp 120129022.0)
            , (Bp 120421639.0)
            , (Bp 124902244.0)
            , (Bp 104043685.0)
            , (Bp 98207768.0)
            , (Bp 94987271.0)
            , (Bp 90702639.0)
            , (Bp 61431566.0)
            , (Bp 17103129.0)
            , (Bp 9174469.0)
            ]

mouseColors :: ChrCtx (color :: Color)
mouseColors = Map.fromFoldable $ map (\x -> Tuple x {color: navy}) mouseChrIds

mouseColors' :: ChrCtx (color :: Color)
mouseColors' = Map.fromFoldable $ Array.zip mouseChrIds $ map (\x -> {color: x})
               [ navy, blue, aqua, teal, olive
               , green, lime, yellow, orange, red
               , maroon, fuchsia, purple, navy, blue
               , aqua, teal, olive, green, lime
               , yellow ]

mouseChrCtx :: ChrCtx (size :: Bp)
mouseChrCtx = map (\size -> { size }) mouseChrs


Fetch gene chrs from ensembl, add gene renderer

module Genetics.Browser.Track.Backend where

import Color.Scheme.Clrs
import Prelude

import Color (Color, black)
import Control.Monad.Aff (Aff, launchAff)
import Control.Monad.Eff (Eff, foreachE)
import Control.Monad.Eff.Class (liftEff)
import Control.Monad.Eff.Console (log)
import Control.Monad.Eff.Exception (error)
import Control.Monad.Error.Class (class MonadError, throwError)
import Control.Monad.Reader (class MonadReader, Reader, ReaderT(..), ask, runReader, runReaderT)
import Control.MonadPlus (guard)
import Data.Argonaut (Json, _Array, _Number, _Object, _String)
import Data.Array ((:))
import Data.Array as Array
import Data.Distributive (distribute)
import Data.Either (Either, Either(..))
import Data.Filterable (class Filterable, filterMap)
import Data.Foldable (class Foldable, fold, foldMap, foldl, for_, length, maximum, minimum, sum)
import Data.Lens (Fold', Prism', Getter', preview, re, (^?))
import Data.Lens.Index (ix)
import Data.List (List(..), mapMaybe)
import Data.List as List
import Data.Map (Map)
import Data.Map as Map
import Data.Maybe (Maybe(..), fromJust, fromMaybe, isJust)
import Data.Monoid (mempty)
import Data.Monoid.Additive (Additive(..))
import Data.Newtype (ala, alaF, over, under, unwrap, wrap)
import Data.NonEmpty (foldl1, fromNonEmpty)
import Data.NonEmpty as NE
import Data.Ord.Max (Max(..))
import Data.Ordering (Ordering(..))
import Data.Record as Record
import Data.Record.Unsafe (unsafeSet)
import Data.Symbol (SProxy(..))
import Data.Traversable (class Traversable, scanl, traverse, traverse_)
import Data.Tuple (Tuple(..), fst, snd, uncurry)
import Data.Unfoldable (singleton, unfoldr)
import Data.Variant (Variant)
import Data.Variant as Variant
import Debug.Trace (trace)
import Debug.Trace as Debug
import Genetics.Browser.DataSource (DataSource)
import Genetics.Browser.Types (Bp(..), BpPerPixel(..), Chr, ChrId(..), Point, Pos, Range, _ChrId, bpToPixels)
import Genetics.Browser.UI.Native (getScreenSize)
import Genetics.Browser.UI.Native.View as View
import Genetics.Browser.View (View, Pixels)
import Global (readFloat, readInt)
import Global.Unsafe (unsafeStringify)
import Graphics.Canvas (CanvasElement, Context2D, getCanvasElementById, getCanvasHeight, getContext2D, setCanvasWidth)
import Graphics.Canvas as C
import Graphics.Drawing (Drawing, circle, fillColor, filled, outlineColor, outlined, rectangle, scale, translate)
import Graphics.Drawing as Drawing
import Graphics.Drawing.Font (bold, font, sansSerif)
import Math as Math
import Network.HTTP.Affjax as Affjax
import Partial.Unsafe (unsafePartial)
import Text.Parsing.CSV as CSV
import Text.Parsing.Parser as CSV
import Type.Prelude (class IsSymbol, class RowLacks)
import Unsafe.Coerce (unsafeCoerce)

type Frame a = { padding :: Pixels, width :: Pixels, contents :: a }

type ChrCtx r = Map ChrId (Record r)

type ChrFrames r a = ChrCtx (frame :: Frame a | r)

totalSize :: forall r.
             ChrCtx (size :: Bp | r)
          -> Bp
totalSize chrCtx = alaF Additive foldMap _.size chrCtx

newtype Normalized a = Normalized a

normalizePoint :: { width :: Number
                  , height :: Number }
               -> Point
               -> Normalized Point
normalizePoint {width, height} {x, y} = Normalized { x: x', y: y' }
  where x' = x / width
        y' = y / height

normPoint :: Point -> Maybe (Normalized Point)
normPoint p@{x, y}
  | x < 0.0 || x > 1.0 = Nothing
  | y < 0.0 || y > 1.0 = Nothing
  | otherwise = Just $ Normalized p

nPointToFrame :: Pixels -> Pixels -> Normalized Point -> Point
nPointToFrame w h (Normalized p) = {x: p.x * w, y: p.y * h}

-- A renderer draws a single value to a normalized "canvas",
-- i.e. the whole chromosome is located in x = (0.0, 1.0),
-- and the whole track height in y = (0.0, 1.0).
-- It can fail if the given feature for some reason cannot be rendered,
-- either because the feature lacks information, or because the context lacks something.
type PureRenderer a = a -> Maybe { drawing :: Drawing
                                 , point   :: Normalized Point
                                 }

-- If both renderers work on exactly the same data point,
-- there's no canonical way of picking between the results,
-- so we just pick the first
catRenderers :: forall a.
                PureRenderer a
             -> PureRenderer a
             -> PureRenderer a
catRenderers r1 r2 = \x -> case r1 x, r2 x of
  Nothing, Nothing -> Nothing
  Just y, _ -> pure y
  _, Just y -> pure y

type GWASFeature r = { score :: Number
                     , pos :: Bp
                     , chrId :: ChrId | r }

mkGwasRenderer :: forall m rf rctx.
                  MonadReader (ChrCtx (size :: Bp, color :: Color | rctx)) m
               -- => MonadError Unit m
               => {min :: Number, max :: Number}
               -> m (PureRenderer (GWASFeature rf))
mkGwasRenderer {min, max} = do
  ctx <- ask
  let renderer f = do
        size <- _.size <$> Map.lookup f.chrId ctx
        color <- _.color <$> Map.lookup f.chrId ctx
        let r = 2.2
            x = unwrap $ f.pos / size
            y = (f.score - min) / (max - min)
            c = circle x y r
            out = outlined (outlineColor color) c
            fill = filled (fillColor color) c
            drawing = out <> fill

        point <- normPoint {x, y}
        pure { drawing, point }

  pure renderer

-- A single Frame that contains all data that is required to draw it
type ReadyFrame a = { width :: Pixels
                    , padding :: Pixels
                    , height :: Pixels
                    , contents :: a }

drawFrame :: forall a.
             PureRenderer a
          -> ReadyFrame (Array a)
          -> Drawing
drawFrame r {width, padding, height, contents} = foldMap draw' contents
  where draw' a = case r a of
          Nothing -> mempty
          Just {drawing, point} ->
            let {x,y} = nPointToFrame (width - padding * 2.0) height $ point
            in translate (x + padding) y drawing

-- Draw a collection of ReadyFrames, lined up one after another
drawFrames :: forall a.
              Array (ReadyFrame (Array a))
           -> PureRenderer a
           -> Drawing
drawFrames frames r = fold drawings
  where os = Array.cons 0.0 (framesOffsets frames) -- add 0.0 since scanl in framesOffsets doesn't include seed
        drawings = Array.zipWith (\o -> translate o 0.0) os
                     $ map (drawFrame r) frames

framesOffsets :: forall f r.
                 Traversable f
              => f { width :: Pixels | r }
              -> f Pixels
framesOffsets fs = scanl (\o {width} -> width + o) 0.0 fs

insertFrame :: forall r a.
               RowLacks "frame" r
            => a
            -> { | r }
            -> { frame :: a | r }
insertFrame = Record.insert (SProxy :: SProxy "frame")

createFrames :: forall r a.
                RowLacks "frame" (size :: Bp | r)
             => Pixels
             -> Pixels
             -> Map ChrId {size :: Bp | r}
             -> Map ChrId a
             -> Map ChrId { size :: Bp, frame :: Frame a | r }
createFrames totalWidth padding sizes as = zipMapsWith zippy sizes as
  where ts = totalSize sizes
        mkWidth size = (unwrap $ size / ts) * totalWidth
        mkFrame' {size} a = { padding, width: mkWidth size, contents: a }
        zippy r a = insertFrame (mkFrame' r a) r

intersection :: forall k a b.
                Ord k
             => Map k a
             -> Map k b
             -> Map k a
intersection a b = Map.filterKeys (flip Map.member b) a

zipMaps :: forall k a b.
           Ord k
        => Map k a
        -> Map k b
        -> Map k (Tuple a b)
zipMaps a b =
  let kas = Array.unzip $ Map.toAscUnfoldable $ a `intersection` b
      kbs = Array.unzip $ Map.toAscUnfoldable $ b `intersection` a
      kabs = uncurry Array.zip $ map (Array.zip (snd kas)) kbs
  in Map.fromFoldable kabs

zipMapsWith :: forall k a b c.
               Ord k
            => (a -> b -> c)
            -> Map k a
            -> Map k b
            -> Map k c
zipMapsWith f a b = uncurry f <$> zipMaps a b

groupToChrs :: forall r.
               List { chrId :: ChrId | r }
            -> Map ChrId (List { chrId :: ChrId | r })
groupToChrs = foldl (\chrs r@{chrId} -> Map.alter (add r) chrId chrs ) mempty
  where add x Nothing   = Just $ singleton x
        add x (Just xs) = Just $ Cons x xs

chrFrames :: forall a.
             Pixels
          -> Pixels
          -> Map ChrId a
          -> Map ChrId { size :: Bp
                       , frame :: { padding :: Pixels
                                  , width :: Pixels
                                  , contents :: a
                                  }
                       }
chrFrames w p as = createFrames w p chrCtx' as
  where chrCtx' = Map.filterKeys (\k -> k `Map.member` as) mouseChrCtx

addHeight :: forall a r.
             Pixels
          -> Map ChrId { size :: Bp
                       , frame :: { padding :: Pixels
                                  , width :: Pixels
                                  , contents :: a
                                  }
                       | r
                       }
          -> Map ChrId { size :: Bp
                       , frame :: (ReadyFrame a)
                       | r }
addHeight h fs = map f fs
  where f entry = entry { frame = Record.insert (SProxy :: SProxy "height") h entry.frame }

fetchJSON :: forall a.
             (Json -> Maybe a)
          -> String
          -> Aff _ (Array a)
fetchJSON p url = do
  json <- _.response <$> Affjax.get url

  case json ^? _Array of
    Nothing -> throwError $ error "Parse error: JSON is not an array"
    Just as -> case traverse p as of
      Nothing -> throwError $ error "Failed to parse JSON features"
      Just fs -> pure $ fs

gemmaJSONParse :: Json -> Maybe (GWASFeature ())
gemmaJSONParse j = do
  obj <- j ^? _Object
  chrId <- ChrId <$> obj ^? ix "chr" <<< _String
  pos   <- Bp    <$> obj ^? ix "ps"  <<< _Number
  score <-           obj ^? ix "af"  <<< _Number
  pure {pos, score, chrId}

fetchGemmaJSON = fetchJSON gemmaJSONParse

type GeneRow' r = ( geneID :: String
                  , desc :: String
                  , start :: Bp
                  , end :: Bp
                  , name :: String | r )

type Gene' r = Record (GeneRow' r)

type GeneRow r = ( geneID :: String
                 , desc :: String
                 , start :: Bp
                 , end :: Bp
                 , name :: String
                 , chrId :: ChrId | r )

type Gene r = Record (GeneRow r)

geneJSONParse :: Json -> Maybe (Gene' ())
geneJSONParse j = do
  obj <- j ^? _Object
  geneID  <- obj ^? ix "Gene stable ID" <<< _String
  desc  <- obj ^? ix "Gene description" <<< _String
  name  <- obj ^? ix "Gene name" <<< _String
  start <- Bp    <$> obj ^? ix "Gene start (bp)"  <<< _Number
  end   <- Bp    <$> obj ^? ix "Gene end (bp)"  <<< _Number
  pure {geneID, desc, name, start, end}

fetchGene'JSON :: String -> Aff _ (Array (Gene' ()))
fetchGene'JSON = fetchJSON geneJSONParse

geneFetchChrId :: forall r.
                  RowLacks "chrId" (GeneRow' r)
               => Gene' r
               -> Aff _ (Gene r)
geneFetchChrId gene = do
  let url = "https://rest.ensembl.org/lookup/id/" <>
            gene.geneID <> "?content-type=application/json"

  res <- _.response <$> Affjax.get url

  let chrId :: _
      chrId = do
        obj <- res ^? _Object
        ChrId <$> obj ^? ix "seq_region_name" <<< _String

  case chrId of
    Nothing  -> throwError $ error $ "Could not find chrId for gene " <> gene.geneID
    Just chr -> pure $ Record.insert (SProxy :: SProxy "chrId") chr gene

fetchGeneJSON :: String -> Aff _ (Array (Gene ()))
fetchGeneJSON url = traverse geneFetchChrId =<< fetchJSON geneJSONParse url

mkGeneRenderer :: forall m rf rctx.
                  MonadReader (ChrCtx (size :: Bp | rctx)) m
               => m (PureRenderer (Gene rf))
mkGeneRenderer = do
  ctx <- ask

  let font' = font sansSerif 12 mempty
  let renderer f = do
        size <- _.size <$> Map.lookup f.chrId ctx
        let rad = 20.0
            lHand = unwrap $ f.start / size
            rHand = unwrap $ f.end / size
            w = (rHand - lHand) * unwrap size
            x = lHand + (rHand - lHand)
            y = 0.5
            rect = rectangle lHand (y - rad/2.0) w (rad/2.0)
            out = outlined (outlineColor red) rect
            fill = filled (fillColor maroon) rect
            -- the canvas is flipped so y-axis increases upward, need to flip the text too
            text' = scale 1.0 (-1.0)
                    $ Drawing.text font' x (y-0.25) (fillColor black) f.desc
            drawing = out <> fill <> text'

        point <- normPoint {x, y}
        pure { drawing, point }

  pure renderer

drawGeneric :: forall f a.
               Foldable f
            => { w :: Pixels, h :: Pixels, p :: Pixels, y :: Pixels }
            -> Array ChrId
            -> PureRenderer a
            -> Map ChrId (f a)
            -> Array ChrId
            -> Drawing
drawGeneric {w,h,p,y} allChrs renderer dat chrs =
  let features :: Map ChrId (Array a)
      features = map Array.fromFoldable
                 $ Map.filterKeys (\k -> k `Array.elem` chrs) dat

      readyFrames :: Map ChrId { size :: Bp, frame :: ReadyFrame (Array a) }
      readyFrames = addHeight h $ chrFrames w p features

      -- needed since the Ord instance on ChrId isn't necessarily correct
      comp (Tuple x _) (Tuple y _ ) = compareChrId allChrs x y

      toDraw :: Array (ReadyFrame (Array a))
      toDraw = _.frame <<< snd <$> (Array.sortBy comp $ Map.toUnfoldable readyFrames)

  in drawFrames toDraw renderer

drawGemma :: Pixels -> Pixels -> Pixels -> Pixels
          -> { min :: Number, max :: Number }
          -> Array ChrId
          -> List (GWASFeature ())
          -> Drawing
drawGemma w h p y s chrs gemma =
  let features :: Map ChrId (Array (GWASFeature ()))
      features = map Array.fromFoldable
                 $ Map.filterKeys (\k -> k `Array.elem` chrs)
                 $ groupToChrs gemma

      readyFrames :: Map ChrId { size :: Bp, frame :: ReadyFrame (Array (GWASFeature ())) }
      readyFrames = addHeight h $ chrFrames w p features

      comp (Tuple x _) (Tuple y _ ) = compareChrId mouseChrIds x y
      toDraw :: Array (ReadyFrame (Array (GWASFeature ())))
      toDraw = _.frame <<< snd <$> (Array.sortBy comp $ Map.toUnfoldable readyFrames)

      colors' = zipMapsWith (\r {color} -> Record.insert (SProxy :: SProxy "color") color r)
                  readyFrames mouseColors

      renderer' :: PureRenderer (GWASFeature ())
      renderer' = mkGwasRenderer s colors'

  -- pure $ translate 0.0 (-y) $ drawFrames toDraw renderer'
  in drawFrames toDraw renderer'

mouseChrIds :: Array ChrId
mouseChrIds =
            [ (ChrId "1")
            , (ChrId "2")
            , (ChrId "3")
            , (ChrId "4")
            , (ChrId "5")
            , (ChrId "6")
            , (ChrId "7")
            , (ChrId "8")
            , (ChrId "9")
            , (ChrId "10")
            , (ChrId "11")
            , (ChrId "12")
            , (ChrId "13")
            , (ChrId "14")
            , (ChrId "15")
            , (ChrId "16")
            , (ChrId "17")
            , (ChrId "18")
            , (ChrId "19")
            , (ChrId "X")
            , (ChrId "Y")
            ]

compareChrId :: Array ChrId -> (ChrId -> ChrId -> Ordering)
compareChrId chrs x y = fromMaybe EQ do
  x <- Array.elemIndex x chrs
  y <- Array.elemIndex y chrs
  pure $ x `compare` y

mouseChrs :: Map ChrId Bp
mouseChrs = Map.fromFoldable $ Array.zip mouseChrIds $
            [ (Bp 195471971.0)
            , (Bp 182113224.0)
            , (Bp 160039680.0)
            , (Bp 156508116.0)
            , (Bp 151834684.0)
            , (Bp 149736546.0)
            , (Bp 145441459.0)
            , (Bp 129401213.0)
            , (Bp 124595110.0)
            , (Bp 130694993.0)
            , (Bp 122082543.0)
            , (Bp 120129022.0)
            , (Bp 120421639.0)
            , (Bp 124902244.0)
            , (Bp 104043685.0)
            , (Bp 98207768.0)
            , (Bp 94987271.0)
            , (Bp 90702639.0)
            , (Bp 61431566.0)
            , (Bp 17103129.0)
            , (Bp 9174469.0)
            ]

mouseColors :: _
mouseColors = Map.fromFoldable $ Array.zip mouseChrIds $ map (\x -> {color: x})
              [ navy, blue, aqua, teal, olive
              , green, lime, yellow, orange, red
              , maroon, fuchsia, purple, navy, blue
              , aqua, teal, olive, green, lime
              , yellow ]

mouseChrCtx :: ChrCtx (size :: Bp)
mouseChrCtx = map (\size -> { size }) mouseChrs

mouseChrSize :: Bp -> ChrId -> Bp
mouseChrSize def chr = fromMaybe def $ Map.lookup chr mouseChrs

module Genetics.Browser.Track.Backend where

import Color.Scheme.Clrs
import Prelude

import Color (Color)
import Control.Monad.Aff (Aff, launchAff)
import Control.Monad.Eff (Eff, foreachE)
import Control.Monad.Eff.Class (liftEff)
import Control.Monad.Eff.Console (log)
import Control.Monad.Eff.Exception (error)
import Control.Monad.Error.Class (class MonadError, throwError)
import Control.Monad.Reader (class MonadReader, Reader, ReaderT(..), ask, runReader, runReaderT)
import Control.MonadPlus (guard)
import Data.Argonaut (Json, _Array, _Number, _Object, _String)
import Data.Array ((:))
import Data.Array as Array
import Data.Distributive (distribute)
import Data.Either (Either, Either(..))
import Data.Filterable (class Filterable, filterMap)
import Data.Foldable (class Foldable, fold, foldMap, foldl, for_, length, maximum, minimum, sum)
import Data.Lens (Fold', Prism', Getter', preview, re, (^?))
import Data.Lens.Index (ix)
import Data.List (List(..), mapMaybe)
import Data.List as List
import Data.Map (Map)
import Data.Map as Map
import Data.Maybe (Maybe(..), fromJust, fromMaybe, isJust)
import Data.Monoid (mempty)
import Data.Monoid.Additive (Additive(..))
import Data.Newtype (ala, alaF, over, under, unwrap, wrap)
import Data.NonEmpty (foldl1, fromNonEmpty)
import Data.NonEmpty as NE
import Data.Ord.Max (Max(..))
import Data.Ordering (Ordering(..))
import Data.Record as Record
import Data.Record.Unsafe (unsafeSet)
import Data.Symbol (SProxy(..))
import Data.Traversable (class Traversable, scanl, traverse, traverse_)
import Data.Tuple (Tuple(..), fst, snd, uncurry)
import Data.Unfoldable (singleton, unfoldr)
import Data.Variant (Variant)
import Data.Variant as Variant
import Debug.Trace (trace)
import Debug.Trace as Debug
import Genetics.Browser.DataSource (DataSource)
import Genetics.Browser.Types (Bp(..), BpPerPixel(..), Chr, ChrId(..), Point, Pos, Range, _ChrId, bpToPixels)
import Genetics.Browser.UI.Native (getScreenSize)
import Genetics.Browser.UI.Native.View as View
import Genetics.Browser.View (View, Pixels)
import Global (readFloat, readInt)
import Global.Unsafe (unsafeStringify)
import Graphics.Canvas (CanvasElement, Context2D, getCanvasElementById, getCanvasHeight, getContext2D, setCanvasWidth)
import Graphics.Canvas as C
import Graphics.Drawing (Drawing, circle, fillColor, filled, outlineColor, outlined, rectangle, translate)
import Graphics.Drawing as Drawing
import Math as Math
import Network.HTTP.Affjax as Affjax
import Partial.Unsafe (unsafePartial)
import Text.Parsing.CSV as CSV
import Text.Parsing.Parser as CSV
import Type.Prelude (class IsSymbol, class RowLacks)
import Unsafe.Coerce (unsafeCoerce)

type Frame a = { padding :: Pixels, width :: Pixels, contents :: a }

type ChrCtx r = Map ChrId (Record r)

type ChrFrames r a = ChrCtx (frame :: Frame a | r)

totalSize :: forall r.
             ChrCtx (size :: Bp | r)
          -> Bp
totalSize chrCtx = alaF Additive foldMap _.size chrCtx

newtype Normalized a = Normalized a

normalizePoint :: { width :: Number
                  , height :: Number }
               -> Point
               -> Normalized Point
normalizePoint {width, height} {x, y} = Normalized { x: x', y: y' }
  where x' = x / width
        y' = y / height

normPoint :: Point -> Maybe (Normalized Point)
normPoint p@{x, y}
  | x < 0.0 || x > 1.0 = Nothing
  | y < 0.0 || y > 1.0 = Nothing
  | otherwise = Just $ Normalized p

nPointToFrame :: Pixels -> Pixels -> Normalized Point -> Point
nPointToFrame w h (Normalized p) = {x: p.x * w, y: p.y * h}

-- A renderer draws a single value to a normalized "canvas",
-- i.e. the whole chromosome is located in x = (0.0, 1.0),
-- and the whole track height in y = (0.0, 1.0).
-- It can fail if the given feature for some reason cannot be rendered,
-- either because the feature lacks information, or because the context lacks something.
type PureRenderer a = a -> Maybe { drawing :: Drawing
                                 , point   :: Normalized Point
                                 }

-- If both renderers work on exactly the same data point,
-- there's no canonical way of picking between the results,
-- so we just pick the first
catRenderers :: forall a.
                PureRenderer a
             -> PureRenderer a
             -> PureRenderer a
catRenderers r1 r2 = \x -> case r1 x, r2 x of
  Nothing, Nothing -> Nothing
  Just y, _ -> pure y
  _, Just y -> pure y

type GWASFeature r = { score :: Number
                     , pos :: Bp
                     , chrId :: ChrId | r }

mkGwasRenderer :: forall m rf rctx.
                  MonadReader (ChrCtx (size :: Bp, color :: Color | rctx)) m
               -- => MonadError Unit m
               => {min :: Number, max :: Number}
               -> m (PureRenderer (GWASFeature rf))
mkGwasRenderer {min, max} = do
  ctx <- ask
  let renderer f = do
        size <- _.size <$> Map.lookup f.chrId ctx
        color <- _.color <$> Map.lookup f.chrId ctx
        let r = 2.2
            x = unwrap $ f.pos / size
            y = 1.0 - ((f.score - min) / (max - min))
            c = circle x y r
            out = outlined (outlineColor color) c
            fill = filled (fillColor color) c
            drawing = out <> fill

        point <- normPoint {x, y}
        pure { drawing, point }

  pure renderer

-- A single Frame that contains all data that is required to draw it
type ReadyFrame a = { width :: Pixels
                    , padding :: Pixels
                    , height :: Pixels
                    , contents :: a }

drawFrame :: forall a.
             PureRenderer a
          -> ReadyFrame (Array a)
          -> Drawing
drawFrame r {width, padding, height, contents} = foldMap draw' contents
  where draw' a = case r a of
          Nothing -> mempty
          Just {drawing, point} ->
            let {x,y} = nPointToFrame (width - padding * 2.0) height $ point
            in translate (x + padding) y drawing

-- Draw a collection of ReadyFrames, lined up one after another
drawFrames :: forall a.
              Array (ReadyFrame (Array a))
           -> PureRenderer a
           -> Drawing
drawFrames frames r = fold drawings
  where os = Array.cons 0.0 (framesOffsets frames) -- add 0.0 since scanl in framesOffsets doesn't include seed
        drawings = Array.zipWith (\o -> translate o 0.0) os
                     $ map (drawFrame r) frames

framesOffsets :: forall f r.
                 Traversable f
              => f { width :: Pixels | r }
              -> f Pixels
framesOffsets fs = scanl (\o {width} -> width + o) 0.0 fs

insertFrame :: forall r a.
               RowLacks "frame" r
            => a
            -> { | r }
            -> { frame :: a | r }
insertFrame = Record.insert (SProxy :: SProxy "frame")

createFrames :: forall r a.
                RowLacks "frame" (size :: Bp | r)
             => Pixels
             -> Pixels
             -> Map ChrId {size :: Bp | r}
             -> Map ChrId a
             -> Map ChrId { size :: Bp, frame :: Frame a | r }
createFrames totalWidth padding sizes as = zipMapsWith zippy sizes as
  where ts = totalSize sizes
        mkWidth size = (unwrap $ size / ts) * totalWidth
        mkFrame' {size} a = { padding, width: mkWidth size, contents: a }
        zippy r a = insertFrame (mkFrame' r a) r

intersection :: forall k a b.
                Ord k
             => Map k a
             -> Map k b
             -> Map k a
intersection a b = Map.filterKeys (flip Map.member b) a

zipMaps :: forall k a b.
           Ord k
        => Map k a
        -> Map k b
        -> Map k (Tuple a b)
zipMaps a b =
  let kas = Array.unzip $ Map.toAscUnfoldable $ a `intersection` b
      kbs = Array.unzip $ Map.toAscUnfoldable $ b `intersection` a
      kabs = uncurry Array.zip $ map (Array.zip (snd kas)) kbs
  in Map.fromFoldable kabs

zipMapsWith :: forall k a b c.
               Ord k
            => (a -> b -> c)
            -> Map k a
            -> Map k b
            -> Map k c
zipMapsWith f a b = uncurry f <$> zipMaps a b

groupToChrs :: forall r.
               List { chrId :: ChrId | r }
            -> Map ChrId (List { chrId :: ChrId | r })
groupToChrs = foldl (\chrs r@{chrId} -> Map.alter (add r) chrId chrs ) mempty
  where add x Nothing   = Just $ singleton x
        add x (Just xs) = Just $ Cons x xs

chrFrames :: forall a.
             Pixels
          -> Pixels
          -> Map ChrId a
          -> Map ChrId { size :: Bp
                       , frame :: { padding :: Pixels
                                  , width :: Pixels
                                  , contents :: a
                                  }
                       }
chrFrames w p as = createFrames w p chrCtx' as
  where chrCtx' = Map.filterKeys (\k -> k `Map.member` as) mouseChrCtx

addHeight :: forall a r.
             Pixels
          -> Map ChrId { size :: Bp
                       , frame :: { padding :: Pixels
                                  , width :: Pixels
                                  , contents :: a
                                  }
                       | r
                       }
          -> Map ChrId { size :: Bp
                       , frame :: (ReadyFrame a)
                       | r }
addHeight h fs = map f fs
  where f entry = entry { frame = Record.insert (SProxy :: SProxy "height") h entry.frame }

fetchJSON :: forall a.
             (Json -> Maybe a)
          -> String
          -> Aff _ (Array a)
fetchJSON p url = do
  json <- _.response <$> Affjax.get url

  case json ^? _Array of
    Nothing -> throwError $ error "Parse error: JSON is not an array"
    Just as -> case traverse p as of
      Nothing -> throwError $ error "Failed to parse JSON features"
      Just fs -> pure $ fs

gemmaJSONParse :: Json -> Maybe (GWASFeature ())
gemmaJSONParse j = do
  obj <- j ^? _Object
  chrId <- ChrId <$> obj ^? ix "chr" <<< _String
  pos   <- Bp    <$> obj ^? ix "ps"  <<< _Number
  score <-           obj ^? ix "af"  <<< _Number
  pure {pos, score, chrId}

fetchGemmaJSON = fetchJSON gemmaJSONParse

type GeneRow' r = ( geneID :: String
                  , desc :: String
                  , start :: Bp
                  , end :: Bp
                  , name :: String | r )

type Gene' r = Record (GeneRow' r)

findRenderer :: forall a.
                Map String (PureRenderer a)
             -> PureRenderer (Tuple String a)
findRenderer rs (Tuple s a) = do
  renderer <- Map.lookup s rs
  renderer a

geneJSONParse :: Json -> Maybe (Gene' ())
geneJSONParse j = do
  obj <- j ^? _Object
  geneID  <- obj ^? ix "Gene stable ID" <<< _String
  desc  <- obj ^? ix "Gene description" <<< _String
  name  <- obj ^? ix "Gene name" <<< _String
  start <- Bp    <$> obj ^? ix "Gene start (bp)"  <<< _Number
  end   <- Bp    <$> obj ^? ix "Gene end (bp)"  <<< _Number
  pure {geneID, desc, name, start, end}

fetchGene'JSON :: String -> Aff _ (Array (Gene' ()))
fetchGene'JSON = fetchJSON geneJSONParse

      -- needed since the Ord instance on ChrId isn't necessarily correct
      comp (Tuple x _) (Tuple y _ ) = compareChrId allChrs x y

      toDraw :: Array (ReadyFrame (Tuple String (Array a)))
      toDraw = _.frame <<< snd <$> (Array.sortBy comp $ Map.toUnfoldable readyFrames)

      -- r :: PureRenderer (Tuple String a)
      -- r = findRenderer renderers
      r :: PureRenderer (Tuple String a)
      r (Tuple s a) =  do
            renderer <- Map.lookup s renderers
            renderer a

  -- TODO fix this; (Tuple String (Array a)) -> Array (Tuple String a)
  -- in unsafeCoerce unit
  in drawFrames toDraw $ findRenderer renderers

drawGeneric :: forall f a.
               Foldable f
            => { w :: Pixels, h :: Pixels, p :: Pixels, y :: Pixels }
            -> Array ChrId
            -> PureRenderer a
            -> Map ChrId (f a)
            -> Array ChrId
            -> Drawing
drawGeneric {w,h,p,y} allChrs renderer dat chrs =
  let features :: Map ChrId (Array a)
      features = map Array.fromFoldable
                 $ Map.filterKeys (\k -> k `Array.elem` chrs) dat

      readyFrames :: Map ChrId { size :: Bp, frame :: ReadyFrame (Array a) }
      readyFrames = addHeight h $ chrFrames w p features

      -- needed since the Ord instance on ChrId isn't necessarily correct
      comp (Tuple x _) (Tuple y _ ) = compareChrId allChrs x y

      toDraw :: Array (ReadyFrame (Array a))
      toDraw = _.frame <<< snd <$> (Array.sortBy comp $ Map.toUnfoldable readyFrames)

  in drawFrames toDraw renderer

drawGemma :: Pixels -> Pixels -> Pixels -> Pixels
          -> { min :: Number, max :: Number }
          -> Array ChrId
          -> List (GWASFeature ())
          -> Drawing
drawGemma w h p y s chrs gemma =
  let features :: Map ChrId (Array (GWASFeature ()))
      features = map Array.fromFoldable
                 $ Map.filterKeys (\k -> k `Array.elem` chrs)
                 $ groupToChrs gemma

      readyFrames :: Map ChrId { size :: Bp, frame :: ReadyFrame (Array (GWASFeature ())) }
      readyFrames = addHeight h $ chrFrames w p features

      comp (Tuple x _) (Tuple y _ ) = compareChrId mouseChrIds x y
      toDraw :: Array (ReadyFrame (Array (GWASFeature ())))
      toDraw = _.frame <<< snd <$> (Array.sortBy comp $ Map.toUnfoldable readyFrames)

      colors' = zipMapsWith (\r {color} -> Record.insert (SProxy :: SProxy "color") color r)
                  readyFrames mouseColors

      renderer' :: PureRenderer (GWASFeature ())
      renderer' = mkGwasRenderer s colors'

  -- pure $ translate 0.0 (-y) $ drawFrames toDraw renderer'
  in drawFrames toDraw renderer'

mouseChrIds :: Array ChrId
mouseChrIds =
            [ (ChrId "1")
            , (ChrId "2")
            , (ChrId "3")
            , (ChrId "4")
            , (ChrId "5")
            , (ChrId "6")
            , (ChrId "7")
            , (ChrId "8")
            , (ChrId "9")
            , (ChrId "10")
            , (ChrId "11")
            , (ChrId "12")
            , (ChrId "13")
            , (ChrId "14")
            , (ChrId "15")
            , (ChrId "16")
            , (ChrId "17")
            , (ChrId "18")
            , (ChrId "19")
            , (ChrId "X")
            , (ChrId "Y")
            ]

compareChrId :: Array ChrId -> (ChrId -> ChrId -> Ordering)
compareChrId chrs x y = fromMaybe EQ do
  x <- Array.elemIndex x chrs
  y <- Array.elemIndex y chrs
  pure $ x `compare` y

mouseChrs :: Map ChrId Bp
mouseChrs = Map.fromFoldable $ Array.zip mouseChrIds $
            [ (Bp 195471971.0)
            , (Bp 182113224.0)
            , (Bp 160039680.0)
            , (Bp 156508116.0)
            , (Bp 151834684.0)
            , (Bp 149736546.0)
            , (Bp 145441459.0)
            , (Bp 129401213.0)
            , (Bp 124595110.0)
            , (Bp 130694993.0)
            , (Bp 122082543.0)
            , (Bp 120129022.0)
            , (Bp 120421639.0)
            , (Bp 124902244.0)
            , (Bp 104043685.0)
            , (Bp 98207768.0)
            , (Bp 94987271.0)
            , (Bp 90702639.0)
            , (Bp 61431566.0)
            , (Bp 17103129.0)
            , (Bp 9174469.0)
            ]

mouseColors :: _
mouseColors = Map.fromFoldable $ Array.zip mouseChrIds $ map (\x -> {color: x})
              [ navy, blue, aqua, teal, olive
              , green, lime, yellow, orange, red
              , maroon, fuchsia, purple, navy, blue
              , aqua, teal, olive, green, lime
              , yellow ]

mouseChrCtx :: ChrCtx (size :: Bp)
mouseChrCtx = map (\size -> { size }) mouseChrs

mouseChrSize :: Bp -> ChrId -> Bp
mouseChrSize def chr = fromMaybe def $ Map.lookup chr mouseChrs