module LinearAlgebra where

import qualified Numeric.LAPACK.Matrix.Triangular as Triangular
import qualified Numeric.LAPACK.Matrix.Layout as Layout
import qualified Numeric.LAPACK.Matrix as Matrix
import qualified Numeric.LAPACK.Singular as Singular
import qualified Numeric.LAPACK.Vector as Vector
import Numeric.LAPACK.Matrix (ShapeInt, (#*|))
import Numeric.LAPACK.Format ((##))

import qualified Combinatorics

import qualified Data.Array.Comfort.Boxed as BoxedArray
import qualified Data.Array.Comfort.Storable as Array
import qualified Data.Array.Comfort.Shape as Shape
import qualified Data.IntSet as IntSet
import qualified Data.Set as Set
import Data.Array.Comfort.Storable (Array)

import Data.Foldable (for_)
import Data.Tuple.HT (mapSnd)



{- $setup
>>> import qualified Data.Array.Comfort.Storable as Array
>>> import Data.Tuple.HT (mapSnd)
>>>
>>> myRound :: Double -> Integer
>>> myRound = round
-}


example0 :: [Double]
example0 = [3,1,4,1,5]

{-
        ...
      ... ...
    ... ... ...
  ... 100 ... .79
.31 .41 ... .26 ...
-}
example2 :: [[Maybe Double]]
example2 =
   let __ = Nothing; d = Just in
                [__] :
             [__,    __] :
          [__,   __,    __] :
       [__,  d 100,  __,  d 79] :
   [d 31, d 41,  __,   d 26,  __] :
   []

pyramid ::
   Array ShapeInt Double ->
   Array (Shape.LowerTriangular ShapeInt) Double
pyramid xs =
   let shape = Array.shape xs
       baseRow = Shape.size shape - 1
       arr =
         fmap (\(i,j) ->
            if i==baseRow
               then xs Array.! j
               else arr BoxedArray.! (i+1,j) + arr BoxedArray.! (i+1,j+1)) $
         BoxedArray.indices $ Shape.lowerTriangular shape
   in Array.fromBoxed arr

basis ::
   ShapeInt -> Matrix.General ShapeInt (Shape.LowerTriangular ShapeInt) Double
basis shape@(Shape.ZeroBased n) =
   Matrix.fromRows (Shape.lowerTriangular shape) $
   map (pyramid . Vector.unit shape) $ take n [0..]

addIndices :: [[Maybe a]] -> [((Int,Int), a)]
addIndices puzzle = do
   (i, xs) <- zip [0..] puzzle
   (j, Just x) <- zip [0..] xs
   return ((i,j),x)

{- |
>>> mapSnd (map myRound . Array.toList) $ solve 3 [((0,0),8), ((2,0),1), ((2,2),3)]
(3,[8,3,5,1,2,3])
-}
solve ::
   Int -> [((Int,Int),Double)] ->
   (Int, Array (Shape.LowerTriangular ShapeInt) Double)
solve n indexed =
   let fullBasis = Matrix.transpose $ basis (Matrix.shapeInt n)
       selected =
         Matrix.takeRowArray
            (BoxedArray.vectorFromList (map fst indexed))
            fullBasis
   in mapSnd ((fullBasis #*|) . Matrix.flattenColumn)
         (Singular.leastSquaresMinimumNormRCond 1e-5 selected $
          Matrix.singleColumn Layout.ColumnMajor $
          Vector.autoFromList (map snd indexed))

solvable :: Int -> [(Int,Int)] -> Bool
solvable n =
   let shape = Matrix.shapeInt n
       fullBasis = basis shape
   in \ixs ->
         (n==) $ length $ takeWhile (1e-5<) $ Vector.toList $
         (#*| Vector.one shape) $ Singular.values $
         Matrix.takeColumnArray (BoxedArray.vectorFromList ixs) fullBasis

solvables :: Int -> [[(Int,Int)]]
solvables n =
   let check = solvable n
   in filter check $
      Combinatorics.tuples n $ Shape.indices $
      Shape.lowerTriangular $ Matrix.shapeInt n

{-
Check, whether a sum pyramid contains a sub-pyramid of size k
with more than k given fields.
If yes, then the pyramid has redundancies in a sub-pyramid
and is not solvable.
-}
wellcrowded :: Int -> [(Int,Int)] -> Bool
wellcrowded n ixs =
   let triShape = Shape.lowerTriangular $ Matrix.shapeInt n
       set = Set.fromList ixs
       countSubTriangle (i,j) k =
         length $
         filter (\(si,sj) -> Set.member (i+si,j+sj) set) $
         Shape.indices $ Shape.lowerTriangular $ Matrix.shapeInt k
   in and $ do
         (i,j) <- Shape.indices triShape
         k <- [0..n-i]
         return $ countSubTriangle (i,j) k <= k

wellcrowdedIntSet :: Int -> [(Int,Int)] -> Bool
wellcrowdedIntSet n ixs =
   let triShape = Shape.lowerTriangular $ Matrix.shapeInt n
       set = IntSet.fromList $ map (Shape.offset triShape) ixs
       countSubTriangle (i,j) k =
         length $
         filter (\(si,sj) ->
            flip IntSet.member set $ Shape.offset triShape (i+si,j+sj)) $
         Shape.indices $ Shape.lowerTriangular $ Matrix.shapeInt k
   in and $ do
         (i,j) <- Shape.indices triShape
         k <- [0..n-i]
         return $ countSubTriangle (i,j) k <= k


{-
Check whether the linear independence criterion matches
the subpyramid criterion 'wellcrowded'.
Well, it does not, the smallest counterexample is:

   *
  . .
 . * .
* . . *
-}
counterexamples :: Int -> [[(Int, Int)]]
counterexamples n =
   let check = solvable n
   in filter (\ixs -> check ixs /= wellcrowded n ixs) $
      Combinatorics.tuples n $ Shape.indices $
      Shape.lowerTriangular $ Matrix.shapeInt n


boolMatrix :: Int -> [(Int, Int)] -> Matrix.Lower ShapeInt Float
boolMatrix n ixs =
   Triangular.fromLowerRowMajor $
   Array.fromAssociations 0
      (Shape.lowerTriangular $ Matrix.shapeInt n)
      (map (\ix -> (ix,1::Float)) ixs)


test :: IO ()
test = do
   Triangular.fromLowerRowMajor (pyramid (Vector.autoFromList example0))
      ## "%.0f"

   let xs = example2
    in mapSnd Triangular.fromLowerRowMajor (solve (length xs) (addIndices xs))
          ## "%.0f"

   putStrLn "\nsolvable:"
   let n = 3
    in for_ (solvables n) $ \ixs -> do
          putStrLn ""
          boolMatrix n ixs ## "%.0f"

   putStrLn "\nunsolvable:"
   for_ [0..5] $ \n ->
      for_ (counterexamples n) $ \ixs -> do
         putStrLn ""
         boolMatrix n ixs ## "%.0f"

{-
https://oeis.org/A014068

map (\n -> Comb.binomial (div (n*(n+1)) 2) n) [1..10::Integer]

Possibilities of choosing n numbers from a n-sized pyramid.
-}