{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE DeriveGeneric #-}
{-# OPTIONS_GHC -Wno-type-defaults -Wno-name-shadowing #-}
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as UM
import qualified Data.Vector.Unboxing as V
import qualified Data.Vector as B
import qualified Data.Vector.Generic as G
import System.Environment (getArgs, getProgName)
import Control.Monad
import Data.List
import qualified Poly as P

modulo :: Int
modulo = 17

newtype IntMod = IntMod Int deriving (Eq)

instance Show IntMod where
  show (IntMod n) = show n

instance V.Unboxable IntMod where
  type Rep IntMod = Int

instance Num IntMod where
  IntMod x + IntMod y = IntMod ((x + y) `rem` modulo)
  IntMod x - IntMod y = IntMod ((x - y) `mod` modulo)
  IntMod x * IntMod y = IntMod ((x * y) `rem` modulo)
  negate (IntMod x) = IntMod (negate x `mod` modulo)
  fromInteger n = IntMod (fromInteger (n `mod` fromIntegral modulo))
  abs = undefined; signum = undefined

instance Fractional IntMod where
  recip x = let x2 = x * x
                x5 = x2 * x2 * x
                x15 = x5 * x5 * x5
            in x15 -- x^15
  fromRational = undefined

trim :: U.Vector Int -> U.Vector Int
trim xs
  | U.null xs = U.empty
  | U.last xs == 0 = trim (U.init xs)
  | otherwise = xs

addMod, subMod, mulMod :: Int -> Int -> Int
addMod x y = (x + y) `rem` modulo
subMod x y = (x - y) `mod` modulo
mulMod x y = (x * y) `rem` modulo

sumMod :: [Int] -> Int
sumMod = foldl' addMod 0

negateMod :: Int -> Int
negateMod x = (negate x) `mod` modulo

recipMod :: Int -> Int
recipMod x = let x2 = x `mulMod` x
                 x5 = (x2 `mulMod` x2) `mulMod` x
                 x15 = (x5 `mulMod` x5) `mulMod` x5
             in x15

polyAdd :: U.Vector Int -> U.Vector Int -> U.Vector Int
polyAdd xs ys
  | n < m = U.create $ do
      v <- UM.new m
      forM_ [0..n-1] $ \i -> UM.write v i ((xs U.! i) `addMod` (ys U.! i))
      forM_ [n..m-1] $ \i -> UM.write v i (ys U.! i)
      return v
  | m < n = U.create $ do
      v <- UM.new n
      forM_ [0..m-1] $ \i -> UM.write v i ((xs U.! i) `addMod` (ys U.! i))
      forM_ [m..n-1] $ \i -> UM.write v i (xs U.! i)
      return v
  | otherwise = trim $ U.zipWith addMod xs ys
  where n = U.length xs
        m = U.length ys

polySub :: U.Vector Int -> U.Vector Int -> U.Vector Int
polySub xs ys
  | n < m = U.create $ do
      v <- UM.new m
      forM_ [0..n-1] $ \i -> UM.write v i ((xs U.! i) `subMod` (ys U.! i))
      forM_ [n..m-1] $ \i -> UM.write v i (negateMod (ys U.! i))
      return v
  | m < n = U.create $ do
      v <- UM.new n
      forM_ [0..m-1] $ \i -> UM.write v i ((xs U.! i) `subMod` (ys U.! i))
      forM_ [m..n-1] $ \i -> UM.write v i (xs U.! i)
      return v
  | otherwise = trim $ U.zipWith subMod xs ys
  where n = U.length xs
        m = U.length ys

polyMul :: U.Vector Int -> U.Vector Int -> U.Vector Int
polyMul xs ys
    | n == 0 || m == 0 = U.empty
    | otherwise = U.generate (n + m - 1) (\i -> sumMod [(xs U.! j) `mulMod` (ys U.! (i - j)) | j <- [max 0 (i - m + 1)..min i (n-1)]])
    where n = U.length xs
          m = U.length ys

scalePoly :: Int -> U.Vector Int -> U.Vector Int
scalePoly a xs
  | a == 0 = U.empty
  | otherwise = U.map (`mulMod` a) xs

toMonicPoly :: U.Vector Int -> U.Vector Int
toMonicPoly xs | U.null xs = U.empty
               | otherwise = U.map (`mulMod` recipMod (U.last xs)) xs

divModPoly :: U.Vector Int -> U.Vector Int -> (U.Vector Int, U.Vector Int)
divModPoly f g
  | U.null g = error "divMod: divide by zero"
  | U.length f < U.length g = (U.empty, f)
  | otherwise = loop U.empty (scalePoly (recipMod b) f)
  where
    g' = toMonicPoly g
    b = U.last g
    -- invariant: f == q * g + scale b p
    loop q !p | U.length p < U.length g = (q, scalePoly b p)
              | otherwise = let !q' = U.drop (U.length g - 1) p
                            in loop (q `polyAdd` q') (p `polySub` (q' `polyMul` g'))

modPoly :: U.Vector Int -> U.Vector Int -> U.Vector Int
modPoly f g = snd (divModPoly f g)

powModPoly :: U.Vector Int -> Int -> U.Vector Int -> U.Vector Int
powModPoly _ 0 _modulo = U.singleton 1
powModPoly f n modulo = loop (n-1) f f
  where loop 0 !_ !acc = acc
        loop 1 !m !acc = (m `polyMul` acc) `modPoly` modulo
        loop i !m !acc
          | even i = loop (i `quot` 2) ((m `polyMul` m) `modPoly` modulo) acc
          | otherwise = loop (i `quot` 2) ((m `polyMul` m) `modPoly` modulo) ((m `polyMul` acc) `modPoly` modulo)

powPoly :: U.Vector Int -> Int -> U.Vector Int
powPoly _ 0 = U.singleton 1
powPoly f n = loop (n-1) f f
  where loop 0 !_ !acc = acc
        loop 1 !m !acc = m `polyMul` acc
        loop i !m !acc
          | even i = loop (i `quot` 2) (m `polyMul` m) acc
          | otherwise = loop (i `quot` 2) (m `polyMul` m) (m `polyMul` acc)

-- Specialization for unboxing vectors + IntMod:
{-# SPECIALIZE P.addPoly :: P.Poly V.Vector IntMod -> P.Poly V.Vector IntMod -> P.Poly V.Vector IntMod #-}
{-# SPECIALIZE P.subPoly :: P.Poly V.Vector IntMod -> P.Poly V.Vector IntMod -> P.Poly V.Vector IntMod #-}
{-# SPECIALIZE P.mulPoly :: P.Poly V.Vector IntMod -> P.Poly V.Vector IntMod -> P.Poly V.Vector IntMod #-}
{-# SPECIALIZE P.divMod :: P.Poly V.Vector IntMod -> P.Poly V.Vector IntMod -> (P.Poly V.Vector IntMod, P.Poly V.Vector IntMod) #-}
{-# SPECIALIZE P.powMod :: P.Poly V.Vector IntMod -> Int -> P.Poly V.Vector IntMod -> P.Poly V.Vector IntMod #-}

-- Specialization for boxed vectors + IntMod:
{-# SPECIALIZE P.addPoly :: P.Poly B.Vector IntMod -> P.Poly B.Vector IntMod -> P.Poly B.Vector IntMod #-}
{-# SPECIALIZE P.subPoly :: P.Poly B.Vector IntMod -> P.Poly B.Vector IntMod -> P.Poly B.Vector IntMod #-}
{-# SPECIALIZE P.mulPoly :: P.Poly B.Vector IntMod -> P.Poly B.Vector IntMod -> P.Poly B.Vector IntMod #-}
{-# SPECIALIZE P.divMod :: P.Poly B.Vector IntMod -> P.Poly B.Vector IntMod -> (P.Poly B.Vector IntMod, P.Poly B.Vector IntMod) #-}
{-# SPECIALIZE P.powMod :: P.Poly B.Vector IntMod -> Int -> P.Poly B.Vector IntMod -> P.Poly B.Vector IntMod #-}

main :: IO ()
main = do
  args <- getArgs
  case args of
    "unboxed":_ -> do
      let f = U.fromList [modulo-1,modulo-1,modulo-1] `polyAdd` (powPoly (U.fromList [0,1]) 2000) {- U.fromList [if k==2000 then 1 else 0 | k<-[0..2000]] -}
      let g = powModPoly (U.fromList [0,1]) 1000000000 f
      print $ sumMod $ U.toList g -- should print '1'
    "unboxing":_ -> do
      let f = P.x^2000 - (P.x^2 + P.x + 1) :: P.Poly V.Vector IntMod
      let g = P.powMod P.x 1000000000 f
      print $ G.sum $ P.coeffAsc g -- should print '1'
    "boxed":_ -> do
      let f = P.x^2000 - (P.x^2 + P.x + 1) :: P.Poly B.Vector IntMod
      let g = P.powMod P.x 1000000000 f
      print $ G.sum $ P.coeffAsc g -- should print '1'
    _ -> do
      progName <- getProgName
      putStrLn $ progName ++ " (unboxed|unboxing|boxed)"
      putStr $ unlines ["This program computes the polynomial x^1000000000 mod (x^2000 - x^2 - x - 1)"
                       ,"and prints its value at x=1 in the finite field F_17 (or GF(17))."
                       ,"Run with '+RTS -t' to show memory usage."
                       ]