{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}

import Math.NumberTheory.Primes.Testing
import Math.NumberTheory.Moduli
import Options.Applicative
import System.Random
import System.IO
import System.Exit
import Data.Bits
import Data.Function
import Data.List
import Data.Maybe
import Test.QuickCheck
import Debug.Trace
import Control.Monad
import Statistics.Sample
import Statistics.Quantile
import Control.Exception (evaluate)
import Data.Bifunctor
import GHC.Stack
import Text.Read
import Text.Printf
import Data.List.Split
import qualified Data.Vector as V
import qualified Data.Sequence as S
import qualified Data.IntMap as IM
import qualified Data.IntSet as IS
import Data.Monoid

-- Some convenient names for some parameters
type BitLength = Int
type Width = Int

-- http://stackoverflow.com/a/13301611/946226
rndPrime :: BitLength -> Gen Integer
rndPrime 1 = error "Cannot create distinct p and q for 1 bit"
rndPrime bits = do
    x <- fmap (.|. 1) $ choose (2^(bits - 1), 2^bits - 1)
    if isPrime x then return x else rndPrime bits

data RSAParams = RSAParams { bits :: Int, p, q, n, e, dp, dq, kp, kq :: Integer } deriving Show

instance Arbitrary RSAParams where
    arbitrary = sized (\s -> genRSAParams (min 1024 (max s 2)))

genRSAParams :: BitLength -> Gen RSAParams
genRSAParams bits = do
    p <- rndPrime bits
    q <- rndPrime bits
    let n = p*q
        e = 65537
        mdp = invertMod e (p-1)
        mdq = invertMod e (q-1)
    case (mdp, mdq) of
        (Just dp, Just dq) -> do
            let kp = (e*dp-1) `div` (p-1)
                kq = (e*dq-1) `div` (q-1)
            return (RSAParams {..})
        _ -> genRSAParams bits -- retry

data Step = S | M deriving (Eq, Ord, Show) -- square or multiply
type Leak = V.Vector Step

leftToRight :: BitLength -> Width -> Integer -> Leak
leftToRight bits w s | w <= 0 = error "width must be positive"
leftToRight bits w s = V.fromList $ go (bits-1)
   where
    go i | i < 0         = []
    go i | s `testBit` i = replicate (slideLength-1) S ++ M : go (i - slideLength)
      where
        Just slideLength = find (\l -> s `testBit` (i-l+1)) $ [w,w-1..1]
    go i | otherwise     = S : go (i-1)

-- Bits learned
leftToRightGMPBits :: BitLength -> Width -> Integer -> Integer
leftToRightGMPBits bits w s | w <= 0 = error "width must be positive"
leftToRightGMPBits bits w s = go (bits -1) 0
   where
    go i !acc | i < 0         = acc
    go i !acc | s `testBit` i = go (i - slideLength) (acc + learned)
      where
        Just slideLength = find (\l -> s `testBit` (i-l+1)) $ [w,w-1..1]
        learned | slideLength == 1 = 1
                | otherwise        = 2
    go i ! acc | otherwise    =  go (i-1) (acc + 1)

rightToLeft :: BitLength -> Width -> Integer -> Leak
rightToLeft bits w s | w <= 0 = error "width must be positive"
rightToLeft bits w s = V.fromList $ go 0 []
   where
    go i acc | i == bits     = acc
    go i acc | s `testBit` i = go (i+slideLength) (replicate (slideLength-1) S ++ M : acc)
      where slideLength = min w (bits - i)
    go i acc | otherwise     = go (i+1) (S : acc)

-- Bits learned
rightToLeftBits :: BitLength -> Width -> Integer -> Integer
rightToLeftBits bits w s | w <= 0 = error "width must be positive"
rightToLeftBits bits w s = go 0 0
   where
    go i !acc | i == bits     = acc
    go i !acc | s `testBit` i = go (i+slideLength) (1+acc)
      where slideLength = min w (bits - i)
    go i !acc | otherwise     = go (i+1) (1+acc)

-- A datatype to track known bits
data BitMask = BM { unknownBits :: !Integer, knownBits :: !Integer }
    deriving Eq

showBitMask bits bm = map go [(bits-1),(bits-2)..0]
  where go i | not (bm `isKnownBit` i) = '_'
             | bm `getKnownBit` i      = '1'
             | otherwise               = '0'

instance Show BitMask where
    show bm = showBitMask bits bm
      where bits = head [i | i <- [0..], 2^i > knownBits bm && 2^i > unknownBits bm]

isKnownBit :: BitMask -> Int -> Bool
bm `isKnownBit` i = not (unknownBits bm `testBit` i)

getKnownBit :: BitMask -> Int -> Bool
bm `getKnownBit` i = knownBits bm `testBit` i

possibleBits :: BitMask -> Int -> [Bool]
possibleBits bm i
    | bm `isKnownBit` i = [bm `getKnownBit` i]
    | otherwise         = [False, True]

matches :: Integer -> BitMask -> Bool
matches p bm = (p .|. unknownBits bm) == (knownBits bm .|. unknownBits bm)

isKnown :: BitMask -> Maybe Integer
isKnown bm | unknownBits bm == 0 = Just (knownBits bm)
           | otherwise           = Nothing

writeBit :: Int -> Bool -> Integer -> Integer
writeBit i True  p = p `setBit` i
writeBit i False p = p `clearBit` i

writeKnownBit :: Int -> Bool -> BitMask -> BitMask
writeKnownBit i b bm = BM (unknownBits bm `clearBit` i) (writeBit i b (knownBits bm))

nothingKnown :: BitLength -> BitMask
nothingKnown bits = BM (2^bits-1) 0

anyOddNumber :: BitLength -> BitMask
anyOddNumber bits = writeKnownBit 0 True (nothingKnown bits)

fromKnownNumber :: Integer -> BitMask
fromKnownNumber i = BM 0 i

-- can be optimized, I guess
allMatching :: BitMask -> [Integer]
allMatching bm = map sum $ mapM go [(bits-1),(bits-2)..0]
  where go i | not (bm `isKnownBit` i) = [2^i,0]
             | bm `getKnownBit` i      = [2^i]
             | otherwise               = [0]
        bits = head [i | i <- [0..], 2^i > knownBits bm && 2^i > unknownBits bm]

maskFromMatching :: [Integer] -> BitMask
maskFromMatching (m:ms) = foldl' go (fromKnownNumber m) ms
  where
    go :: BitMask -> Integer -> BitMask
    go bm i = BM (unknownBits bm .|. disagree) (knownBits bm `clear` disagree)
      where disagree = (unknownBits bm .|. knownBits bm) `xor` (unknownBits bm .|. i)

x `clear` y = complement (complement x .|. y)

countKnownBits :: BitLength -> BitMask -> Int
countKnownBits bits bm = bits - popCount (unknownBits bm)

