module Quox.Syntax.Shift

import public Quox.Syntax.Var
import public Quox.Thin

import Data.Nat
import Data.So
import Data.DPair

%default total


||| represents the difference between a smaller scope and a larger one.
public export
data Shift : (from, to : Nat) -> Type where
  SZ : Shift from from
  SS : Shift from to -> Shift from (S to)
%name Shift by, bz
%builtin Natural Shift

public export
(.nat) : Shift from to -> Nat
(SZ).nat    = Z
(SS by).nat = S by.nat
%transform "Shift.(.nat)" Shift.(.nat) = believe_me

public export Cast (Shift from to) Nat     where cast = (.nat)
public export Cast (Shift from to) Integer where cast = cast . cast {to = Nat}

export Eq  (Shift from to) where (==) = (==) `on` (.nat)
export Ord (Shift from to) where compare = compare `on` (.nat)


||| shift equivalence, ignoring indices
public export
data Eqv : Shift from1 to1 -> Shift from2 to2 -> Type where
  EqSZ : SZ `Eqv` SZ
  EqSS : by `Eqv` bz -> SS by `Eqv` SS bz
%name Shift.Eqv e

using (by : Shift from to, bz : Shift from to)
  ||| two equivalent shifts are equal if they have the same indices.
  export
  0 fromEqv : by `Eqv` bz -> by = bz
  fromEqv EqSZ     = Refl
  fromEqv (EqSS e) = cong SS $ fromEqv e

  ||| two equal shifts are equivalent.
  export
  0 toEqv : by = bz -> by `Eqv` bz
  toEqv Refl {by = SZ}      = EqSZ
  toEqv Refl {by = (SS by)} = EqSS $ toEqv Refl


export
cmpLen : Shift from1 to -> Shift from2 to -> Either Ordering (from1 = from2)
cmpLen SZ      SZ      = Right Refl
cmpLen SZ      (SS by) = Left LT
cmpLen (SS by) SZ      = Left GT
cmpLen (SS by) (SS bz) = cmpLen by bz

export
0 shiftDiff : (by : Shift from to) -> to = by.nat + from
shiftDiff SZ      = Refl
shiftDiff (SS by) = cong S $ shiftDiff by

export
0 shiftVarLT : (by : Shift from to) -> (i : Var from) ->
               by.nat + i.nat `LT` to
shiftVarLT by i =
  rewrite plusSuccRightSucc by.nat i.nat in
  transitive
    (plusLteMonotoneLeft by.nat (S i.nat) from (toNatLT i))
    (replace {p=(`LTE` to)} (shiftDiff by) reflexive)


public export
fromNat : (by : Nat) -> Shift from (by + from)
fromNat Z      = SZ
fromNat (S by) = SS $ fromNat by
%transform "Shift.fromNat" Shift.fromNat x = believe_me x

public export
fromNat0 : (by : Nat) -> Shift 0 by
fromNat0 by = rewrite sym $ plusZeroRightNeutral by in fromNat by

export
0 fromToNat : (by : Shift from to) -> by `Eqv` fromNat by.nat {from}
fromToNat SZ      = EqSZ
fromToNat (SS by) = EqSS $ fromToNat by

export
0 toFromNat : (from, by : Nat) -> by = (fromNat by {from}).nat
toFromNat from 0     = Refl
toFromNat from (S k) = cong S $ toFromNat from k

export
0 toNatInj' : (by : Shift from1 to1) -> (bz : Shift from2 to2) ->
              by.nat = bz.nat -> by `Eqv` bz
toNatInj' SZ      SZ      prf = EqSZ
toNatInj' (SS by) (SS bz) prf = EqSS $ toNatInj' by bz $ injective prf
toNatInj' (SS by) SZ      Refl impossible

export
0 toNatInj : {by, bz : Shift from to} -> by.nat = bz.nat -> by = bz
toNatInj {by, bz} e = fromEqv $ toNatInj' by bz e

export %inline
Injective Shift.(.nat) where injective eq = irrelevantEq $ toNatInj eq


public export
ssDown : Shift (S from) to -> Shift from to
ssDown SZ      = SS SZ
ssDown (SS by) = SS (ssDown by)

