{-# LANGUAGE TupleSections #-}

module Lib where

import qualified Data.List as L
import qualified Data.Map as M
import Data.Maybe (fromJust)
import qualified Data.Set as S
import Parser
  ( Addend (..),
    Matrix (..),
    QN (Neg, Qua),
    QNMFormula (..),
    Quantifier (Exists, ForAll),
  )

mget :: (Ord a) => a -> M.Map a b -> b
mget a m = fromJust $ M.lookup a m

automatizem :: M.Map Addend Int -> Matrix -> DFA [Int] Int
automatizem dir (MEq as bs) =
  eq (length dir) (map (`mget` dir) as) (map (`mget` dir) bs)
automatizem dir (MLe a b) = le (mget a dir) (mget b dir)
automatizem dir (MLt a b) = lt (mget a dir) (mget b dir)
automatizem dir (MAnd x y) = conj (automatizem dir x) (automatizem dir y)
automatizem dir (MOr x y) = disj (automatizem dir x) (automatizem dir y)
automatizem dir (MNot x) = compl (automatizem dir x)

automatize :: M.Map Addend Int -> QNMFormula -> NFA [Int] Int
automatize dir (QNMFormula [] m) =
  let dfa = automatizem dir m
      mdfa = minDistinguishability dfa (length dir)
   in minReachability (length dir) $ nondeterminize mdfa
automatize dir (QNMFormula (Neg : quas) m) =
  let nfa = automatize dir (QNMFormula quas m)
   in minReachability (length dir) (ncompl nfa)
  where
    ncompl = nondeterminize . compl . determinize
automatize dir (QNMFormula (Qua (Exists, v) : quas) m) =
  let nfa = automatize dir (QNMFormula quas m)
   in case M.lookup (Var v) dir of
        Nothing -> nfa
        (Just i) -> minReachability (length dir) $ existentialize i nfa
automatize _ (QNMFormula (Qua (ForAll, _) : _) _) =
  error "should be universal-free"

assign :: S.Set Addend -> M.Map Addend Int
assign xs = M.fromList (zip (S.toList xs) [0 ..])

collect :: Matrix -> S.Set Addend
collect (MEq as bs) = S.fromList (as ++ bs)
collect (MLe a b) = S.fromList [a, b]
collect (MLt a b) = S.fromList [a, b]
collect (MAnd x y) = S.union (collect x) (collect y)
collect (MOr x y) = S.union (collect x) (collect y)
collect (MNot x) = collect x

literals :: M.Map Addend Int -> [[Int]]
literals m =
  let addends = L.sortOn snd (M.toList m)
      reversedLiterals = map (\(a, _) -> addendReverseBinary a) addends
      padTo = L.maximum (map length reversedLiterals)
      paddedReversed =
        map (\x -> x ++ replicate (padTo - length x) 0) reversedLiterals
      padded = map reverse paddedReversed
   in L.transpose padded

addendReverseBinary :: Addend -> [Int]
addendReverseBinary (Var _) = []
addendReverseBinary (Con n) = reverseBinary n

reverseBinary :: Int -> [Int]
reverseBinary 0 = []
reverseBinary n = mod n 2 : reverseBinary (div n 2)

eval :: QNMFormula -> Bool
eval f@(QNMFormula _ m) =
  let dir = (assign $ collect m)
      nfa = automatize dir f
      input = literals (assign $ collect m)
   in runNFA nfa input

data State a = S a | Double (State a, State a) | Multi [State a] deriving (Eq, Ord, Show)

data DFA c a = DFA [State a] (State a) [State a] (State a -> c -> State a)

data NFA c a = NFA [State a] [State a] [State a] (State a -> c -> [State a])

runDFA :: (Ord a) => DFA c a -> [c] -> Bool
runDFA (DFA _ start accepts f) cs = foldl f start cs `elem` accepts

minReachability :: (Ord a) => Int -> NFA [Int] a -> NFA [Int] a
minReachability n nfa@(NFA _ starts accepts f) = NFA states' starts' accepts' f
  where
    states' = closure nfa (chars n) starts
    starts' = starts `L.intersect` states'
    accepts' = accepts `L.intersect` states'

minDistinguishability :: (Ord a, Show a) => DFA [Int] a -> Int -> DFA [Int] a
minDistinguishability dfa@(DFA _ start accepts f) n =
  let partition = hopcroftPartitionRefinement dfa n
      states' = map head partition
      start' = head $ head $ filter (L.elem start) partition
      accepts' = L.intersect accepts states'
      f' s c = let s' = f s c in head $ head $ filter (L.elem s') partition
   in DFA states' start' accepts' f'

hopcroftPartitionRefinement ::
  (Ord a, Show a) => DFA [Int] a -> Int -> [[State a]]
hopcroftPartitionRefinement (DFA states _ accepts f) n =
  let p0 = [states L.\\ accepts, accepts] in recurse p0
  where
    sigma = chars n
    normalizePartition partition = S.fromList (map S.fromList partition)
    findInPartition xs x = fromJust $ L.find (L.elem x) xs
    equalUnder partition x y =
      findInPartition partition x == findInPartition partition y
    indistinguishable partition x y =
      all (\c -> equalUnder partition (f x c) (f y c)) sigma
    partialRefine _ pk [] = pk
    partialRefine pkPrev [] (x : xs) = partialRefine pkPrev [[x]] xs
    partialRefine pkPrev pk (x : xs) =
      case filter (\s -> indistinguishable pkPrev (head s) x) pk of
        [] -> partialRefine pkPrev ([x] : pk) xs
        [match] -> partialRefine pkPrev ((x : match) : L.delete match pk) xs
        _ -> error "should be unreachable"
    recurse pkPrev =
      let pk = concatMap (partialRefine pkPrev []) pkPrev
       in if normalizePartition pk == normalizePartition pkPrev
            then pk
            else recurse pk

runNFA :: (Ord a) => NFA c a -> [c] -> Bool
runNFA (NFA _ starts accepts f) cs =
  foldl (\xs c -> L.nub $ concatMap (`f` c) xs) starts cs
    `L.intersect` accepts
    /= []

reversal :: (Ord a) => NFA c a -> NFA c a
reversal (NFA states starts accepts f) = NFA states accepts starts f'
  where
    f' s c = filter (\state -> s `elem` f state c) states

eq :: Int -> [Int] -> [Int] -> DFA [Int] Int
eq n is js = determinize $ minReachability n $ reversal $ nondeterminize dfa
  where
    states = S <$> [- (length js - 1) .. length is - 1 + 1]
    start = S 0
    accepts = [S 0]
    rejector = last states
    f :: State Int -> [Int] -> State Int
    f carrystate@(S carry) c =
      if carrystate == rejector
        then rejector
        else
          let si = sum (map (c !!) is)
              sj = sum (map (c !!) js)
              parityok = mod (carry + si) 2 == mod sj 2
              newcarry = div (carry + si - sj) 2
           in if parityok then S newcarry else rejector
    dfa = DFA states start accepts f

le :: Int -> Int -> DFA [Int] Int
le = less LessEqual

lt :: Int -> Int -> DFA [Int] Int
lt = less LessThan

data LessType = LessEqual | LessThan deriving (Eq, Ord, Show)

less :: LessType -> Int -> Int -> DFA [Int] Int
less op i j = DFA [S 0, S 1, S 2] (S 0) accepts f
  where
    accepts = if op == LessEqual then [S 0, S 2] else [S 2]
    f s c = case (s, (c !! i, c !! j)) of
      (S 0, (1, 1)) -> S 0
      (S 0, (0, 0)) -> S 0
      (S 0, (1, 0)) -> S 1
      (S 0, _) -> S 2
      (S 1, _) -> S 1
      (_, _) -> S 2

prod :: [a] -> [b] -> [(a, b)]
prod _ [] = []
prod [] _ = []
prod (x : xs) ys = fmap (x,) ys ++ prod xs ys

data JunctionType = Conj | Disj deriving (Eq, Ord, Show)

junction :: (Ord a) => JunctionType -> DFA c a -> DFA c a -> DFA c a
junction jt (DFA states1 start1 accepts1 f1) (DFA states2 start2 accepts2 f2) =
  DFA states' start' accepts' f'
  where
    newStates = prod states1 states2
    states' = Double <$> newStates
    start' = Double (start1, start2)
    accepts' =
      if jt == Conj
        then Double <$> prod accepts1 accepts2
        else
          Double
            <$> filter (\(s, t) -> s `elem` accepts1 || t `elem` accepts2) newStates
    f' (Double (s, t)) c = Double (f1 s c, f2 t c)

conj :: (Ord a) => DFA c a -> DFA c a -> DFA c a
conj = junction Conj

disj :: (Ord a) => DFA c a -> DFA c a -> DFA c a
disj = junction Disj

compl :: (Ord a) => DFA c a -> DFA c a
compl (DFA states start accepts f) = DFA states start (states L.\\ accepts) f

nondeterminize :: (Ord a) => DFA c a -> NFA c a
nondeterminize (DFA states start accepts f) = NFA states [start] accepts f'
  where
    f' s c = [f s c]

change :: [a] -> Int -> a -> [a]
change xs idx b = take idx xs ++ [b] ++ drop (idx + 1) xs

closure :: (Ord a) => NFA c a -> [c] -> [State a] -> [State a]
closure nfa@(NFA _ _ _ f) cs initstates =
  let new = concatMap (\state -> concatMap (f state) cs) initstates
   in if L.nub new L.\\ L.nub initstates /= []
        then closure nfa cs (L.nub $ new ++ initstates)
        else L.nub initstates

existentialize :: (Ord a) => Int -> NFA [Int] a -> NFA [Int] a
existentialize idx nfa@(NFA states starts accepts f) =
  NFA
    states
    starts'
    accepts
    f'
  where
    zeroer = replicate 50 0
    oneer = change zeroer idx 1
    starts' = closure nfa [zeroer, oneer] starts
    f' s c = f s (change c idx 0) ++ f s (change c idx 1)

powerset :: [a] -> [[a]]
powerset [] = [[]]
powerset (x : xs) = let rest = powerset xs in map (x :) rest ++ rest

determinize :: (Ord a) => NFA c a -> DFA c a
determinize (NFA states start accepts f) = DFA states' start' accepts' f'
  where
    newStates = map L.sort $ powerset states
    states' = Multi <$> newStates
    start' = Multi $ L.sort start
    accepts' =
      Multi <$> filter (\state' -> state' `L.intersect` accepts /= []) newStates
    f' (Multi s) c = Multi $ L.nub $ L.sort $ concatMap (`f` c) s

chars :: Int -> [[Int]]
chars 0 = [[]]
chars n = let r = chars (n - 1) in map (1 :) r ++ map (0 :) r