module Generators where

import qualified Solution as Sol
import qualified Types_Efficient as SolT

import Control.Monad (foldM)
import Test.QuickCheck as QC

instance Arbitrary SolT.Polarity where
  arbitrary = oneof $ map return [SolT.Pos, SolT.Neg]

-- mV: maximum variable index
genEName :: Int -> Gen SolT.EName
genEName mV = do
  n <- choose (1, mV)
  return $ show n

instance Arbitrary SolT.ELiteral where
  arbitrary = sized $ \n -> do
    p <- arbitrary
    v <- genEName n
    return $ SolT.Literal p v
  shrink (SolT.Literal p n) = map (SolT.Literal p . show) [1..(read n - 1)]

-- fP: frequency of positive literals
-- fN: frequency of negative literals
genELiteral :: Int -> Int -> Int -> Gen SolT.ELiteral
genELiteral mV fP fN = do
  v <- genEName mV
  frequency [(fP, return $ SolT.Literal SolT.Pos v), (fN, return $ SolT.Literal SolT.Neg v)]

-- lC: maximum size of clauses
genEClause :: Int -> Int -> Int -> Int -> Gen SolT.EClause
genEClause mV fP fN lC = resize lC $ listOf $ genELiteral mV fP fN

-- same as genClause but generates only non-empty clauses
genEClause1 :: Int -> Int -> Int -> Int -> Gen SolT.EClause
genEClause1 mV fP fN lC = resize lC $ listOf1 $ genELiteral mV fP fN

-- lC: maximum number of clauses
genEConjForm :: Int -> Int -> Int -> Int -> Int -> Gen SolT.EConjForm
genEConjForm mV fP fN lC l = resize l $ listOf $ genEClause mV fP fN lC

-- same as genConjForm but generates only non-empty conjunctions
genEConjForm1 :: Int -> Int -> Int -> Int -> Int -> Gen SolT.EConjForm
genEConjForm1 mV fP fN lC l = resize l $ listOf1 $ genEClause1 mV fP fN lC


-- genEConjForm1 + at least one positive literal per clause
genEConjForm1Pos :: Int -> Int -> Int -> Int -> Int -> Gen SolT.EConjForm
genEConjForm1Pos mV fP fN lC l = do
  cs <- genEConjForm1 mV fP fN lC l
  mapM addPos cs
  where
    addPos c = do
      l <- genELiteral mV 0 1
      i <- choose (0,length c-1)
      return $ insertAt i l c
    insertAt n a xs = let (ys,zs) = splitAt n xs in ys ++ [a] ++ zs

genEConjForm1Neg :: Int -> Int -> Int -> Int -> Int -> Gen SolT.EConjForm
genEConjForm1Neg mV fP fN lC l = do
  cs <- genEConjForm1 mV fP fN lC l
  mapM addPos cs
  where
    addPos c = do
      l <- genELiteral mV 1 0
      i <- choose (0,length c-1)
      return $ insertAt i l c

insertAt :: Int -> a -> [a] -> [a]
insertAt n a xs = let (ys,zs) = splitAt n xs in ys ++ [a] ++ zs

-- strictly negative below lN and positive above lP
genELiteralPure :: Int -> Int -> Int -> Int -> Int -> Gen SolT.ELiteral
genELiteralPure mV fP fN bN bP = do
  v <- genEName mV
  let i = read v
  if i < bN then return $ SolT.Literal SolT.Neg v
  else if i > bP then return $ SolT.Literal SolT.Pos v
  else frequency [(fP, return $ SolT.Literal SolT.Pos v), (fN, return $ SolT.Literal SolT.Neg v)]

-- non-empty + strictly negative below lN and positive above lP
genEClause1Pure :: Int -> Int -> Int -> Int -> Int -> Int -> Gen SolT.EClause
genEClause1Pure mV fP fN lC bN bP = resize lC $ listOf1 (genELiteralPure mV fP fN bN bP)

-- non-empty + strictly negative below lN and positive above lP
genEConjForm1Pure :: Int -> Int -> Int -> Int -> Int -> Int -> Int -> Gen SolT.EConjForm
genEConjForm1Pure mV fP fN lC l bN bP = resize l $ listOf1 (genEClause1Pure mV fP fN lC bN bP)

-- non-empty + olr applies
genEConjForm1Olr :: Int -> Int -> Gen SolT.EConjForm
genEConjForm1Olr l m = run 1 []
  where
    -- 2020: I cried a bit - why didn't I just shuffle??
    -- Update 2021: whelp, I'm not gonna rewrite this.. just re-use it from last year
    run i cs | i > l = do
      (c,_,_) <- foldM buildClause ([],l-1,False) [1..l-1] 
      let nl = SolT.Literal (if l `mod` m == 0 then SolT.Pos else SolT.Neg) $ show l
      j <- choose (0,l-1)
      let c' = insertAt j nl c
      j' <- choose (0,l)
      return $ insertAt j' c' cs
    run i cs = do
      (c,_,_) <- foldM buildClause ([],i,True) [1..i]
      let lc = length cs
      j <- choose (0,lc-1)
      run (i+1) (insertAt j c cs)
    buildClause (acc,b,normal) i = do
      let p1 = if normal && i == b then SolT.Neg else SolT.Pos
      let p2 = if normal && i == b then SolT.Pos else SolT.Neg
      let l = SolT.Literal (if i `mod` m == 0 then p1 else p2) $ show i
      j <- choose (0,i-1)
      return (insertAt j l acc,b,normal)

genEConjFormPigeonhole :: Int -> Gen SolT.EConjForm
genEConjFormPigeonhole s = do
  let c  = [ [SolT.Literal SolT.Pos (show $ pair i j) | j<-[1..s-1]] | i<-[1..s]]
  let c' = [ [SolT.Literal SolT.Neg (show $ pair i k), SolT.Literal SolT.Neg (show $ pair j k)] | k<-[1..s-1],i<-[1..s],j<-[1..s],i/=j]
  shuffle $ c ++ c'
  where
    pair a b = ((a + b)*(a + b + 1) + b) `div` 2

genEKeyConjFormSaturated :: Int -> Int -> Int -> Int -> Int -> Gen (SolT.EProof, SolT.EKeyConjForm)
genEKeyConjFormSaturated mV fP fN lC l = do
  cs <- genEConjForm1 mV fP fN lC l 
  let (p, kcs) = Sol.resolution cs
  return (SolT.toEProof p, SolT.toEKeyConjForm kcs)

genEKeyConjFormSaturatedSatisfiable :: Int -> Int -> Int -> Int -> Int -> Gen SolT.EKeyConjForm
genEKeyConjFormSaturatedSatisfiable mV fP fN lC l = genEKeyConjFormSaturated mV fP fN lC l `suchThatMap` satisfiableKeyConjForm
  where
    satisfiableKeyConjForm (SolT.Model _, kcs) = Just kcs
    satisfiableKeyConjForm _ = Nothing