module Data.Fin.Strong

import Data.Fin
import Decidable.Equality.Core
import Syntax.PreorderReasoning

-- Definition ------------------------------------------------------------------

public export
data SFin : Nat -> Type where
  FZ : SFin (S n)
  FS : SFin (S n) -> SFin (S (S n))

-- Properties

export
Injective (FS {n}) where
  injective Refl = Refl

export
Uninhabited (SFin 0) where uninhabited x impossible

export
Uninhabited (Strong.FS x = FZ) where uninhabited prf impossible

-- Conversion ------------------------------------------------------------------

-- Utlities

public export
0 indexPred : (0 x : Fin n) -> Nat
indexPred {n = S n} x = n

public export
0 indexPredPrf : (0 x : Fin n) -> n = S (indexPred x)
indexPredPrf {n = S n} x = Refl

public export
0 indexPred' : (0 x : SFin n) -> Nat
indexPred' {n = S n} x = n

public export
0 indexPred'Prf : (0 x : SFin n) -> n = S (indexPred' x)
indexPred'Prf {n = S n} x = Refl

public export
finToStrong : Fin n -> SFin n
finToStrong FZ = FZ
finToStrong (FS x) =
  rewrite indexPredPrf x in
  FS (replace {p = SFin} (indexPredPrf x) (finToStrong x))

public export
strongToFin : SFin n -> Fin n
strongToFin FZ = FZ
strongToFin (FS x) = FS (strongToFin x)

-- Properties

export
finToStrongToFin : (x : Fin n) -> strongToFin (finToStrong x) = x
finToStrongToFin FZ = Refl
finToStrongToFin (FS FZ) = Refl
finToStrongToFin (FS (FS x)) = cong FS $ finToStrongToFin (FS x)

export
strongToFinToStrong : (x : SFin n) -> finToStrong (strongToFin x) = x
strongToFinToStrong FZ = Refl
strongToFinToStrong (FS x) = cong FS (strongToFinToStrong x)

export
Injective (strongToFin {n}) where
  injective {x, y} prf = Calc $
    |~ x
    ~~ finToStrong (strongToFin x) ..<(strongToFinToStrong x)
    ~~ finToStrong (strongToFin y) ...(cong finToStrong prf)
    ~~ y                           ...(strongToFinToStrong y)

-- Decidable Equality ----------------------------------------------------------

export
DecEq (SFin n) where
  decEq x y = viaEquivalence
    (MkEquivalence (injective {f = strongToFin}) (\prf => cong strongToFin prf))
    (decEq (strongToFin x) (strongToFin y))

-- Useful Constructor ----------------------------------------------------------

%inline
export
FS' : SFin n -> SFin (S n)
FS' = finToStrong . FS . strongToFin

export
Injective (FS' {n}) where
  injective {x = FZ, y = FZ} prf = Refl
  injective {x = FS x, y = FS y} prf = cong FS $ injective {f = FS'} $ injective prf

export
sucNonZero : (x : SFin n) -> Not (FS' x = FZ)
sucNonZero FZ prf = absurd prf
sucNonZero (FS x) prf = absurd prf

export
sucIsFS : (x : SFin (S n)) -> FS' x = FS x
sucIsFS x = rewrite strongToFinToStrong x in Refl

-- Num Interface ---------------------------------------------------------------

public export
{n : Nat} -> Num (SFin (S n)) where
  x + y = finToStrong (strongToFin x + strongToFin y)
  x * y = finToStrong (strongToFin x * strongToFin y)
  fromInteger = finToStrong . restrict n

-- Weaken and Shift ------------------------------------------------------------

export
weakenN : (0 k : Nat) -> SFin n -> SFin (n + k)
weakenN k = finToStrong . weakenN k . strongToFin

export
shift : (k : Nat) -> SFin n -> SFin (k + n)
shift k = finToStrong . shift k . strongToFin