-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Lens
-- Copyright   :  (c) 2010 University of Minho
-- License     :  BSD3
--
-- Maintainer  :  hpacheco@di.uminho.pt
-- Stability   :  experimental
-- Portability :  non-portable
--
-- Pointless Rewrite:
-- automatic transformation system for point-free programs
-- 
-- Evaluation of point-free lens representations.
--
-----------------------------------------------------------------------------

module Data.Lens where

import Data.Type
import Data.Pf
import Data.Spine
import Data.Equal
import Data.Default

import Prelude hiding (Functor(..))
import Control.Monad hiding (Functor(..))

import Generics.Pointless.Functors hiding (rep)
import Generics.Pointless.Lenses

-- | Computes the inverse lens for isomorphic lenses.
inv :: MonadPlus m => Type (Lens a b) -> Pf (Lens a b) -> m (Pf (Lens b a))
inv _ ID_LNS = return ID_LNS
inv (Lns c a) (COMP_LNS b f g) = do
    g' <- inv (Lns c b) g
    f' <- inv (Lns b a) f
    return $ COMP_LNS b g' f'
inv (Lns (Prod a b) (Prod c d)) (f `PROD_LNS` g) = do
    f' <- inv (Lns a c) f
    g' <- inv (Lns b d) g
    return $ f' ><<< g'
inv (Lns (Either a b) (Either c d)) (f `SUM_LNS` g) = do
    f' <- inv (Lns a c) f
    g' <- inv (Lns b d) g
    return $ f' -|-<< g'
inv (Lns c (Prod One _)) BANGL_LNS =
    return $ SND_LNS BANG
inv (Lns c (Prod _ One)) (BANGR_LNS) =
    return $ FST_LNS BANG
inv _ SWAP_LNS = return SWAP_LNS
inv _ COSWAP_LNS = return COSWAP_LNS
inv _ DISTL_LNS = return UNDISTL_LNS
inv _ UNDISTL_LNS = return DISTL_LNS
inv _ DISTR_LNS = return UNDISTR_LNS
inv _ UNDISTR_LNS = return DISTR_LNS
inv _ ASSOCL_LNS = return ASSOCR_LNS
inv _ ASSOCR_LNS = return ASSOCL_LNS
inv _ COASSOCL_LNS = return COASSOCR_LNS
inv _ COASSOCR_LNS = return COASSOCL_LNS
inv (Lns c a@(dataFctr -> Just fctr)) INN_LNS = case teq c (rep fctr a) of
    { Just Eq -> return OUT_LNS
    ; otherwise -> fail "inv INN_LNS" }
inv (Lns a@(dataFctr -> Just fctr) c) OUT_LNS = case teq c (rep fctr a) of
    { Just Eq -> return INN_LNS
    ; otherwise -> fail "inv OUT_LNS" }
inv _ _ = mzero 

-- | Lifts a point-free function into a lens (unsafe).
lensify :: MonadPlus m => Type (a -> b) -> Pf (a -> b) -> m (Pf (Lens a b))
lensify (Fun _ _) (GET l) = return l

lensify (Fun a c) (COMP b f g) = do
    f' <- lensify (Fun b c) f
    g' <- lensify (Fun a b) g
    return $ COMP_LNS b f' g'
lensify (Fun (Prod a b) _) FST = return $ FST_LNS (constPf $ defvalue b)
lensify (Fun (Prod a b) _) SND = return $ SND_LNS (constPf $ defvalue a)
lensify (Fun (Prod a b) (Prod c d)) (f `PROD` g) = do
    f' <- lensify (Fun a c) f
    g' <- lensify (Fun b d) g
    return $ f' ><<< g'
lensify (Fun (Either a b) (Either c d)) ((COMP _ INL f) `EITHER` (COMP _ INR g)) = do
    f' <- lensify (Fun a c) f
    g' <- lensify (Fun b d) g
    return $ f' `SUM_LNS` g'
    
