module Quox.Thin.Coprod

import public Quox.Thin.Base
import public Quox.Thin.Comp
import public Quox.Thin.View
import public Quox.Thin.List
import public Quox.Thin.Cover
import Data.DPair
import Data.Nat
import Control.Function

%default total


namespace Coprod
  public export
  data Comps : OPE scope n scopeMask ->
               OPEList scope -> OPEList n -> Type where
    Nil  : Comps sub [] []
    (::) : Comp  sub inner  full ->
           Comps sub inners fulls ->
           Comps sub (inner :: inners) (full :: fulls)
  %name Comps comps

  public export
  record Coprod (fulls : OPEList n) where
    constructor MkCoprod
    {scopeMask : Nat}
    {0 sub : OPE scope n scopeMask}
    inners : OPEList scope
    0 comps : Comps sub inners fulls
    0 cov : Cover inners
  %name Coprod cop

export
0 compsLength : Comps s ts us -> length ts = length us
compsLength []           = Refl
compsLength (_ :: comps) = cong S $ compsLength comps


export
coprodNil : Coprod []
coprodNil = MkCoprod [] [] [] {sub = zero}

private
coprodHead : {n : Nat} -> (opes : OPEList (S n)) ->
             Either (Cover1 opes) (All IsDrop opes)
coprodHead [] = Right []
coprodHead (ope :: opes) = case view ope of
  DropV {} => case coprodHead opes of
    Left  cov1  => Left  $ There cov1
    Right drops => Right $ ItIsDrop :: drops
  KeepV {} => Left Here


private
0 compsConsDrop : (opes : OPEList (S n)) ->
                  All IsDrop opes ->
                  All2 IsTail opes tails ->
                  Comps sub inners tails -> Comps (drop sub) inners opes
compsConsDrop [] [] [] [] = []
compsConsDrop (Drop ope Refl :: opes) (ItIsDrop :: ds) (DropT :: ts) (c :: cs) =
  DropZ c :: compsConsDrop opes ds ts cs
compsConsDrop (_ :: _) [] _ _ impossible

private
coprodConsDrop : (0 opes : OPEList (S n)) ->
                 (0 ds : All IsDrop opes) ->
                 (0 ts : All2 IsTail opes tails) ->
                 Coprod tails -> Coprod opes
coprodConsDrop opes ds ts (MkCoprod inners comps cov) =
  MkCoprod inners (compsConsDrop opes ds ts comps) cov


private
copyHeads : {m : Nat} ->
            (src : OPEList (S m)) -> (tgt : OPEList n) ->
            (0 eq : length src = length tgt) => OPEList (S n)
copyHeads [] [] = []
copyHeads (s :: ss) (t :: ts) =
  case view s of
    DropV mask ope => drop t :: copyHeads ss ts @{inj S eq}
    KeepV mask ope => keep t :: copyHeads ss ts @{inj S eq}

private
0 copyHeadsComps : (eq : length outers = length inners) ->
                   All2 IsTail outers tails ->
                   Comps sub inners tails ->
                   Comps (keep sub) (copyHeads outers inners) outers
copyHeadsComps _ [] [] = []
copyHeadsComps eq (DropT {eq = eq2} :: ps) ((c :: cs) {full}) =
  let (Refl) = eq2 in -- coverage checker weirdness
  rewrite viewDrop full Refl in
  KDZ c :: copyHeadsComps (inj S eq) ps cs
copyHeadsComps eq (KeepT {eq = eq2} :: ps) ((c :: cs) {full}) =
  let (Refl) = eq2 in
  rewrite viewKeep full Refl in
  KeepZ c :: copyHeadsComps (inj S eq) ps cs

-- should be erased (coverage checker weirdness)
-- it is possibly https://github.com/idris-lang/Idris2/issues/1417 that keeps
-- happening. not 100% sure
private
cover1CopyHeads : {m : Nat} ->
                  (ss : OPEList (S m)) -> (ts : OPEList n) ->
                  (eq : length ss = length ts) ->
                  (cov1 : Cover1 ss) -> Cover1 (copyHeads ss ts)
cover1CopyHeads (Keep s Refl :: ss) (t :: ts) eq Here =
  rewrite viewKeep s Refl in Here
cover1CopyHeads (s :: ss) (t :: ts) eq (There c) with (view s)
  cover1CopyHeads (Drop {} :: ss) (t :: ts) eq (There c) | DropV {} =
    There $ cover1CopyHeads ss ts (inj S eq) c
  cover1CopyHeads (Keep {} :: ss) (t :: ts) eq (There c) | KeepV {} =
    Here

private
copyHeadsTails : {m : Nat} ->
                 (ss : OPEList (S m)) -> (ts : OPEList n) ->
                 (eq : length ss = length ts) ->
                 All2 IsTail (copyHeads ss ts) ts
copyHeadsTails [] [] eq = []
copyHeadsTails (s :: ss) (t :: ts) eq with (view s)
  copyHeadsTails (Drop ope Refl :: ss) (t :: ts) eq | DropV mask ope =
    DropT :: copyHeadsTails ss ts (inj S eq)
  copyHeadsTails (Keep ope Refl :: ss) (t :: ts) eq | KeepV mask ope =
    KeepT :: copyHeadsTails ss ts (inj S eq)

private
coprodConsKeep : {n : Nat} ->
                 (opes : OPEList (S n)) ->
                 {0 tails : OPEList n} ->
                 (cov1 : Cover1 opes) ->
                 (0 ts : All2 IsTail opes tails) ->
                 Coprod tails -> Coprod opes
coprodConsKeep opes cov1 ts (MkCoprod inners comps cov) =
  MkCoprod
    (copyHeads opes inners @{all2Length ts `trans` sym (compsLength comps)})
    (copyHeadsComps _ ts comps)
    ((cover1CopyHeads {cov1, _} :: cov) @{copyHeadsTails {}})


export
coprod : {n : Nat} -> (opes : OPEList n) -> Coprod opes

private
coprod0 : (opes : OPEList 0) -> Coprod opes

private
coprodS : {n : Nat} -> (opes : OPEList (S n)) -> Coprod opes

coprod {n = 0}   opes = coprod0 opes
coprod {n = S n} opes = coprodS opes

coprod0 [] = coprodNil
coprod0 (ope :: opes) with %syntactic 0 (zeroIsStop ope) | (coprod opes)
  coprod0 (Stop :: opes)
      | ItIsStop | MkCoprod {sub} inners comps cov
      with %syntactic 0 (zeroIsStop sub)
    coprod0 (Stop :: opes)
      | ItIsStop | MkCoprod {sub = Stop} inners comps cov | ItIsStop
      = MkCoprod (Stop :: inners) (StopZ :: comps) []

coprodS [] = coprodNil
coprodS opes =
  let hs              = heads opes
      Element ts tprf = tails_ opes
      tcop            = coprod $ assert_smaller opes ts
  in
  case coprodHead opes of
    Left  cov1  => coprodConsKeep opes cov1 tprf tcop
    Right drops => coprodConsDrop opes drops tprf tcop