{-# LANGUAGE GADTs #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}

import Control.Applicative
import Control.Arrow
import Control.Category
import Prelude hiding ((.),id)

class Coalgebraic (k :: * -> * -> *) (x :: *) where
  type A k x :: *
  type B k x :: *

type family ArgOf (f_b :: *) :: *
type instance ArgOf (f b) = b

instance (Functor f, f_b ~ f b) => Coalgebraic (->) (a -> f_b) where
  type A (->) (a -> f_b) = a
  type B (->) (a -> f_b) = ArgOf f_b

class (Category k, Coalgebraic k x, Coalgebraic k y) => Isomorphic k x y where
  iso   :: (A k y -> A k x) -> (B k x -> B k y) -> k x y
  isoid :: (A k x ~ A k y, B k x ~ B k y) => k x y

instance (Functor f, x ~ (a -> f b), y ~ (s -> f t)) => Isomorphic (->) x y where
  iso sa bt afb s = bt <$> afb (sa s)
  isoid = id

instance Coalgebraic Isoid (a,b) where
  type A Isoid (a,b) = a
  type B Isoid (a,b) = b

instance (x ~ (a,b), y ~ (s,t)) => Isomorphic Isoid x y where
  iso   = Iso
  isoid = Isoid

data Isoid x y where
  Isoid :: Isoid x x
  Iso   :: (s -> a) -> (b -> t) -> Isoid (a,b) (s,t)

instance Category Isoid where
  id = Isoid
  Isoid . x = x
  x . Isoid = x
  Iso xs ty . Iso sa bt = Iso (sa.xs) (ty.bt)

type AnIso s t a b = Isoid (a,b) (s,t)

type Iso s t a b = forall k x y. (Isomorphic k x y, A k x ~ a, B k x ~ b, A k y ~ s, B k y ~ t) => k x y

plus1 :: (Num a, Num b) => Iso a b a b
plus1 = iso (+1) (subtract 1)

from :: AnIso s t a b -> Iso b a t s
from (Iso sa bt) = iso bt sa
from Isoid       = isoid

class Isomorphic k x y => Prismatic k x y where
  prism :: (B k x -> B k y) -> (A k y -> Either (B k y) (A k x)) -> k x y

instance (Applicative f, x ~ (a -> f b), y ~ (s -> f t)) => Prismatic (->) x y where
  prism bt seta afb = either pure (fmap bt . afb) . seta

data Prismoid x y where
  Prismoid :: Prismoid x x
  Prism :: (b -> t) -> (s -> Either t a) -> Prismoid (a,b) (s,t)

instance Coalgebraic Prismoid (a,b) where
  type A Prismoid (a,b) = a
  type B Prismoid (a,b) = b

instance Category Prismoid where
  id = Prismoid
  x . Prismoid = x
  Prismoid . x = x
  Prism ty xeys . Prism bt seta = Prism (ty.bt) $ \x ->
    case xeys x of
      Left y  -> Left y
      Right s -> case seta s of
        Left t  -> Left (ty t)
        Right a -> Right a

instance (x ~ (a,b), y ~ (s,t)) => Isomorphic Prismoid x y where
  iso sa bt = Prism bt (Right . sa)
  isoid = Prismoid

instance (x ~ (a,b), y ~ (s,t)) => Prismatic Prismoid x y where
  prism = Prism

type Prism s t a b = forall k x y. (Prismatic k x y, A k x ~ a, B k x ~ b, A k y ~ s, B k y ~ t) => k x y

_right :: Prism (Either c a) (Either c b) a b
_right = prism Right (left Left)