-- Special either expressions
lensify (Fun (Either _ a) e@(Either x y)) (INL `EITHER` f) = do
    f' <- lensify (Fun a e) f
    return $ COMP_LNS (Either (Either x x) y) ((ID_LNS .\/<< ID_LNS) -|-<< ID_LNS)
           $ COMP_LNS (Either x e) COASSOCL_LNS $ ID_LNS -|-<< f'
lensify (Fun (Either a _) e@(Either x y)) (f `EITHER` INR) = do
    f' <- lensify (Fun a e) f
    return $ COMP_LNS (Either x (Either y y)) (ID_LNS -|-<< (ID_LNS \/.<< ID_LNS))
           $ COMP_LNS (Either e y) COASSOCR_LNS $ f' -|-<< ID_LNS
-- Regular either expression
lensify fun@(Fun _ l@(List a)) (ZERO `EITHER` f) = lensify fun
    $ COMP (Either One (Prod a l)) INN
    $ COMP (Either One l) (INL `EITHER` OUT) (BANG `SUM` f)
lensify fun@(Fun _ a@(isNat -> Just Eq)) (ZERO `EITHER` f) = lensify fun
    $ COMP (Either One a) INN
    $ COMP (Either One a) (INL `EITHER` OUT) (BANG `SUM` f)
lensify fun@(Fun _ l@(List a)) (f `EITHER` ZERO) = lensify fun
    $ COMP (Either One (Prod a l)) INN
    $ COMP (Either One l) (INL `EITHER` OUT)
    $ COMP (Either l One) COSWAP (f `SUM` BANG)
lensify fun@(Fun _ a@(isNat -> Just Eq)) (f `EITHER` ZERO) = lensify fun
    $ COMP (Either One a) INN
    $ COMP (Either One a) (INL `EITHER` OUT)
    $ COMP (Either a One) COSWAP (f `SUM` BANG)
    
lensify (Fun (Either a b) c) (f `EITHER` g) = do
    f' <- lensify (Fun a c) f
    g' <- lensify (Fun b c) g
    return $ EITHER_LNS (COMP One INL BANG) f' g'
lensify (Fun (Either a b) (Either c d)) (f `SUM` g) = do
    f' <- lensify (Fun a c) f
    g' <- lensify (Fun b d) g
    return $ f' `SUM_LNS` g'
lensify (Fun a _) BANG = return $ BANG_LNS (constPf $ defvalue a)

lensify (Fun _ a@(isList -> Just Eq)) PLUS = return CAT_LNS
lensify (Fun _ a@(isList -> Just Eq)) FOLD = return CONCAT_LNS
lensify (Fun _ a@(isNat -> Just Eq)) PLUS = return PLUSN_LNS
lensify (Fun _ a@(isNat -> Just Eq)) FOLD = return SUMN_LNS
lensify (Fun (List a) _) LENGTH = return $ LENGTH_LNS (defvalue a)
lensify (Fun (List a) (List b)) (MAP f) = do
    f' <- lensify (Fun a b) f
    return $ MAP_LNS f'
    
lensify (Fun _ _) ID = return ID_LNS
lensify (Fun _ _) SWAP = return SWAP_LNS
lensify (Fun _ _) COSWAP = return COSWAP_LNS
lensify (Fun _ _) DISTL = return DISTL_LNS
lensify (Fun _ _) UNDISTL = return UNDISTL_LNS
lensify (Fun _ _) DISTR = return DISTR_LNS
lensify (Fun _ _) UNDISTR = return UNDISTR_LNS
lensify (Fun _ _) ASSOCL = return ASSOCL_LNS
lensify (Fun _ _) ASSOCR = return ASSOCR_LNS
lensify (Fun _ _) COASSOCL = return COASSOCL_LNS
lensify (Fun _ _) COASSOCR = return COASSOCR_LNS

