{- |
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.Triangular as Triangular
import qualified Numeric.LAPACK.Matrix as Matrix
import qualified Numeric.LAPACK.Vector as Vector
import Numeric.LAPACK.Format ((##))

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.Enumeration, Shape.Enumeration, Shape.Enumeration)

edgeShape :: EdgeShape
edgeShape = (Shape.Enumeration, Shape.Enumeration, Shape.Enumeration)


{-
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.fromRowMajor $
   Array.sample
      (edgeShape, cornerShape)
      (\((ed,ex,ey), c) ->
         let ((cx, cy), cz) = selectCornerCoords ed c
         in  if ex==cx && ey==cy
               then
                  case cz of
                     C0 ->  1
                     C1 -> -1
               else 0)


selectCornerCoords :: Dim -> Corner -> ((Coord, Coord), Coord)
selectCornerCoords ed (cx,cy,cz) =
   case ed of
      D0 -> ((cy, cz), cx)
      D1 -> ((cx, cz), cy)
      D2 -> ((cx, cy), cz)

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

resistances :: Vector.Vector EdgeShape Double
resistances = Vector.constant edgeShape 1


fullMatrix :: Triangular.Symmetric (():+:(EdgeShape:+:CornerShape)) Double
fullMatrix =
   Triangular.stackSymmetric
      (Triangular.symmetricFromList MatrixShape.RowMajor () [0])
      (Matrix.singleRow MatrixShape.RowMajor $
         Vector.unit (edgeShape:+:cornerShape) (Right sourceCorner)) $
   Triangular.stackSymmetric
      (Triangular.diagonal MatrixShape.RowMajor resistances)
      voltageMatrix
      (Vector.constant
         (MatrixShape.symmetric MatrixShape.RowMajor cornerShape) 0)


main :: IO ()
main = do
   fullMatrix ## "%2.f"
   let solutionVec =
         Matrix.unliftColumn MatrixShape.ColumnMajor
            (Triangular.solve fullMatrix) $
         Vector.unit (():+:(edgeShape:+:cornerShape)) (Right (Right destCorner))
   print $ - solutionVec ! Right (Right destCorner)

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