-- A breadth first search limiting the count by n
bfSearch :: Int -> a -> (a -> [a]) -> (a -> Bool) -> Maybe Int
bfSearch limit s0 next done = go 0 (S.singleton s0)
  where
    go n _ | n >= limit = Nothing
    go n todo = case S.viewl todo of
        S.EmptyL -> Nothing
        s S.:< todo' | done s    -> Just n
                     | otherwise -> go (n+1) (todo' S.>< S.fromList (next s))

dfSearch :: Int -> a -> (a -> [a]) -> (a -> Bool) -> Maybe Int
dfSearch limit s0 next done = go 0 [s0]
  where
    go n _ | n >= limit = Nothing
    go n [] = Nothing
    go n (s:todo') | done s    = Just n
                   | otherwise = go (n+1) (next s ++ todo')

-- Currently unused
linearWorkList :: forall a. Eq a =>
    IS.IntSet -> IM.IntMap a -> (Int -> IM.IntMap a -> a) -> IM.IntMap a
linearWorkList todo s0 recalc = go todo s0
  where
    go :: IS.IntSet -> IM.IntMap a -> IM.IntMap a
    go todo s | IS.null todo         = s
    go todo s | oldValue == newValue = go todo' s
              | otherwise            = go newTodo newS
      where
        (k, todo') = IS.deleteFindMin todo
        oldValue = s IM.! k
        newValue = recalc k s
        newS = IM.insert k newValue s
        newTodo = (if (k-1) `IM.member` s then IS.insert (k-1) else id) $
                  (if (k+1) `IM.member` s then IS.insert (k+1) else id) $ todo'

type WindowLengthRanges = [(Int, (Int,Int))]

possibleWindowLengths :: BitLength -> Width -> Leak -> WindowLengthRanges
possibleWindowLengths bits w leak = IM.elems result
  where
    clamp width = max 1 (min w width)

    -- positions of the window end (from the back, i.e. bits)
    windowEnds = IM.fromList $ zip [0..] $ map ((bits - 1) - ) $
                 V.toList $ V.findIndices (== M) leak

    windows = IM.keysSet windowEnds
    maxKey = IS.findMax windows

    zipIM = IM.intersectionWith (,)

    windowMax = IM.mapWithKey (\k _ -> w_max k) windowEnds
    windowMin = IM.mapWithKey (\k _ -> w_min k) windowEnds
    result = windowEnds `zipIM` (windowMin `zipIM` windowMax)

    msbIsOne = False

    -- The first window spans precisely from the MSB to the first M
    firstWinLength =  clamp (bits - pos)
      where pos = windowEnds IM.! 0

    -- minimum length. computed from right to left
    w_min k | k == 0, msbIsOne = firstWinLength
            | k == maxKey      = 1
            | otherwise        = clamp (j + w_min_j + w -i)
        where
            i       = windowEnds IM.! k
            j       = windowEnds IM.! (k+1)
            w_min_j = windowMin IM.! (k+1)

    -- maximum length. computed from left to right
    w_max k | k == 0    = if msbIsOne then firstWinLength else w
            | otherwise = clamp (i + w_max_i - w -j)
        where
            i       = windowEnds IM.! (k-1)
            j       = windowEnds IM.! k
            w_max_i = windowMax IM.! (k-1)

countPossibilities :: BitLength -> Width -> WindowLengthRanges -> Integer
countPossibilities bits w lengths = go (bits-1) 0
  where
    vec = V.fromList lengths

    pos k = fst (vec V.! k)
    windowLengthRange k = [minWin..maxWin]
      where
        (_, (minWin, maxWin)) =  vec V.! k

    go soonestOne k | k == V.length vec = 2^0
    go soonestOne k = sum
        [ countMemo k winLength
        | winLength <- windowLengthRange k
        , let firstOneOfWindow = pos k + winLength -1
        , firstOneOfWindow <= soonestOne
        ]

    -- simple memoization
    countTable = V.fromList
        [ count k winLength | k <- [0..V.length vec -1] , winLength <- [1..w] ]
    countMemo k winLength = countTable V.!  (k * w + (winLength - 1))

    count k winLength = 2^unknown * go nextPossibleOne (k+1)
      where
        firstOneOfWindow = pos k + winLength -1
        unknown = max 0 (winLength - 2)
        nextPossibleOne = firstOneOfWindow - w

testCountPossibilities :: BitLength -> Width -> Integer -> Property
testCountPossibilities bits w s = fromIntegral (length candidates) === poss
  where
    leak = leftToRight bits w s
    windowLengths = possibleWindowLengths bits w leak
    poss = countPossibilities bits w windowLengths

    mask = getHardBits bits w windowLengths
    candidates = realCandidates bits w leak mask

getWindowStarts :: WindowLengthRanges -> IS.IntSet
getWindowStarts list = IS.fromList
    [ pos + maxWin - 1
    | (pos, (minWin, maxWin)) <- list
    ]

getHardBits :: BitLength -> Width -> WindowLengthRanges -> BitMask
getHardBits bits w final_list = go 0 (nothingKnown bits)
      where
        go i !bm | i == bits = bm
                 | i `IS.member` ones    = go (i+1) (writeKnownBit i True bm)
                 | i `IS.member` zeros   = go (i+1) (writeKnownBit i False bm)
                 | i `IS.member` inRange = go (i+1) bm
                 | otherwise             = go (i+1) (writeKnownBit i False bm)

        onesOn =     IS.fromList [ pos
                                 | (pos,(minWin, maxWin)) <- final_list ]
        onesStart =  IS.fromList [ pos + minWin - 1
                                 | (pos,(minWin, maxWin)) <- final_list
                                 , minWin == maxWin ]
        zerosAfter = IS.fromList [ j
                                 | (pos,(minWin, maxWin)) <- final_list
                                 , j <- [pos + maxWin - w .. pos - 1] ]
        inRange =    IS.fromList [ j
                                 | (pos,(minWin, maxWin)) <- final_list
                                 , j <- [pos .. pos + maxWin - 1] ]

        ones  = IS.unions [ onesOn, onesStart ]
        zeros = IS.unions [ zerosAfter ]

countHardBits :: BitLength -> Width -> WindowLengthRanges -> Int
countHardBits bits w final_list = bits - sum (map uncertain final_list)
  where uncertain (_, (minWin, maxWin)) | minWin == maxWin = max 0 (maxWin - 2)
                                        | otherwise        = max 0 (maxWin - 1)

testCountHardBits bits w final_list =
    countKnownBits bits (getHardBits bits w final_list) == countHardBits bits w final_list

knownBitsFromLeak :: BitLength -> Width -> Leak -> BitMask
knownBitsFromLeak bits w leak = getHardBits bits w $ possibleWindowLengths bits w leak

testKnownBitsFromLeak :: BitLength -> Positive Integer -> Bool
testKnownBitsFromLeak bits (Positive p) = p `matches` bm
  where
    leak = leftToRight bits 4 p
    bm = knownBitsFromLeak bits 4 leak

testKnownBitsFromLeak2 :: RSAParams -> Bool
testKnownBitsFromLeak2 params
    = dp params `matches` bitmask_dp && dq params `matches` bitmask_dq
  where
    leak_dp    = leftToRight       (bits params) 4 (dp params)
    bitmask_dp = knownBitsFromLeak (bits params) 4 leak_dp
    leak_dq    = leftToRight       (bits params) 4 (dq params)
    bitmask_dq = knownBitsFromLeak (bits params) 4 leak_dq

type Candidate = (Int, Integer, Integer, Integer, Integer)

-- | An oracle is a function that takes a possible suffix (given with length
-- and value) and indicates if this suffix is possible.
type Oracle = Int -> Integer -> Bool

selectPortion :: Double -> Int -> [Int]
selectPortion h to = go 0 0
  where
    go chosen i | i > to    = []
                | f < h     = i : go (chosen + 1) (i+1)
                | otherwise =     go  chosen      (i+1)
      where f = fromIntegral chosen / fromIntegral i

selectRandomPortion :: StdGen -> Double -> Int -> [Int]
selectRandomPortion rng h to = go rng 0
  where
    go rng i | i > to    = []
             | x < h     = i : go rng' (i+1)
             | otherwise =     go rng' (i+1)
      where (x, rng') = random rng

mkHardBitOracle :: BitLength -> Double -> Integer -> Oracle
mkHardBitOracle bits h input = oracle
  where
    rng = mkStdGen (fromIntegral (input+3))
    knownBits = IS.fromList $ selectRandomPortion rng h bits
    oracle i guess
        | i `IS.member` knownBits
        = guess `testBit` i == input `testBit` i
        | otherwise
        = True

mkBlockSumOracle :: Int -> Integer -> Oracle
mkBlockSumOracle bs input = oracle
  where
    oracle i guess | j `mod` bs == 0 = guess ==? input
                   | otherwise       = True
      where
        j = i + 1
        select k = (k `div` 2^(j-bs)) `mod` 2^bs
        (==?) = ((==) `on` popCount) `on` select

testOracle :: (Integer -> Oracle) -> Positive Integer -> Positive Int ->  Bool
testOracle mkO (Positive input) (Positive i) = o i (input `mod` 2^(i+1))
  where o = mkO input

mkLeftToRightOracle bits w input = oracle
  where
   s = leftToRight bits w input

   windowLengths = possibleWindowLengths bits w s
   bitmask = getHardBits bits w windowLengths
   windowStarts = getWindowStarts windowLengths

   shouldCheck i = (i-1) `IS.member` windowStarts

   oracle i guess
        | (guess `testBit` i) `notElem` possibleBits bitmask i
        = False
        | shouldCheck i
        , let leak' = leftToRight i w guess
        , V.drop (bits - i) s /= leak'
        = False
        | otherwise
        = True

-- Recovers the bits from the leaks for p and q
lsbRecovery :: Int -> RSAParams -> (Oracle, Oracle) -> Maybe Int
lsbRecovery limit params (op, oq) = dfSearch limit s0 next done
  where
    s0 :: Candidate
    s0 = ( 0 -- last bit guessed/set; this and all later bits are known
         , 1 -- known suffix of dp
         , 1 -- known suffix of dq
         , 1 -- known suffix of p
         , 1 -- known suffix of q
         )

    next :: Candidate -> [Candidate]
    next (i, dp, dq, p, q)
        | j == (bits params) = []
        | otherwise =
            filter prune
            [ (j, dp', dq', p', q')
            | (p',dp') <- filter (pruneP j) $ (,) <$> nextP <*> nextDP
            , (q',dq') <- filter (pruneQ j) $ (,) <$> nextQ <*> nextDQ
            ]
      where
        j = i+1
        nextDP = [dp, dp `setBit` j]
        nextDQ = [dq, dq `setBit` j]
        nextP  = [p, p `setBit` j]
        nextQ  = [q, q `setBit` j]

    pruneP i (p, dp) = and
        [ (e params*dp - 1) `mod` modulus == (kp params * (p-1)) `mod` modulus
        , op i dp
        ]
      where modulus = 2^(i+1)

    pruneQ i (q, dq) = and
        [ (e params*dq - 1) `mod` modulus == (kq params * (q-1)) `mod` modulus
        , oq i dq
        ]
      where modulus = 2^(i+1)

    prune :: Candidate -> Bool
    prune (i, dp, dq, p, q)
        = and [ (p * q) `mod` modulus == n params `mod` modulus
              ]
      where modulus = 2^(i+1)

    done :: Candidate -> Bool
    done (i, dp', dq', p', q')
        = i + 1 == bits params
        && p' * q' == n params
        {-
        && Just dp' == invertMod (e params) (p'-1)
        && Just dq' == invertMod (e params) (q'-1)
        && (dp' == dp params || traceShow ('P',p',q', dp', dq', params) True)
        && (dq' == dq params || traceShow ('Q',p',q', dp', dq', params) True)
        && (dp' `matches` bitmask_dp || traceShow (bits params, dp params) True)
        && (dq' `matches` bitmask_dq || traceShow (bits params, dq params) True)
        -}

runLSBRecovery :: Int -> Width -> RSAParams -> Maybe Int
runLSBRecovery limit w params =
    let op = mkLeftToRightOracle (bits params) w (dp params)
        oq = mkLeftToRightOracle (bits params) w (dq params)
    in lsbRecovery limit params (op, oq)

main :: IO ()
main = join . customExecParser (prefs showHelpOnEmpty) $
  info (helper <*> parser)
  (  fullDesc
  <> header "Sliding right experiment"
  <> progDesc "Slides stuff around"
  <> footer "Some parameter understand range syntax in the form 1,2,3,4 or 4,8,..,2048"
  )
  where
    parser :: Parser (IO ())
    parser =
      hsubparser (mconcat
        [ command "run_recover"        recoverCommand
        , command "left_to_right"      printLTRCommand
        , command "test_hard_rules"    testHardRulesCommand
        , command "measure_leakage"    measureLeakageCommand
        , command "collision_entropy"  measureCollisionEntropyCommand
        , command "hard_soft_scatter"  hardSoftScatterCommand
        , command "leak_preinput_size" leakPreinputCommand
        , command "entropy_scatter"    entropyScatterCommand
        ])

    bitsArg = (`div` 2) <$> -- internally, we want the key length of p and q
      option auto
      (  long "bits"
         <> metavar "BITS"
         <> help "Bit length of p and q combined (usually 1024 or 2048)"
         <> value 1024
         <> showDefault
      )

    bitsRangeArg = (map (`div` 2)) <$> -- internally, we want the key length of p and q
      option (eitherReader parseRange)
      (  long "bits"
         <> metavar "BITS,.."
         <> help "Bit length of p and q combined. Range syntax allowed"
         <> value [1024]
         <> showDefaultWith showRange
      )

    knownBitsRangeArg =
      option (eitherReader parseRange)
      (  long "known-bits"
         <> metavar "PROB,.."
         <> help "Proportion of known bits in range 0...1000. Range syntax allowed"
         <> value [0,100..1000]
         <> showDefaultWith showRange
      )

    blockSumRangeArg =
      option (eitherReader parseRange)
      (  long "sum-length"
         <> metavar "INT,.."
         <> help "Lenght of blocks of which sum is known"
         <> value [1,2,3,4]
         <> showDefaultWith showRange
      )

    widthArg = option auto
         (  long "width"
         <> metavar "WIDTH"
         <> help "Width of the window (e.g. 4 or 5)"
         <> value 4
         <> showDefault
         )

    widthRangeArg = option (eitherReader parseRange)
         (  long "width"
         <> metavar "WIDTH,.."
         <> help "Width of the window (e.g. 4 or 5). Range syntax allowed."
         <> value [4]
         <> showDefaultWith showRange
         )

    iterationsArg = option auto
         (  long "iterations"
         <> metavar "N"
         <> help "Number of iterations"
         <> value 100
         <> showDefault
         )

    limitArg = option auto
         (  long "limit"
         <> metavar "N"
         <> help "Search limit"
         <> value 100000
         <> showDefault
         )

    tsvArg = switch
         (  long "tsv"
         <> help "Tab separated output (no variance)"
         )

    printLTRCommand = info
        (pure printLTR <*> bitsArg <*> widthArg<*> bitSequenceArg)
        (progDesc "Run left-to-right and print sequence")
      where bitSequenceArg = argument auto (metavar "input_sequence")

    testHardRulesCommand = info
        (pure testHardRules <*> bitsArg <*> widthArg<*> verboseArg)
        (progDesc "Verify hard bit calculation (no output = all good)")
      where verboseArg = switch (long "verbose" <> help "Show all tested data")

    recoverCommand = info
        (pure runRecover <*> bitsArg <*> widthArg<*> iterationsArg <*> limitArg )
        (progDesc "Run key recovery, and calculate search tree size")

    measureLeakageCommand = info
        (pure measureLeakage <*> bitsRangeArg <*> widthRangeArg<*> tsvArg <*> iterationsArg)
        (progDesc "Calculate the Shannon entropy rate and hard bits")

    measureCollisionEntropyCommand = info
        (pure measureCollisionEntropy <*> bitsRangeArg <*> widthRangeArg<*> tsvArg <*> iterationsArg)
        (progDesc "Calculate the collision entropy rate")

    hardSoftScatterCommand = info
        (pure hardSoftScatter <*> bitsRangeArg <*> widthRangeArg<*> iterationsArg)
        (progDesc "Draw a hard/soft bit scatter plot")

    leakPreinputCommand = info
        (pure leakPreinput <*> widthArg)
        (progDesc "Number of input sequences for given square-and-mulitply sequence")

    entropyScatterCommand = info
        (pure runScatter <*> bitsArg <*> knownBitsRangeArg <*> widthRangeArg <*> blockSumRangeArg <*>iterationsArg <*> limitArg )
        (progDesc "Simulates recovery with various oracles of given entropy, measuing tree size")

runRecover :: BitLength -> Width -> Int -> Int -> IO ()
runRecover  bits w iter limit = do
    samples <- replicateM iter $ do
       params <- generate (genRSAParams bits)
       let res = runLSBRecovery limit w params
       putStrLn $ maybe "-" show res
       return res
    let (mean, variance) = meanVariance (V.fromList (map fromIntegral (catMaybes samples)))
    putStrLn $ "-- Mean: " ++ show mean ++ " Variance: " ++ show variance

popCountCollEntropyRate :: Int -> Double
popCountCollEntropyRate w = (/fromIntegral w) $ negate $ logBase 2 $ sum [
    (fromIntegral (binom w i) / 2^w)^2
    | i <- [0..w]
    ]

binom :: Int -> Int -> Int
binom n m = fac n `div` (fac m * fac (n-m))

fac :: Int -> Int
fac n = product [1..n]

runScatter :: BitLength -> [Int] -> [Width] -> [Int] -> Int -> Int -> IO ()
runScatter  bits knownBitsRatios widths blockSums iter limit = do
    printf "leak\th\tmean\tsolved\tcalc_mean\n"
    let todo = [ (p, o, "hard_bits")
               | ratio <- sort knownBitsRatios
               , let p = fromIntegral ratio / 1000 :: Double
               , let o = mkHardBitOracle bits p
               ] ++
               [ (h, o, "ltr")
               | w <- sortBy (flip compare) widths
               , let h = [undefined, 1, 0.786, 0.652, 0.545, 0.461, 0.395, 0.345] !! w
               , let o = mkLeftToRightOracle bits w
               ] ++
               [ (h, o, "popcnt")
               | bs <- sortBy (flip compare) blockSums
               , let h = popCountCollEntropyRate bs
               , let o = mkBlockSumOracle bs
               ]
    forM_ todo $ \(h, mkO, desc) -> do
        samples <- replicateM iter $ do
           params <- generate (genRSAParams bits)
           let op = mkO (dp params)
               oq = mkO (dq params)
               res = lsbRecovery limit params (op, oq)
           -- putStrLn $ maybe "-" show res
           return res
        let (mean, variance) = meanVariance (V.fromList (map fromIntegral (catMaybes samples)))
        let solved = fromIntegral (length (filter isJust samples)) / fromIntegral iter :: Double
        printf "%s\t%.3f\t%.1f\t%.3f\t%.1g\n"
            desc h mean solved (treeSize bits h)

treeSize :: BitLength -> Double -> Double
treeSize k h = k' * (1 + (1 - 2**(k'*(1-2*h)))/(1-2**(1-2*h)))
  where k' = 2*fromIntegral k

printLargePrime :: BitLength -> Width -> IO ()
printLargePrime bits w = generate (rndPrime bits) >>= print

printRSA :: BitLength -> Width -> IO ()
printRSA bits w = generate (genRSAParams bits) >>= print

showLeak :: Leak -> String
showLeak = map go . V.toList
  where go S = '_'
        go M = 'M'

printLTR :: BitLength -> Width -> Integer -> IO ()
printLTR bits w s = putStrLn $ showLeak $ leftToRight bits w s

realCandidates :: BitLength -> Width -> Leak -> BitMask -> [Integer]
realCandidates bits w leak mask = ms'
   where ms    = allMatching mask
         ms'   = filter (\m -> leftToRight bits w m == leak) ms

reallyKnownBits :: BitLength -> Width -> Leak -> BitMask -> BitMask
reallyKnownBits bits w leak mask = mask'
   where ms'   = realCandidates bits w leak mask
         mask' = maskFromMatching ms'

testKnownBits bits w n = do
    putStrLn (showBitMask bits (fromKnownNumber n))
    let leak = leftToRight bits w n
    putStrLn (showLeak leak)
    let mask = knownBitsFromLeak bits w leak
    putStrLn (showBitMask bits mask)
    let mask' = reallyKnownBits bits w leak mask
    putStrLn (showBitMask bits mask')

testHardRules :: BitLength -> Width -> Bool -> IO ()
testHardRules bits w verbose = forM_ [2^(bits-1)..(2^bits)-1] $ \n -> do
    let leak  = leftToRight bits w n
    let mask  = knownBitsFromLeak bits w leak
    let ms    = allMatching mask
    let ms'   = filter (\m -> leftToRight bits w m == leak) ms
    let mask' = maskFromMatching ms'

    when (verbose || mask /= mask' || n `notElem` ms) $ do
        putStr $ showBitMask bits (fromKnownNumber n) ++ ": "
        putStr $ showLeak leak ++ " "
        putStr $ showBitMask bits mask ++ " "
        putStr $ showBitMask bits mask' ++ " "
        putStr $ if mask == mask' then "✓" else "☹"
        unless (n `elem` ms) $ putStr " HORROR"
        putStr "\n"

measureCollisionEntropy :: [BitLength] -> [Width] -> Bool -> Int -> IO ()
measureCollisionEntropy bitsRange wRange tsv iter= do
    when tsv $ putStrLn $ intercalate "\t" $
        words "bits width entropy"
    sequence_ [ measureCollisionEntropyOne bits w tsv iter
              | bits <- bitsRange, w <- wRange ]

measureCollisionEntropyOne :: BitLength -> Width -> Bool -> Int -> IO ()
measureCollisionEntropyOne bits w tsv iter = do
    probs <- replicateM iter doOne
    -- There are serious problems here with numerical precision
    -- and the results seem to vary a log depending on how the logarithm
    -- is calculated
    let entropy = - (logBase 2.0 (sum probs) - fromIntegral (2*bits) - logBase 2.0 (fromIntegral iter)) / fromIntegral (2*bits)
    if tsv then
        putStrLn $ intercalate "\t" $ map show [2*bits, w] ++ map show [entropy]
    else putStrLn $ "Collision entropy rate:" ++ show entropy
  where
    stats :: V.Vector Double -> [Double]
    stats samples =
        [ minimum samples ] ++
        [ weightedAvg n 4 samples | n <- [1..3]] ++
        [ maximum samples ]
    doOne = do
        dp <- fmap (.|. 1) $ randomRIO (2^(bits - 1), 2^bits - 1)
        dq <- fmap (.|. 1) $ randomRIO (2^(bits - 1), 2^bits - 1)

        let sp = leftToRight bits w dp
            wlp = possibleWindowLengths bits w sp
            hardBitsp = countHardBits bits w wlp
            possp = countPossibilities bits w wlp

            sq = leftToRight bits w dq
            wlq = possibleWindowLengths bits w sq
            hardBitsq = countHardBits bits w wlq
            possq = countPossibilities bits w wlq

            totalBits = fromIntegral (2*bits)
            prob = fromIntegral possp * fromIntegral possq
        return prob

measureLeakage :: [BitLength] -> [Width] -> Bool -> Int -> IO ()
measureLeakage bitsRange wRange tsv iter= do
    when tsv $ putStrLn $ intercalate "\t" $
        words "bits width" ++
        concatMap (\n -> [n ++ "_q" ++ show i | i <- [0..4]])
            (words "rightToLeft hardBitsLearned informationGained GMP")
    sequence_ [ measureLeakageOne bits w tsv iter
              | bits <- bitsRange, w <- wRange ]

measureLeakageOne :: BitLength -> Width -> Bool -> Int -> IO ()
measureLeakageOne bits w tsv iter= do
    (samplesRTL, samplesHard, samplesReal, samplesGMP) <- V.unzip4 . V.fromList <$> replicateM iter doOne
    if tsv then do
        putStrLn $ intercalate "\t" $
            map show [2*bits, w] ++
            map show (concatMap stats [ samplesRTL, samplesHard, samplesReal, samplesGMP])
    else do
        let (mean, variance) = meanVariance samplesHard
        let good = V.sum (V.map (\n -> if n > 0.5 then 1 else 0) samplesHard) / fromIntegral iter
        putStrLn $ "Hard bits mean:  " ++ show mean ++
                   " variance: " ++ show variance ++
                   " >50%: " ++ show good
        let (mean, variance) = meanVariance samplesReal
        let good = V.sum (V.map (\n -> if n > 0.5 then 1 else 0) samplesReal) / fromIntegral iter
        putStrLn $ "Total data mean: " ++ show mean ++
                   " variance: " ++ show variance ++
                   " >50%: " ++ show good
  where
    stats :: V.Vector Double -> [Double]
    stats samples =
        [ minimum samples ] ++
        [ weightedAvg n 4 samples | n <- [1..3]] ++
        [ maximum samples ]
    doOne = do
        dp <- fmap (.|. 1) $ randomRIO (2^(bits - 1), 2^bits - 1)
        dq <- fmap (.|. 1) $ randomRIO (2^(bits - 1), 2^bits - 1)

        let sp = leftToRight bits w dp
            wlp = possibleWindowLengths bits w sp
            hardBitsp = countHardBits bits w wlp
            possp = countPossibilities bits w wlp

            sq = leftToRight bits w dq
            wlq = possibleWindowLengths bits w sq
            hardBitsq = countHardBits bits w wlq
            possq = countPossibilities bits w wlq

            totalBits = fromIntegral (2*bits)
            hardBitsFrac = fromIntegral (hardBitsp + hardBitsq) / totalBits

            realBitsFrac = ( totalBits - logBase 2.0 (fromIntegral possp) - logBase 2.0 (fromIntegral possq) ) / totalBits

            rtlBitsFrac = ( fromIntegral $ rightToLeftBits bits w dp + rightToLeftBits bits w dq
                          ) / totalBits

            gmpBitsFrac = ( fromIntegral $ leftToRightGMPBits bits w dp + leftToRightGMPBits bits w dq
                          ) / totalBits

        return (rtlBitsFrac, hardBitsFrac, realBitsFrac, gmpBitsFrac)

hardSoftScatter :: [BitLength] -> [Width] -> Int -> IO ()
hardSoftScatter bitsRange wRange iter = do
    putStrLn $ intercalate "\t" $ words "bits width hard soft"
    forM_ bitsRange $ \bits -> forM_ wRange $ \w -> do
        pairs <- replicateM iter (doOne bits w)
        forM_ pairs $ \(h,s) ->
            putStrLn $ intercalate "\t" (map show [bits, w] ++ map show [h,s])
  where
    doOne bits w = do
        dp <- fmap (.|. 1) $ randomRIO (2^(bits - 1), 2^bits - 1)
        dq <- fmap (.|. 1) $ randomRIO (2^(bits - 1), 2^bits - 1)

        let sp = leftToRight bits w dp
            wlp = possibleWindowLengths bits w sp
            hardBitsp = countHardBits bits w wlp
            possp = countPossibilities bits w wlp

            sq = leftToRight bits w dq
            wlq = possibleWindowLengths bits w sq
            hardBitsq = countHardBits bits w wlq
            possq = countPossibilities bits w wlq

            totalBits = fromIntegral (2*bits) :: Double
            hardBitsFrac = fromIntegral (hardBitsp + hardBitsq) / totalBits

            realBitsFrac = ( totalBits - logBase 2.0 (fromIntegral possp) - logBase 2.0 (fromIntegral possq) ) / totalBits

            rtlBitsFrac = ( fromIntegral $ rightToLeftBits bits w dp + rightToLeftBits bits w dq
                          ) / totalBits

            gmpBitsFrac = ( fromIntegral $ leftToRightGMPBits bits w dp + leftToRightGMPBits bits w dq
                          ) / totalBits

        return (hardBitsFrac, realBitsFrac)

leakPreinput ::  Width -> IO ()
leakPreinput w = do
    line <- getLine
    let line' = fixFirst line
    unless (all (`elem` "SM_") line') $ do
        hPutStrLn stderr "Error: Input is not a sequence of 'S', '_' or 'M'"
        exitFailure
    let bits = length line'
        sp = V.fromList (map toStep line')
        wlp = possibleWindowLengths bits w sp
            --hardBitsp = countHardBits bits w wlp
        possp = countPossibilities bits w wlp
    print possp
  where
    toStep 'M' = M
    toStep 'S' = S
    toStep '_' = S

    -- The input may explicitly say that the first character is a 1.
    -- This code knows that, but still expects there to be a S or M.
    -- Recover whether this is a S or M from looking at the following characters.
    fixFirst ('1':xs) = if 'M' `elem` take (w - 1) xs then 'S':xs else 'M':xs
    fixFirst xs = xs

showRange :: [Int] -> String
showRange = intercalate "," . map show

parseRange :: String -> Either String [Int]
parseRange s = case splitOn "," s of
    [n1,n2,"..",n3] -> enumFromThenTo <$> readEither n1 <*> readEither n2 <*> readEither n3
    [n1,"..",n3] -> enumFromTo <$> readEither n1 <*> readEither n3
    ns -> mapM readEither ns


{-# LANGUAGE FunctionalDependencies, FlexibleInstances, FlexibleContexts, BangPatterns #-}

module AI.Learning.LinearRegression where

import Data.List (foldl')
import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Util (ones)

import AI.Util.Matrix

{-

A linear model should contain the following information:

* Regression coefficients
* Covariance matrix of regression coefficients
* Residuals
* Confidence intervals for regression coefficients
* Confidence intervals for residuals
* T-Statistic for regression coefficients
* P-Values for regression coefficients
* F-Statistic for regression
* P-Value for regression
* MSE
* SST
* SSE
* PRESS Statistic
* R squared
* adjusted r squared
*

-}

-------------
-- Options --
-------------

-- |Type recording what kind of linear model we want to build. Choose from
--  ordinary least squares (OLS), ridge regression with parameter lambda
--  and LASSO with parameter lambda.
data LMType = OLS
            | Ridge Double
            | LASSO Double
            deriving (Show)

---- |Data type for linear regression options.
data LMOpts = LMOpts { fitIntercept :: Bool
                     , standardizeRegressors :: Bool }
                     deriving (Show)

---- |Standard options for linear regression (used by 'regress').
stdLMOpts = LMOpts { fitIntercept = True
                   , standardizeRegressors = False }

---------------------
---- Linear Models --
---------------------

-- |Data type for a linear model. It consists of the coefficient vector,
--  together with options that specify how to transform any data that is used
--  to make new predictions.
data LinearModel = LM { coefs   :: Vector Double
                      , lmIntercept :: Bool
                      , lmStandardize :: Bool
                      , lmMean :: Vector Double
                      , lmStd  :: Vector Double }
                      deriving (Show)

-- |Statistics structure for a linear regression.
data LMStats = LMStats { covBeta :: Maybe (Matrix Double)
                       , betaCI :: Maybe (Matrix Double)
                       , sst :: Double
                       , sse :: Double
                       , mse :: Double
                       , rSquare :: Double
                       , tBeta :: Maybe (Vector Double)
                       , pBeta :: Maybe (Vector Double)
                       , fRegression :: Maybe Double
                       , pRegression :: Maybe Double }
                       deriving (Show)

type Predictor = Matrix Double
type Response = Vector Double

-- |Fit an ordinary least squares linear model to data.
lm :: Matrix Double -> Vector Double -> LinearModel
lm = lmWith OLS stdLMOpts

-- |Fit a ridge regression (least squares with quadratic penalty) to data.
lmRidge :: Matrix Double -> Vector Double -> Double -> LinearModel
lmRidge x y lambda = lmWith (Ridge lambda) stdLMOpts x y

-- |Fit a linear model to data. The exact model fitted is specified by the
--  'LMType' argument, and the options are specified in 'LMOpts'.
lmWith :: LMType -> LMOpts -> Matrix Double -> Vector Double -> LinearModel
lmWith kind opts x y = LM { coefs = beta
                          , lmIntercept = fitIntercept opts
                          , lmStandardize = standardizeRegressors opts
                          , lmMean = mu
                          , lmStd = sigma }
    where
        (xx,mu,sigma) = lmPrepare opts x
        beta = case kind of
            OLS     -> regress xx y
            Ridge a -> ridgeRegress xx y (fitIntercept opts) a
            LASSO a -> error "LASSO not implemented."

-- |Prepare data according to the specified options structure. This may involve
--  centering and standardizing the data, or adding a column of constants.
lmPrepare :: LMOpts
          -> Matrix Double
          -> (Matrix Double, Vector Double, Vector Double)
lmPrepare opts x = (x3,mu,sigma)
    where
        (x1,mu,sigma) = standardize x
        x2            = if standardizeRegressors opts then x1 else x
        x3            = if fitIntercept opts then addOnes x2 else x2

---------------------------------
---- Predict from linear model --
---------------------------------

-- |Make predictions from a linear model.
lmPredict :: LinearModel -> Matrix Double -> Vector Double
lmPredict model x = x2 <> beta
    where
        beta    = coefs model
        xbar    = lmMean model
        sigma   = lmStd model
        x1      = if lmStandardize model
                    then eachRow (\x -> (x - xbar) / sigma) x
                    else x
        x2      = if lmIntercept model
                    then addOnes x1
                    else x1

-- |Calculate statistics for a linear regression.
lmStats :: LinearModel -> Matrix Double -> Vector Double -> LMStats
lmStats model x y =
    let ybar      = constant (mean y) (dim y)
        yhat      = lmPredict model x
        residuals = y - yhat
        covBeta   = Nothing
        betaCI    = Nothing
        sst       = sumVector $ (y - ybar) ^ 2
        sse       = sumVector $ residuals ^ 2
        mse       = sse / fromIntegral (rows x)
        rSq       = 1 - sse / sst
        tBeta     = Nothing
        pBeta     = Nothing
        fReg      = Nothing
        pReg      = Nothing
    in LMStats covBeta betaCI sst sse mse rSq tBeta pBeta fReg pReg

mseEvalFun :: Response -> Response -> Double
mseEvalFun actual predicted = mean $ (actual - predicted) ^ 2

lmPredFun :: Predictor -> Response -> Predictor -> Response -> Double
lmPredFun xtrain ytrain xtest ytest = mseEvalFun ytest ypred
    where
        ypred = lmPredict model xtest
        model = lm xtrain ytrain

--------------------------
---- Perform Regression --
--------------------------

---- |Regress a vector y against a matrix of predictors x, with the specified
----  options.
regress :: Matrix Double -- X
        -> Vector Double -- y
        -> Vector Double -- beta
regress x y
    | rows x /= dim y = error "Inconsistent dimensions -- REGRESS"
    | otherwise = let (_,n) = size x
                      (_,r) = qr x
                      rr = takeRows n r
                   in (trans rr <> rr) <\> trans x <> y

-- |Ridge regression. This adds a penalty term to OLS regression, which
--  discourages large coefficient values to prevent overfitting.
ridgeRegress :: Matrix Double  -- X
             -> Vector Double  -- y
             -> Bool           -- useConst?
             -> Double         -- lambda
             -> Vector Double  -- beta
ridgeRegress x y useConst lambda
    | rows x /= dim y = error "Inconsistent dimensions -- RIDGEREGRESS"
    | otherwise = let (_,n) = size x
                      (_,r) = qr x
                      rr = takeRows n r
                      ww = if useConst
                            then diag $ join [0, constant 1 (n-1)]
                            else ident n
                   in (trans rr <> rr + lambda `scale` ww) <\> trans x <> y

-----------------
---- Utilities --
-----------------

-- |De-mean a sample.
demean :: Matrix Double -> Matrix Double
demean x = eachRow (subtract $ mean x) x

-- |Standardize a sample to have zero mean and unit variance. Returns a triple
--  consisting of the standardized sample, the sample mean and the sample
--  standard deviation.
standardize :: Matrix Double -> (Matrix Double, Vector Double, Vector Double)
standardize m = (eachRow (\x -> (x - mu) / sigma) m, mu, sigma)
    where mu    = mean m
          sigma = std m

-- |Standardize a sample to have zero mean and unit variance. This variant of
--  the function discards the mean and standard deviation vectors, only
--  returning the standardized sample.
standardize_ :: Matrix Double -> Matrix Double
standardize_ x = a where (a,_,_) = standardize x

-------------------------
-- Mean, Variance etc. --
-------------------------

class Mean a b | a -> b where
    mean :: a -> b

class Floating b => Variance a b | a -> b where
    var :: a -> b
    var x = y * y where y = std x

    std :: a -> b
    std x = sqrt $ var x

instance Mean (Vector Double) Double where
    mean v = sumVector v / fromIntegral (dim v)

instance Variance (Vector Double) Double where
    var v  = mean $ (v - constant vbar (dim v)) ^ 2
        where vbar = mean v

instance Mean (Matrix Double) (Vector Double) where
    mean m = fromList $ mapCols mean m

instance Variance (Matrix Double) (Vector Double) where
    var m  = fromList $ mapCols var m


import Data.Complex

direction :: Char -> Complex Float
direction 'U' = 0 :+ 1
direction 'D' = 0 :+ (-1)
direction 'L' = (-1) :+ 0
direction 'R' = 1 :+ 0

clamp :: Float -> Float
clamp = max (-1) . min 1

advance :: Complex Float -> Complex Float -> Complex Float
advance pos dir =
  let (x :+ y) = pos + dir
  in clamp x :+ clamp y

complexToButton :: Complex Float -> Char
complexToButton ((-1) :+   1)  = '1'
complexToButton (  0  :+   1)  = '2'
complexToButton (  1  :+   1)  = '3'
complexToButton ((-1) :+   0)  = '4'
complexToButton (  0  :+   0)  = '5'
complexToButton (  1  :+   0)  = '6'
complexToButton ((-1) :+ (-1)) = '7'
complexToButton (  0  :+ (-1)) = '8'
complexToButton (  1  :+ (-1)) = '9'

part1 :: String -> String
part1 input = do
  line <- lines input
  let button = complexToButton $ foldl advance 0 $ map direction line
  return button

complexToButton2 :: Complex Float -> Char
complexToButton2 (  0  :+   2)  = '1'
complexToButton2 ((-1) :+   1)  = '2'
complexToButton2 (  0  :+   1)  = '3'
complexToButton2 (  1  :+   1)  = '4'
complexToButton2 ((-2) :+   0)  = '5'
complexToButton2 ((-1) :+   0)  = '6'
complexToButton2 (  0  :+   0)  = '7'
complexToButton2 (  1  :+   0)  = '8'
complexToButton2 (  2  :+   0)  = '9'
complexToButton2 ((-1) :+ (-1)) = 'A'
complexToButton2 (  0  :+ (-1)) = 'B'
complexToButton2 (  1  :+ (-1)) = 'C'
complexToButton2 (  0  :+ (-2)) = 'D'
complexToButton2 _              = '0' -- default

advance2 :: Complex Float -> Complex Float -> Complex Float
advance2 pos dir =
  let pos' = pos + dir
  in if complexToButton2 pos' == '0'
  then pos
  else pos'

part2 :: String -> String
part2 input = do
  line <- lines input
  let button = complexToButton2 $ foldl advance2 0 $ map direction line
  return button


module Corfunc (corfunc,fitdata) where

import Data.Complex
import Numeric.FFT

porod :: Double -> Double -> Double -> Double
porod k sigma q = (k*q**(-4))*exp (-q**2*sigma)

guinier :: Double -> Double -> Double -> Double
guinier a b q = a * exp (b*q**2)

beta :: [Double] -> [Double] -> Double
beta x y = rval x y * std y / std x

alpha :: [Double] -> [Double] -> Double
alpha x y = mean y - beta x y * mean x

mean :: [Double] -> Double
mean xs = sum xs / fromIntegral (length xs)

std :: [Double] -> Double
std xs = sqrt $ mean (map (**2) xs) - mean xs**2

rval :: [Double] -> [Double] -> Double
rval xs ys = sum (zipWith (*) xdiff ydiff) / sqrt (sum (map (**2) xdiff) * sum (map (**2) ydiff))
    where
      diff qs = map (\z -> z - mean qs) qs
      xdiff = diff xs
      ydiff = diff ys

fitguinier :: [Double] -> [Double] -> Double -> Double
fitguinier qs is = guinier a b
    where
      xs = map (**2) qs
      ys = map log is
      a = exp $ alpha xs ys
      b = beta xs ys

fitporod :: [Double] -> [Double] -> Double -> Double
fitporod qs is = porod k sigma
    where
      is' = filter (>0) is
      qs' = filterKey (>0) is qs
      xs = map (**2) qs'
      ys = map log . zipWith (*) is' . map (**4) $ qs'
      k = exp $ alpha xs ys
      sigma = 0 * beta xs ys

filterKey :: (a->Bool) -> [a] -> [b] -> [b]
filterKey f ks = map snd . filter (f.fst) . zip ks

takeKey :: (a->Bool) -> [a] -> [b] -> [b]
takeKey f ks = map snd . takeWhile (f . fst) . zip ks

dropKey :: (a->Bool) -> [a] -> [b] -> [b]
dropKey f ks = map snd . dropWhile (f . fst) . zip ks

fitdata :: [Double] -> [Double] -> Double -> Double
fitdata qs is = let d = interp qs is
                    maxq = 0.04
                    highq = dropWhile (< maxq) qs
                    highi = dropKey (< maxq) qs is
                    fitp = fitporod highq highi
                    minq = 0.0065*3
                    lowq = takeWhile (< minq) qs
                    lowi = takeKey (< minq) qs is
                    fitg = fitguinier lowq lowi
                    s1 = smooth d fitp (maxq,maximum qs)
                    s2 = smooth fitg s1 (minimum qs,minq)
                in
                  s2

-- | Smoothly transition between two functions over a range of functions
smooth :: (Double -> Double) -> (Double -> Double) -> (Double, Double) -> Double -> Double
smooth f g (start, stop) x
    | x <= start = f x
    | x >= stop = g x
    | otherwise = h*g x +(1-h)*f x
    where
      h = 1/(1+(x-stop)**2/(start-x)**2)

interpHead :: ((Double, Double), (Double, Double)) -> Double -> Double
interpHead ((x1,x2),(y1,y2)) q = (x2-q)/(x2-x1)*y1 + (q-x1)/(x2-x1)*y2

interp :: [Double] -> [Double] -> Double -> Double
interp xs ys x =
    let ranges = zip (zip xs (tail xs)) (zip ys (tail ys))
    in
      case dropWhile (\q -> snd (fst q) < x) ranges of
        [] -> last ys
        p:_ -> interpHead p x

dct :: [Double] -> [Double]
dct = map realPart . fft . map (:+ 0)

corfunc :: [Double] -> [Double] -> [(Double,Double)]
corfunc qs is = let fullData = fitdata qs is
                    factor = 16
                    deltaq = (qs!!1-head qs)/factor
                    qgrid = map (*deltaq) [0..factor*2^(floor . log . (/ deltaq) . maximum $ qs)]
                    igrid = zipWith (*) (map (**2) qgrid) $ map fullData qgrid
                    ys = dct igrid
                    xs = map (\q -> 2 * pi * q / deltaq /fromIntegral (length qgrid)) [0..fromIntegral $ length qgrid]
                in
                  zip xs ys


module Main where

import AI.HNN.FF.Network
import Numeric.LinearAlgebra

samples :: Samples Double
samples = [ (fromList [0, 0], fromList [0])
          , (fromList [0, 1], fromList [1])
          , (fromList [1, 0], fromList [1])
          , (fromList [1, 1], fromList [0])
          ]

main :: IO ()
main = do
  n <- createNetwork 2 [2] 1
  mapM_ (print . output n tanh . fst) samples
  putStrLn "------------------"
  let n' = trainNTimes 1000 0.8 tanh tanh' n samples
  mapM_ (print . output n' tanh . fst) samples

{-

OUTPUT :
[0.5835034660982111]
[0.7161238877902711]
[0.7476745582239942]
[0.7959844102953423]
------------------
[4.427877031872795e-2]
[0.9735228418859667]
[0.9544892359343744]
[-2.7925107071558753e-2]
-}


module FET where

import qualified Data.ByteString.Lazy.Char8             as BS
import           Data.Eq
import           Data.Function
import           Data.Functor
import           Data.Int
import           Data.List
import           Data.Ord
import           Prelude                                (Double, fromIntegral,
                                                         (*), (+), (-), (/),
                                                         (^), otherwise, (&&))
import           Statistics.Distribution                (complCumulative,
                                                         cumulative,
                                                         probability)
import           Statistics.Distribution.ChiSquared
import           Statistics.Distribution.Hypergeometric
import           Text.Show

--the Fisher's exact test can be computed from the hypergeometric distribution
--eg. test whether the presence / absence of a gene is statistically significant
--between two bacterial groups
--           Group 1 | Group 2
--          -------------------
-- has gene |   2    |    7
-- no gene  |   5    |    3
--          -------------------
-- the hypergeometric distribution is calculated in Statistics.Distribution.Hypergeometric
-- hypergeometric = m l k
-- From the package description: The parameters of the distribution describe
-- k elements chosen from a population of l, with m elements of one type,
-- and l-m of the other (all are positive integers).
-- In our case:
-- m is the sum of "has factor" instances, so 9
-- l is the total population size, so 17
-- k is the sum of a particular group, so 7 for Group 1
-- let d = distribution 9 17 7
-- This gives us the distribution, but we need to test the probability
-- that the presence of this gene is statistically different, using the two-tail
-- approach.
-- The one-tail probability in our example of Group 1 having 2 members with the
-- gene can be tested using the cumulative distribution function from
-- Statistics.Distribution, our previously calculated distribution "d", and
-- the number of "successes" (x), which is 2.
-- let oneTail = cumulative d 2 = 0.11703002879473468
-- However, we have no previous expectation for this group to have or not have
-- this gene, so what we really want is the two-tailed test.
-- Please see http://www.pmean.com/11/Fishers.html for an excellent,
-- detailed description.
-- In short, we need to sum all of the probabilities that are smaller than or
-- equal to the probability of the observed table (eg. x = 2 in our example).
-- The twoTailValue is calculated by generating all of the probabilities for
-- our sample, filtering the ones that have a probability too large, and
-- summing the remaining values.
--
-- For a total number of observations greater than 1000, we use the Chi Square
-- test, as the FET returns NaN.
-- ChiSqr statistic  for
--           Group 1 | Group 2
--          -------------------
-- has gene |   2    |    7
-- no gene  |   5    |    3
--          -------------------
-- is [(2*3)-(5*7)]^2 * (2+7+5+3) / (2+7)(5+3)(7+3)(2+5)
-- With 1 degree of freedom, we can calculate the Chi Square distribution
-- with chiSquared 1 from Statistics.Distribution
-- To determine the P value of a Chi Square statistic, we use the
-- complCumulative function and the
-- Chi Square distribution with 1 degree of freedom.
-- For Chi Square, the one-tail is simply half the two-tail
-- Statistics.Test.ChiSquared only return "Significant" or "NotSignificant"
-- and we wanted the actual P-value, hence the function.
-- Likewise there is no Statistics.Test.FET

newtype GroupOneA = GroupOneA { unGroupOneA :: Int } deriving (Eq, Show)
newtype GroupOneB = GroupOneB { unGroupOneB :: Int } deriving (Eq, Show)
newtype GroupTwoA = GroupTwoA { unGroupTwoA :: Int } deriving (Eq, Show)
newtype GroupTwoB = GroupTwoB { unGroupTwoB :: Int } deriving (Eq, Show)
newtype PValue = PValue {unPValue :: Double} deriving (Eq, Show)
newtype FETName = FETName {unFETName :: BS.ByteString} deriving(Eq, Show)
data FETMode = OneTail | TwoTail deriving (Eq,Show)

data FETResult = FETResult{groupOneA :: Int
                          ,groupOneB :: Int
                          ,groupTwoA :: Int
                          ,groupTwoB :: Int
                          ,pvalue    :: Double
                          ,fetName   :: BS.ByteString} deriving (Eq, Show)

fet :: FETName
    -> GroupOneA
    -> GroupOneB
    -> GroupTwoA
    -> GroupTwoB
    -> FETMode
    -> FETResult
fet (FETName fName) (GroupOneA goa) (GroupOneB gob) (GroupTwoA gta) (GroupTwoB gtb) mode
  | l > 0 && goa >= 0 && gob >= 0 && gta >= 0 && gtb >= 0 && k > 0 && gt > 0 =
    case mode of
        OneTail -> theResult{pvalue = cumD}
        TwoTail -> theResult{pvalue = twoTailValue}
  | otherwise = theResult
  where
    cumD = if l >= 1000
            then chiP / 2
            else cumulative d $ fromIntegral goa
    probX = probability d goa
    twoTailValue = if l >= 1000
                      then chiP
                      else sum . filter (<= probX) . fmap (probability d) $ [0..k]
    d = hypergeometric m l k
    chiP = complCumulative d' $ chiSquare goa gob gta gtb
    d' = chiSquared 1
    m = goa + gta
    l = m + gob + gtb
    k = goa + gob
    gt = gta + gtb
    theResult = FETResult{groupOneA = goa
                         ,groupOneB = gob
                         ,groupTwoA = gta
                         ,groupTwoB = gtb
                         ,pvalue = 1
                         ,fetName = fName
                         }

chiSquare :: Int
          -> Int
          -> Int
          -> Int
          -> Double
chiSquare a' b' c' d' = (x' * x' * (a + b + c + d)) /
                        ((a+b) * (c+d) * (b+d) * (a+c))
  where
    x' = (a * d) - (b * c)
    a = fromIntegral a'
    b = fromIntegral b'
    c = fromIntegral c'
    d = fromIntegral d'


module Algoritmo where

import Prelude as P
import Data.Complex
import Control.Monad.State
import Numeric.Statistics (percentile)
import Numeric.LinearAlgebra
import Data.Vector.Storable as VS hiding (modify, null)

import FFT
import ArrayPoking
import Bases

type L2  = Vector (Complex Double)
type Dim = Int
type Niveles = Int

secAnalitica :: Niveles -> L2 -> State [L2] ()
secAnalitica n zF =
  do let dim = VS.length zF
         uF  = conj . fft $ haarU dim
         vF  = conj . fft $ haarV dim
         x1  = downSampF (zF * vF) dim
         y1  = downSampF (zF * uF) dim

     -- [y3, x3, x2, x1]
     -- modify (x1 :)
     if n <= 1
       -- Haz nada, este es el ultimo nivel
       then modify (zF :)
       --
       else modify (x1 :) >> secAnalitica (n - 1) y1

secSintetica :: Niveles -> State [L2] L2
secSintetica n =
  do y2 : x2 : rest <- get
     let dim = VS.length y2
         dim_n = dim * 2
         x2' = upSampF x2 dim
         y2' = upSampF y2 dim
         y1  = x2' * (fft (haarV dim_n)) +
               y2' * (fft (haarU dim_n))
     put (y1 : rest)

     if null rest || n <= 1
       then return y1
       else secSintetica (n - 1)

{-
  Idea: Mantener los @k terminos mayores y lo demas setearlos a 0. Es
  molesto contar, asi que veremos un porcentaje.
  rate = 90.0, es decir 90% de calidad, debo de matar 10% de los
  candidatos.
-}
compresion :: Double -> State [L2] ()
compresion rate =
  do vec <- get
     let cut  = percentile (100 - rate) . sortIP
                . cmap magnitude $ VS.concat vec

         go c = if magnitude c < cut then 0 else c
     put $ P.map (VS.map go) vec

waveletAlgo :: Double -> Niveles -> Niveles -> L2 -> L2
waveletAlgo calidad div recon color = evalState go []
  where
    go :: State [L2] L2
    go = return (fft color)
         >>= secAnalitica div
         -- >> modify (P.map ifft)
         >> compresion calidad
         -- >> modify (P.map fft)
         >> secSintetica recon
         >>= return . ifft


{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}
module Marvin.API.Algorithms.LogisticRegression (
  LogisticRegression (..)
  , defaultLogisticRegression

  , LogisticRegressionModel
  , threshold
  , setThreshold
  , coefficients
  , intercept
) where

import qualified Numeric.LinearAlgebra as LA (toList)

import qualified Marvin.API.Algorithms.Internal.GradientDescent as GD
import Marvin.API.Algorithms.Internal.GradientDescent hiding (learningRate, lambda, addIntercept)

import Marvin.API.Meta.Model
import Marvin.API.Table.Internal
import Marvin.API.Fallible

-- * Parameters

-- | Parameters.
data LogisticRegression = LogisticRegression {
  learningRate :: Double -- ^ Coefficient of the gradient when updating the model.
  , numberOfIterations :: Int -- ^ Number of iterations.
  , lambda :: Double -- ^ Regularization coefficient.
  , addIntercept :: Bool -- ^ Whether to add a constant 1 column before training.
}

-- | Default parameters.
defaultLogisticRegression = LogisticRegression {
  learningRate = 0.001
  , numberOfIterations = 10
  , lambda = 0.1
  , addIntercept = True
}

toGradDesc :: LogisticRegression -> GradientDescent
toGradDesc logReg = defaultGradientDescent {
  numIter = numberOfIterations logReg
  , GD.cost = logisticRegressionCost
  , GD.learningRate = learningRate logReg
  , GD.addIntercept = addIntercept logReg
  , GD.lambda = lambda logReg
}

-- * Model

-- | Resulting model.
newtype LogisticRegressionModel = LogRegModel (Double, LinearModel)

-- | Retrieves the threshold to decide 'True' or 'False' based on probability.
threshold (LogRegModel (t, _)) = t
-- | Sets the threshold that decides 'True' or 'False' based on probability.
setThreshold t (LogRegModel (_, linear)) = LogRegModel (t, linear)

-- | Retrieves a list of the coefficient.
coefficients :: LogisticRegressionModel -> [Double]
coefficients (LogRegModel (_, linearModel)) = LA.toList $ coefs linearModel

-- | Retrieves the constant to be added.
intercept :: LogisticRegressionModel -> Double
intercept (LogRegModel (_, linearModel)) = intercept_ linearModel

instance Show LogisticRegressionModel where
  show (LogRegModel linear) = "Logistic regression model:\n\n" ++ show linear

instance Estimator LogisticRegression (NumericTable, BinaryColumn) where
  type ResultingModel LogisticRegression (NumericTable, BinaryColumn) =
    LogisticRegressionModel
  fit logReg (train, target) = do
    ensurePositive
      (InvalidAlgorithmParameter "LogisticRegression" "numberOfIterations must be positive")
        $ numberOfIterations logReg
    ensurePositive
      (InvalidAlgorithmParameter "LogisticRegression" "learning rate must be positive")
        $ learningRate logReg
    ensureNonNegative
      (InvalidAlgorithmParameter "LogisticRegression" "lambda must be non-negative")
        $ lambda logReg
    let numTarget = binaryToNumericColumn target
    ensureTableColSizeMatch "while training logistic regression" train numTarget
    let linearModel = fit' (toGradDesc logReg) (train, numTarget)
    let defaultThreshold = 0.5
    return $ LogRegModel (defaultThreshold, linearModel)

instance Predictor LogisticRegressionModel where
  type Testing LogisticRegressionModel = NumericTable
  type Prediction LogisticRegressionModel = BinaryColumn
  predict (LogRegModel (thresh, model)) testingTable = do
    numericPredictions <- probaPredict model testingTable
    return $ mapColumnAtoB (\x -> if x < thresh then False else True) numericPredictions

probaPredict model testingTable = do
  numericPredictions <- predict' model testingTable
  return $ mapColumnAtoB sigmoid numericPredictions


-- | <<docs/apollonian_gasket/spaced_gasket.svg>>
module Geometry.Processes.ApollonianGasket (
    -- * The classical shape
      createGasket
    , fifthCircle
    , Circle(..)

    -- * For playing around
    , ApoCircle(..)
    , newCircle
    , toCircle
    , toApoCircle
) where

import Data.Complex
import Data.Tree

import Geometry.Core as G

-- | A circle, but in a format more suitable for Descartes’ Theorem, which the
-- gasket algorithm is based on.
--
-- It interprets the center @'Vec2' x y@ as a complex number \(x + \mathrm i\,y\),
-- and stores the curvature \(k=\frac1r\) instead of the radius \(r\).
--
-- 'Circle' and 'ApoCircle' can be freely converted between using 'toCircle' and
-- 'toApoCircle'.
data ApoCircle = ApoCircle (Complex Double) Double
    deriving (Eq, Show)

instance HasBoundingBox ApoCircle where
    boundingBox = boundingBox . toCircle

-- | Simple conversion function. Inverse of 'toApoCircle'.
toCircle :: ApoCircle -> Circle
toCircle (ApoCircle (x :+ y) k) = Circle (Vec2 x y) (1/k)

-- | Simple conversion function. Inverse of 'toCircle'.
toApoCircle :: Circle -> ApoCircle
toApoCircle (Circle (Vec2 x y) r) = ApoCircle (x :+ y) (1/r)

-- | Curvature of the new Apollonian circle.
newCurvature
    :: (Double -> Double -> Double)
    -> Double -- ^ Curvature of first circle
    -> Double -- ^ Second circle
    -> Double -- ^ Third circle
    -> Double -- ^ New circle
newCurvature (+-) k1 k2 k3 =
    -- Taken straight from Wikipedia: https://en.wikipedia.org/wiki/Descartes%27_theorem
    (k1 + k2 + k3) +- (2 * sqrt (k1*k2 + k2*k3 + k3*k1))

-- | Center of the new Apollonian circle.
newCenter
    :: (Complex Double -> Complex Double -> Complex Double) -- ^ Plus or minus to yield both circles
    -> ApoCircle
    -> ApoCircle
    -> ApoCircle
    -> Double         -- ^ Curvature of the new circle
    -> Complex Double -- ^ Center of the new circle
newCenter (+-) (ApoCircle c1 k1) (ApoCircle c2 k2) (ApoCircle c3 k3) k4 =
    -- Taken straight from Wikipedia: https://en.wikipedia.org/wiki/Descartes%27_theorem
    1/k4' * ( (c1*k1' + c2*k2' + c3*k3')
              +- (2*sqrt (k1'*k2'*c1*c2 + k2'*k3'*c2*c3 + k1'*k3'*c1*c3)))
  where
    [k1', k2', k3', k4'] = map (:+ 0) [k1, k2, k3, k4]

-- | Create a new Apollonian circle, based on three existing ones. This is the raw
-- application of Descartes’ theorem. For something with less pitfalls, use
-- 'fifthCircle'.
--
-- The choice of sign is a mystery to me. For the Apollonian Gasket, the rule is:
--
--   1. The large circle is obtained using \((-,-)\) on the three initial cycles
--   2. The bottom-left circle is obtained using \((+,-)\) on the initial top, initial bottom, and large circle
--   3. All other recursive steps use \((+,+)\)
--   4. If you can explain the \((+,-)\), please contact me, I’m dying to know
--
-- If you’re wondering, this is what you get when you don’t do the \((+,-)\) trick:
--
-- <<docs/apollonian_gasket/missing_the_minus.svg>>
newCircle
    :: (Double -> Double -> Double)
        -- ^ Plus (larger curvature, smaller circle) or minus (smaller curvature, larger circle).
    -> (Complex Double -> Complex Double -> Complex Double) -- ^ Plus or minus; see this function’s doc.
    -> ApoCircle
    -> ApoCircle
    -> ApoCircle
    -> ApoCircle
newCircle plusMinusCurvature plusMinusCenter circ1@(ApoCircle _ k1) circ2@(ApoCircle _ k2) circ3@(ApoCircle _ k3) =
    let k4 = newCurvature plusMinusCurvature k1 k2 k3
        c4 = newCenter plusMinusCenter circ1 circ2 circ3 k4
    in ApoCircle c4 k4

-- | Given four mutually tangent circles, find »the other one« than the first argument.
-- This resolves the »\((+,-)\)« ambiguity of 'newCircle'.
fifthCircle :: ApoCircle -> ApoCircle -> ApoCircle -> ApoCircle -> ApoCircle
fifthCircle (ApoCircle c1 k1) (ApoCircle c2 k2) (ApoCircle c3 k3) (ApoCircle c4 k4) =
    let [k1', k2', k3', k4', k5'] = map (:+ 0) [k1, k2, k3, k4, k5]
        k5 = 2*(k2+k3+k4) - k1
        c5 = (2*(k2'*c2+k3'*c3+k4'*c4) - k1'*c1) / k5'
    in ApoCircle c5 k5

-- | The simple workhorse function. Given three mutually touching, equally sized
-- circles, it calculates the classical Apollonian Gasket:
--
--
-- <<docs/apollonian_gasket/classical_gasket.svg>>
--
-- Be warned though, the number of circles grows exponentially with recursion depth
-- (controlled with the radius parameter). This can easily melt your RAM, CPU or
-- harddrive if you aren’t careful.
--
-- It’s quite fun to play with the parameters though! Here is what happens when the three initial circles
-- aren’t touching:
--
-- <<docs/apollonian_gasket/forgetting_gen0.svg>>
--
-- To get even more creative freedom, I like to copy this algorithm and see what I
-- get, e.g. by commenting out one of the recursions.
createGasket
    :: Double -- ^ Radius threshold: do not recurse to circles smaller than this. Smaller values yield more detail.
    -> Circle -- ^ Initial left circle
    -> Circle -- ^ Initial right circle
    -> Circle -- ^ Initial bottom circle
    -> Tree Circle
createGasket minRadius gen0LCirc gen0RCirc gen0BCirc =
    let gen0L = toApoCircle gen0LCirc
        gen0R = toApoCircle gen0RCirc
        gen0B = toApoCircle gen0BCirc

        large = newCircle (-) (-) gen0L gen0R gen0B
        small = newCircle (+) (+) gen0L gen0R gen0B

        recurse :: ApoCircle -> ApoCircle -> ApoCircle -> ApoCircle -> Tree Circle
        recurse cA c1 c2 c3 =
            let cB@(ApoCircle _ kB) = fifthCircle cA c1 c2 c3
            in Node (toCircle cB)
                (if abs kB < 1/minRadius
                    then
                        [ recurse c3 cB c1 c2
                        , recurse c2 c3 cB c1
                        , recurse c1 c2 c3 cB
                        ]
                    else [])

        circles :: Tree Circle
        circles = Node (toCircle large)
            [ recurse small gen0L gen0R gen0B
            , recurse large gen0L gen0R gen0B
            ]
    in circles


module Models where

import Layer
import Numeric.LinearAlgebra
import Data.List
import Loss

data Model a = Model [a]

addLayer :: Model a -> a -> Model a
addLayer (Model []) layer = Model [layer]
addLayer (Model l) layer = Model (l++[layer])

genParamsList :: Model VLayer -> [Matrix R]
genParamsList (Model []) = []
genParamsList m =
  let genModelParams' res (Model []) = reverse res
      genModelParams' res (Model (mx:mxs)) =
        genModelParams' ((weight mx):res) (Model mxs) in
  genModelParams' [] m

-- update the model's params from the params list
updateModelParams :: Model VLayer -> [Matrix R] -> Model VLayer
updateModelParams (Model []) _ = Model []
updateModelParams m p  =
  let updateModelParams' (Model res) (Model []) _ = Model (reverse res)
      updateModelParams' (Model res) (Model (mx:mxs)) (px:pxs) =
        let newParams = mx {weight=px} in
        updateModelParams' (Model (newParams:res)) (Model mxs) pxs
  in updateModelParams' (Model []) m p

genGParamsList :: Model VLayer -> [Matrix R]
genGParamsList (Model []) = []
genGParamsList l =
  let genModelGParams' res (Model []) = reverse res
      genModelGParams' res (Model (mx:mxs)) =
        let gp = ( (rowNum mx)  >< (colNum mx)) [0..]::(Matrix R) in
        genModelGParams' (gp:res) (Model mxs) in
  genModelGParams' [] l

--updateModelGParams :: Model VLayer -> [Matrix R] -> Model VLayer
-- update the gparams list from the model
updateGParamsList :: Model VLayer -> [Matrix R] -> [Matrix R]
updateGParamsList (Model []) _ = []
updateGParamsList l gpl =
  let updateModelGParams' res (Model []) _ = reverse res
      updateModelGParams res (Model (mx:mxs)) (gpx:gpxs) =
        let deltaUp = delta mx in
        updateModelGParams' ((deltaUp+gpx):res) (Model mxs) gpxs in
  updateModelGParams' [] l gpl

forwardInterface :: Matrix R -> VLayer -> Matrix R
forwardInterface input layer =
  case (layerType layer) of
    LinearLayer -> forwardLinearLayer input layer
    SigmoidLayer -> forwardSigmoidLayer input layer
    --SoftmaxLayer -> forwardSoftmaxLayer input layer
    ReluLayer -> forwardReluLayer input layer

backwardInterface :: VLayer -> Matrix R -> Matrix R
backwardInterface layer output =
  case (layerType layer) of
    LinearLayer -> backwardLinearLayer output layer
    SigmoidLayer -> backwardSigmoidLayer output layer
    ReluLayer -> backwardReluLayer output layer

-- forward process of the neural network model
forward :: Matrix R -> Model VLayer -> Matrix R
forward input (Model []) = input
forward input (Model xs) = foldl forwardInterface input xs

-- the backpropagate algorithm of Neural Network
-- update the delta of each layer from end to start
backward :: String -> Matrix R -> Matrix R -> Model VLayer -> Model VLayer
backward _ _ _ (Model []) = Model []
backward lossType yMat tMat l =
    let loss = getLoss lossType yMat tMat in
    let backward' (Model res) (Model []) _ = Model (reverse res)
        backward' (Model res) (Model (mx:mxs)) dw =
          let dwUp = backwardInterface mx dw
              layerUp = mx {delta=dwUp} in
          backward' (Model (layerUp:res)) (Model mxs) dwUp in
    backward' (Model []) l loss


module LinearRegressionExercise where

import qualified Graphics.Gnuplot.Advanced as GP

import qualified Graphics.Gnuplot.Graph as Graph

import qualified Graphics.Gnuplot.Plot.ThreeDimensional as Plot3D
import qualified Graphics.Gnuplot.Graph.ThreeDimensional as Graph3D

import qualified Graphics.Gnuplot.Plot.TwoDimensional as Plot2D
import qualified Graphics.Gnuplot.Graph.TwoDimensional as Graph2D
import qualified Graphics.Gnuplot.Terminal.Default as DefaultTerm
import Graphics.Gnuplot.Plot.TwoDimensional (linearScale, )

import LinearRegression
import Numeric.LinearAlgebra
import Prelude hiding ((<>))
import System.Directory
import Util

-- replace this with the path to your data.
exerciseOneDataOne :: IO String
exerciseOneDataOne = do
  home <- getHomeDirectory
  return (home ++ "/code/haskell/CourseraMachineLearning/app/ex1data1.txt")

exerciseOneDataTwo :: IO String
exerciseOneDataTwo = do
  home <- getHomeDirectory
  return (home ++ "/code/haskell/CourseraMachineLearning/app/ex1data2.txt")

exerciseOne = do
  -- Part 1: Basic function
  putStrLn "Running warmUpExercise ..."
  putStrLn "5x5 Identity Matrix: "
  let a = ident 5 :: Matrix Double
  print a

  -- Part 2: Plotting
  putStrLn "Plotting Data ..."
  dataOne <- exerciseOneDataOne
  [x', y'] <- fmap toColumns (loadMatrix dataOne)
  GP.plotSync DefaultTerm.cons (plotPoints (toList x') (toList y'))

  -- Part 3: Cost and Gradient Descent
  let m = size y'
  let x = ones m 1 ||| reshape 1 x'
  let y = asColumn y'
  let xPoints = toList x'
  let yPoints = toList y'
  let theta = zeros 2 1
  let iterations = 1500
  let alpha = 0.01
  putStrLn "Testing the cost function ..."
  let j = computeCost x y theta
  putStrLn ("With theta = [0, 0]\n Cost computed = " ++ show j)
  putStrLn "Expected cost value (approx) 32.07"
  let k = computeCost x y (asColumn $ 2 |> [-1, 2])
  putStrLn ("With theta = [-1, 2]\n Cost computed = " ++ show k)
  putStrLn "Expected cost value (approx) 54.24"

  putStrLn "Running Gradient Descent ..."
  let (costHist, thetaOpt) = gradientDescent  x y theta alpha iterations
  putStrLn ("Theta found by gradient descent: " ++ show thetaOpt)
  putStrLn "Expected theta values (approx)\n"
  putStrLn " -3.6303\n  1.1664"
  let yOptPoints = (toList (flatten $ x <> thetaOpt))
  putStrLn " -3.6303\n  1.1664"
  let yOptPoints = (toList (flatten $ x <> thetaOpt))
  GP.plotSync DefaultTerm.cons $ plotLine xPoints yOptPoints `mappend` plotPoints xPoints yPoints

  -- Predicting values
  let predict1 = row [1, 3.5] <> thetaOpt
  putStrLn "For population = 35,000 we predict a profit of "
  print ((predict1 * 1000) `atIndex` (0,0))

  let predict2 = row [1, 20] <> thetaOpt
  putStrLn "For population = 70,000 we predict a profit of "
  print ((predict2 * 1000) `atIndex` (0,0))

  -- Visualising J(theta_0, theta_1)
  let theta0 = toList (linspace 100 (-10 :: Double, 10 :: Double))
  let theta1 = toList (linspace 100 (-1  :: Double, 4  :: Double))
  let jVals = [computeCost x y (col [i, j]) | i <- theta0, j <- theta1]
  -- TODO(mchav): Draw contour plot as well.
  GP.plotSync DefaultTerm.cons $ Plot3D.surface theta0 theta1 (\i j -> computeCost x y (col [i, j]))

exerciseOneOptional = do
  dataTwo <- exerciseOneDataTwo
  columns <- fmap toColumns (loadMatrix dataTwo)
  let y' = last columns
  let y = asColumn y'
  let x' = fromColumns (init columns)
  let m = size y'

  let (xNorm, mu, sigma) = featureNormalize x'
  let x = ones m 1 ||| xNorm

  let theta = zeros 3 1
  let iterations = 50 :: Int
  let alpha = 1

  -- TODO(mchav): Include cost history in gradient descent function and then plot it.
  let (costHist, thetaOpt) = gradientDescent x y theta alpha iterations
  putStrLn ("Theta found by gradient descent: " ++ show thetaOpt)
  -- Plot the graph ingoring the first (zeroth) iteration.
  GP.plotSync DefaultTerm.cons $ plotLine [1..(fromIntegral iterations)] ((tail . reverse) costHist)
  putStrLn "Predicted price of a 1650 sq-ft, 3 br house (using gradient descent)"
  let xTest = ones 1 1 ||| ((row [1650, 3] - mu) / sigma)
  print (xTest <> thetaOpt)
  let x'' = ones m 1 ||| x'
  let normTheta = normalEquation x'' y
  print normTheta
  let xTest' = ones 1 1 ||| row [1650, 3]
  print (xTest' <> normTheta)


module FFT.FFTControl(
    fftbflySPipeline,
    fftdc,
    fftdcPipeline
) where
import Data.Complex
import FFT.Samples
import System.Environment
import Control.Parallel
import Strategies
import Criterion.Main
import FFT.Orig

-- twiddle factors
tw :: Int -> Int -> Complex Float
tw n k = cis (-2 * pi * fromIntegral k / fromIntegral n)

fftdc as = dc split threshold combine worker as
	where
                combine [xs] = xs
		combine [xs,ys] = interleave xs ys
		threshold = (\x -> length x <= 1) --(floor $ sqrt $ fromIntegral as))
		worker a = [a]
		split xs = [cs, ds]
			where (cs, ds) = bflyS xs

fftdcPipeline as = dc split threshold combine worker as
	where
                combine [xs] = xs
		combine [xs,ys] = interleave xs ys
		threshold = (\x -> length x <= 1) --(floor $ sqrt $ fromIntegral as))
		worker a = [a]
		split xs = [cs, ds]
			where (cs, ds) = bflySPipeline xs

fftbflySPipeline as = interleave ls rs
    where
        (cs,ds) = bflySPipeline as
        ls = fft cs
        rs = fft ds

interleave :: [Complex Float] -> [Complex Float] -> [Complex Float]
interleave [] bs = bs
interleave (a:as) bs = a : interleave bs as

bflyS :: [Complex Float] -> ([Complex Float], [Complex Float])
bflyS as = (los,rts)
  where
    (ls,rs) = halve as
    los = zipWith (+) ls rs
    ros = zipWith (-) ls rs
    rts = zipWith (*) ros [tw (length as) i | i <- [0..(length ros) - 1]]

bflySPipeline :: [Complex Float] -> ([Complex Float], [Complex Float])
bflySPipeline as = (los,rts)
    where
        (ls, rs) = halve as
        los = zipWith (+) ls rs
        rts = parpipeline2 (zipWith (-) ls) (zipWith (*) [tw (length as) i | i <- [0..(length rs) - 1]]) rs

-- split the input into two halves
halve as = splitAt n' as
  where
    n' = div (length as + 1) 2


module Main where

import Criterion.Main (defaultMain, bgroup, bench, nf, nfIO, whnf)
import qualified Numeric.LinearAlgebra as LA
import qualified Numeric.Morpheus as M
import qualified Naive as N

main :: IO ()
main = do
  !ar <- LA.rand 1000 1000
  let !ac = LA.tr ar

  defaultMain [
    bgroup "Activation" [
        bench "sigmoid"            $ nf M.sigmoid ar
        , bench "sigmoid gradient" $ nf M.sigmoidGradient ar
        ]
    , bgroup "Matrix Reduce" [
        bench "sum by columns - row major"          $ nf M.columnSum ar
        , bench "sum by columns - column major"       $ nf M.columnSum ac
        , bench "sum by rows - row major"               $ nf M.rowSum ar
        , bench "sum by rows - column major"          $ nf M.rowSum ac
        , bench "max index in columns - row major"    $ nf M.columnMaxIndex ar
        , bench "max index in columns - column major" $ nf M.columnMaxIndex ac
        , bench "max index in rows - row major"       $ nf M.rowMaxIndex ar
        , bench "max index in rows - column major"    $ nf M.rowMaxIndex ac
        ]
    , bgroup "Matrix Reduce (naive)" [
        bench "max index in columns - row major"    $ nf (N.reduceByColumnsV (fromIntegral . LA.maxIndex) ) ar
        , bench "max index in columns - column major" $ nf (N.reduceByColumnsV (fromIntegral . LA.maxIndex) ) ac
        , bench "max index in rows - row major"         $ nf (N.reduceByRowsV (fromIntegral . LA.maxIndex) ) ar
        , bench "max index in rows - column major"    $ nf (N.reduceByRowsV (fromIntegral . LA.maxIndex) ) ac
        ]
    ]


{-# OPTIONS_GHC -Wall #-}

import           Numeric.Sundials.ARKode.ODE
import           Numeric.LinearAlgebra

import           Data.List (intercalate)

import           Text.PrettyPrint.HughesPJClass

butcherTableauTex :: ButcherTable -> String
butcherTableauTex (ButcherTable m c b b2) =
  render $
  vcat [ text ("\n\\begin{array}{c|" ++ (concat $ replicate n "c") ++ "}")
       , us
       , text "\\hline"
       , text bs <+> text "\\\\"
       , text b2s <+> text "\\\\"
       , text "\\end{array}"
       ]
  where
    n = rows m
    rs = toLists m
    ss = map (\r -> intercalate " & " $ map show r) rs
    ts = zipWith (\i r -> show i ++ " & " ++ r) (toList c) ss
    us = vcat $ map (\r -> text r <+> text "\\\\") ts
    bs  = " & " ++ (intercalate " & " $ map show $ toList b)
    b2s = " & " ++ (intercalate " & " $ map show $ toList b2)

main :: IO ()
main = do

  let res = butcherTable (SDIRK_2_1_2 undefined)
  putStrLn $ show res
  putStrLn $ butcherTableauTex res

  let resA = butcherTable (KVAERNO_4_2_3 undefined)
  putStrLn $ show resA
  putStrLn $ butcherTableauTex resA

  let resB = butcherTable (SDIRK_5_3_4 undefined)
  putStrLn $ show resB
  putStrLn $ butcherTableauTex resB

  let resC = butcherTable (FEHLBERG_6_4_5 undefined)
  putStrLn $ show resC
  putStrLn $ butcherTableauTex resC


{-#LANGUAGE ScopedTypeVariables#-}
module Main where

import CV.ColourUtils
import CV.DFT
import CV.Drawing
import CV.Edges
import CV.Filters
import CV.Image
import CV.ImageOp
import CV.Pixelwise
import Control.Applicative hiding ((<**>))
import qualified CV.ImageMath as IM
import qualified CV.Transforms as T
import Data.Complex

fi = fromIntegral

main = do
    Just x <- loadImage "smallLena.jpg"
    let  d :: Pixelwise (Complex D32)
         d  = fromImage $ dft x
         mpper f = \(x,y) -> f (x,y) + (sin (7*fi x))
         d2' = remapImage mpper x
         d2  = dft d2'
         d2c = d2 <# circleOp (0:+0) (130,102) 13 Filled
                  <# circleOp (0:+0) (71,102)  13 Filled

         dftd1 = (toImage d) <# circleOp (0:+0) (101,101) 30 Filled
    saveImage "dftmanip.png" $ montage (2,3) 10
        [
         logarithmicCompression $ fst $ dftSplit $ dftd1
        ,stretchHistogram $ idft dftd1
        ,stretchHistogram $ d2'
        ,logarithmicCompression $ fst $ dftSplit $ d2
        ,stretchHistogram $ fst $ dftSplit $ d2c
        ,stretchHistogram $ idft $ d2c
        ]


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TupleSections #-}
module Main where

import qualified Numeric.LinearAlgebra as Lina
import Numeric.LinearAlgebra hiding( Matrix, Vector )
import qualified Numeric.LinearAlgebra as Lina

import qualified Data.Csv as CSV
import Control.Monad.Except
import qualified Data.ByteString.Lazy as BS
import qualified Data.Vector as Vec
import Data.Char
import Foreign.C.Types( CInt )

type Matrix = Lina.Matrix Double
type Vector = Lina.Vector Double
type Label = CInt

type ErrT m a = ExceptT String m a

trainingDataFormat =
	CSV.defaultDecodeOptions

inputDataFormat =
	CSV.defaultDecodeOptions{
		CSV.decDelimiter = fromIntegral (ord ' ')
	}

-----------------------------------------------------------------
-- IO stuff:
-----------------------------------------------------------------

main :: IO ()
main =
	handleErrors $
	do
		mapM_
			(uncurry4 testWithData . uncurry testParamsFromLabels) $
			allPairs [3,5,7,8]
	where
		handleErrors x =
			do
				valOrErr <- runExceptT x
				either
					(\err -> putStrLn $ "ERROR: " ++ err)
					return
					valOrErr

testWithData :: FilePath -> FilePath -> Label -> Label -> ErrT IO ()
testWithData trainingFile1 trainingFile2 label1 label2 =
	do
		liftIO $ putStrLn $ concat $
			[ "----------------------------------------------\n"
			, "classifying to labels ", show [label1, label2]
			, " in files ", show [trainingFile1, trainingFile2]
			]
		[trainingSet1, trainingSet2] <-
			mapM (readData trainingDataFormat) $
			[ trainingFile1
			, trainingFile2
			]
		(inputLabels, inputData) <-
			prepareInputData (`elem` [fromIntegral label1, fromIntegral label2]) <$>
			readData inputDataFormat "resource/zip.test"
		let beta = calcBeta trainingSet1 trainingSet2
		let classified = classify (label1,label2) beta inputData
		let result =  calcClassificationQuality (cmap round $ inputLabels) classified
		liftIO $ putStrLn $
			descriptionString
				trainingSet1
				trainingSet2
				beta
				inputData
				classified
				result

readData :: CSV.DecodeOptions -> FilePath -> ErrT IO Matrix
readData fmtOpts path =
	(fromRows . Vec.toList) <$>
	(
		ExceptT $
		fmap (CSV.decodeWith fmtOpts CSV.NoHeader) $
		BS.readFile path
	)

-- (helpers: )
-----------------------------------------------------------------

prepareInputData :: (Double -> Bool) -> Matrix -> (Vector, Matrix)
prepareInputData filterRowsBy =
	(\rawInput ->
		(
			flatten $ (rawInput ?? (All, Take 1)),
			rawInput ?? (All, Drop 1)
		)
	)
	.
	fromLists
	.
	filter (filterRowsBy . head)
	.
	toLists

testParamsFromLabels x y =
	let
		[filePath1, filePath2] = map (("resource/train." ++) . show) [x,y]
	in
		(filePath1, filePath2, x, y)

descriptionString set1 set2 beta inputData classified result =
	unlines $
	[ concat $ ["set1 size:", show $ size set1]
	, concat $ ["set2 size:", show $ size set2]
	, concat $ ["beta size:", show $ size beta]
	, concat $ ["inputData size:", show $ size inputData]
	, concat $ ["result: --------------------"]
	, concat $ ["classification quality:", show $ result]
	]

-----------------------------------------------------------------
-- actual code:
-----------------------------------------------------------------

calcBeta :: Matrix -> Matrix -> Vector
calcBeta set1 set2 =
	flatten $ -- matrix to vector
	let
		x =
			prependOnes $
				set1
				===
				set2
		y =
			konst (-1) (count1,1)
			===
			konst 1 (count2,1)
		count1 = rows set1
		count2 = rows set2
		xTr_times_x_SAFE =
			let
				xTr_times_x = tr' x <> x
			in
				if det xTr_times_x > 0.01
				then xTr_times_x
				else
					xTr_times_x + (alpha * ident (rows xTr_times_x))
					where alpha = 0.01
	in
		inv xTr_times_x_SAFE <> tr' x <> y

calcClassificationQuality :: Lina.Vector Label -> Lina.Vector Label -> Double
calcClassificationQuality expected res =
	(/ fromIntegral (size res)) $
	sum $
	map (\x -> if x/=0 then 0 else 1) $
	toList $
	expected - res

classify :: (Label, Label) -> Vector -> Matrix -> Lina.Vector Label
classify (labelNeg, labelPos) beta input =
	cmap (assignLabels . sgn) $
		prependOnes input #> beta
	where
		sgn :: Double -> Double
		sgn x =
			let temp = signum x in
				if temp /= 0 then temp else 1
		assignLabels x
			| x < 0 = labelNeg
			| otherwise = labelPos

prependOnes m =
	konst 1 (rows m,1) ||| m

-----------------------------------------------------------------
-- utils:
-----------------------------------------------------------------

instance CSV.FromRecord Vector where
	parseRecord v =
		fmap fromList $ CSV.parseRecord v

-- return all pairs in a list
allPairs l =
	case l of
		(x:rest) ->
			map (x,) rest ++ allPairs rest
		_ -> []

uncurry4 f (a,b,c,d) = f a b c d


module FFT.DFTData(
    dftInnerMap,
    dftFoldMap,
    dftOuterMap,
    dftInnerMapReduce,
    dftInnerParMapReduce
) where
import Data.Complex
import FFT.Samples
import Control.Parallel
import Strategies

-- twiddle factors
tw :: Int -> Int -> Complex Float
tw n k = cis (-2 * pi * fromIntegral k / fromIntegral n)

dftInnerMap xs = [ sum (innerMap k) | k <- [0..n']]
  where
    innerMap k = parmap (\j -> xs!!j * tw n (j*k)) [0..n']
    n = length xs
    n' = n-1

dftFoldMap xs = [ (parfold (+) (0 :+ 0) (innerMap k)) | k <- [0..n']]
  where
    innerMap k = parmap (\j -> xs!!j * tw n (j*k)) [0..n']
    n = length xs
    n' = n-1

dftOuterMap xs = parmap innerFunc [0..n']
    where
        innerFunc k = sum [ xs!!j * tw n (j*k) | j <- [0..n']]
        n = length xs
        n' = n-1

dftInnerMapReduce xs = map (\k -> parMapReduceSimple rdeepseq (mapper k) rdeepseq reducer [0..n']) [0..n']
    where
        mapper k j = xs!!j * tw n (j*k)
        reducer = sum
        n = length xs
        n' = n-1

dftInnerParMapReduce xs = parmap (\k -> parMapReduceSimple rdeepseq (mapper k) rdeepseq reducer [0..n']) [0..n']
    where
        mapper k j = xs!!j * tw n (j*k)
        reducer = sum
        n = length xs
        n' = n-1


-- |
-- Module:      Data.Euclidean
-- Copyright:   (c) 2019 Andrew Lelechenko
-- Licence:     BSD3
-- Maintainer:  Andrew Lelechenko <andrew.lelechenko@gmail.com>
--

{-# LANGUAGE CPP                        #-}
{-# LANGUAGE DefaultSignatures          #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MagicHash                  #-}

module Data.Euclidean
  ( Euclidean(..)
  , Field
  , GcdDomain(..)
  , WrappedIntegral(..)
  , WrappedFractional(..)
  , gcdExt
  ) where

import Prelude hiding (quotRem, quot, rem, divMod, div, mod, gcd, lcm, negate, (*), Int, Word)
import qualified Prelude as P
import Control.Exception (throw, ArithException(..))
import Data.Bits (Bits)
import Data.Complex (Complex(..))
import Data.Int (Int, Int8, Int16, Int32, Int64)
import Data.Maybe (isJust)
import Data.Ratio (Ratio)
import Data.Semiring (Semiring(..), Ring(..), (*), minus, isZero, Mod2)
import Data.Word (Word, Word8, Word16, Word32, Word64)
import Foreign.C.Types (CFloat, CDouble)
import GHC.Exts (Int(..), Word(..))
import GHC.Integer.GMP.Internals (gcdInt, gcdWord, gcdInteger, lcmInteger)

import Numeric.Natural

---------------------------------------------------------------------
-- Classes
---------------------------------------------------------------------

-- | 'GcdDomain' represents a
-- <https://en.wikipedia.org/wiki/GCD_domain GCD domain>.
-- This is a domain, where GCD can be defined,
-- but which does not necessarily allow a well-behaved
-- division with remainder (as in 'Euclidean' domains).
--
-- For example, there is no way to define 'rem' over
-- polynomials with integer coefficients such that
-- remainder is always "smaller" than divisor. However,
-- 'gcd' is still definable, just not by means of
-- Euclidean algorithm.
--
-- All methods of 'GcdDomain' have default implementations
-- in terms of 'Euclidean'. So most of the time
-- it is enough to write:
--
-- > instance GcdDomain Foo
-- > instance Euclidean Foo where
-- >   quotRem = ...
-- >   degree  = ...
class Semiring a => GcdDomain a where
  -- | Division without remainder.
  --
  -- prop> \x y -> (x * y) `divide` y == Just x
  -- prop> \x y -> maybe True (\z -> x == z * y) (x `divide` y)
  divide :: a -> a -> Maybe a

  default divide :: (Eq a, Euclidean a) => a -> a -> Maybe a
  divide x y = let (q, r) = quotRem x y in
    if isZero r then Just q else Nothing

  -- | Greatest common divisor. Must satisfy
  --
  -- prop> \x y -> isJust (x `divide` gcd x y) && isJust (y `divide` gcd x y)
  -- prop> \x y z -> isJust (gcd (x * z) (y * z) `divide` z)
  gcd :: a -> a -> a

  default gcd :: (Eq a, Euclidean a) => a -> a -> a
  gcd a b
    | isZero b  = a
    | otherwise = gcd b (a `rem` b)

  -- | Lowest common multiple. Must satisfy
  --
  -- prop> \x y -> isJust (lcm x y `divide` x) && isJust (lcm x y `divide` y)
  -- prop> \x y z -> isNothing (z `divide` x) || isNothing (z `divide` y) || isJust (z `divide` lcm x y)
  lcm :: a -> a -> a

  default lcm :: Eq a => a -> a -> a
  lcm a b
    | isZero a || isZero b = zero
    | otherwise = case a `divide` gcd a b of
      Nothing -> error "lcm: violated gcd invariant"
      Just c  -> c * b

  -- | Test whether two arguments are
  -- <https://en.wikipedia.org/wiki/Coprime_integers coprime>.
  -- Must match its default definition:
  --
  -- prop> \x y -> coprime x y == isJust (1 `divide` gcd x y)
  coprime :: a -> a -> Bool

  default coprime :: Eq a => a -> a -> Bool
  coprime x y = isJust (one `divide` gcd x y)

infixl 7 `divide`

-- | Informally speaking, 'Euclidean' is a superclass of 'Integral',
-- lacking 'toInteger', which allows to define division with remainder
-- for a wider range of types, e. g., complex integers
-- and polynomials with rational coefficients.
--
-- 'Euclidean' represents a
-- <https://en.wikipedia.org/wiki/Euclidean_domain Euclidean domain>
-- endowed by a given Euclidean function 'degree'.
--
-- No particular rounding behaviour is expected of 'quotRem'. E. g.,
-- it is not guaranteed to truncate towards zero or towards negative
-- infinity (cf. 'P.divMod'), and remainders are not guaranteed to be non-negative.
-- For a faithful representation of residue classes one can use
-- <http://hackage.haskell.org/package/mod mod> package instead.
class GcdDomain a => Euclidean a where
  {-# MINIMAL (quotRem | quot, rem), degree #-}
  -- | Division with remainder.
  --
  -- prop> \x y -> y == 0 || let (q, r) = x `quotRem` y in x == q * y + r
  quotRem :: a -> a -> (a, a)
  quotRem x y = (quot x y, rem x y)

  -- | Division. Must match its default definition:
  --
  -- prop> \x y -> quot x y == fst (quotRem x y)
  quot :: a -> a -> a
  quot x y = fst (quotRem x y)

  -- | Remainder. Must match its default definition:
  --
  -- prop> \x y -> rem x y == snd (quotRem x y)
  rem :: a -> a -> a
  rem x y = snd (quotRem x y)

  -- | Euclidean (aka degree, valuation, gauge, norm) function on @a@. Usually @'fromIntegral' '.' 'abs'@.
  --
  -- 'degree' is rarely used by itself. Its purpose
  -- is to provide an evidence of soundness of 'quotRem'
  -- by testing the following property:
  --
  -- prop> \x y -> y == 0 || let (q, r) = x `quotRem` y in (r == 0 || degree r < degree y)
  degree :: a -> Natural

infixl 7 `quot`
infixl 7 `rem`

coprimeIntegral :: Integral a => a -> a -> Bool
coprimeIntegral x y = (odd x || odd y) && P.gcd x y == 1

-- | Execute the extended Euclidean algorithm.
-- For elements @a@ and @b@, compute their greatest common divisor @g@
-- and the coefficient @s@ satisfying @as + bt = g@ for some @t@.
gcdExt :: (Eq a, Euclidean a, Ring a) => a -> a -> (a, a)
gcdExt = go one zero
  where
    go s s' r r'
      | r' == zero = (r, s)
      | otherwise  = case quotRem r r' of
        (q, r'') -> go s' (minus s (times q s')) r' r''
{-# INLINABLE gcdExt #-}

-- | 'Field' represents a
-- <https://en.wikipedia.org/wiki/Field_(mathematics) field>,
-- a ring with a multiplicative inverse for any non-zero element.
class (Euclidean a, Ring a) => Field a

---------------------------------------------------------------------
-- Instances
---------------------------------------------------------------------

instance GcdDomain () where
  divide  = const $ const (Just ())
  gcd     = const $ const ()
  lcm     = const $ const ()
  coprime = const $ const True

instance Euclidean () where
  degree  = const 0
  quotRem = const $ const ((), ())
  quot    = const $ const ()
  rem     = const $ const ()

instance Field ()

instance GcdDomain Mod2 where

instance Euclidean Mod2 where
  degree = const 0
  quotRem x y
    | isZero y  = throw DivideByZero
    | otherwise = (x, zero)

instance Field Mod2

-- | Wrapper around 'Integral' with 'GcdDomain'
-- and 'Euclidean' instances.
newtype WrappedIntegral a = WrapIntegral { unwrapIntegral :: a }
  deriving (Eq, Ord, Show, Num, Integral, Real, Enum, Bits)

instance Num a => Semiring (WrappedIntegral a) where
  plus  = (P.+)
  zero  = 0
  times = (P.*)
  one   = 1
  fromNatural = P.fromIntegral

instance Num a => Ring (WrappedIntegral a) where
  negate = P.negate

instance Integral a => GcdDomain (WrappedIntegral a) where
  divide x y = case x `P.quotRem` y of (q, 0) -> Just q; _ -> Nothing
  gcd     = P.gcd
  lcm     = P.lcm
  coprime = coprimeIntegral

instance Integral a => Euclidean (WrappedIntegral a) where
  degree  = P.fromIntegral . abs . unwrapIntegral
  quotRem = P.quotRem
  quot    = P.quot
  rem     = P.rem

instance GcdDomain Int where
  divide x y = case x `P.quotRem` y of (q, 0) -> Just q; _ -> Nothing
#if MIN_VERSION_integer_gmp(0,5,1)
  gcd (I# x) (I# y) = I# (gcdInt x y)
#else
  gcd     = P.gcd
#endif
  lcm     = P.lcm
  coprime = coprimeIntegral

instance GcdDomain Word where
  divide x y = case x `P.quotRem` y of (q, 0) -> Just q; _ -> Nothing
#if MIN_VERSION_integer_gmp(1,0,0)
  gcd (W# x) (W# y) = W# (gcdWord x y)
#else
  gcd     = P.gcd
#endif
  lcm     = P.lcm
  coprime = coprimeIntegral

instance GcdDomain Integer where
  divide x y = case x `P.quotRem` y of (q, 0) -> Just q; _ -> Nothing
  gcd     = gcdInteger
  lcm     = lcmInteger
  coprime = coprimeIntegral

#define deriveGcdDomain(ty)     \
instance GcdDomain (ty) where { \
;  divide x y = case x `P.quotRem` y of { (q, 0) -> Just q; _ -> Nothing } \
;  gcd     = P.gcd              \
;  lcm     = P.lcm              \
;  coprime = coprimeIntegral    \
}

deriveGcdDomain(Int8)
deriveGcdDomain(Int16)
deriveGcdDomain(Int32)
deriveGcdDomain(Int64)
deriveGcdDomain(Word8)
deriveGcdDomain(Word16)
deriveGcdDomain(Word32)
deriveGcdDomain(Word64)
deriveGcdDomain(Natural)

#define deriveEuclidean(ty)       \
instance Euclidean (ty) where {   \
;  degree  = P.fromIntegral . abs \
;  quotRem = P.quotRem            \
;  quot    = P.quot               \
;  rem     = P.rem                \
}

deriveEuclidean(Int)
deriveEuclidean(Int8)
deriveEuclidean(Int16)
deriveEuclidean(Int32)
deriveEuclidean(Int64)
deriveEuclidean(Integer)
deriveEuclidean(Word)
deriveEuclidean(Word8)
deriveEuclidean(Word16)
deriveEuclidean(Word32)
deriveEuclidean(Word64)
deriveEuclidean(Natural)

-- | Wrapper around 'Fractional'
-- with trivial 'GcdDomain'
-- and 'Euclidean' instances.
newtype WrappedFractional a = WrapFractional { unwrapFractional :: a }
  deriving (Eq, Ord, Show, Num, Fractional)

instance Num a => Semiring (WrappedFractional a) where
  plus  = (P.+)
  zero  = 0
  times = (P.*)
  one   = 1
  fromNatural = P.fromIntegral

instance Num a => Ring (WrappedFractional a) where
  negate = P.negate

instance Fractional a => GcdDomain (WrappedFractional a) where
  divide x y = Just (x / y)
  gcd        = const $ const 1
  lcm        = const $ const 1
  coprime    = const $ const True

instance Fractional a => Euclidean (WrappedFractional a) where
  degree      = const 0
  quotRem x y = (x / y, 0)
  quot        = (/)
  rem         = const $ const 0

instance Fractional a => Field (WrappedFractional a)

instance Integral a => GcdDomain (Ratio a) where
  divide x y = Just (x / y)
  gcd        = const $ const 1
  lcm        = const $ const 1
  coprime    = const $ const True

instance Integral a => Euclidean (Ratio a) where
  degree      = const 0
  quotRem x y = (x / y, 0)
  quot        = (/)
  rem         = const $ const 0

instance Integral a => Field (Ratio a)

instance GcdDomain Float where
  divide x y = Just (x / y)
  gcd        = const $ const 1
  lcm        = const $ const 1
  coprime    = const $ const True

instance Euclidean Float where
  degree      = const 0
  quotRem x y = (x / y, 0)
  quot        = (/)
  rem         = const $ const 0

instance Field Float

instance GcdDomain Double where
  divide x y = Just (x / y)
  gcd        = const $ const 1
  lcm        = const $ const 1
  coprime    = const $ const True

instance Euclidean Double where
  degree      = const 0
  quotRem x y = (x / y, 0)
  quot        = (/)
  rem         = const $ const 0

instance Field Double

instance GcdDomain CFloat where
  divide x y = Just (x / y)
  gcd        = const $ const 1
  lcm        = const $ const 1
  coprime    = const $ const True

instance Euclidean CFloat where
  degree      = const 0
  quotRem x y = (x / y, 0)
  quot        = (/)
  rem         = const $ const 0

instance Field CFloat

instance GcdDomain CDouble where
  divide x y = Just (x / y)
  gcd        = const $ const 1
  lcm        = const $ const 1
  coprime    = const $ const True

instance Euclidean CDouble where
  degree      = const 0
  quotRem x y = (x / y, 0)
  quot        = (/)
  rem         = const $ const 0

instance Field CDouble

conjQuotAbs :: Field a => Complex a -> Complex a
conjQuotAbs (x :+ y) = x `quot` norm :+ (negate y) `quot` norm
  where
    norm = (x `times` x) `plus` (y `times` y)

instance Field a => GcdDomain (Complex a) where
  divide x y = Just (x `times` conjQuotAbs y)
  gcd        = const $ const one
  lcm        = const $ const one
  coprime    = const $ const True

instance Field a => Euclidean (Complex a) where
  degree      = const 0
  quotRem x y = (quot x y, zero)
  quot x y    = x `times` conjQuotAbs y
  rem         = const $ const zero

instance Field a => Field (Complex a)