lensify (Fun _ _) INN = return INN_LNS
lensify (Fun _ _) OUT = return OUT_LNS
lensify (Fun _ _) (FMAP fctr t f) = do
    f' <- lensify t f
    return $ FMAP_LNS fctr t f'
lensify (Fun a b@(Data s fctr)) (ANA f) = do
    f' <- lensify (Fun a (rep fctr a)) f
    return $ ANA_LNS f'
lensify (Fun a@(Data s fctr) b) (CATA f) = do
    f' <- lensify (Fun (rep fctr b) b) f
    return $ CATA_LNS f'
lensify t v = fail $ "lensify " ++ gshow (Pf t) v

getof :: Type (Lens c a) -> Pf (Lens c a) -> Pf (c -> a)
getof (Lns _ _) (LNS s l) = FUN (showL ["get",s]) $ get l

getof (Lns c a) (COMP_LNS b f g) = COMP b (getof (Lns b a) f) (getof (Lns c b) g)
getof (Lns _ _) (FST_LNS f) = FST
getof (Lns _ _) (SND_LNS f) = SND
getof (Lns (Prod a b) (Prod c d)) (PROD_LNS f g) = getof (Lns a c) f ><= getof (Lns b d) g
getof (Lns (Either a b) c) (EITHER_LNS x f g) = getof (Lns a c) f \/= getof (Lns b c) g
getof (Lns (Either c d) (Either a b)) (SUM_LNS f g) = getof (Lns c a) f -|-= getof (Lns d b) g
getof (Lns (Either c d) (Either a b)) (SUMW_LNS h i f g) = getof (Lns c a) f -|-= getof (Lns d b) g
getof (Lns c One) (BANG_LNS f) = BANG
getof (Lns c (Prod One _)) (BANGL_LNS) = BANG /\= ID
getof (Lns c (Prod _ One)) (BANGR_LNS) = ID /\= BANG

getof (Lns _ _) ID_LNS = ID
getof (Lns _ _) SWAP_LNS = SWAP
getof (Lns _ _) COSWAP_LNS = COSWAP
getof (Lns _ _) DISTL_LNS = DISTL
getof (Lns _ _) UNDISTL_LNS = UNDISTL
getof (Lns _ _) DISTR_LNS = DISTR
getof (Lns _ _) UNDISTR_LNS = UNDISTR
getof (Lns _ _) ASSOCL_LNS = ASSOCL
getof (Lns _ _) ASSOCR_LNS = ASSOCR
getof (Lns _ _) COASSOCL_LNS = COASSOCL
getof (Lns _ _) COASSOCR_LNS = COASSOCR

getof (Lns _ _) INN_LNS = INN
getof (Lns a@(dataFctr -> Just fctr) c) OUT_LNS = case teq c (rep fctr a) of
    { Just Eq   -> OUT
    ; otherwise -> error "getof OUT" }
getof (Lns _ _) (FMAP_LNS fctr (Fun c a) f) = FMAP fctr (Fun c a) $ getof (Lns c a) f
getof (Lns a b@(dataFctr -> Just fctr)) (ANA_LNS f) = ANA $ getof (Lns a (rep fctr a)) f
getof (Lns a@(dataFctr -> Just fctr) b) (CATA_LNS f) = CATA $ getof (Lns (rep fctr b) b) f
getof (Lns _ _) BOT = BOT
getof _ f = GET f

createof :: Type (Lens c a) -> Pf (Lens c a) -> Pf (a -> c)
createof (Lns _ _) (LNS s l) = FUN (showL ["create",s]) $ create l

createof (Lns c a) (COMP_LNS b f g) = COMP b (createof (Lns c b) g) (createof (Lns b a) f)
createof (Lns _ _) (FST_LNS f) = ID /\= f
createof (Lns _ _) (SND_LNS f) = f /\= ID
createof (Lns (Prod a b) (Prod c d)) (PROD_LNS f g) = createof (Lns a c) f ><= createof (Lns b d) g
createof (Lns (Either a b) c) (EITHER_LNS x f g) = COMP t' (l -|-= r) $ COMP t DISTL (x /\= ID)
    where l = COMP c (createof (Lns a c) f) SND
          r = COMP c (createof (Lns b c) g) SND
          t = Prod (Either One One) c
          t' = Either (Prod One c) (Prod One c)
