Functions for running demo pipeline, a bunch of dumb notes

module Genetics.Browser.Track.Backend where

import Prelude

import Color (Color, black)
import Color.Scheme.Clrs (aqua, blue, fuchsia, green, lime, maroon, navy, olive, orange, purple, red, teal, yellow)
import Control.Alt ((<|>))
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 as Array
import Data.Filterable (class Filterable, filterMap)
import Data.Foldable (class Foldable, fold, foldMap, foldl, maximum)
import Data.Lens ((^?))
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.Monoid.Additive (Additive(..))
import Data.Newtype (class Newtype, alaF, unwrap)
import Data.Ord.Max (Max(..))
import Data.Record as Record
import Data.Semigroup.Foldable (foldMap1)
import Data.Symbol (SProxy(..))
import Data.Traversable (class Traversable, scanl, traverse)
import Data.Tuple (Tuple(Tuple), fst, snd, uncurry)
import Data.Unfoldable (class Unfoldable, singleton)
import Data.Unfoldable as Unfoldable
import Data.Variant (class Contractable, Variant)
import Data.Variant as Variant
import Genetics.Browser.Types (Bp(Bp), ChrId(ChrId), Point)
import Genetics.Browser.View (Pixels)
import Graphics.Drawing (Drawing, circle, fillColor, filled, 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 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

derive instance newtypeNormalized :: Newtype (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
                                 }

tagRenderer :: forall a sym r1 r2.
               RowCons sym a r1 r2
            => IsSymbol sym
            => SProxy sym
            -> PureRenderer a
            -> PureRenderer (Variant r2)
tagRenderer sym renderer var = Variant.prj sym var >>= renderer

combineRenderers :: forall r1 r2 r3.
                    Union r1 r2 r3
                 => Contractable r3 r1
                 => Contractable r3 r2
                 => PureRenderer (Variant r1)
                 -> PureRenderer (Variant r2)
                 -> PureRenderer (Variant r3)
combineRenderers r1 r2 var = (Variant.contract var >>= r1) <|>
                             (Variant.contract var >>= r2)

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 }

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

-- TODO simplify these two functions by extracting `padding` and `height` to another argument;
-- after all, only `width` changes per frame
drawPureFrame :: forall f r.
                 Foldable f
              => f { drawing :: Drawing, point :: Normalized Point | r }
              -> { width :: Pixels, padding :: Pixels, height :: Pixels }
              -> Drawing
drawPureFrame fs { width, padding, height } = foldMap render' fs
  where render' {drawing, point} =
            let {x,y} = nPointToFrame (width - padding * 2.0) height $ point
            in translate (x + padding) y drawing

-- Draw a collection of , lined up one after another
drawPureFrames :: forall a.
                  Array (Tuple (Array { drawing :: Drawing, point :: Normalized Point })
                               { width :: Pixels, padding :: Pixels, height :: Pixels })
               -> Drawing
drawPureFrames frames = fold $ Array.zipWith (\o d -> translate o 0.0 d) os drawings'
  where os = Array.cons 0.0 (framesOffsets $ snd <$> frames) -- add 0.0 since scanl in framesOffsets doesn't include seed
        drawings' = map (uncurry drawPureFrame) 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 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 :: { chrId :: ChrId | rData } -> Maybe _ -> Maybe _
        add x Nothing   = Just $ pure x
        add x (Just xs) = Just $ pure x <> xs

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

mouseChrFrames = chrFrames 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 :: String -> Aff _ (Array (GWASFeature ()))
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)

gene'JSONParse :: Json -> Maybe (Gene' ())
gene'JSONParse 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}

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}

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

