-- SPDX-FileCopyrightText: 2022 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA

module Morley.Util.Bimap
  ( Bimap(..)
  , empty
    -- * Optics
  , flipped
  ) where

import Prelude hiding (empty)

import Control.Lens (At(..), Index, Iso, IxValue, Ixed(..), iso)
import Data.Aeson (FromJSON(..), ToJSON(..))
import Data.Bimap qualified as Bimap
import Data.Coerce (coerce)
import GHC.Exts (IsList)

newtype Bimap a b = Bimap { unBimap :: Bimap.Bimap a b }
  deriving newtype (Show, Eq, Ord, IsList)

empty :: Bimap a b
empty = coerce Bimap.empty

type instance Index (Bimap k _) = k
type instance IxValue (Bimap _ v) = v

-- | Left-biased 'Ixed' instance.
-- It assumes the left value @a@ is the key (just like the @Ix (Map k v)@ instance).
--
-- To flip this assumption, use the 'flipped' optic.
instance (Ord k, Ord v) => Ixed (Bimap k v) where
  ix :: k -> Traversal' (Bimap k v) v
  ix k handler (Bimap bmap) =
    case Bimap.lookup k bmap of
     Just v -> handler v <&> \v' -> Bimap $ Bimap.insert k v' bmap
     Nothing -> pure $ Bimap bmap

-- | Left-biased 'At' instance.
-- It assumes the left value @a@ is the key (just like the @At (Map k v)@ instance).
--
-- To flip this assumption, use the 'flipped' optic.
instance (Ord k, Ord v) => At (Bimap k v) where
  at :: k -> Lens' (Bimap k v) (Maybe v)
  at k handler (Bimap bmap) =
    handler currentValueMaybe <&> \newValueMaybe ->
      Bimap $
        case (currentValueMaybe, newValueMaybe) of
          (Nothing, Nothing) -> bmap
          (Just _, Nothing) -> Bimap.delete k bmap
          (_, Just newValue) -> Bimap.insert k newValue bmap
    where
      currentValueMaybe = Bimap.lookup k bmap

-- | Isomorphism between @Bimap a b@ and @Bimap b a@.
flipped :: Iso (Bimap a1 b1) (Bimap a2 b2) (Bimap b1 a1) (Bimap b2 a2)
flipped = iso (coerce Bimap.twist) (coerce Bimap.twist)

instance (Ord a, Ord b, FromJSON a, FromJSON b) => FromJSON (Bimap a b) where
  parseJSON = fmap (Bimap . Bimap.fromList) . parseJSON

instance (ToJSON a, ToJSON b) => ToJSON (Bimap a b) where
  toJSON = toJSON . Bimap.toList . unBimap