createof (Lns (Either c d) (Either a b)) (SUM_LNS f g) = createof (Lns c a) f -|-= createof (Lns d b) g
createof (Lns (Either c d) (Either a b)) (SUMW_LNS h i f g) = createof (Lns c a) f -|-= createof (Lns d b) g
createof (Lns c One) (BANG_LNS f) = f
createof (Lns c (Prod One _)) (BANGL_LNS) = SND
createof (Lns c (Prod _ One)) (BANGR_LNS) = FST

createof (Lns _ _) ID_LNS = ID
createof (Lns _ _) SWAP_LNS = SWAP
createof (Lns _ _) COSWAP_LNS = COSWAP
createof (Lns _ _) DISTL_LNS = UNDISTL
createof (Lns _ _) UNDISTL_LNS = DISTL
createof (Lns _ _) DISTR_LNS = UNDISTR
createof (Lns _ _) UNDISTR_LNS = DISTR
createof (Lns _ _) ASSOCL_LNS = ASSOCR
createof (Lns _ _) ASSOCR_LNS = ASSOCL
createof (Lns _ _) COASSOCL_LNS = COASSOCR
createof (Lns _ _) COASSOCR_LNS = COASSOCL

createof (Lns _ _) INN_LNS = OUT
createof (Lns a@(dataFctr -> Just fctr) c) OUT_LNS = case teq c (rep fctr a) of
    { Just Eq   -> INN
    ; otherwise -> error "createof OUT" }
createof (Lns _ _) (FMAP_LNS fctr (Fun c a) f) = FMAP fctr (Fun a c) $ createof (Lns c a) f
createof (Lns a b@(dataFctr -> Just fctr)) (ANA_LNS f) = CATA $ createof (Lns a (rep fctr a)) f
createof (Lns a@(dataFctr -> Just fctr) b) (CATA_LNS f) = ANA $ createof (Lns (rep fctr b) b) f

createof (Lns (List a) (List b)) (MAP_LNS f) = MAP $ createof (Lns a b) f

createof (Lns _ _) BOT = BOT
createof t f = CREATE f

putof :: Type (Lens c a) -> Pf (Lens c a) -> Pf ((a,c) -> c)
putof (Lns _ _) (LNS s l) = FUN (showL ["put",s]) $ put l

