{- |
Consider a cube of resistors of equal resistance.
What is the overall resistance from one corner to the opposite one?
-}
{-# LANGUAGE TypeOperators #-}
module Main where

import qualified Numeric.LAPACK.Matrix.Shape as MatrixShape
import qualified Numeric.LAPACK.Matrix.Symmetric as Symmetric
import qualified Numeric.LAPACK.Matrix as Matrix
import qualified Numeric.LAPACK.Vector as Vector
import Numeric.LAPACK.Matrix.Symmetric ((#%%%#))
import Numeric.LAPACK.Matrix ((#\|))
import Numeric.LAPACK.Format ((##))

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 Data.Array.Comfort.Storable ((!))
import Data.Array.Comfort.Shape ((::+)((::+)))


data Coord = C0 | C1 deriving (Eq, Ord, Show, Enum, Bounded)
data Dim = D0 | D1 | D2 deriving (Eq, Ord, Show, Enum, Bounded)
type Corner = (Coord,Coord,Coord)
type Edge = (Dim,Coord,Coord)

type ShapeEnum = Shape.Enumeration
type CornerShape = (ShapeEnum Coord, ShapeEnum Coord, ShapeEnum Coord)
type EdgeShape = (ShapeEnum Dim, ShapeEnum Coord, ShapeEnum Coord)

type Matrix height width = Matrix.General height width


cornerShape :: CornerShape
cornerShape = Shape.static

edgeShape :: EdgeShape
edgeShape = Shape.static


{-
We also need a currentMatrix that implements Kirchhoff's current law
but it turns out to equal 'transpose voltageMatrix'.
This makes fullMatrix symmetric.
-}
voltageMatrix :: Matrix EdgeShape CornerShape Double
voltageMatrix =
   Matrix.fromRowArray cornerShape $
   fmap
      (\e ->
         Array.fromAssociations 0 cornerShape
            [(edgeCorner e C0, 1), (edgeCorner e C1, -1)]) $
   BoxedArray.indices edgeShape

edgeCorner :: Edge -> Coord -> Corner
edgeCorner (ed,ex,ey) coord =
   case ed of
      D0 -> (coord,ex,ey)
      D1 -> (ex,coord,ey)
      D2 -> (ex,ey,coord)

sourceCorner, destCorner :: Corner
sourceCorner = (C0,C0,C0)
destCorner = (C1,C1,C1)

resistances :: Vector.Vector EdgeShape Double
resistances = Vector.one edgeShape


fullMatrix :: Matrix.Symmetric (()::+EdgeShape::+CornerShape) Double
fullMatrix =
   let order = MatrixShape.RowMajor
       symmetricZero sh = Matrix.zero $ MatrixShape.symmetric order sh
       unit = Vector.unit (edgeShape::+cornerShape) (Right sourceCorner)
   in (symmetricZero (), Matrix.singleRow order unit)
      #%%%#
      (Symmetric.diagonal order resistances, voltageMatrix)
      #%%%#
      symmetricZero cornerShape


main :: IO ()
main = do
   fullMatrix ## "%2.f"
   let ix = Right (Right destCorner)
   let solutionVec = fullMatrix #\| Vector.unit (Symmetric.size fullMatrix) ix
   print $ - solutionVec ! ix

{-
result: total resistance is 5/(2+2+2)
-}