{-# LANGUAGE OverloadedStrings, DeriveGeneric, DeriveFunctor, DeriveAnyClass, TemplateHaskell, RankNTypes, FlexibleInstances, MultiParamTypeClasses, ScopedTypeVariables, ImpredicativeTypes #-}
module Linden.Types (
   BranchSym(..), Branch(..), X, Y, Angle, Scale
 , offspring
 , LState(..), Light(..), BranchPossition
 , step, EditCommand(..)
 , LEnvT, LEnv, runLEnvT, runLEnv, Supply(..)
 , Rule, AxiomSource
 , PrettyJSTree(..)
 , GardenStore(..), GardenSave, GardenCAS, GardenExists
 , UserCommand(..)
 , filterTree
 , Nullable(..)
 , makePossitions
 ) where

import GHC.Generics
import qualified Data.Text as T
import qualified Data.Aeson as JS
import Data.Aeson.TH
import Data.Bifunctor
import Data.Biapplicative
import Data.Word
import Data.Tree
import Data.Maybe
import Data.Time
import Data.UUID (UUID)
import Data.Tree.Zipper
import Data.Random.RVar
import Control.Monad.Identity
import Control.Monad.State.Lazy
import Control.Monad.Supply (MonadSupply(..))
import qualified Data.Map.Lazy as Map
import Control.DeepSeq

import Linden.TH

data BranchSym =
  BranchSym {
      bsImg :: Maybe T.Text
    , bsRoot :: !(Int, Int)
    , bsAttach :: ![(Int, Int)]
    , bsImmutable :: !Bool
    , bsRigid :: !Bool
    , bsClasses :: ![Int]
    }
  deriving (Read, Show, Eq, Ord, Generic, NFData)

type Angle = Double
type Scale = Double

data Branch =
  Branch {
      bId :: {-# UNPACK #-} !Word32
    , bImg :: Maybe T.Text
    , bRoot :: !(Int, Int)
    , bAttach :: !(Int, Int)
    , bAngle :: {-# UNPACK #-} !Angle
    , bScale ::  {-# UNPACK #-}!Scale
    , bImmutable :: !Bool
    , bRigid :: !Bool
    , bClass :: ![Int]
    , bProxyFor :: Maybe Word32
    , bWaviness :: {-# UNPACK #-}!Double
    }
  deriving (Read, Show, Eq, Ord, Generic, NFData)

$(deriveJSON defaultOptions{fieldLabelModifier = makeCamel 1} ''Branch)

data Light =
    Light {
      lId ::  {-# UNPACK #-}!Word32
    , lPos :: !(Int, Int)
    , lPointAngle ::  {-# UNPACK #-}!Angle
    , lBeamAngle ::  {-# UNPACK #-}!Angle
    , lTemp :: {-# UNPACK #-}!Double
    }
  deriving (Read, Show, Eq, Ord, Generic, NFData)

$(deriveJSON defaultOptions{fieldLabelModifier = makeCamel 1} ''Light)

newtype PrettyJSTree a =
    PrettyJSTree (Tree a)
  deriving (Read, Show, Eq, Generic, NFData)

instance JS.ToJSON a => JS.ToJSON (PrettyJSTree a) where
    toJSON (PrettyJSTree (Node a c)) =
        JS.object [ ("node", JS.toJSON a)
                  , ("children", JS.toJSON . map PrettyJSTree $ c)]

instance JS.FromJSON a => JS.FromJSON (PrettyJSTree a) where
    parseJSON (JS.Object v) = do
      n <- v JS..: "node"
      c <- v JS..: "children"
      return . PrettyJSTree . Node n . map (\(PrettyJSTree t) -> t) $ c
    parseJSON _ = mzero

data LState =
  LS {
    lsGarden :: {-# UNPACK #-} !UUID
  , lsLastUpdate :: {-# UNPACK #-} !UTCTime
  , lsSupply ::  {-# UNPACK #-}!Word32
  , lsLights :: ![Light]
  , lsTree :: Maybe (PrettyJSTree Branch)
  }
  deriving (Read, Show, Eq, Generic, NFData)

$(deriveJSON defaultOptions{fieldLabelModifier = makeCamel 2} ''LState)

type GardenSave = LState -> IO ()
-- Returns Nothing if the UUID doesn't exist in the store.
-- Action returns Nothing if there is no update.
type GardenCAS = forall a. UUID -> (LState -> RVar (Maybe LState, a)) -> IO (Maybe a)
type GardenExists = UUID -> IO Bool

data GardenStore =
    GardenStore GardenSave GardenCAS GardenExists

data Nullable a =
    IsNull
  | Exists a
  deriving (Read, Show, Eq, Ord, Generic, NFData)

instance JS.ToJSON a => JS.ToJSON (Nullable a) where
    toJSON (Exists a) = JS.toJSON a
    toJSON IsNull = JS.Null

instance JS.FromJSON a => JS.FromJSON (Nullable a) where
    parseJSON JS.Null = return IsNull
    parseJSON v = Exists <$> JS.parseJSON v

data UserCommand =
    UserDel { ucGarden :: UUID, ucEditKey :: String, ugNode :: Word32 }
  | UserEditLight { ucGarden :: UUID, ucEditKey :: String
                  , ucLight :: Nullable Word32, ucPos :: (Int, Int)
                  , ucPointAngle :: Angle, ucBeamAngle :: Angle
                  , ucTemp :: Double }
  deriving (Read, Show, Eq, Ord, Generic, NFData)

$(deriveJSON defaultOptions {
                   fieldLabelModifier = makeCamel 2
                 , constructorTagModifier = drop 4} ''UserCommand)

filterTree :: (Branch -> Bool) -> Tree Branch -> Maybe (Tree Branch)
filterTree f (Node b sf) =
    if f b
    then Just (Node b (mapMaybe (filterTree f) sf))
    else Nothing

data EditCommand = DoNothing | AddChildren [Tree Branch] | DeleteMe

newtype Supply = Supply Word32

type X = Double
type Y = Double

newtype LEnvT m a =
    LEnvT { runLEnvT :: Supply -> m (a, Supply) }

type LEnv = LEnvT Identity

runLEnv :: LEnv a -> Supply -> (a, Supply)
runLEnv a = runIdentity . runLEnvT a

instance (Functor m) => Functor (LEnvT m) where
    fmap f m = LEnvT $ \s ->
               fmap (\ ~(a, s') -> (f a, s')) $ runLEnvT m s

instance (Functor m, Monad m) => Applicative (LEnvT m) where
    pure a = LEnvT $ \s -> return (a, s)
    LEnvT mf <*> LEnvT mx = LEnvT $ \s -> do
        ~(f, s') <- mf s
        ~(x, s'') <- mx s'
        return (f x, s'')
    {-# INLINE (<*>) #-}

instance (Monad m) => Monad (LEnvT m) where
    m >>= k  = LEnvT $ \s -> do
        ~(a, s') <- runLEnvT m s
        runLEnvT (k a) s'
    fail str = LEnvT $ \_ -> fail str

instance (Monad m) => MonadState Supply (LEnvT m) where
    state f = LEnvT (return . f)

instance MonadTrans LEnvT where
    lift m = LEnvT $ \s -> do
               a <- m
               return (a, s)

instance Monad m => MonadSupply Word32 (LEnvT m) where
    supply = state (\(Supply s) -> (s, Supply $ s+1))
    peek = (\(Supply s) -> s) <$> get
    exhausted = (\(Supply s) -> s == maxBound) <$> get

type BranchPossition = Map.Map Word32 ((X, Y), Angle)

makePossitions :: Tree Branch -> BranchPossition
makePossitions (Node b' sf') =
    go (Map.singleton (bId b') (attachOffset b', bAngle b')) (bId b') sf'
  where
    doublefy :: (Int, Int) -> (Double, Double)
    doublefy = bimap fromIntegral fromIntegral
    attachOffset :: Branch -> (X, Y)
    attachOffset b =
        ((-), (-)) <<*>> (doublefy $ bAttach b) <<*>> (doublefy $ bRoot b)
    go :: BranchPossition -> Word32 -> Forest Branch -> BranchPossition
    go memo pid cl = foldl (addChild pid) memo cl
    addChild :: Word32 -> BranchPossition -> Tree Branch -> BranchPossition
    addChild pid memo (Node c sf) =
      go (Map.insert (bId c) (attachPos pid memo c) memo) (bId c) sf
    attachPos :: Word32 -> BranchPossition -> Branch -> ((X, Y), Angle)
    attachPos pid memo c =
      let ((parX, parY), pAng) = memo Map.! pid
          ang = pAng + (bAngle c)
          (offX, offY) = attachOffset c
          (xosq, yosq) = ((^(2::Int)), (^(2::Int))) <<*>> attachOffset c
          cLen = (bScale c) * sqrt(xosq+yosq)
          rang = ang*(pi/180) + atan2 offY offX
          cX = cLen*(cos rang)
          cY = cLen*(sin rang)
      in ((parX+cX, parY+cY), ang)

type Rule = BranchPossition -> [Light] -> TreePos Full Branch -> RVar (LEnv EditCommand)

-- Axiom generator
type AxiomSource = UUID -> UTCTime -> RVar LState

offspring :: TreePos Full Branch -> [TreePos Full Branch]
offspring tp =
    go . children $ tp
  where
    go :: TreePos Empty Branch -> [TreePos Full Branch]
    go cf =
      case nextTree cf of
        Nothing -> []
        Just c -> c:(go . nextSpace $ c)

step :: Rule -> UTCTime -> LState -> RVar LState
step _ _ l@(LS _ _ _ _ Nothing) = return l
step rule n (LS i _ s0 lights (Just (PrettyJSTree t0))) =
    fmap (\(t1, Supply s1) -> force $ LS i n s1 lights . fmap PrettyJSTree $ t1) .
    (`runLEnvT` (Supply s0)) . applyProductions .
    fromTree $ (t0::Tree Branch)
  where
    bps = makePossitions t0
    applyProductions :: TreePos Full Branch -> LEnvT RVar (Maybe (Tree Branch))
    applyProductions tp = do
      let sym = label tp
      cuped <- fmap catMaybes . forM (offspring tp) $ \c -> applyProductions c
      com <- (lift . rule bps lights $ tp) >>= (\a -> state (runLEnv a))
      case com of
        DoNothing -> return . Just $ Node sym cuped
        AddChildren ts -> return . Just $ Node sym (ts++cuped)
        DeleteMe -> return Nothing