{-# LANGUAGE OverloadedStrings #-}
-- | Rule-based linguistic feature encoder (POS, lemma, dependency).
-- No ML training. Deterministic, auditable.
-- Author: Ahmad Ali Parr · Trust: Bel Esprit D'Accord Irrevocable Trust
module LinguisticEncoder
  ( encodeLinguisticFeatures
  , combineWithBERT
  , ruleBasedPOS
  , ruleBasedLemma
  ) where

import qualified Data.Vector      as V
import           Data.Vector      (Vector)
import           Data.List        (isSuffixOf)
import           NeurologicalRegex (Token(..), Embedding, bertEmbeddings)

type LinguisticEmbedding = Vector Double

ruleBasedPOS :: String -> String
ruleBasedPOS w
  | "ed"  `isSuffixOf` w = "VERB"
  | "ing" `isSuffixOf` w = "VERB"
  | "ly"  `isSuffixOf` w = "ADV"
  | "s"   `isSuffixOf` w && length w > 1 = "NOUN"
  | otherwise             = "NOUN"

ruleBasedLemma :: String -> String
ruleBasedLemma w
  | "ing" `isSuffixOf` w && length w > 3 = take (length w - 3) w
  | "ed"  `isSuffixOf` w && length w > 2 = take (length w - 2) w
  | "s"   `isSuffixOf` w && length w > 1 = init w
  | otherwise = w

ruleBasedDepRel :: [String] -> Int -> String
ruleBasedDepRel ws i
  | i == 0    = "ROOT"
  | i > 0 && ruleBasedPOS (ws !! (i-1)) == "VERB"
             && ruleBasedPOS (ws !!  i)  == "NOUN" = "nsubj"
  | otherwise = "dep"

encodeLinguisticFeatures :: [String] -> [LinguisticEmbedding]
encodeLinguisticFeatures ws =
  [ featureToEmb (ruleBasedPOS w) (ruleBasedDepRel ws i) (ruleBasedLemma w)
  | (i, w) <- zip [0..] ws ]
  where
    featureToEmb pos dep lemma =
      V.generate 16 $ \j ->
        let base = fromIntegral (sum (map fromEnum (pos++dep++lemma)) `mod` 1024)
        in (base + fromIntegral j * 17.0) / 2048.0

combineWithBERT :: [Token] -> [LinguisticEmbedding] -> [Embedding]
combineWithBERT toks lingEmbs =
  let bertEmbs = map (\t -> bertEmbeddings 8 t) toks
      padLing  = map (\e -> V.generate 8 (\i -> V.unsafeIndex e (i `mod` V.length e))) lingEmbs
      pairs    = zip bertEmbs (padLing ++ repeat (V.replicate 8 0))
  in map (\(b,l) -> b V.++ l) pairs