putof (Lns c a) (COMP_LNS b f g) = COMP t (putof (Lns c b) g)
    ((COMP t' (putof (Lns b a) f) (ID ><= (getof (Lns c b) g))) /\= SND)
    where t  = Prod b c
          t' = Prod a b
putof (Lns _ _) (FST_LNS f) = ID ><= SND
putof (Lns (Prod a b) _) (SND_LNS va) = COMP (Prod b a) SWAP (ID ><= FST)
putof (Lns (Prod a b) (Prod c d)) (PROD_LNS f g) = COMP t (putof (Lns a c) f ><= putof (Lns b d) g) distp_pf
    where t = Prod (Prod c a) (Prod d b)
putof (Lns (Either a b) c) (EITHER_LNS x f g) = COMP t (putof (Lns a c) f -|-= putof (Lns b c) g) DISTR
    where t = Either (Prod c a) (Prod c b)
putof (Lns (Either c d) (Either a b)) (SUM_LNS f g) = COMP t (l -|-= r) (dists_pf (Prod (Either a b) (Either c d)))
    where l = putof (Lns c a) f \/= COMP a (createof (Lns c a) f) FST
          r = COMP b (createof (Lns d b) g) FST \/= putof (Lns d b) g
          t = Either (Either (Prod a c) (Prod a d)) (Either (Prod b c) (Prod b d))
putof (Lns (Either c d) (Either a b)) (SUMW_LNS h i f g) = COMP t (putof (Lns c a) f -|-= putof (Lns d b) g) $
    COMP t' (l -|-= r) (dists_pf (Prod (Either a b) (Either c d)))
    where l  = ID \/= (FST /\= h)
          r  = (FST /\= i) \/= ID
          t  = Either (Prod a c) (Prod b d)
          t' = Either (Either (Prod a c) (Prod a d)) (Either (Prod b c) (Prod b d))
putof (Lns c One) (BANG_LNS f) = SND
putof (Lns c (Prod One _)) (BANGL_LNS) = COMP (Prod One c) SND FST
putof (Lns c (Prod _ One)) (BANGR_LNS) = COMP (Prod c One) FST FST

putof (Lns c a) ID_LNS = COMP a (createof (Lns c a) ID_LNS) FST
putof (Lns c a) SWAP_LNS = COMP a (createof (Lns c a) SWAP_LNS) FST
putof (Lns c a) COSWAP_LNS = COMP a (createof (Lns c a) COSWAP_LNS) FST
putof (Lns c a) DISTL_LNS = COMP a (createof (Lns c a) DISTL_LNS) FST
putof (Lns c a) UNDISTL_LNS = COMP a (createof (Lns c a) UNDISTL_LNS) FST
putof (Lns c a) DISTR_LNS = COMP a (createof (Lns c a) DISTR_LNS) FST
putof (Lns c a) UNDISTR_LNS = COMP a (createof (Lns c a) UNDISTR_LNS) FST
putof (Lns c a) ASSOCL_LNS = COMP a (createof (Lns c a) ASSOCL_LNS) FST
putof (Lns c a) ASSOCR_LNS = COMP a (createof (Lns c a) ASSOCR_LNS) FST
putof (Lns c a) COASSOCL_LNS = COMP a (createof (Lns c a) COASSOCL_LNS) FST
putof (Lns c a) COASSOCR_LNS = COMP a (createof (Lns c a) COASSOCR_LNS) FST

putof (Lns c a) INN_LNS = COMP a OUT FST
putof (Lns c@(dataFctr -> Just fctr) a) OUT_LNS = case teq a (rep fctr c) of
    { Just Eq   -> COMP (rep fctr c) INN FST
    ; otherwise -> error "putof OUT" }
putof (Lns _ _) (FMAP_LNS fctr (Fun c a) f) = COMP (rep fctr (Prod a c)) (FMAP fctr (Fun (Prod a c) c) (putof (Lns c a) f)) $
    FZIP fctr (Fun a c) (createof (Lns c a) f)
putof x@(Lns a b@(dataFctr -> Just fctr)) (ANA_LNS f) = COMP (fixof kfctr) g h
    where g = CATA $ putof (Lns a (rep fctr a)) f
          h = ANA $ ((COMP t (FZIP fctr (Fun b a) $ createof x (ANA_LNS f)) (OUT ><= getof (Lns a (rep fctr a)) f)) /\= SND)
          kfctr = fctr :*!: K a
          t = Prod (rep fctr b) (rep fctr a)
putof x@(Lns b@(dataFctr -> Just fctr) a) (CATA_LNS f) = ANA $ COMP t aux1 $ COMP t' aux2 (ID ><= OUT)
    where aux1 = FZIP fctr (Fun a b) $ createof x (CATA_LNS f)
          aux2 = COMP t'' (putof (Lns (rep fctr a) a) f) (ID ><= aux3) /\= SND
          aux3 = FMAP fctr (Fun b a) $ getof x (CATA_LNS f)
          t = Prod (rep fctr a) (rep fctr b)
          t' = Prod a (rep fctr b)
          t'' = Prod a (rep fctr a)
putof (Lns _ _) BOT = BOT
putof (Lns _ _) f = PUT f