{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}

module Transpose (Transpose, transpose) where

import Data.Proxy
import GHC.TypeLits

data N = Z | S N

type family Dec (x :: N) :: N where
  Dec 'Z = 'Z
  Dec ('S n) = n

type family DecIfBigger (x :: N) (y :: N) (z :: N) :: N where
  DecIfBigger Z _ z = Dec z
  DecIfBigger (S _) Z z = z
  DecIfBigger (S n) (S m) z = DecIfBigger n m z

type family DecBigger (x :: N) (xs :: [N]) :: [N] where
  DecBigger _ '[] = '[]
  DecBigger x (n ': ys) = (DecIfBigger x n n) ': (DecBigger x ys)

class Transpose (ixs :: [Nat]) f g | ixs f -> g where
  transpose :: proxy ixs -> f -> g

class FromNats (xs :: [Nat]) (ys :: [N]) | xs -> ys
instance FromNats '[] '[]
instance (FromNat x y, FromNats xs ys) => FromNats (x ':xs) (y ': ys)

class FromNat (x :: Nat) (y :: N) | x -> y
instance FromNat_ x y (x <=? 0) => FromNat x y

class FromNat_ (x :: Nat) (y :: N) (b :: Bool) | x b -> y
instance FromNat_ x Z 'True
instance FromNat_ (x-1) y (x-1 <=? 0) => FromNat_ x (S y) 'False

instance (FromNats ixs ixs', Transpose_ ixs' f g) => Transpose ixs f g where
  transpose _ = transpose_ (Proxy @ixs')

class Transpose_ (ixs :: [N]) f g | ixs f -> g where
  transpose_ :: proxy ixs -> f -> g

instance Transpose_ '[] f f where
  transpose_ _ = id

instance
  ( PromoteAxis x f (h a)
  , Functor h
  , Transpose_ (DecBigger x xs) a g
  ) =>  Transpose_ (x ': xs) f (h g) where
  transpose_ _ = fmap (transpose_ (Proxy @(DecBigger x xs))) . promoteAxis (Proxy @x)

class PromoteAxis (ix :: N) f g | ix f -> g where
  promoteAxis :: proxy ix -> f -> g

instance PromoteAxis 'Z f f where
  promoteAxis _ = id

instance (Applicative g, Traversable g, Traversable f, Functor f, PromoteAxis n a (g b)) => PromoteAxis ('S n) (f a) (g (f b)) where
  promoteAxis _ = sequenceA . fmap (promoteAxis (Proxy @n))

example ::
  ( Applicative f1, Traversable f1
  , Applicative f2, Traversable f2
  , Applicative f3, Traversable f3
  , Applicative f4, Traversable f4
  ) => f1 (f2 (f3 (f4 a))) -> f2 (f1 (f4 (f3 a)))
example = transpose (Proxy @[1,0,3,2])