{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
module Main where

import qualified Numeric.LAPACK.Orthogonal as Ortho
import qualified Numeric.LAPACK.Matrix.Shape as MatrixShape
import qualified Numeric.LAPACK.Matrix.Layout as Layout
import qualified Numeric.LAPACK.Matrix.Square as Square
import qualified Numeric.LAPACK.Matrix.Array as ArrMatrix
import qualified Numeric.LAPACK.Matrix as Matrix
import qualified Numeric.LAPACK.Vector as Vector
import qualified Numeric.LAPACK.Shape as ExtShape
import qualified Numeric.Netlib.Class as Class
import Numeric.LAPACK.Matrix
         (ShapeInt,
          (#-#), (#*#), (#\##), (#*##), (##*#), (#*|), (#*\), (\*#), (|||))

import qualified Data.Array.Comfort.Storable as Array
import qualified Data.Array.Comfort.Boxed as BoxedArray
import qualified Data.Array.Comfort.Shape as Shape
import Data.Array.Comfort.Shape ((::+))
import Data.Array.Comfort.Storable ((!))

import qualified Data.Complex as Complex
import Data.Complex (Complex)

import Text.Printf (printf)

import qualified Data.List.Key as Key
import qualified Data.List as List
import Data.Semigroup ((<>))
import Data.Tuple.HT (swap)


type ShapeInt2 = ShapeInt ::+ ShapeInt

data Score = P1 | P2 | P3 | P4 | P5 | P6 deriving (Eq, Ord, Enum, Bounded, Show)
type ScoreShape = Shape.Enumeration Score
type TriShape = Shape.UpperTriangular ScoreShape
type TriShapeInt = ExtShape.IntIndexed TriShape
type Transition = Matrix.Square TriShape
type Transformation = Matrix.Square TriShapeInt
type TallMatrix = Matrix.Tall TriShape ShapeInt
type Matrix = Matrix.General TriShape TriShape
type Vector = Vector.Vector TriShape
type Eigenvalues = Vector.Vector TriShapeInt


triShape :: TriShape
triShape = Shape.upperTriangular Shape.Enumeration

transition :: Transition Double
transition =
   Square.fromFull $ Matrix.scale (1/6) $
   Matrix.fromRowArray triShape $ BoxedArray.fromList triShape $
   map (Vector.fromList triShape) $
      [6,1,0,0,0,0,6,0,0,0,0,0,0,0,0,0,0,0,0,0,0] :
      [0,1,1,0,0,0,0,6,0,0,0,0,0,0,0,0,0,0,0,0,0] :
      [0,1,0,1,0,0,0,0,6,0,0,0,0,0,0,0,0,0,0,0,0] :
      [0,1,0,0,1,0,0,0,0,6,0,0,0,0,0,0,0,0,0,0,0] :
      [0,1,0,0,0,1,0,0,0,0,6,0,0,0,0,0,0,0,0,0,0] :
      [0,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0] :
      [0,0,1,0,0,0,0,0,0,0,0,6,0,0,0,0,0,0,0,0,0] :
      [0,0,1,1,0,0,0,0,0,0,0,0,6,0,0,0,0,0,0,0,0] :
      [0,0,1,0,1,0,0,0,0,0,0,0,0,6,0,0,0,0,0,0,0] :
      [0,0,1,0,0,1,0,0,0,0,0,0,0,0,6,0,0,0,0,0,0] :
      [0,0,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0] :
      [0,0,0,1,0,0,0,0,0,0,0,0,0,0,0,6,0,0,0,0,0] :
      [0,0,0,1,1,0,0,0,0,0,0,0,0,0,0,0,6,0,0,0,0] :
      [0,0,0,1,0,1,0,0,0,0,0,0,0,0,0,0,0,6,0,0,0] :
      [0,0,0,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0] :
      [0,0,0,0,1,0,0,0,0,0,0,0,0,0,0,0,0,0,6,0,0] :
      [0,0,0,0,1,1,0,0,0,0,0,0,0,0,0,0,0,0,0,6,0] :
      [0,0,0,0,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0] :
      [0,0,0,0,0,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,6] :
      [0,0,0,0,0,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0] :
      [0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0] :
      []

-- cf. Graph.move, roll
transitionAlt :: Transition Double
transitionAlt =
   Matrix.scale (1/6) $ Matrix.transpose $ ArrMatrix.fromVector $
   Array.fromAssociations 0 (Layout.square MatrixShape.RowMajor triShape) $
      ((((P1,P1),(P1,P1)), 6) :
       [(((a,b),(pred a, pred b)), 6) | a<-[P2 .. P6], b<-[a .. P6]] ++
       [((from,to), 1) |
         b<-[P2 .. P6], let b' = pred b, a<-[P1 .. P6],
         let from = (P1,b); to = (min a b', max a b')])

initial1 :: Vector Double
initial1 =
   Vector.scale (1/6) $
   Vector.fromList triShape [1,1,1,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]

initial2 :: Vector Double
initial2 =
   Vector.scale (1/36) $
   Vector.fromList triShape [1,2,2,2,2,2,1,2,2,2,2,1,2,2,2,1,2,2,1,2,1]

eigenSystem ::
   (a ~ Complex Double) =>
   (Transformation a, Eigenvalues a, Transformation a)
eigenSystem =
   let (vr,d,vlAdj) =
         Square.eigensystem $ Square.mapSize ExtShape.IntIndexed transition
       scal = Square.takeDiagonal $ vlAdj <> vr
   in (vr, d, Vector.recip scal \*# vlAdj)

similarityTransformation ::
   (Class.Floating a, Vector.RealOf a ~ ar, Class.Real ar) =>
   (Transformation a, Eigenvalues a, Transformation a) -> Transition a
similarityTransformation (trafo, diag, trafoInv) =
   Square.mapSize ExtShape.deconsIntIndexed $ trafo #*\ diag ##*# trafoInv

reconstruct :: Transition (Complex Double)
reconstruct = similarityTransformation eigenSystem

reconstructionError :: (Complex Double, ((Score,Score), (Score,Score)))
reconstructionError =
   swap $ Vector.argAbsMaximum $ ArrMatrix.toVector $
   Matrix.sub (Matrix.fromReal transition) reconstruct


unwrapEigenvector :: Vector.Vector TriShapeInt a -> Vector a
unwrapEigenvector = Array.mapShape ExtShape.deconsIntIndexed

eigenvectors :: [(Complex Double, Vector (Complex Double))]
eigenvectors =
   let (m,v,_minv) = eigenSystem
   in  Key.sort (negate . Complex.magnitude . fst) $
       zip (Vector.toList v) (map unwrapEigenvector $ Matrix.toColumns m)

tallMatrixFromColumns :: [Vector Double] -> TallMatrix Double
tallMatrixFromColumns =
   either id (error "matrix not tall") . Matrix.caseTallWide .
   Matrix.fromColumns (Matrix.height transition)

dominantEigenmatrix :: TallMatrix Double
dominantEigenmatrix =
   tallMatrixFromColumns $
   map (Array.map Complex.realPart . snd) $ take 2 eigenvectors

{-
This is like the first step of the Schur decomposition.

https://cs.stackexchange.com/questions/33068/partially-diagonalizing-a-matrix
https://math.stackexchange.com/questions/3967036/partial-eigendecomposition-of-a-positive-semi-definite-matrix
-}
partialEigenSystem ::
   (Matrix.LiberalSquare TriShape ShapeInt2 Double,
    Matrix.LiberalSquare TriShape ShapeInt2 Double)
partialEigenSystem =
   let (ml,mr) = dominantEigenmatrixPairs
       scal = Square.takeDiagonal $ Square.fromFull $ Matrix.adjoint ml #*# mr
       mlScal = Matrix.generalizeTall $ ml #*\ Vector.recip scal
       ortho = Matrix.generalizeTall $ Ortho.complement ml
   in (Square.liberalFromFull $
         mlScal  |||  ortho #-# mlScal #*# (Matrix.adjoint mr #*# ortho),
       Square.liberalFromFull $ Matrix.generalizeTall mr ||| ortho)

partiallyDiagonalized :: Matrix.LiberalSquare ShapeInt2 ShapeInt2 Double
partiallyDiagonalized =
   case fromInteger 1 :: Int of
      0 -> let mr = snd partialEigenSystem in mr #\## transition #*## mr
      _ ->
         let (ml,mr) = partialEigenSystem
         in Matrix.adjoint ml #*## transition #*## mr

{-
U are the given eigenvectors.
Uorth is the basis for their orthogonal complement.

U _|_ Uorth
M*U _|_ M*Uorth - does not apply

Thus, (U | Uorth) cannot diagonalize the transition matrix.
-}


filterEig :: (Double -> Bool) -> Transition (Complex Double)
filterEig p =
   let (trafo, diag, trafoInv) = eigenSystem
   in similarityTransformation
         (trafo,
          Array.map (\x -> if p (Complex.magnitude x) then 1 else 0) diag,
          trafoInv)

powerCoeff :: Complex Double
powerCoeff =
   (filterEig (\x->0.99<x && x<1.01) #*| Vector.fromReal initial2) ! (P1,P1)


outerProduct ::
   (Eq height, Eq width, Shape.C height, Shape.C width, Class.Floating a) =>
   Vector.Vector height a ->
   Vector.Vector width a ->
   Matrix.General height width a
outerProduct = Matrix.outer MatrixShape.RowMajor

rank2matrix :: Matrix.General ShapeInt ShapeInt Double
rank2matrix =
   Matrix.add
      (Vector.autoFromList [38,37,20,80,38]
       `outerProduct`
       Vector.autoFromList [97,30,19,49,37])
      (Vector.autoFromList [7,43,89,64,20]
       `outerProduct`
       Vector.autoFromList [4,96,44,49,80])

eigenvectorPairs ::
   [(Complex Double, (Vector (Complex Double), Vector (Complex Double)))]
eigenvectorPairs =
   Key.sort (negate . Complex.magnitude . fst) $
   (\(vr,d,vlAdj) ->
      zip (Vector.toList d) $
      zip (map unwrapEigenvector $ Matrix.toRows vlAdj)
          (map unwrapEigenvector $ Matrix.toColumns vr)) $
   Square.eigensystem $ Square.mapSize ExtShape.IntIndexed transition

dominantEigenmatrixPairs :: (TallMatrix Double, TallMatrix Double)
dominantEigenmatrixPairs =
   (tallMatrixFromColumns $ map (Array.map Complex.realPart . fst . snd) $
      take 2 eigenvectorPairs,
    tallMatrixFromColumns $ map (Array.map Complex.realPart . snd . snd) $
      take 2 eigenvectorPairs)

rankReduced :: Matrix (Complex Double)
rankReduced =
   foldl1 Matrix.add $
   map
      (\(l,(vl,vr)) ->
         Matrix.scale (l / Vector.inner vl vr) $ outerProduct vr vl)
      (take 2 eigenvectorPairs)

powerParameters :: Vector Double -> [(Complex Double, Complex Double)]
powerParameters initial =
   let initialC = Vector.fromReal initial
   in map
         (\(l,(vl,vr)) ->
            (l, (vr!(P1,P1)) * Vector.inner vl initialC / Vector.inner vl vr))
         eigenvectorPairs


probabilities :: Vector Double -> [Double]
probabilities = map (! (P1,P1)) . iterate (transition #*|)

probabilitiesApprox1 :: [Double]
probabilitiesApprox1 =
   map (1-) $ iterate (0.9083578427694454*) 0.9237265161044081

probabilitiesApprox2 :: [Double]
probabilitiesApprox2 =
   map (1-) $ iterate (0.9083578427694454*) 1.025929108036849

formatProbabilities :: [String]
formatProbabilities =
   List.zipWith5
      (printf "%3d & %1.5f & %1.5f & %1.5f & %1.5f \\\\")
      [0::Int ..]
      (probabilities initial1)
      probabilitiesApprox1
      (probabilities initial2)
      probabilitiesApprox2


main :: IO ()
main = mapM_ putStrLn $ take 101 formatProbabilities