module Previous where

import Control.Applicative (liftA2)
import Data.Array
import Data.List (unfoldr, transpose)
import Data.Tuple (swap)

zipAdd :: Num a => [a] -> [a] -> [a]
zipAdd []     ds     = ds
zipAdd cs     []     = cs
zipAdd (c:cs) (d:ds) = (c + d):zipAdd cs ds

{------------------------------------------------------------------------------
 -  ___     _                     _      _
 - | _ \___| |_  _ _ _  ___ _ __ (_)__ _| |___
 - |  _/ _ \ | || | ' \/ _ \ '  \| / _` | (_-<
 - |_| \___/_|\_, |_||_\___/_|_|_|_\__,_|_/__/
 -            |__/
 ------------------------------------------------------------------------------}

-- A polynomial is its ascending list of coefficients (of type a)
newtype Polynomial a = Poly { coeffs :: [a] } deriving (Functor)

instance (Eq a, Num a) => Num (Polynomial a) where
  (+) x@(Poly as) y@(Poly bs) = Poly $ zipAdd as bs

  --convolution
  (*) (Poly []) (Poly bs) = Poly []
  (*) (Poly as) (Poly []) = Poly []
  (*) (Poly as) (Poly bs) = Poly $ zeroize $ finish $ foldl convolve' ([], []) as
    where convolve' (xs, ys) a = (a:xs, sum (zipWith (*) (a:xs) bs):ys)
          finish (xs, ys) = (reverse ys ++) $ finish' xs $ tail bs
          finish' xs [] = []
          finish' xs ys = sum (zipWith (*) xs ys):finish' xs (tail ys)
          zeroize xs    = if all (==0) xs then [] else xs

  -- this definition works
  negate = fmap negate
  -- but these two might run into some problems
  abs    = fmap abs
  signum = fmap signum
  fromInteger 0    = Poly []
  fromInteger a    = Poly [fromInteger a]

instance (Eq a, Num a) => Eq (Polynomial a) where
  (==) as bs = all (==0) $ coeffs $ as - bs


-- Interpret a number's base-b expansion as a polynomial
-- Build a list with f, which returns either Nothing
-- or Just (next element of list, next argument to f)
asPoly :: Int -> Int -> Polynomial Int
asPoly b = Poly . unfoldr f where
  -- Divide x by b. Emit the remainder and recurse with the quotient.
  f x | x /= 0    = Just $ swap $ divMod x b
  -- If there's nothing left to divide out, terminate
      | otherwise = Nothing

-- Horner evaluation of a polynomial at the integer b
-- Start with the highest coefficient
-- Multiply by b at each step and add the coefficient of the next term
evalPoly b (Poly p) = foldr (\y acc -> b*acc + y) 0 p

-- Divide the polynomial ps by qs (coefficients in descending degree order)
synthDiv' :: (Eq a, Num a) => [a] -> [a] -> ([a], [a])
synthDiv' ps qs
  | head qs /= 1 = error "Cannot divide by non-monic polynomial"
  | otherwise   = splitAt deg $ doDiv ps deg
  where
    -- Negate the denominator and ignore leading term
    qNeg = map negate $ tail qs
    -- The degree of the result, based on the degrees of the numerator and denominator
    deg   = max 0 (length ps - length qs + 1)
    -- Pluck off the head of the list and add a shifted and scaled version of
    -- qs to the tail of the list. Repeat this d times
    doDiv xs    0  = xs
    doDiv (x:xs) d = x:doDiv (zipAdd xs $ map (*x) qNeg) (d - 1)

-- Use Polynomial (coefficients in ascending degree order) instead of lists
synthDiv :: (Eq a, Num a) => Polynomial a -> Polynomial a
  -> (Polynomial a, Polynomial a)
synthDiv (Poly p) (Poly q) = (Poly $ reverse quot, Poly $ reverse rem)
  where (quot, rem)        = synthDiv' (reverse p) (reverse q)


{------------------------------------------------------------------------------
 -  ___     _                     _      _    ___
 - | _ \___| |_  _ _ _  ___ _ __ (_)__ _| |  / __| ___ __ _ _  _ ___ _ _  __ ___ ___
 - |  _/ _ \ | || | ' \/ _ \ '  \| / _` | |  \__ \/ -_) _` | || / -_) ' \/ _/ -_|_-<
 - |_| \___/_|\_, |_||_\___/_|_|_|_\__,_|_|  |___/\___\__, |\_,_\___|_||_\__\___/__/
 -            |__/                                       |_|
 ------------------------------------------------------------------------------}

-- All monic polynomials of degree d with coefficients mod n
monics :: Int -> Int -> [Polynomial Int]
monics n d = map (asPoly n) [n^d..2*(n^d) - 1]

-- All monic polynomials with coefficients mod n, ordered by degree
allMonics :: Int -> [Polynomial Int]
allMonics n = concat [monics n d | d <- [1..]]

-- All irreducible monic polynomials with coefficients mod n
irreducibles :: Int -> [Polynomial Int]
irreducibles n = go [] $ allMonics n where
  -- Divide the polynomial x by i, then take the remainder mod n
  remModN x i      = fmap (`mod` n) $ snd $ synthDiv x i
  -- Find remainders of x divided by every irreducible in "is".
  -- If any give the zero polynomial, then x is a multiple of an irreducible
  notMultiple x is = and [not $ all (==0) $ coeffs $ remModN x i | i <- is]
  -- Sieve out by notMultiple
  go is (x:xs)
    | notMultiple x is = x:go (x:is) xs
    | otherwise        = go is xs


{------------------------------------------------------------------------------
 -  __  __      _       _
 - |  \/  |__ _| |_ _ _(_)__ ___ ___
 - | |\/| / _` |  _| '_| / _/ -_|_-<
 - |_|  |_\__,_|\__|_| |_\__\___/__/
 -
 ------------------------------------------------------------------------------}

