{-# LANGUAGE OverloadedStrings, DeriveGeneric #-}
module Linden.RuleDSL (
   RuleContext(..), RuleDSL, compileRules
 , addOption
 , doNothing, deleteMe, addBranch, addBranches
 , alreadyExists, noDup, nodeCount, maxNodes
 , curBranch, gaiaPathSyms, angleSum, angleToVertical
 , distance, distanceSq, curChildren
 , lightCones, lightHits, lightAng, lightIntensity
 , colorTemp
 , materialChildrenSyms, directSiblings
 , weighted, readBranchSym, attachLocation, allLights
 -- To suppress warnings about constructors used for TH naming
 , pX, pY, gRoot, gAttach
 , Point(..), Grips(..)
 ) where

import Data.Char
import Data.Maybe
import Data.Tree
import Data.Tree.Zipper
import Data.Bifunctor
import Data.Biapplicative
import Data.Random.RVar
import Data.Random.Distribution.Categorical
import Control.Monad.Identity
import Control.Monad.Reader
import Control.Monad.Writer.Strict
import Control.Monad.Supply (MonadSupply(..))
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as BSL
import Data.Aeson ((.:))
import qualified Data.Aeson as JS
import qualified Data.Text as T
import GHC.Generics
import System.FilePath
import Data.Digest.Pure.SHA (sha1, showDigest)
import qualified Data.Map.Lazy as Map
import Data.Hashable (hash)

import Debug.Trace

import Linden.Types

data Point =
    Point { pX :: Int, pY :: Int }
  deriving (Read, Show, Eq, Ord, Generic)

instance JS.ToJSON Point where
  toJSON (Point x y) = JS.object [("x", JS.toJSON x), ("y", JS.toJSON y)]

instance JS.FromJSON Point where
  parseJSON (JS.Object o) =
    Point <$> o .: "x" <*> o .: "y"
  parseJSON _ = mzero

data Grips =
  Grips {
    gRoot :: Point
  , gAttach :: [Point]
  }
  deriving (Read, Show, Eq, Ord, Generic)

instance JS.ToJSON Grips where
  toJSON (Grips r a) = JS.object [("attach", JS.toJSON a), ("root", JS.toJSON r)]

instance JS.FromJSON Grips where
  parseJSON (JS.Object o) =
    Grips <$> o .: "root" <*> o .: "attach"
  parseJSON _ = mzero

unPoint :: Point -> (Int, Int)
unPoint (Point x y) = (x, y)

readBranchSym :: [Int] -> FilePath -> ReaderT FilePath IO BranchSym
readBranchSym = readBranchSym'

readBranchSym' :: [Int] -> FilePath -> ReaderT FilePath IO BranchSym
readBranchSym' cls fp = do
  bp <- ask
  imfl <- liftIO $ BSL.readFile (bp </> "tiff" </> (addExtension fp "tif"))
  let imhash = take 16 . showDigest . sha1 $ imfl
  let gfl = bp </> "assets" </> (take 16 imhash++".json")
  gjson <- liftIO $ BS.readFile gfl
  case JS.eitherDecodeStrict' gjson of
    Left err -> fail $ "Failed to parse json of "++fp++" because "++show err
    Right (Grips (Point rX rY) as) ->
        return $ BranchSym (Just . T.pack $ addExtension fp "png")
                           (rX, rY) (map unPoint as) False False (map hash cls)

weighted :: [(Double, LEnv EditCommand)] -> RVar (LEnv EditCommand)
weighted = weightedCategorical

data RuleContext =
    RC {
      rcLights :: [Light]
    , rcAttachLocs :: BranchPossition
    , rcTreePos :: TreePos Full Branch
    }
  deriving (Read, Show, Eq)

type RuleDSL a = ReaderT RuleContext (WriterT [(Double, LEnv EditCommand)] RVar) a

compileRules :: RuleDSL () -> Rule
compileRules act bps lights tp =
    weighted =<< (execWriterT . runReaderT act $ RC lights bps tp)

curBranch :: RuleDSL Branch
curBranch = label <$> asks rcTreePos

curChildren :: RuleDSL Int
curChildren = (length . offspring) <$> asks rcTreePos

directSiblings :: RuleDSL Int
directSiblings =
  (fromMaybe 0 . fmap ((1-) . length . subForest . tree) . parent) <$> asks rcTreePos

-- | Finds if a command adds an edge that already exists.
alreadyExists :: LEnv EditCommand -> RuleDSL Bool
alreadyExists mc = do
  let c = runLEnv mc (Supply 0) -- We're an ID-independant comparison.
  case fst c of
    DeleteMe -> return False
    DoNothing -> return False
    AddChildren nsf -> do
      esf <- offspring <$> asks rcTreePos
      return . or $ [ ((bAngle nc) == (bAngle ec)) && ((bImg nc) == (bImg ec))
                    | nc <- rootLabel <$> nsf, ec <- fmap label esf]

nodeCount :: RuleDSL Int
nodeCount = (length . flatten . toTree) <$> asks rcTreePos

maxNodes :: Int -> RuleDSL () -> RuleDSL ()
maxNodes mx act = do
  nc <- nodeCount
  case nc of
    l | l >= mx -> doNothing 1
    _ -> act

isMaterial :: RuleDSL Bool
isMaterial = (not . isJust . bProxyFor . label) <$> asks rcTreePos

materialChildrenSyms :: RuleDSL [Branch]
materialChildrenSyms = do
    tp <- asks rcTreePos
    im <- isMaterial
    case im of
      -- We're material, so only our children can be immaterial.
      True -> return $ searchDown . tree $ tp
      False -> return $ maybe [] searchDown . findMaterial $ tp
  where
    findMaterial :: TreePos Full Branch -> Maybe (Tree Branch)
    findMaterial tp =
      case parent tp of
        Nothing -> Nothing
        Just tp' | (isJust . bProxyFor . label $ tp') -> findMaterial tp'
        Just tp' -> Just . tree $ tp'
    searchDown :: Tree Branch -> [Branch]
    searchDown (Node b sf) =
      case isJust . bProxyFor $ b of
        True -> concatMap searchDown sf
        False -> [b]