export
0 ssDownEqv : (by : Shift (S from) to) -> ssDown by `Eqv` SS by
ssDownEqv SZ      = EqSS EqSZ
ssDownEqv (SS by) = EqSS $ ssDownEqv by

private %inline
ssDownViaNat : Shift (S from) to -> Shift from to
ssDownViaNat by =
  rewrite shiftDiff by in
  rewrite sym $ plusSuccRightSucc by.nat from in
  fromNat $ S by.nat

%transform "Shift.ssDown" ssDown = ssDownViaNat


public export
weak : (s : Nat) -> Shift from to -> Shift (s + from) (s + to)
weak s SZ = SZ
weak s (SS by) {to = S to} =
  rewrite sym $ plusSuccRightSucc s to in
  SS $ weak s by

private
weakViaNat : (s : Nat) -> Shift from to -> Shift (s + from) (s + to)
weakViaNat s by =
  rewrite shiftDiff by in
  rewrite plusAssociative s by.nat from in
  rewrite plusCommutative s by.nat in
  rewrite sym $ plusAssociative by.nat s from in
  fromNat by.nat

%transform "Shift.weak" Shift.weak = weakViaNat


public export
shift : Shift from to -> Var from -> Var to
shift SZ      i = i
shift (SS by) i = VS $ shift by i

private %inline
shiftViaNat' : (by : Shift from to) -> (i : Var from) ->
               (0 p : by.nat + i.nat `LT` to) -> Var to
shiftViaNat' by i p = V $ by.nat + i.nat

private %inline
shiftViaNat : Shift from to -> Var from -> Var to
shiftViaNat by i = shiftViaNat' by i $ shiftVarLT by i

private
0 shiftViaNatCorrect : (by : Shift from to) -> (i : Var from) ->
                       (0 p : by.nat + i.nat `LT` to) ->
                       shiftViaNat' by i p = shift by i
shiftViaNatCorrect SZ      i (LTESucc p) = fromToNat i _
shiftViaNatCorrect (SS by) i (LTESucc p) = cong VS $ shiftViaNatCorrect by i p

%transform "Shift.shift" shift = shiftViaNat


public export
(.) : Shift from mid -> Shift mid to -> Shift from to
by . SZ    = by
by . SS bz = SS $ by . bz

private
0 compNatProof : (by : Shift from mid) -> (bz : Shift mid to) ->
                 to = by.nat + bz.nat + from
compNatProof by bz =
  trans (shiftDiff bz) $
  trans (cong (bz.nat +) (shiftDiff by)) $
  trans (plusAssociative bz.nat by.nat from) $
  cong (+ from) (plusCommutative bz.nat by.nat)

private %inline
compViaNat' : (by : Shift from mid) -> (bz : Shift mid to) ->
              Shift from (by.nat + bz.nat + from)
compViaNat' by bz = fromNat $ by.nat + bz.nat

private %inline
compViaNat : (by : Shift from mid) -> (bz : Shift mid to) -> Shift from to
compViaNat by bz = rewrite compNatProof by bz in compViaNat' by bz

private
0 compViaNatCorrect : (by : Shift from mid) -> (bz : Shift mid to) ->
                      by . bz `Eqv` compViaNat' by bz
compViaNatCorrect by SZ =
  rewrite plusZeroRightNeutral by.nat in fromToNat by
compViaNatCorrect by (SS bz) =
  rewrite sym $ plusSuccRightSucc by.nat bz.nat in
  EqSS $ compViaNatCorrect by bz

%transform "Shift.(.)" Shift.(.) = compViaNat


infixl 8 //
public export
interface CanShift f where
  (//) : f from -> Shift from to -> f to

export %inline
CanShift Var where i // by = shift by i

namespace CanShift
  public export %inline
  [Map] (Functor f, CanShift tm) => CanShift (f . tm) where
    x // by = map (// by) x

  public export %inline
  [Const] CanShift (\_ => a) where x // _ = x


export
shiftOPE : {mask : Nat} -> (0 ope : OPE m n mask) ->
           Shift n n' -> Subset Nat (OPE m n')
shiftOPE ope SZ      = Element _ ope
shiftOPE ope (SS by) =
  let Element _ ope = shiftOPE ope by in Element _ $ drop ope

export
CanShift (Thinned f) where
  Th ope tm // by = Th (shiftOPE ope by).snd tm