newtype Matrix a = Mat { unMat :: Array (Int, Int) a } deriving (Functor)

asScalar = Mat . listArray ((0,0),(0,0)) . pure

transposeM :: Matrix a -> Matrix a
transposeM (Mat m) = Mat $ ixmap (bounds m) swap m

dotMul :: Num a => [a] -> [a] -> a
dotMul = (sum .) . zipWith (*)

matMul' xs ys = reshape (length ys) $ liftA2 dotMul xs ys
matMul  xs ys = matMul' xs $ transpose ys

-- Pointwise application
zipWithArr :: Ix i => (a -> b -> c)
  -> Array i a -> Array i b -> Array i c
zipWithArr f a b
  | ab == bb  = array ab $ map (\x -> (x, f (a!x) (b!x))) $ indices a
  | otherwise = error "Array dimension mismatch" where
    ab = bounds a
    bb = bounds b

-- Maps
mapRange :: Ix i => (i -> e) -> (i, i) -> [(i, e)]
mapRange g r = map (\x -> (x, g x)) $ range r

instance (Num a, Eq a) => Eq (Matrix a) where
  (==) (Mat x) (Mat y)
    | (0,0) == (snd $ bounds x) = and [x!(0,0) == if m == n then u else 0 | ((m,n), u) <- assocs y]
    | (0,0) == (snd $ bounds y) = and [x!(0,0) == if m == n then u else 0 | ((m,n), u) <- assocs x]
    | otherwise                 = x == y


