content: {-# LANGUAGE DeriveAnyClass #-}
  {-# LANGUAGE DeriveGeneric #-}
  
  -- | Rank-1 Constraint
  module Keelung.Constraint.R1C (R1C (..), satisfy, freeVars) where
  
  import Control.DeepSeq (NFData)
  import Data.Field.Galois (GaloisField)
  import Data.IntMap (IntMap)
  import Data.IntMap qualified as IntMap
  import Data.IntSet (IntSet)
  import Data.Serialize (Serialize)
  import GHC.Generics (Generic)
  import Keelung.Data.Polynomial (Poly)
  import Keelung.Data.Polynomial qualified as Poly
  
  --------------------------------------------------------------------------------
  
  -- | A Rank-1 Constraint is a relation between 3 polynomials
  --      Ax * Bx = Cx
  data R1C n = R1C (Either n (Poly n)) (Either n (Poly n)) (Either n (Poly n))
    deriving (Generic, NFData)
  
  instance (Eq n, Num n) => Eq (R1C n) where
    R1C a b c == R1C a' b' c' =
      -- if the RHS are the same
      (c == c' && (a == a' && b == b' || a == b' && b == a'))
        -- if the RHS are the negation of each other
        || ( negate' c == c'
               && ( negate' a == a' && b == b'
                      || a == a' && negate' b == b'
                      || negate' a == b' && b == a'
                      || a == b' && negate' b == a'
                  )
           )
      where
        negate' (Left n) = Left (-n)
        negate' (Right p) = Right (Poly.negate p)
  
  instance Functor R1C where
    fmap f (R1C a b c) = R1C (fmapE a) (fmapE b) (fmapE c)
      where
        fmapE (Left x) = Left (f x)
        fmapE (Right xs) = Right (fmap f xs)
  
  instance Serialize n => Serialize (R1C n)
  
  instance (Num n, Eq n, Ord n) => Ord (R1C n) where
    compare x@(R1C a b c) y@(R1C e f g) = case (isRank1 x, isRank1 y) of
      (True, False) -> LT
      (False, True) -> GT
      (True, True) ->
        -- both are of rank 1
        -- the one with a constant term on the RHS is considered smaller
        case (c, g) of
          (Left _, Right _) -> LT
          (Right _, Left _) -> GT
          (Left c0, Left c1) -> compare c0 c1 -- compare the constant terms
          (Right v0, Right v1) -> compare v0 v1 -- compare the polynomials
      (False, False) ->
        -- both are of rank 2
        compare (a, b, c) (e, f, g)
  
  instance (Show n, Ord n, Eq n, Num n) => Show (R1C n) where
    show (R1C aX bX cX) = case (aX, bX, cX) of
      (Left 0, _, _) -> "0 = " ++ showVec cX
      (_, Left 0, _) -> "0 = " ++ showVec cX
      (Left 1, _, _) -> showVec bX ++ " = " ++ showVec cX
      (_, Left 1, _) -> showVec aX ++ " = " ++ showVec cX
      (_, _, _) -> showVecWithParen aX ++ " * " ++ showVecWithParen bX ++ " = " ++ showVec cX
      where
        showVec :: (Show n, Ord n, Eq n, Num n) => Either n (Poly n) -> String
        showVec (Left c) = show c
        showVec (Right xs) = show xs
  
        -- wrap the string with parenthesis if it has more than 1 term
        showVecWithParen :: (Show n, Ord n, Eq n, Num n) => Either n (Poly n) -> String
        showVecWithParen (Left c) = showVec (Left c) -- no parenthesis
        showVecWithParen (Right xs) =
          let termNumber =
                IntMap.size (Poly.coeffs xs)
                  + if Poly.constant xs == 0
                    then 0
                    else 1
           in if termNumber < 2
                then showVec (Right xs)
                else "(" ++ showVec (Right xs) ++ ")"
  
  -- | See if a R1C is satified by a given assignment
  satisfy :: GaloisField a => R1C a -> IntMap a -> Bool
  satisfy constraint assignment
    | R1C aV bV cV <- constraint =
        evaluate aV assignment * evaluate bV assignment == evaluate cV assignment
    where
      evaluate :: GaloisField a => Either a (Poly a) -> IntMap a -> a
      evaluate (Left x) _ = x
      evaluate (Right p) w = Poly.evaluate p w
  
  -- | Free variables in a R1C
  freeVars :: R1C n -> IntSet
  freeVars (R1C a b c) = freeVarsE a <> freeVarsE b <> freeVarsE c
    where
      freeVarsE (Left _) = mempty
      freeVarsE (Right p) = Poly.vars p
  
  -- | An R1C is of rank 1 if either side of the multiplication is a constant
  isRank1 :: R1C n -> Bool
  isRank1 (R1C a b _) = isConstant a || isConstant b
    where
      isConstant (Left _) = True
      isConstant (Right _) = False