geneFetchChrId :: forall r.
                  RowLacks "chrId" (GeneRow' r)
               => Gene' r
               -> Aff _ (Gene r)
geneFetchChrId gene = do
  let url = "http://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

fetchGeneJSONEnsembl :: String -> Aff _ (Array (Gene ()))
fetchGeneJSONEnsembl url = traverse geneFetchChrId =<< fetchJSON gene'JSONParse url

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

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

{-
What is needed is a way to get the score for some other thing
at or close to the annotation.

Maybe we can use the `point` property of the renderers.
In fact, that is the only thing we need to modify.
However, modifying that becomes tricky -- where do we do it,
and how do we keep track of it so that it's reversible?

THIS IS HOW:
"Render" everything once, at start. Then just use those Drawings and Normalized Points
when the chr view changes, zooms, etc. For now we can assume that we can just wait
a few seconds at start; later we can show the mid-steps.
Using this data (plus some sort of Shape usable for collision detection and/or some
provided "bumping" rules) we can iterate the track until everything is where it
should be.
That could even use Event's `fix` function to show the bumping take place...

There's one problem -- the horizontal interval a glyph covers is
a function of the BpPerPixels scale. I.e. even though each feature is always
drawn using the same Drawing and Normalized Point, the /relative horizontal/
positioning of glyphs is a function of:

the Drawings, Normalized Points, Collision Shapes in the neighborhood, and
the /current/ BpPerPixels in the chromosome.

However, it should be a pretty basic function. Linear, even.

Anyway, this is going to be tricky enough to not attack as a first attempt;
get something that works with what we've got right now, first.
-}
shiftRendererMinY :: forall rf.
                      RowLacks "minY" rf
                   => Number
                   -> { min :: Number, max :: 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 = (s.max - f.minY) / (s.max - s.min)
      y' = min (y + normY - dist) 1.0
  point' <- normPoint {x, y: y'}
  pure {drawing, point: point'}

bumpedGeneRenderer :: forall rf.
                      RowLacks "minY" (GeneRow rf)
                   => { min :: Number, max :: Number }
                   -> PureRenderer (Gene rf)
                   -> PureRenderer (Gene (minY :: Number | rf))
bumpedGeneRenderer s = shiftRendererMinY 0.05 s

bumpGenes :: forall rGene rGWAS.
             RowLacks "minY" (GeneRow rGene)
          => Map ChrId (List (GWASFeature rGWAS))
          -> Map ChrId (List (Gene rGene))
          -> Map ChrId (List (Gene (minY :: Number | rGene)))
bumpGenes = zipMapsWith (map <<< bumpGene)
  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
        bumpGene :: List (GWASFeature rGWAS) -> Gene rGene -> Gene (minY :: Number | rGene)
        bumpGene l g = let minY = maxScoreIn g.start g.end l
                       in Record.insert (SProxy :: SProxy "minY") minY g

getDataDemo :: { gwas :: String
               , genes :: String }
            -> Aff _ { gwas  :: Map ChrId (List (GWASFeature ()       ))
                     , genes :: Map ChrId (List (Gene (minY :: Number)))
                     }
getDataDemo urls = do
  gwas <- fetchGemmaJSON urls.gwas
  genes' <- fetchGeneJSON urls.genes
  -- Divide data by chromosomes
  let gwasChr :: Map ChrId (List (GWASFeature ()))
      gwasChr = List.fromFoldable <$> groupToChrs gwas
  let genesChr :: Map ChrId (List (Gene ()))
      genesChr = List.fromFoldable <$> groupToChrs genes'
  let bumpedGenes :: Map ChrId (List (Gene (minY :: Number)))
      bumpedGenes = bumpGenes gwasChr genesChr

  pure { gwas: gwasChr, genes: bumpedGenes }

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

        genes :: PureRenderer (Gene (minY :: Number))
        genes = bumpedGeneRenderer s $ mkGeneRenderer mouseChrCtx

drawDemo :: forall f.
            Foldable f
         => Filterable f
         => { width :: Pixels, height :: Pixels, padding :: Pixels, yOffset :: Pixels }
         -> { min :: Number, max :: Number }
         -> { gwas :: Map ChrId (f (GWASFeature ()))
            , genes :: Map ChrId (f (Gene (minY :: Number))) }
         -> Array ChrId
         -> Drawing
drawDemo f s {gwas, genes} chrs = (drawData' f mouseChrCtx renderers.gwas gwas chrs) <>
                                  (drawData' f mouseChrCtx renderers.genes genes chrs)
  where renderers = renderersDemo s

{-
The only way _this_ function would work (without a ridiculous number
of type annotations, if then) would be if the `r` row contains *all*
intermediate types.

But then it could work. No idea if it's in any way reasonable
or useful even if it does work, tho :) :) :)

Actually, it might work kind of nice if Variant.onMatch is used.
Then we'd just have to provide a record where each field corresponds
to a function modifying the respective data depending on one of the
possible data types; the function would itself use Variant.onMatch
since each type of data should be able to interact with every other
type of data.
-}

pipeline :: forall r.
            List (List (Variant r) -> List (Variant r))
         -> Map ChrId (List (Variant r))
         -> Map ChrId (List (Variant r))
pipeline = unsafeCoerce unit

-- Draw data by rendering it only once
drawData' :: forall f a r.
             Foldable f
          => Filterable f
          -- NOTE! the width is the canvas width, not frame. Same w/ yOffset
          => { width :: Pixels, height :: Pixels, padding :: Pixels, yOffset :: Pixels }
          -> Map ChrId { size :: Bp | r}
          -> PureRenderer a
          -> Map ChrId (f a)
          -> Array ChrId
          -> Drawing
drawData' frameBox chrCtx renderer dat chrs = translate 0.0 (frameBox.height - frameBox.yOffset)
                                                $ scale 1.0 (-1.0)
                                                $ drawPureFrames (frames chrs)
  -- TODO add a traceAny statement here to make sure it only runs once
  where drawings :: Map ChrId (f {drawing :: Drawing, point :: _})
        drawings = map (pureRender renderer) dat

        mkFrame :: Bp -> ChrId -> Maybe {width :: _, height :: _, padding :: _}
        mkFrame total chr = do
          {size} <- Map.lookup chr chrCtx
          let width = (unwrap $ size / total) * frameBox.width
          pure { height: frameBox.height, padding: frameBox.padding, width }

        mkFrames :: Array ChrId -> Array { width :: _, height :: _, padding :: _ }
        mkFrames = filterMap (mkFrame $ totalSize chrCtx)

        frames :: Array ChrId -> Array (Tuple (Array _) (_))
        frames chrs' = Array.zip (filterMap (\c -> Array.fromFoldable <$> Map.lookup c drawings) chrs')
                                 (mkFrames chrs')
{-
drawGenes :: forall r.
             { w :: Number, h :: Number
             , p :: Number, y :: Number }
          -> List (Gene r)
          -> Array ChrId
          -> Drawing
drawGenes pos dat chrs = drawData pos renderer (groupToChrs dat) chrs
  where renderer = mkGeneRenderer mouseChrCtx

drawGemma :: forall r.
             { w :: Number, h :: Number
             , p :: Number, y :: Number }
          -> { min :: Number, max :: Number }
          -> List (GWASFeature r)
          -> Array ChrId
          -> Drawing
drawGemma pos s dat chrs = drawData pos renderer (groupToChrs dat) chrs
  where colorsCtx = zipMapsWith (\r {color} -> Record.insert (SProxy :: SProxy "color") color r)
                    mouseChrCtx mouseColors
        renderer :: PureRenderer (GWASFeature r)
        renderer = mkGwasRenderer s colorsCtx

drawBoth :: forall r.
            { w :: Number, h :: Number
            , p :: Number, y :: Number }
         -> { min :: Number, max :: Number }
         -> Map ChrId (List (Variant ("Gene" :: Gene (), "GWAS" :: GWASFeature ())))
         -> Array ChrId
         -> Drawing
drawBoth pos s dat chrs = drawData pos renderer dat chrs
  where geneR :: PureRenderer (Variant ("Gene" :: Gene ()))
        geneR = tagRenderer (SProxy :: SProxy "Gene") $ mkGeneRenderer mouseChrCtx
        gemmaR :: PureRenderer (Variant ("GWAS" :: GWASFeature ()))
        gemmaR = tagRenderer (SProxy :: SProxy "GWAS") $ mkGwasRenderer s colorsCtx
        colorsCtx = zipMapsWith (\r {color} -> Record.insert (SProxy :: SProxy "color") color r)
                    mouseChrCtx mouseColors
        renderer = combineRenderers geneR gemmaR
-}

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 $ 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 Prelude

import Color (Color, black)
import Color.Scheme.Clrs (aqua, blue, fuchsia, green, lime, maroon, navy, olive, orange, purple, red, teal, yellow)
import Control.Alt ((<|>))
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 as Array
import Data.Filterable (class Filterable, filterMap)
import Data.Foldable (class Foldable, fold, foldMap, foldl, maximum)
import Data.Lens ((^?))
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.Monoid.Additive (Additive(..))
import Data.Newtype (class Newtype, alaF, unwrap)
import Data.Ord.Max (Max(..))
import Data.Record as Record
import Data.Semigroup.Foldable (foldMap1)
import Data.Symbol (SProxy(..))
import Data.Traversable (class Traversable, scanl, traverse)
import Data.Tuple (Tuple(Tuple), fst, snd, uncurry)
import Data.Unfoldable (class Unfoldable, singleton)
import Data.Unfoldable as Unfoldable
import Data.Variant (class Contractable, Variant)
import Data.Variant as Variant
import Genetics.Browser.Types (Bp(Bp), ChrId(ChrId), Point)
import Genetics.Browser.View (Pixels)
import Graphics.Drawing (Drawing, circle, fillColor, filled, 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 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

derive instance newtypeNormalized :: Newtype (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
                                 }

tagRenderer :: forall a sym r1 r2.
               RowCons sym a r1 r2
            => IsSymbol sym
            => SProxy sym
            -> PureRenderer a
            -> PureRenderer (Variant r2)
tagRenderer sym renderer var = Variant.prj sym var >>= renderer

combineRenderers :: forall r1 r2 r3.
                    Union r1 r2 r3
                 => Contractable r3 r1
                 => Contractable r3 r2
                 => PureRenderer (Variant r1)
                 -> PureRenderer (Variant r2)
                 -> PureRenderer (Variant r3)
combineRenderers r1 r2 var = (Variant.contract var >>= r1) <|>
                             (Variant.contract var >>= r2)

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 }

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

-- TODO simplify these two functions by extracting `padding` and `height` to another argument;
-- after all, only `width` changes per frame
drawPureFrame :: forall f r.
                 Foldable f
              => f { drawing :: Drawing, point :: Normalized Point | r }
              -> { width :: Pixels, padding :: Pixels, height :: Pixels }
              -> Drawing
drawPureFrame fs { width, padding, height } = foldMap render' fs
  where render' {drawing, point} =
            let {x,y} = nPointToFrame (width - padding * 2.0) height $ point
            in translate (x + padding) y drawing

-- Draw a collection of , lined up one after another
drawPureFrames :: forall a.
                  Array (Tuple (Array { drawing :: Drawing, point :: Normalized Point })
                               { width :: Pixels, padding :: Pixels, height :: Pixels })
               -> Drawing
drawPureFrames frames = fold $ Array.zipWith (\o d -> translate o 0.0 d) os drawings'
  where os = Array.cons 0.0 (framesOffsets $ snd <$> frames) -- add 0.0 since scanl in framesOffsets doesn't include seed
        drawings' = map (uncurry drawPureFrame) 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 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 :: { chrId :: ChrId | rData } -> Maybe _ -> Maybe _
        add x Nothing   = Just $ pure x
        add x (Just xs) = Just $ pure x <> xs

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

mouseChrFrames = chrFrames 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 :: String -> Aff _ (Array (GWASFeature ()))
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)

gene'JSONParse :: Json -> Maybe (Gene' ())
gene'JSONParse 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}

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}

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