instance Num a => Num (Matrix a) where
  (+) (Mat x) (Mat y)
    | (0,0) == (snd $ bounds x) && (0,0) == (snd $ bounds y) = Mat $ zipWithArr (+) x y
    | (0,0) == (snd $ bounds x) = Mat y + (((x!(0,0)) *) <$> eye ((+1) $ snd $ snd $ bounds y)) --adding scalars
    | (0,0) == (snd $ bounds y) = Mat x + (((y!(0,0)) *) <$> eye ((+1) $ snd $ snd $ bounds x)) --adding scalars
    | otherwise                 = Mat $ zipWithArr (+) x y
  (*) x y
    | (0,0) == (snd $ bounds $ unMat x) = fmap (((unMat x)!(0,0))*) y --multiplying scalars
    | (0,0) == (snd $ bounds $ unMat y) = fmap (((unMat y)!(0,0))*) x --multiplying scalars
    | otherwise                         = toMatrix $ matMul' (fromMatrix x) (fromMatrix $ transposeM y)
  abs = fmap abs
  negate = fmap negate
  signum = undefined
  fromInteger = asScalar . fromInteger --"scalars"

reshape :: Int -> [a] -> [[a]]
reshape n = unfoldr (reshape' n) where
  reshape' n x = if null x then Nothing else Just $ splitAt n x

-- Simple function for building a Matrix from lists
toMatrix :: [[a]] -> Matrix a
toMatrix l = Mat $ listArray ((0,0),(n-1,m-1)) $ concat l where
  m = length $ head l
  n = length l

-- ...and its inverse
fromMatrix :: Matrix a -> [[a]]
fromMatrix (Mat m) = let (_,(_,n)) = bounds m in reshape (n+1) $ elems m

-- Zero matrix
zero :: Num a => Int -> Matrix a
zero n = Mat $ listArray ((0,0),(n-1,n-1)) $ repeat 0

-- Identity matrix
eye :: Num a => Int -> Matrix a
eye n = Mat $ unMat (zero n) // take n [((i,i), 1) | i <- [0..]]

-- Companion matrix
companion :: (Eq a, Num a) => Polynomial a -> Matrix a
companion (Poly ps)
  | last ps' /= 1 = error "Cannot find companion matrix of non-monic polynomial"
  | otherwise     = Mat $ array ((0,0), (n-1,n-1)) $ lastRow ++ shiftI where
    -- The degree of the polynomial, as well as the size of the matrix
    n       = length ps' - 1
    -- Remove trailing 0s from ps
    ps'     = reverse $ dropWhile (==0) $ reverse ps
    -- Address/value tuples for a shifted identity matrix:
    -- 1s on the diagonal just above the main diagonal, 0s elsewhere
    shiftI  = map (\p@(x,y) -> (p, if y == x + 1 then 1 else 0)) $
      range ((0,0),(n-2,n-1))
    -- Address/value tuples for the last row of the companion matrix:
    -- ascending powers of the polynomial
    lastRow = zipWith (\x y -> ((n-1, x), y)) [0..n-1] $ map negate ps'

mDims = snd . bounds . unMat

-- trace of a matrix: sum of diagonal entries
mTrace :: (Num a) => Matrix a -> a
mTrace a = sum [(unMat a)!(i,i) | i <- [0..(uncurry min $ mDims a)]]

-- Faddeev-Leverrier algorithm for computing the determinant, characteristic polynomial, and inverse matrix
-- this relies on dividing the trace by the current iteration count, so it is only suitable for Integral types
fadeev :: (Num a, Integral a) => Matrix a -> ([a], a, Matrix a)
fadeev a = fadeev' [1] (eye $ n+1) 1 where
  n = snd $ mDims a
  fadeev' rs m i
    | all (==0) (unMat nextm) = (replicate (n - (fromIntegral i) + 1) 0 ++ tram:rs , -tram, m)
    | otherwise               = fadeev' (tram:rs) nextm (i+1) where
      am    = a * m
      tram  = -mTrace am `div` i
      nextm = am + (fmap (tram *) $ eye (n+1))

-- determinant from Faddeev-Leverrier algorithm
determinant :: (Num a, Integral a) => Matrix a -> a
determinant = (\(_,x,_) -> x) . fadeev
-- characteristic polynomial from Faddeev-Leverrier algorithm
charpoly :: (Num a, Integral a) => Matrix a -> Polynomial a
charpoly    = Poly . (\(x,_,_) -> x) . fadeev