content: module Stan.Data where
  
  import Data.List
  import qualified Data.Vector as V
  import qualified Data.Map.Strict as Map
  import Data.Random.Source.PureMT
  
  class Dump1 a where
    dump1 :: a -> ([String], [Int])
  
  instance Dump1 Double where
      dump1 x = ([show x], [])
  instance Dump1 Int where
      dump1 x = ([show x], [])
  instance Dump1 Bool where
      dump1 True = (["1"], [])
      dump1 False = (["0"], [])
  
  --dump1c xs = concat $ "c(" : intersperse "," (map dump1 xs) ++[")"]
  
  instance Dump1 a => Dump1 [a] where
      dump1 xs = let xss = map dump1 xs
                 in (concat $ transpose $ map fst xss, length xss:(snd $ head xss) )
  
  instance Dump1 a => Dump1 (V.Vector a) where
      dump1 xs = dump1 $ V.toList xs
  
  dumpAll :: ([String], [Int]) -> String
  dumpAll ([s], []) = s
  dumpAll (ss, [_]) = concat $ "c(" : intersperse "," ss ++[")"]
  dumpAll (ss, ns) =
      concat $ "structure(c(" : intersperse "," ss ++["), .Dim = c("]
                               ++intersperse "," (map show ns)++["))"]
  
  dumpAs :: ToStanData a => String -> a -> StanData
  dumpAs nm x = Map.singleton nm $ toStanData x
  
  (<~) :: ToStanData a => String -> a -> StanData
  (<~) = dumpAs
  
  type StanData = StanEnv
  
  inGroupsOf :: Int -> [a] -> [[a]]
  inGroupsOf _ [] = []
  inGroupsOf n xs = let (a,b) = splitAt n xs in a : inGroupsOf n b
  
  data StanValue = VDouble Double
                 | VInt Int
                 | VArray (V.Vector StanValue)
                 | VSamples (V.Vector StanValue)
                 | VSeed PureMT
                 deriving (Show)
  
  instance Dump1 StanValue where
      dump1 (VDouble x) = dump1 x
      dump1 (VInt x) = dump1 x
      dump1 (VArray x) = dump1 x
      dump1 (VSamples _) = error "Dump1 VSamples ¯\\_(ツ)_/¯"
  
  type StanEnv = Map.Map String StanValue
  
  dumpEnv :: StanEnv -> [String]
  dumpEnv = map f . Map.toList where
      f (nm,sv) = nm++"<-"++(dumpAll $ dump1 sv)
  
  class ToStanData a where
      toStanData :: a -> StanValue
  
  instance ToStanData Double where
      toStanData x = VDouble x
  instance ToStanData Int where
      toStanData x = VInt x
  instance ToStanData Bool where
      toStanData True = VInt 1
      toStanData False = VInt 0
  instance ToStanData a => ToStanData [a] where
      toStanData xs = VArray $ V.fromList $ map toStanData xs
  instance ToStanData a => ToStanData (V.Vector a) where
      toStanData xs = VArray $ V.map toStanData xs
  
  --TODO class Reify a where reify :: String -> StanEnv -> Maybe a
  
  unDouble :: StanValue -> Double
  unDouble (VDouble x) = x
  unDouble (VInt x) = realToFrac x
  unDouble v = error $ "unDouble: "++show v
  
  unDoubles :: StanValue -> [Double]
  unDoubles (VArray vs) = V.toList $ V.map unDouble vs
  
  unPairDoubles :: StanValue -> [(Double,Double)]
  unPairDoubles (VArray vs)
     = V.toList $ V.map (\(VArray vs) -> (unDouble $ vs V.! 0, unDouble $ vs V.! 1)) vs
  
  sumValue :: StanValue -> StanValue -> StanValue
  sumValue (VDouble x) (VDouble y) = VDouble (x+y)
  sumValue (VArray v1) (VArray v2) = VArray $ V.zipWith sumValue v1 v2
  
  scaleValue :: Double -> StanValue -> StanValue
  scaleValue x (VDouble y) = VDouble (y*x)
  scaleValue x (VArray v) = VArray $ V.map (scaleValue x) v