geneFetchChrId :: forall r.
                  RowLacks "chrId" (GeneRow' r)
               => Gene' r
               -> Aff _ (Gene r)
geneFetchChrId gene = do
  let url = "http://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

fetchGeneJSONEnsembl :: String -> Aff _ (Array (Gene ()))
fetchGeneJSONEnsembl url = traverse geneFetchChrId =<< fetchJSON gene'JSONParse url

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

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

shiftRendererMinY :: forall rf.
                      RowLacks "minY" rf
                   => Number
                   -> { min :: Number, max :: 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 = (s.max - f.minY) / (s.max - s.min)
      y' = min (y + normY - dist) 1.0
  point' <- normPoint {x, y: y'}
  pure {drawing, point: point'}

bumpedGeneRenderer :: forall rf.
                      RowLacks "minY" (GeneRow rf)
                   => { min :: Number, max :: Number }
                   -> PureRenderer (Gene rf)
                   -> PureRenderer (Gene (minY :: Number | rf))
bumpedGeneRenderer s = shiftRendererMinY 0.05 s

bumpGenes :: forall rGene rGWAS.
             RowLacks "minY" (GeneRow rGene)
          => Map ChrId (List (GWASFeature rGWAS))
          -> Map ChrId (List (Gene rGene))
          -> Map ChrId (List (Gene (minY :: Number | rGene)))
bumpGenes = zipMapsWith (map <<< bumpGene)
  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
        bumpGene :: List (GWASFeature rGWAS) -> Gene rGene -> Gene (minY :: Number | rGene)
        bumpGene l g = let minY = maxScoreIn g.start g.end l
                       in Record.insert (SProxy :: SProxy "minY") minY g

drawData :: forall f a.
            Foldable f
         => { w :: Pixels, h :: Pixels, p :: Pixels, y :: Pixels }
         -> PureRenderer a
         -> Map ChrId (f a)
         -> Array ChrId
         -> Drawing
drawData {w,h,p,y} 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 $ mouseChrFrames w p features

      toDraw :: Array (ReadyFrame (Array a))
      toDraw = _.frame  <$> (filterMap (\chr -> Map.lookup chr readyFrames) chrs)

  in translate 0.0 (h-y) $ scale 1.0 (-1.0) $ drawFrames toDraw renderer

-- Draw data by rendering it only once
drawData' :: forall f a r.
             Foldable f
          => Filterable f
          -- NOTE! the width is the canvas width, not frame. Same w/ yOffset
          => { width :: Pixels, height :: Pixels, padding :: Pixels, yOffset :: Pixels }
          -> Map ChrId { size :: Bp | r}
          -> PureRenderer a
          -> Map ChrId (f a)
          -> Array ChrId
          -> Drawing
drawData' frameBox chrCtx renderer dat chrs = translate 0.0 (frameBox.height - frameBox.yOffset)
                                                $ scale 1.0 (-1.0)
                                                $ drawPureFrames (frames chrs)
  -- TODO add a traceAny statement here to make sure it only runs once
  where drawings :: Map ChrId (f {drawing :: Drawing, point :: _})
        drawings = map (pureRender renderer) dat

        mkFrame :: Bp -> ChrId -> Maybe {width :: _, height :: _, padding :: _}
        mkFrame total chr = do
          {size} <- Map.lookup chr chrCtx
          let width = (unwrap $ size / total) * frameBox.width
          pure { height: frameBox.height, padding: frameBox.padding, width }

        mkFrames :: Array ChrId -> Array { width :: _, height :: _, padding :: _ }
        mkFrames = filterMap (mkFrame $ totalSize chrCtx)

        frames :: Array ChrId -> Array (Tuple (Array _) (_))
        frames chrs' = Array.zip (filterMap (\c -> Array.fromFoldable <$> Map.lookup c drawings) chrs')
                                 (mkFrames chrs')
{-
drawGenes :: forall r.
             { w :: Number, h :: Number
             , p :: Number, y :: Number }
          -> List (Gene r)
          -> Array ChrId
          -> Drawing
drawGenes pos dat chrs = drawData pos renderer (groupToChrs dat) chrs
  where renderer = mkGeneRenderer mouseChrCtx

drawGemma :: forall r.
             { w :: Number, h :: Number
             , p :: Number, y :: Number }
          -> { min :: Number, max :: Number }
          -> List (GWASFeature r)
          -> Array ChrId
          -> Drawing
drawGemma pos s dat chrs = drawData pos renderer (groupToChrs dat) chrs
  where colorsCtx = zipMapsWith (\r {color} -> Record.insert (SProxy :: SProxy "color") color r)
                    mouseChrCtx mouseColors
        renderer :: PureRenderer (GWASFeature r)
        renderer = mkGwasRenderer s colorsCtx

drawBoth :: forall r.
            { w :: Number, h :: Number
            , p :: Number, y :: Number }
         -> { min :: Number, max :: Number }
         -> Map ChrId (List (Variant ("Gene" :: Gene (), "GWAS" :: GWASFeature ())))
         -> Array ChrId
         -> Drawing
drawBoth pos s dat chrs = drawData pos renderer dat chrs
  where geneR :: PureRenderer (Variant ("Gene" :: Gene ()))
        geneR = tagRenderer (SProxy :: SProxy "Gene") $ mkGeneRenderer mouseChrCtx
        gemmaR :: PureRenderer (Variant ("GWAS" :: GWASFeature ()))
        gemmaR = tagRenderer (SProxy :: SProxy "GWAS") $ mkGwasRenderer s colorsCtx
        colorsCtx = zipMapsWith (\r {color} -> Record.insert (SProxy :: SProxy "color") color r)
                    mouseChrCtx mouseColors
        renderer = combineRenderers geneR gemmaR
-}

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 $ 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


Add MonadRec instance

Add MonadRec instance

module Test.StrongCheck.Gen
  ( GenT(..)
  , Gen(..)
  , GenState(..)
  , GenOut(..)
  , Size
  , allInArray
  , allInRange
  , applyGen
  , arrayOf
  , arrayOf1
  , charGen
  , choose
  , chooseInt
  , chunked
  , collectAll
  , decorateSeed
  , dropGen
  , elements
  , extend
  , foldGen
  , foldGen'
  , frequency
  , infinite
  , interleave
  , nChooseK
  , oneOf
  , perms
  , perturbGen
  , repeatable
  , resize
  , runGen
  , sample
  , sample'
  , showSample
  , showSample'
  , shuffle
  , shuffle'
  , shuffleArray
  , sized
  , stateful
  , suchThat
  , suchThatMaybe
  , takeGen
  , toLazyList
  , transGen
  , uniform
  , unGen
  , unGenOut
  , unGenState
  , updateSeedState
  , variant
  , vectorOf
  , wrapEffect
  ) where

import Prelude

import Control.Alt (class Alt, (<|>))
import Control.Alternative (class Alternative)
import Control.Lazy as CL
import Control.Monad.Eff (Eff)
import Control.Monad.Eff.Console (logShow, CONSOLE)
import Control.Monad.List.Trans as ListT
import Control.Monad.Trampoline (runTrampoline, Trampoline)
import Control.Monad.Gen as CMG
import Control.Monad.Rec.Class (class MonadRec, Step(..), tailRecM)
import Control.MonadPlus (class MonadPlus)
import Control.MonadZero (class MonadZero)
import Control.Plus (class Plus)

import Data.Array as A
import Data.Array.Partial as AP
import Data.Char (fromCharCode)
import Data.Foldable (fold)
import Data.Int (toNumber)
import Data.Lazy (Lazy, defer)
import Data.List as L
import Data.Machine.Mealy as Mealy
import Data.Maybe (fromMaybe, maybe, Maybe(..), fromJust)
import Data.Monoid (mempty, class Monoid)
import Data.Monoid.Additive (Additive(..))
import Data.Newtype (unwrap)
import Data.Profunctor (arr, lmap)
import Data.Profunctor.Strong ((&&&))
import Data.Tuple (Tuple(..), snd, fst)

import Math as M

import Partial.Unsafe (unsafePartial)

import Test.StrongCheck.LCG (Seed, mkSeed, runSeed, lcgPerturb, lcgNext, lcgN)

type Size = Int

newtype GenState = GenState { seed :: Seed, size :: Size }

unGenState :: GenState -> { seed :: Seed, size :: Size }
unGenState (GenState s) = s

stateM
  :: ({ seed :: Seed, size :: Size } -> { seed :: Seed, size :: Size })
  -> GenState
  -> GenState
stateM f = GenState <<< f <<< unGenState

newtype GenOut a = GenOut { state :: GenState, value :: a }

unGenOut :: forall a. GenOut a -> { state :: GenState, value :: a }
unGenOut (GenOut v) = v

newtype GenT f a = GenT (Mealy.MealyT f GenState (GenOut a))

type Gen a = GenT Trampoline a

unGen :: forall f a. GenT f a -> Mealy.MealyT f GenState (GenOut a)
unGen (GenT m) = m

decorateSeed :: forall f a. Monad f => GenT f a -> GenT f (Tuple Seed a)
decorateSeed gen = stateful \sIn -> Tuple (unGenState sIn).seed <$> gen

lcgStep :: forall f. Monad f => GenT f Int
lcgStep = GenT $ arr \s ->
  GenOut { state: updateSeedState s, value: runSeed (unGenState s).seed }

uniform :: forall f. Monad f => GenT f Number
uniform = (\n -> toNumber n / toNumber lcgN) <$> lcgStep

stepGen
  :: forall f a
   . Monad f
  => GenState
  -> GenT f a
  -> f (Maybe (GenOut (Tuple a (GenT f a))))
stepGen st (GenT m) = h <$> Mealy.stepMealy st m
  where
  h Mealy.Halt = Nothing
  h (Mealy.Emit a m') = Just $ flip Tuple (GenT m') <$> a

evalGen :: forall f a. Monad f => GenT f a -> GenState -> f (Maybe a)
evalGen g st = h <$> stepGen st g
  where
  h Nothing = Nothing
  h (Just (GenOut { value: Tuple a _ })) = Just a

pureGen :: forall f a. Monad f => (GenState -> GenOut a) -> GenT f a
pureGen f = GenT $ arr f

repeatable' :: forall f a b. Monad f => (a -> GenT f b) -> GenT f (a -> f b)
repeatable' = unsafePartial \f -> GenT $
  Mealy.pureMealy \s ->
    Mealy.Emit
      (GenOut { state: s, value: \a -> fromJust <$> evalGen (f a) s })
      Mealy.halt

-- | Creates a function generator that will always generate the same output
-- | for the same input.
repeatable :: forall a b. (a -> Gen b) -> Gen (a -> b)
repeatable f = g <$> repeatable' f
  where
  g f' = \a -> runTrampoline $ f' a

-- | Creates a generator that depends on access to the generator state.
stateful :: forall f a. Monad f => (GenState -> GenT f a) -> GenT f a
stateful f = GenT $ unGen <<< f =<< Mealy.take 1 id

-- | Fixes a generator on a certain variant, given by the specified seed.
variant :: forall f a. Monad f => Seed -> GenT f a -> GenT f a
variant n g = GenT $ lmap (stateM (\s -> s { seed = n })) (unGen g)

-- | Creates a generator that depends on the size parameter.
sized :: forall f a. Monad f => (Size -> GenT f a) -> GenT f a
sized f = stateful \s -> f (unGenState s).size

-- | Resizes the generator so the size parameter passed into the generator
-- | will be equal to the specified size.
resize :: forall f a. Monad f => Size -> GenT f a -> GenT f a
resize sz g = GenT $ lmap (stateM (\s -> s { size = sz })) (unGen g)

-- | A generator for characters.
charGen :: forall f. Monad f => GenT f Char
charGen = fromCharCode <$> chooseInt 0 65535

-- | Creates a generator that generates real numbers between the specified
-- | inclusive range.
choose :: forall f. Monad f => Number -> Number -> GenT f Number
choose a b = (*) (max - min) >>> (+) min <$> uniform
  where
  min = M.min a b
  max = M.max a b

-- | Creates a generator that generates integers between the specified
-- | inclusive range.
chooseInt :: forall f. Monad f => Int -> Int -> GenT f Int
chooseInt a b = clamp <$> lcgStep
  where
  clamp :: Int -> Int
  clamp x = case x `mod` (b - a + one) of
              r | r >= 0 -> a + r
                | otherwise -> b + r + one

-- | Creates a generator that chooses another generator from the specified list
-- | at random, and then generates a value with that generator.
oneOf :: forall f a. Monad f => GenT f a -> Array (GenT f a) -> GenT f a
oneOf x xs = do
  n <- chooseInt 0 (A.length xs)
  if n == 0 then x else fromMaybe x (xs A.!! (n - 1))

-- | Generates elements by the specified frequencies (which will be normalized).
frequency
  :: forall f a
   . Monad f
  => Tuple Number (GenT f a)
  -> L.List (Tuple Number (GenT f a))
  -> GenT f a
frequency x xs = do
  let xxs :: L.List (Tuple Number (GenT f a))
      xxs = L.Cons x xs

      total :: Number
      total = unwrap $ fold ((Additive <<< fst) <$> xxs)

      pick :: Number -> GenT f a -> L.List (Tuple Number (GenT f a)) -> GenT f a
      pick _ d L.Nil = d
      pick n d (L.Cons (Tuple k x') xs') = if n <= k then x' else pick (n - k) d xs'

  n <- choose 1.0 total
  pick n (snd x) xxs

-- | Creates a generator of elements ranging from 0 to the maximum size.
arrayOf :: forall f a. Monad f => GenT f a -> GenT f (Array a)
arrayOf g = sized \n -> do
  k <- chooseInt 0 n
  vectorOf k g

-- | Creates a generator of elements ranging from 1 to the maximum size.
arrayOf1 :: forall f a. Monad f => GenT f a -> GenT f (Tuple a (Array a))
arrayOf1 g = sized \n -> do
  k <- chooseInt 0 n
  x <- g
  xs <- vectorOf (k - 1) g
  pure $ Tuple x xs

-- | Creates a generator that generates arrays of some specified size.
vectorOf :: forall f a. Monad f => Int -> GenT f a -> GenT f (Array a)
vectorOf n g = transGen f [] (extend n g)
  where
  f b a =
    let b' = b <> [a]
    in if A.length b' >= n then Tuple [] (Just b') else Tuple b' Nothing

-- | Creates a generator that chooses an element from among a set of elements.
elements :: forall f a. Monad f => a -> L.List a -> GenT f a
elements x xs = do
  n <- chooseInt 0 (L.length xs)
  pure if n == 0 then x else fromMaybe x (xs L.!! (n - 1))

foreign import float32ToInt32 :: Number -> Int

perturbGen :: forall f a. Monad f => Number -> GenT f a -> GenT f a
perturbGen n (GenT m) =
  GenT $ lmap (stateM (\s -> s { seed = lcgPerturb (toNumber (float32ToInt32 n)) s.seed })) m

updateSeedState :: GenState -> GenState
updateSeedState (GenState s) = GenState { seed: lcgNext s.seed, size: s.size }

updateSeedGen :: forall f a. Monad f => GenT f a -> GenT f a
updateSeedGen (GenT m) = GenT $ lmap updateSeedState m

liftMealy
  :: forall f a
   . Monad f
   => (Mealy.MealyT f GenState (GenOut a) -> Mealy.MealyT f GenState (GenOut a))
   -> GenT f a
   -> GenT f a
liftMealy f = \g -> GenT $ f (unGen g)

-- | Takes the first number of values from the generator. Will turn an infinite
-- | generator into a finite generator.
takeGen :: forall f a. Monad f => Int -> GenT f a -> GenT f a
takeGen n = liftMealy $ Mealy.take n

-- | Drops a certain number of values from the generator. May produce
-- | an empty generator if called on a finite generator.
dropGen :: forall f a. Monad f => Int -> GenT f a -> GenT f a
dropGen n = liftMealy $ Mealy.drop n

-- | Extends a generator to produce *at least* the specified number of values.
-- | Will not turn a finite generator into an infinite one.
extend :: forall f a. Monad f => Int -> GenT f a -> GenT f a
extend n (GenT m) = (GenT $ loop 0 m) <> (GenT m)
  where
  loop i m' = Mealy.mealy \st ->
    let
      f (Mealy.Emit s m'') = pure $ Mealy.Emit s (loop (i + 1) m'')
      f Mealy.Halt =
        if i >= n then pure Mealy.Halt else Mealy.stepMealy st (loop i m)
    in Mealy.stepMealy st m' >>= f

-- | Fairly interleaves two generators.
interleave :: forall f a. Monad f => GenT f a -> GenT f a -> GenT f a
interleave (GenT g1) (GenT g2) = GenT $ Mealy.interleave g1 g2

-- | Ensures that a given generator can produce an infinite number of values,
-- | assuming it can produce at least one.
infinite :: forall f a. Monad f => GenT f a -> GenT f a
infinite = liftMealy $ Mealy.loop

-- | Folds over a generator to produce a value. Either the generator or the
-- | user-defined function may halt the fold. Returns not just the value
-- | created through folding, but also the successor generator.
foldGen'
  :: forall f a b
   . Monad f
  => (b -> a -> Maybe b)
  -> b
  -> GenState
  -> GenT f a -> f (Tuple b (GenT f a))
foldGen' f b s (GenT m) = loop s m b
  where
  loop st m' b' = Mealy.stepMealy st m' >>= g
    where
    g Mealy.Halt = pure $ Tuple b' (GenT Mealy.halt)
    g (Mealy.Emit (GenOut { value: a, state: st' }) m'') =
      let b'' = f b' a in maybe (pure $ Tuple b' (GenT m'')) (loop st' m'') b''

-- | Folds over a generator to produce a value. Either the generator or the
-- | user-defined function may halt the fold.
foldGen
  :: forall f a b
   . Monad f
  => (b -> a -> Maybe b)
  -> b
  -> GenState
  -> GenT f a
  -> f b
foldGen f b s g = fst <$> foldGen' f b s g

-- | Transforms one gen into another, passing along user-supplied state.
-- | Either the generator being transformed or the transforming function may
-- | halt the transformation.
transGen
  :: forall f a b c
   . Monad f
  => (b -> a -> Tuple b (Maybe c))
  -> b
  -> GenT f a
  -> GenT f c
transGen f b (GenT m) = GenT $ loop m b
  where
  loop m' b' = Mealy.mealy \st -> Mealy.stepMealy st m' >>= g
    where
    g Mealy.Halt = pure Mealy.Halt
    g (Mealy.Emit (GenOut { value: a, state: st }) m'') =
      case f b' a of
        Tuple b'' Nothing  -> Mealy.stepMealy st (loop m'' b'')
        Tuple b'' (Just c) ->
          let c' = GenOut { value: c, state: st }
          in  pure $ Mealy.Emit c' (loop m'' b'')

-- | A deterministic generator that produces all possible permutations of
-- | the specified array.
perms :: forall f a. Monad f => Array a -> GenT f (Array a)
perms = unsafePartial \arr ->
  if A.length arr == 0
  then pure [ ]
  else do
    let hd = AP.head arr
        tl = AP.tail arr
    xs <- perms tl
    let f n = let prefix = A.take n xs
                  suffix = A.drop n xs
              in prefix <> [hd] <> suffix
    allInArray $ f <$> A.range 0 (A.length xs)

-- | A deterministic generator that produces all possible combinations of
-- | choosing exactly k elements from the specified array.
nChooseK :: forall f a. Monad f => Int -> Array a -> GenT f (Array a)
nChooseK = unsafePartial \k arr ->
  if k == 0 then pure []
  else if A.length arr == 0
       then mempty
       else
         let tl = A.drop 1 arr
             hd = A.singleton $ A.unsafeIndex arr 0
         in ((hd <> _) <$> (nChooseK (k - 1) tl)) <> nChooseK k tl

-- | Same as `nChooseK` but for `List`
nChooseKL :: forall f a. Monad f => Int -> L.List a -> GenT f (L.List a)
nChooseKL 0 _ = pure L.Nil
nChooseKL _ L.Nil = mempty
nChooseKL k (L.Cons x xs) =
  (((L.singleton x) <> _) <$> (nChooseKL (k - 1) xs)) <> nChooseKL k xs

-- | Filters a generator to produce only values satisfying the specified
-- | predicate.
suchThat :: forall f a. Monad f => GenT f a -> (a -> Boolean) -> GenT f a
suchThat g p = g `suchThatMaybe` p >>= maybe (sized next) pure
  where
  next n = resize (n + 1) (g `suchThat` p)

-- | Tries to filter a generator such that it only produces values satisfying
-- | the specified predicate.
suchThatMaybe
  :: forall f a
   . Monad f
  => GenT f a
  -> (a -> Boolean)
  -> GenT f (Maybe a)
suchThatMaybe g p = sized $ try 0
  where
  try _ 0 = pure Nothing
  try k n = do
    x <- resize (2 * k + n) g
    if p x then pure (Just x) else try (k + 1) (n - 1)

-- | A deterministic generator that produces integers from the specified
-- | inclusive range, in sequence.
allInRange :: forall f. Monad f => Int -> Int -> GenT f Int
allInRange min max = GenT $ go min
  where
  go cur = Mealy.pureMealy \s ->
    if cur > max
    then Mealy.Halt
    else Mealy.Emit (GenOut {state:s, value:cur}) (go (cur + 1))

-- | A deterministic generator that produces values from the specified array,
-- | in sequence.
allInArray :: forall f a. Monad f => Array a -> GenT f a
allInArray a = GenT $ go 0
  where
  go i = Mealy.pureMealy \s ->
    maybe
      Mealy.Halt
      (\a' -> Mealy.Emit (GenOut { state: s, value: a' }) (go (i + 1)))
      (a A.!! i)

-- | Drains a finite generator of all values. Or blows up if you called it on
-- | an infinite generator.
collectAll :: forall f a. Monad f => GenState -> GenT f a -> f (Array a)
collectAll = foldGen f []
  where
  f v a = Just $ v <> [a]

-- | Applies a state to a generator to possibly produce the next state,
-- | a value, and the next generator.
applyGen
  :: forall f a
   . Monad f
  => GenState
  -> GenT f a
  -> f (Maybe (GenOut (Tuple a (GenT f a))))
applyGen s (GenT m) = f <$> Mealy.stepMealy s m
  where
  f Mealy.Halt = Nothing
  f (Mealy.Emit (GenOut { state, value: a }) m') =
    Just $ GenOut { state, value: Tuple a (GenT m')}

-- | Samples a generator, producing the specified number of values.
sample' :: forall f a. Monad f => Int -> GenState -> GenT f a -> f (Array a)
sample' n st g = fst <$> runGen n st g

-- | Samples a generator, producing the specified number of values. Uses
-- | default settings for the initial generator state.
sample :: forall f a. Monad f => Int -> GenT f a -> f (Array a)
sample n = sample' n (GenState { size: 10, seed: mkSeed 445645874 })

-- | Shows a sample of values generated from the specified generator.
showSample'
  :: forall r a
   . Show a
  => Int
  -> Gen a
  -> Eff (console :: CONSOLE | r) Unit
showSample' n g = logShow $ runTrampoline $ sample n g

-- | Shows a sample of values generated from the specified generator.
showSample :: forall r a. Show a => Gen a -> Eff (console :: CONSOLE | r) Unit
showSample = showSample' 10

-- | Runs a generator to produce a specified number of values, returning both
-- | an array containing the values and the successor Gen that can be used to
-- | continue the generation process at a later time.
runGen
  :: forall f a
   . Monad f
  => Int
  -> GenState
  -> GenT f a -> f (Tuple (Array a) (GenT f a))
runGen n st g = foldGen' f [] st (extend n g)
  where
  f v a = if A.length v < n then Just $ v <> [a] else Nothing

-- | Creates a generator that produces chunks of values in the specified size.
-- | Will extend the generator if necessary to produce a chunk of the specified
-- | size, but will not turn a finite generator into an infinite generator.
chunked :: forall f a. Monad f => Int -> GenT f a -> GenT f (Array a)
chunked n g = transGen f [] (extend n g)
  where
  f xs a =
    let xs' = a A.: xs
    in if A.length xs' >= n then Tuple [] (Just xs') else Tuple xs' Nothing

-- | Wraps an effect in a generator that ignores the input state.
wrapEffect :: forall f a. Monad f => f a -> GenT f a
wrapEffect fa = GenT $ g <$> (id &&& (Mealy.wrapEffect fa))
  where
  g (Tuple l r) = GenOut { state: l, value: r }

-- | Creates a generator that mixes up the order of the specified generator.
-- | This is achieved by chunking the generator with the specified size
-- | and then shuffling each chunk before continuing to the next.
shuffle' :: forall f a. Monad f => Int -> GenT f a -> GenT f a
shuffle' n g = do
  chunks <- chunked n g
  shuffled <- shuffleArray chunks
  allInArray shuffled

-- | Same as shuffle' but with default for the chunk size.
shuffle :: forall f a. Monad f => GenT f a -> GenT f a
shuffle = shuffle' 100

-- | Creates a generator that produces shuffled versions of the provided array.
shuffleArray :: forall f a. Monad f => Array a -> GenT f (Array a)
shuffleArray = shuffle0 []
  where
  shuffle0 acc [] = pure $ acc
  shuffle0 acc xs = do
    i <- chooseInt 0 (A.length xs - 1)
    let acc' = acc <> (maybe [] A.singleton (xs A.!! i))
        xs' = fromMaybe xs $ A.deleteAt i xs
    shuffle0 acc' xs'

-- | Converts the generator into a function that, given the initial state,
-- | returns a lazy list.
toLazyList :: forall a. Gen a -> GenState -> ListT.ListT Lazy a
toLazyList (GenT m) s = ListT.wrapLazy $ defer \_ -> loop m s
  where
  loop m' s' =
    case runTrampoline (Mealy.stepMealy s' m') of
      Mealy.Halt -> ListT.nil
      Mealy.Emit (GenOut { value, state }) m'' ->
        ListT.prepend' value (defer \_ -> loop m'' state)

instance semigroupGenState :: Semigroup GenState where
  append (GenState a) (GenState b) =
    GenState { seed: lcgPerturb (toNumber (runSeed a.seed)) b.seed, size: b.size }

instance monoidGenState :: Monoid GenState where
  mempty = GenState { seed: mkSeed 85734629, size: 10 }

instance semigroupGenOut :: Semigroup a => Semigroup (GenOut a) where
  append (GenOut a) (GenOut b) =
    GenOut { state: a.state <> b.state, value: a.value <> b.value }

instance monoidGenOut :: Monoid a => Monoid (GenOut a) where
  mempty = GenOut { state: mempty, value: mempty }

instance functorGenOut :: Functor GenOut where
  map f (GenOut m) = GenOut { state: m.state, value: f m.value }

instance applyGenOut :: Apply GenOut where
  apply (GenOut f) (GenOut x) =
    GenOut { state: x.state, value: f.value x.value }

-- GenT instances
instance functorGenT :: Monad f => Functor (GenT f) where
  map f (GenT m) = GenT $ map f <$> m

instance applyGenT :: Monad f => Apply (GenT f) where
  apply f x = GenT $ apply <$> unGen f <*> unGen (updateSeedGen x)

instance applicativeGenT :: Monad f => Applicative (GenT f) where
  pure t = GenT $ arr (\s -> GenOut { state: updateSeedState s, value: t })

instance semigroupGenT :: Monad f => Semigroup (GenT f a) where
  append (GenT a) (GenT b) = GenT (a <> b)

instance monoidGenT :: Monad f => Monoid (GenT f a) where
  mempty = GenT mempty

instance bindGenT :: Monad f => Bind (GenT f) where
  bind (GenT m) f = GenT do
    a <- m
    unGen $ updateSeedGen (f (unGenOut a).value)

instance monadGenT :: Monad f => Monad (GenT f)

instance altGenT :: Monad f => Alt (GenT f) where
  alt (GenT x) (GenT y) = GenT $ x <|> y

instance plusGenT :: Monad f => Plus (GenT f) where
  empty = mempty

instance alternativeGenT :: Monad f => Alternative (GenT f)

instance monadZeroGenT :: Monad f => MonadZero (GenT f)

instance monadPlusGenT :: Monad f => MonadPlus (GenT f)

instance monadRecGenT :: MonadRec f => MonadRec (GenT f) where
  tailRecM k a = k a >>= case _ of
    Loop b -> tailRecM k b
    Done r -> pure r

instance lazyGenT :: Monad f => CL.Lazy (GenT f a) where
  defer f = GenT $ CL.defer (unGen <<< f)

instance monadGenGenT :: Monad f => CMG.MonadGen (GenT f) where
  chooseInt = chooseInt
  chooseFloat = choose
  chooseBool = (_ < 0.5) <$> uniform
  resize f g = stateful \(GenState state) -> resize (f state.size) g
  sized = sized

module Test.StrongCheck.Gen
  ( GenT(..)
  , Gen(..)
  , GenState(..)
  , GenOut(..)
  , Size
  , allInArray
  , allInRange
  , applyGen
  , arrayOf
  , arrayOf1
  , charGen
  , choose
  , chooseInt
  , chunked
  , collectAll
  , decorateSeed
  , dropGen
  , elements
  , extend
  , foldGen
  , foldGen'
  , frequency
  , infinite
  , interleave
  , nChooseK
  , oneOf
  , perms
  , perturbGen
  , repeatable
  , resize
  , runGen
  , sample
  , sample'
  , showSample
  , showSample'
  , shuffle
  , shuffle'
  , shuffleArray
  , sized
  , stateful
  , suchThat
  , suchThatMaybe
  , takeGen
  , toLazyList
  , transGen
  , uniform
  , unGen
  , unGenOut
  , unGenState
  , updateSeedState
  , variant
  , vectorOf
  , wrapEffect
  ) where

import Prelude

import Control.Alt (class Alt, (<|>))
import Control.Alternative (class Alternative)
import Control.Lazy as CL
import Control.Monad.Eff (Eff)
import Control.Monad.Eff.Console (logShow, CONSOLE)
import Control.Monad.List.Trans as ListT
import Control.Monad.Trampoline (runTrampoline, Trampoline)
import Control.Monad.Gen as CMG
import Control.MonadPlus (class MonadPlus)
import Control.MonadZero (class MonadZero)
import Control.Plus (class Plus)

import Data.Array as A
import Data.Array.Partial as AP
import Data.Char (fromCharCode)
import Data.Foldable (fold)
import Data.Int (toNumber)
import Data.Lazy (Lazy, defer)
import Data.List as L
import Data.Machine.Mealy as Mealy
import Data.Maybe (fromMaybe, maybe, Maybe(..), fromJust)
import Data.Monoid (mempty, class Monoid)
import Data.Monoid.Additive (Additive(..))
import Data.Newtype (unwrap)
import Data.Profunctor (arr, lmap)
import Data.Profunctor.Strong ((&&&))
import Data.Tuple (Tuple(..), snd, fst)

import Math as M

import Partial.Unsafe (unsafePartial)

import Test.StrongCheck.LCG (Seed, mkSeed, runSeed, lcgPerturb, lcgNext, lcgN)

type Size = Int

newtype GenState = GenState { seed :: Seed, size :: Size }

unGenState :: GenState -> { seed :: Seed, size :: Size }
unGenState (GenState s) = s

stateM
  :: ({ seed :: Seed, size :: Size } -> { seed :: Seed, size :: Size })
  -> GenState
  -> GenState
stateM f = GenState <<< f <<< unGenState

newtype GenOut a = GenOut { state :: GenState, value :: a }

unGenOut :: forall a. GenOut a -> { state :: GenState, value :: a }
unGenOut (GenOut v) = v

newtype GenT f a = GenT (Mealy.MealyT f GenState (GenOut a))

type Gen a = GenT Trampoline a

unGen :: forall f a. GenT f a -> Mealy.MealyT f GenState (GenOut a)
unGen (GenT m) = m

decorateSeed :: forall f a. Monad f => GenT f a -> GenT f (Tuple Seed a)
decorateSeed gen = stateful \sIn -> Tuple (unGenState sIn).seed <$> gen

lcgStep :: forall f. Monad f => GenT f Int
lcgStep = GenT $ arr \s ->
  GenOut { state: updateSeedState s, value: runSeed (unGenState s).seed }

uniform :: forall f. Monad f => GenT f Number
uniform = (\n -> toNumber n / toNumber lcgN) <$> lcgStep

stepGen
  :: forall f a
   . Monad f
  => GenState
  -> GenT f a
  -> f (Maybe (GenOut (Tuple a (GenT f a))))
stepGen st (GenT m) = h <$> Mealy.stepMealy st m
  where
  h Mealy.Halt = Nothing
  h (Mealy.Emit a m') = Just $ flip Tuple (GenT m') <$> a

evalGen :: forall f a. Monad f => GenT f a -> GenState -> f (Maybe a)
evalGen g st = h <$> stepGen st g
  where
  h Nothing = Nothing
  h (Just (GenOut { value: Tuple a _ })) = Just a

pureGen :: forall f a. Monad f => (GenState -> GenOut a) -> GenT f a
pureGen f = GenT $ arr f

repeatable' :: forall f a b. Monad f => (a -> GenT f b) -> GenT f (a -> f b)
repeatable' = unsafePartial \f -> GenT $
  Mealy.pureMealy \s ->
    Mealy.Emit
      (GenOut { state: s, value: \a -> fromJust <$> evalGen (f a) s })
      Mealy.halt

-- | Creates a function generator that will always generate the same output
-- | for the same input.
repeatable :: forall a b. (a -> Gen b) -> Gen (a -> b)
repeatable f = g <$> repeatable' f
  where
  g f' = \a -> runTrampoline $ f' a

-- | Creates a generator that depends on access to the generator state.
stateful :: forall f a. Monad f => (GenState -> GenT f a) -> GenT f a
stateful f = GenT $ unGen <<< f =<< Mealy.take 1 id

-- | Fixes a generator on a certain variant, given by the specified seed.
variant :: forall f a. Monad f => Seed -> GenT f a -> GenT f a
variant n g = GenT $ lmap (stateM (\s -> s { seed = n })) (unGen g)

-- | Creates a generator that depends on the size parameter.
sized :: forall f a. Monad f => (Size -> GenT f a) -> GenT f a
sized f = stateful \s -> f (unGenState s).size

-- | Resizes the generator so the size parameter passed into the generator
-- | will be equal to the specified size.
resize :: forall f a. Monad f => Size -> GenT f a -> GenT f a
resize sz g = GenT $ lmap (stateM (\s -> s { size = sz })) (unGen g)

-- | A generator for characters.
charGen :: forall f. Monad f => GenT f Char
charGen = fromCharCode <$> chooseInt 0 65535

-- | Creates a generator that generates real numbers between the specified
-- | inclusive range.
choose :: forall f. Monad f => Number -> Number -> GenT f Number
choose a b = (*) (max - min) >>> (+) min <$> uniform
  where
  min = M.min a b
  max = M.max a b

-- | Creates a generator that generates integers between the specified
-- | inclusive range.
chooseInt :: forall f. Monad f => Int -> Int -> GenT f Int
chooseInt a b = clamp <$> lcgStep
  where
  clamp :: Int -> Int
  clamp x = case x `mod` (b - a + one) of
              r | r >= 0 -> a + r
                | otherwise -> b + r + one

-- | Creates a generator that chooses another generator from the specified list
-- | at random, and then generates a value with that generator.
oneOf :: forall f a. Monad f => GenT f a -> Array (GenT f a) -> GenT f a
oneOf x xs = do
  n <- chooseInt 0 (A.length xs)
  if n == 0 then x else fromMaybe x (xs A.!! (n - 1))

-- | Generates elements by the specified frequencies (which will be normalized).
frequency
  :: forall f a
   . Monad f
  => Tuple Number (GenT f a)
  -> L.List (Tuple Number (GenT f a))
  -> GenT f a
frequency x xs = do
  let xxs :: L.List (Tuple Number (GenT f a))
      xxs = L.Cons x xs

      total :: Number
      total = unwrap $ fold ((Additive <<< fst) <$> xxs)

      pick :: Number -> GenT f a -> L.List (Tuple Number (GenT f a)) -> GenT f a
      pick _ d L.Nil = d
      pick n d (L.Cons (Tuple k x') xs') = if n <= k then x' else pick (n - k) d xs'

  n <- choose 1.0 total
  pick n (snd x) xxs

-- | Creates a generator of elements ranging from 0 to the maximum size.
arrayOf :: forall f a. Monad f => GenT f a -> GenT f (Array a)
arrayOf g = sized \n -> do
  k <- chooseInt 0 n
  vectorOf k g

-- | Creates a generator of elements ranging from 1 to the maximum size.
arrayOf1 :: forall f a. Monad f => GenT f a -> GenT f (Tuple a (Array a))
arrayOf1 g = sized \n -> do
  k <- chooseInt 0 n
  x <- g
  xs <- vectorOf (k - 1) g
  pure $ Tuple x xs

-- | Creates a generator that generates arrays of some specified size.
vectorOf :: forall f a. Monad f => Int -> GenT f a -> GenT f (Array a)
vectorOf n g = transGen f [] (extend n g)
  where
  f b a =
    let b' = b <> [a]
    in if A.length b' >= n then Tuple [] (Just b') else Tuple b' Nothing

-- | Creates a generator that chooses an element from among a set of elements.
elements :: forall f a. Monad f => a -> L.List a -> GenT f a
elements x xs = do
  n <- chooseInt 0 (L.length xs)
  pure if n == 0 then x else fromMaybe x (xs L.!! (n - 1))

foreign import float32ToInt32 :: Number -> Int

perturbGen :: forall f a. Monad f => Number -> GenT f a -> GenT f a
perturbGen n (GenT m) =
  GenT $ lmap (stateM (\s -> s { seed = lcgPerturb (toNumber (float32ToInt32 n)) s.seed })) m

updateSeedState :: GenState -> GenState
updateSeedState (GenState s) = GenState { seed: lcgNext s.seed, size: s.size }

updateSeedGen :: forall f a. Monad f => GenT f a -> GenT f a
updateSeedGen (GenT m) = GenT $ lmap updateSeedState m

liftMealy
  :: forall f a
   . Monad f
   => (Mealy.MealyT f GenState (GenOut a) -> Mealy.MealyT f GenState (GenOut a))
   -> GenT f a
   -> GenT f a
liftMealy f = \g -> GenT $ f (unGen g)

-- | Takes the first number of values from the generator. Will turn an infinite
-- | generator into a finite generator.
takeGen :: forall f a. Monad f => Int -> GenT f a -> GenT f a
takeGen n = liftMealy $ Mealy.take n

-- | Drops a certain number of values from the generator. May produce
-- | an empty generator if called on a finite generator.
dropGen :: forall f a. Monad f => Int -> GenT f a -> GenT f a
dropGen n = liftMealy $ Mealy.drop n

-- | Extends a generator to produce *at least* the specified number of values.
-- | Will not turn a finite generator into an infinite one.
extend :: forall f a. Monad f => Int -> GenT f a -> GenT f a
extend n (GenT m) = (GenT $ loop 0 m) <> (GenT m)
  where
  loop i m' = Mealy.mealy \st ->
    let
      f (Mealy.Emit s m'') = pure $ Mealy.Emit s (loop (i + 1) m'')
      f Mealy.Halt =
        if i >= n then pure Mealy.Halt else Mealy.stepMealy st (loop i m)
    in Mealy.stepMealy st m' >>= f

-- | Fairly interleaves two generators.
interleave :: forall f a. Monad f => GenT f a -> GenT f a -> GenT f a
interleave (GenT g1) (GenT g2) = GenT $ Mealy.interleave g1 g2

-- | Ensures that a given generator can produce an infinite number of values,
-- | assuming it can produce at least one.
infinite :: forall f a. Monad f => GenT f a -> GenT f a
infinite = liftMealy $ Mealy.loop

-- | Folds over a generator to produce a value. Either the generator or the
-- | user-defined function may halt the fold. Returns not just the value
-- | created through folding, but also the successor generator.
foldGen'
  :: forall f a b
   . Monad f
  => (b -> a -> Maybe b)
  -> b
  -> GenState
  -> GenT f a -> f (Tuple b (GenT f a))
foldGen' f b s (GenT m) = loop s m b
  where
  loop st m' b' = Mealy.stepMealy st m' >>= g
    where
    g Mealy.Halt = pure $ Tuple b' (GenT Mealy.halt)
    g (Mealy.Emit (GenOut { value: a, state: st' }) m'') =
      let b'' = f b' a in maybe (pure $ Tuple b' (GenT m'')) (loop st' m'') b''

-- | Folds over a generator to produce a value. Either the generator or the
-- | user-defined function may halt the fold.
foldGen
  :: forall f a b
   . Monad f
  => (b -> a -> Maybe b)
  -> b
  -> GenState
  -> GenT f a
  -> f b
foldGen f b s g = fst <$> foldGen' f b s g

-- | Transforms one gen into another, passing along user-supplied state.
-- | Either the generator being transformed or the transforming function may
-- | halt the transformation.
transGen
  :: forall f a b c
   . Monad f
  => (b -> a -> Tuple b (Maybe c))
  -> b
  -> GenT f a
  -> GenT f c
transGen f b (GenT m) = GenT $ loop m b
  where
  loop m' b' = Mealy.mealy \st -> Mealy.stepMealy st m' >>= g
    where
    g Mealy.Halt = pure Mealy.Halt
    g (Mealy.Emit (GenOut { value: a, state: st }) m'') =
      case f b' a of
        Tuple b'' Nothing  -> Mealy.stepMealy st (loop m'' b'')
        Tuple b'' (Just c) ->
          let c' = GenOut { value: c, state: st }
          in  pure $ Mealy.Emit c' (loop m'' b'')

-- | A deterministic generator that produces all possible permutations of
-- | the specified array.
perms :: forall f a. Monad f => Array a -> GenT f (Array a)
perms = unsafePartial \arr ->
  if A.length arr == 0
  then pure [ ]
  else do
    let hd = AP.head arr
        tl = AP.tail arr
    xs <- perms tl
    let f n = let prefix = A.take n xs
                  suffix = A.drop n xs
              in prefix <> [hd] <> suffix
    allInArray $ f <$> A.range 0 (A.length xs)

-- | A deterministic generator that produces all possible combinations of
-- | choosing exactly k elements from the specified array.
nChooseK :: forall f a. Monad f => Int -> Array a -> GenT f (Array a)
nChooseK = unsafePartial \k arr ->
  if k == 0 then pure []
  else if A.length arr == 0
       then mempty
       else
         let tl = A.drop 1 arr
             hd = A.singleton $ A.unsafeIndex arr 0
         in ((hd <> _) <$> (nChooseK (k - 1) tl)) <> nChooseK k tl

-- | Same as `nChooseK` but for `List`
nChooseKL :: forall f a. Monad f => Int -> L.List a -> GenT f (L.List a)
nChooseKL 0 _ = pure L.Nil
nChooseKL _ L.Nil = mempty
nChooseKL k (L.Cons x xs) =
  (((L.singleton x) <> _) <$> (nChooseKL (k - 1) xs)) <> nChooseKL k xs

-- | Filters a generator to produce only values satisfying the specified
-- | predicate.
suchThat :: forall f a. Monad f => GenT f a -> (a -> Boolean) -> GenT f a
suchThat g p = g `suchThatMaybe` p >>= maybe (sized next) pure
  where
  next n = resize (n + 1) (g `suchThat` p)

-- | Tries to filter a generator such that it only produces values satisfying
-- | the specified predicate.
suchThatMaybe
  :: forall f a
   . Monad f
  => GenT f a
  -> (a -> Boolean)
  -> GenT f (Maybe a)
suchThatMaybe g p = sized $ try 0
  where
  try _ 0 = pure Nothing
  try k n = do
    x <- resize (2 * k + n) g
    if p x then pure (Just x) else try (k + 1) (n - 1)

-- | A deterministic generator that produces integers from the specified
-- | inclusive range, in sequence.
allInRange :: forall f. Monad f => Int -> Int -> GenT f Int
allInRange min max = GenT $ go min
  where
  go cur = Mealy.pureMealy \s ->
    if cur > max
    then Mealy.Halt
    else Mealy.Emit (GenOut {state:s, value:cur}) (go (cur + 1))

-- | A deterministic generator that produces values from the specified array,
-- | in sequence.
allInArray :: forall f a. Monad f => Array a -> GenT f a
allInArray a = GenT $ go 0
  where
  go i = Mealy.pureMealy \s ->
    maybe
      Mealy.Halt
      (\a' -> Mealy.Emit (GenOut { state: s, value: a' }) (go (i + 1)))
      (a A.!! i)

-- | Drains a finite generator of all values. Or blows up if you called it on
-- | an infinite generator.
collectAll :: forall f a. Monad f => GenState -> GenT f a -> f (Array a)
collectAll = foldGen f []
  where
  f v a = Just $ v <> [a]

-- | Applies a state to a generator to possibly produce the next state,
-- | a value, and the next generator.
applyGen
  :: forall f a
   . Monad f
  => GenState
  -> GenT f a
  -> f (Maybe (GenOut (Tuple a (GenT f a))))
applyGen s (GenT m) = f <$> Mealy.stepMealy s m
  where
  f Mealy.Halt = Nothing
  f (Mealy.Emit (GenOut { state, value: a }) m') =
    Just $ GenOut { state, value: Tuple a (GenT m')}

-- | Samples a generator, producing the specified number of values.
sample' :: forall f a. Monad f => Int -> GenState -> GenT f a -> f (Array a)
sample' n st g = fst <$> runGen n st g

-- | Samples a generator, producing the specified number of values. Uses
-- | default settings for the initial generator state.
sample :: forall f a. Monad f => Int -> GenT f a -> f (Array a)
sample n = sample' n (GenState { size: 10, seed: mkSeed 445645874 })

-- | Shows a sample of values generated from the specified generator.
showSample'
  :: forall r a
   . Show a
  => Int
  -> Gen a
  -> Eff (console :: CONSOLE | r) Unit
showSample' n g = logShow $ runTrampoline $ sample n g

-- | Shows a sample of values generated from the specified generator.
showSample :: forall r a. Show a => Gen a -> Eff (console :: CONSOLE | r) Unit
showSample = showSample' 10

-- | Runs a generator to produce a specified number of values, returning both
-- | an array containing the values and the successor Gen that can be used to
-- | continue the generation process at a later time.
runGen
  :: forall f a
   . Monad f
  => Int
  -> GenState
  -> GenT f a -> f (Tuple (Array a) (GenT f a))
runGen n st g = foldGen' f [] st (extend n g)
  where
  f v a = if A.length v < n then Just $ v <> [a] else Nothing

-- | Creates a generator that produces chunks of values in the specified size.
-- | Will extend the generator if necessary to produce a chunk of the specified
-- | size, but will not turn a finite generator into an infinite generator.
chunked :: forall f a. Monad f => Int -> GenT f a -> GenT f (Array a)
chunked n g = transGen f [] (extend n g)
  where
  f xs a =
    let xs' = a A.: xs
    in if A.length xs' >= n then Tuple [] (Just xs') else Tuple xs' Nothing

-- | Wraps an effect in a generator that ignores the input state.
wrapEffect :: forall f a. Monad f => f a -> GenT f a
wrapEffect fa = GenT $ g <$> (id &&& (Mealy.wrapEffect fa))
  where
  g (Tuple l r) = GenOut { state: l, value: r }

-- | Creates a generator that mixes up the order of the specified generator.
-- | This is achieved by chunking the generator with the specified size
-- | and then shuffling each chunk before continuing to the next.
shuffle' :: forall f a. Monad f => Int -> GenT f a -> GenT f a
shuffle' n g = do
  chunks <- chunked n g
  shuffled <- shuffleArray chunks
  allInArray shuffled

-- | Same as shuffle' but with default for the chunk size.
shuffle :: forall f a. Monad f => GenT f a -> GenT f a
shuffle = shuffle' 100

-- | Creates a generator that produces shuffled versions of the provided array.
shuffleArray :: forall f a. Monad f => Array a -> GenT f (Array a)
shuffleArray = shuffle0 []
  where
  shuffle0 acc [] = pure $ acc
  shuffle0 acc xs = do
    i <- chooseInt 0 (A.length xs - 1)
    let acc' = acc <> (maybe [] A.singleton (xs A.!! i))
        xs' = fromMaybe xs $ A.deleteAt i xs
    shuffle0 acc' xs'

-- | Converts the generator into a function that, given the initial state,
-- | returns a lazy list.
toLazyList :: forall a. Gen a -> GenState -> ListT.ListT Lazy a
toLazyList (GenT m) s = ListT.wrapLazy $ defer \_ -> loop m s
  where
  loop m' s' =
    case runTrampoline (Mealy.stepMealy s' m') of
      Mealy.Halt -> ListT.nil
      Mealy.Emit (GenOut { value, state }) m'' ->
        ListT.prepend' value (defer \_ -> loop m'' state)

instance semigroupGenState :: Semigroup GenState where
  append (GenState a) (GenState b) =
    GenState { seed: lcgPerturb (toNumber (runSeed a.seed)) b.seed, size: b.size }

instance monoidGenState :: Monoid GenState where
  mempty = GenState { seed: mkSeed 85734629, size: 10 }

instance semigroupGenOut :: Semigroup a => Semigroup (GenOut a) where
  append (GenOut a) (GenOut b) =
    GenOut { state: a.state <> b.state, value: a.value <> b.value }

instance monoidGenOut :: Monoid a => Monoid (GenOut a) where
  mempty = GenOut { state: mempty, value: mempty }

instance functorGenOut :: Functor GenOut where
  map f (GenOut m) = GenOut { state: m.state, value: f m.value }

instance applyGenOut :: Apply GenOut where
  apply (GenOut f) (GenOut x) =
    GenOut { state: x.state, value: f.value x.value }

-- GenT instances
instance functorGenT :: Monad f => Functor (GenT f) where
  map f (GenT m) = GenT $ map f <$> m

instance applyGenT :: Monad f => Apply (GenT f) where
  apply f x = GenT $ apply <$> unGen f <*> unGen (updateSeedGen x)

instance applicativeGenT :: Monad f => Applicative (GenT f) where
  pure t = GenT $ arr (\s -> GenOut { state: updateSeedState s, value: t })

instance semigroupGenT :: Monad f => Semigroup (GenT f a) where
  append (GenT a) (GenT b) = GenT (a <> b)

instance monoidGenT :: Monad f => Monoid (GenT f a) where
  mempty = GenT mempty

instance bindGenT :: Monad f => Bind (GenT f) where
  bind (GenT m) f = GenT do
    a <- m
    unGen $ updateSeedGen (f (unGenOut a).value)

instance monadGenT :: Monad f => Monad (GenT f)

instance altGenT :: Monad f => Alt (GenT f) where
  alt (GenT x) (GenT y) = GenT $ x <|> y

instance plusGenT :: Monad f => Plus (GenT f) where
  empty = mempty

instance alternativeGenT :: Monad f => Alternative (GenT f)

instance monadZeroGenT :: Monad f => MonadZero (GenT f)

instance monadPlusGenT :: Monad f => MonadPlus (GenT f)

instance lazyGenT :: Monad f => CL.Lazy (GenT f a) where
  defer f = GenT $ CL.defer (unGen <<< f)

instance monadGenGenT :: Monad f => CMG.MonadGen (GenT f) where
  chooseInt = chooseInt
  chooseFloat = choose
  chooseBool = (_ < 0.5) <$> uniform
  resize f g = stateful \(GenState state) -> resize (f state.size) g
  sized = sized