noDup :: RuleDSL a -> RuleDSL a
noDup act = do
    rc <- ask
    (a, w) <- lift . lift . runWriterT . runReaderT act $ rc
    forM_ w $ \(weight, ec) -> do
      d <- alreadyExists ec
      unless d $ tell [(weight, ec)]
    return a

addOption :: Double -> LEnv EditCommand -> RuleDSL ()
addOption w ec = tell [(w, ec)]

doNothing :: Double -> RuleDSL ()
doNothing w = addOption w . return $ DoNothing

deleteMe :: Double -> RuleDSL ()
deleteMe w = do
  im <- (not . bImmutable) <$> curBranch
  when im $ addOption w . return $ DeleteMe

explodeSym :: BranchSym -> Angle -> Scale -> Double -> LEnv (Tree Branch)
explodeSym (BranchSym img r [attach] im rgd cls) a s wv = do
  i <- supply
  return $ Node (Branch i img r attach a s im rgd cls Nothing wv) []
explodeSym (BranchSym img r attachs im rgd cls) a s wv = do
  i <- supply
  sf <- forM attachs $ \at -> do
     i' <- supply
     return $ Node (Branch i' Nothing (0, 0) at 0 1 im True cls (Just i) 0) []
  return $ Node (Branch i img r (0, 0) a s im rgd cls Nothing wv) sf

addBranch :: Double -> BranchSym -> Angle -> Scale -> Double -> RuleDSL ()
addBranch w bs a s wv =
  addOption w ((AddChildren . pure) <$> explodeSym bs a s wv)

addBranches :: Double -> [(BranchSym, Angle, Scale, Double)] -> RuleDSL ()
addBranches w as = addOption w $ do
  cs <- forM as $ \(bs, a, s, wv) -> do
         explodeSym bs a s wv
  return . AddChildren $ cs

gaiaPathSyms :: RuleDSL [Branch]
gaiaPathSyms = (map (\(_, p, _) -> p) . parents) <$> asks rcTreePos

angleToVertical :: RuleDSL Angle
angleToVertical = negate <$> angleSum

angleSum :: RuleDSL Angle
angleSum = (absAngle . sum . map bAngle) <$> gaiaPathSyms

absAngle :: Angle -> Angle
absAngle d = d + (fromIntegral $ (ceiling ((negate d) / 360.0)::Integer)) * 360.0

lightAng :: (X, Y) -> Light -> Angle
lightAng p l =
  let (dx, dy) = ((-), (-)) <<*>> p <<*>> (bimap fromIntegral fromIntegral $ lPos l)
  in absAngle $ (180/pi) * atan2 dy dx

lightHits :: (X, Y) -> Light -> Bool
lightHits p l =
    let ang = lightAng p l
    in (abs $ ang-(absAngle $ lPointAngle l)) <= (lBeamAngle l)/2

attachLocation :: RuleDSL (X, Y)
attachLocation =
    (fst . fromJust) <$> (Map.lookup <$> (bId <$> curBranch) <*> (asks rcAttachLocs))

allLights :: RuleDSL [Light]
allLights = asks rcLights

lightCones :: RuleDSL [Light]
lightCones = do
  p <- attachLocation
  filter (lightHits p) <$> asks rcLights

distance :: (X, Y) -> (Int, Int) -> Double
distance a b = sqrt . distanceSq a $ b

distanceSq :: (X, Y) -> (Int, Int) -> Double
distanceSq (x0, y0) (x1', y1') =
    ((x0-x1)^(2::Int))+((y0-y1)^(2::Int))
  where
    x1 = fromIntegral x1'
    y1 = fromIntegral y1'

colorTemp :: RuleDSL Double
colorTemp = do
  lcs <- lightCones
  return $ if (null lcs)
           then -10
           else (sum . map lTemp $ lcs) / (fromIntegral $ length lcs)

lightIntensity :: RuleDSL Double
lightIntensity = do
  p <- attachLocation
  hitting <- lightCones
  let intenses = map (\l -> (10^(4::Int))*(360/(lBeamAngle l)) / (distanceSq p (lPos l))) hitting
  return . sum $ 0:intenses