module Data.Nimber where

module Data.Nimber
  ( Nimber (..),
    floorLog,
    sqr,
    pow,
    artinSchreierRoot,
  )
where

twoPowers :: (Num a, Bits a, Integral b) => a -> [b]
twoPowers 0 = []
twoPowers m =
  (if m `testBit` 0 then (0 :) else id) . fmap (+ 1) $ twoPowers (m `shiftR` 1)

-- | Index of highest-order set bit, or -1 if there are none.

floorLog 0 = error "Logarithm of 0"

floorLog 0 = -1

mult' :: Int -> Nimber -> Nimber -> Nimber
mult' _ 0 _ = 0
mult' _ _ 0 = 0
mult' _ 1 b = b
mult' _ a 1 = a
mult' m a b =
  let semiD = bit (bit m - 1) -- semimultiple of D
      s1 = a `shiftR` bit m -- a = a1D+a2
      s2 = a .^. (s1 `shiftL` bit m)
      t1 = b `shiftR` bit m -- b = b1D+b2
      t2 = b .^. (t1 `shiftL` bit m)
      c = mult' (m - 1) s2 t2
   in ((mult' (m - 1) (s1 + s2) (t1 + t2) - c) `shiftL` bit m) + mult' (m - 1) (mult' (m - 1) s1 t1) semiD + c

  0 * _ = 0
  _ * 0 = 0
  1 * b = b
  a * 1 = a

        mult' _ 0 _ = 0
        mult' _ _ 0 = 0
        mult' _ 1 t = t
        mult' _ s 1 = s
        mult' k s t =
          let semiD = bit (bit k - 1) -- semimultiple of D
              s1 = s `shiftR` bit k -- a = a1D+a2
              s2 = s .^. (s1 `shiftL` bit k)
              t1 = t `shiftR` bit k -- b = b1D+b2
              t2 = t .^. (t1 `shiftL` bit k)
              c = mult' (k-1) s2 t2
           in ((mult' (k-1) (s1 + s2) (t1 + t2) - c) `shiftL` bit k) + mult' (k-1) (mult' (k-1) s1 t1) semiD + c

sqr' :: Int -> Nimber -> Nimber
sqr' _ 0 = 0
sqr' _ 1 = 1
sqr' m n =
  let a = n `shiftR` bit m -- n = aD+b
      aD = a `shiftL` bit m
      b = n .^. aD
      semiD = bit (bit m - 1) -- semimultiple of D
      sqra = sqr' (m - 1) a
   in sqra `shiftL` bit m + mult' (m - 1) sqra semiD + sqr' (m - 1) b

sqr 0 = 0
sqr 1 = 1

      sqr' _ 0 = 0
      sqr' _ 1 = 1
      sqr' k l =
        let a = l `shiftR` bit k -- n = aD+b
            aD = a `shiftL` bit k
            b = l .^. aD
            semiD = bit (bit k - 1) -- semimultiple of D
            sqra = sqr' (k - 1) a
         in sqra `shiftL` bit k + sqra * semiD + sqr' (k - 1) b

  | otherwise = product . fmap snd . filter (testBit n . fst) . zip [0 ..] . take (1 + floorLog (n + 1)) $ iterate sqr x

  | otherwise =
      let m = floorLog @Int $ floorLog x
       in product . fmap snd . filter (testBit n . fst) . zip [0 ..] . take (1 + floorLog (n + 1)) $ iterate (sqr' m) x

  recip 0 = error "Divide by zero"
  recip 1 = 1

        recip' _ 0 = error "Divide by zero"

           in (aD + a + b) * recip' (k - 1) (semiD * sqr a + b * (a + b))

           in mult' k (aD + a + b) $ recip' (k - 1) (mult' (k - 1) semiD (sqr' (k - 1) a) + mult' (k - 1) b (a + b))

  sqrt 0 = 0
  sqrt 1 = 1

           in sqrta `shiftL` bit k + sqrta * sqrt semiD + sqrt' (k - 1) b

           in sqrta `shiftL` bit k + mult' (k - 1) sqrta (sqrt' (k - 1) semiD) + sqrt' (k - 1) b

      squares = iterate sqr n

      squares = iterate (sqr' m) n

      t4k = sum $ take (bit (m' - 1)) quarts -- trace of the Artin-Screier root of n
   in if t4k == 1

      semiTrace = sum $ take (bit (m' - 1)) quarts -- trace of the Artin-Screier root of n
   in if semiTrace == 1

                pure $ (squares !! (shiftL i 1 - 1 .^. bit (m' - 1))) * (squares !! (shiftL j 1 - 2))

                pure $ mult' m' (squares !! (shiftL i 1 - 1 .^. bit (m' - 1))) (squares !! (shiftL j 1 - 2))