(.module:
  [lux #*
   [abstract
    [codec (#+ Codec)]
    [monad (#+ do)]]
   [control
    ["." try (#+ Try)]
    ["ex" exception (#+ exception:)]
    ["p" parser
     ["l" text (#+ Parser)]
     ["s" code]]]
   [data
    [number
     ["i" int ("#@." decimal)]]
    ["." text ("#@." monoid)]]
   [type
    abstract]
   [macro
    ["." code]
    [syntax (#+ syntax:)]]])

(exception: #export zero-cannot-be-a-modulus)

(abstract: #export (Modulus m)
  Int

  {#.doc (doc "A number used as a modulus in modular arithmetic."
              "It cannot be 0.")}

  (def: #export (from-int value)
    (Ex [m] (-> Int (Try (Modulus m))))
    (if (i.= +0 value)
      (ex.throw zero-cannot-be-a-modulus [])
      (#try.Success (:abstraction value))))

  (def: #export (to-int modulus)
    (All [m] (-> (Modulus m) Int))
    (|> modulus :representation))
  )

(exception: #export [m] (incorrect-modulus {modulus (Modulus m)}
                                           {parsed Int})
  (ex.report ["Expected" (i@encode (to-int modulus))]
             ["Actual" (i@encode parsed)]))

(exception: #export [rm sm] (cannot-equalize-moduli {reference (Modulus rm)}
                                                    {sample (Modulus sm)})
  (ex.report ["Reference" (i@encode (to-int reference))]
             ["Sample" (i@encode (to-int sample))]))

(def: #export (congruent? modulus reference sample)
  (All [m] (-> (Modulus m) Int Int Bit))
  (|> sample
      (i.- reference)
      (i.% (to-int modulus))
      (i.= +0)))

(syntax: #export (modulus {modulus s.int})
  (case (from-int modulus)
    (#try.Success _)
    (wrap (list (` (try.assume (..from-int (~ (code.int modulus)))))))
    
    (#try.Failure error)
    (p.fail error)))

(def: intL
  (Parser Int)
  (p.codec i.decimal
           (l.and (l.one-of "-+") (l.many l.decimal))))

(abstract: #export (Mod m)
  {#remainder Int
   #modulus (Modulus m)}

  {#.doc "A number under a modulus."}

  (def: #export (mod modulus)
    (All [m] (-> (Modulus m) (-> Int (Mod m))))
    (function (_ value)
      (:abstraction {#remainder (i.mod (to-int modulus) value)
                     #modulus modulus})))

  (def: #export (un-mod modular)
    (All [m] (-> (Mod m) [Int (Modulus m)]))
    (:representation modular))

  (def: separator " mod ")

  (structure: #export (codec modulus)
    (All [m] (-> (Modulus m) (Codec Text (Mod m))))

    (def: (encode modular)
      (let [[remainder modulus] (:representation modular)]
        ($_ text@compose
            (i@encode remainder)
            separator
            (i@encode (to-int modulus)))))

    (def: decode
      (l.run (do p.monad
               [[remainder _ _modulus] ($_ p.and intL (l.this separator) intL)
                _ (p.assert (ex.construct incorrect-modulus [modulus _modulus])
                            (i.= (to-int modulus) _modulus))]
               (wrap (mod modulus remainder))))))

  (def: #export (equalize reference sample)
    (All [r s] (-> (Mod r) (Mod s) (Try (Mod r))))
    (let [[reference reference-modulus] (:representation reference)
          [sample sample-modulus] (:representation sample)]
      (if (i.= (to-int reference-modulus)
               (to-int sample-modulus))
        (#try.Success (:abstraction {#remainder sample
                                     #modulus reference-modulus}))
        (ex.throw cannot-equalize-moduli [reference-modulus sample-modulus]))))

  (template [<name> <op>]
    [(def: #export (<name> reference sample)
       (All [m] (-> (Mod m) (Mod m) Bit))
       (let [[reference _] (:representation reference)
             [sample _] (:representation sample)]
         (<op> reference sample)))]

    [= i.=]
    [< i.<]
    [<= i.<=]
    [> i.>]
    [>= i.>=]
    )

  (template [<name> <op>]
    [(def: #export (<name> param subject)
       (All [m] (-> (Mod m) (Mod m) (Mod m)))
       (let [[param modulus] (:representation param)
             [subject _] (:representation subject)]
         (:abstraction {#remainder (|> subject
                                       (<op> param)
                                       (i.mod (to-int modulus)))
                        #modulus modulus})))]

    [+ i.+]
    [- i.-]
    [* i.*]
    )
  
  (def: (gcd+ a b)
    (-> Int Int [Int Int Int])
    (if (i.= +0 a)
      [+0 +1 b]
      (let [[ak bk gcd] (gcd+ (i.% a b) a)]
        [(i.- (i.* ak
                   (i./ a b))
              bk)
         ak
         gcd])))

  (def: #export (inverse modular)
    (All [m] (-> (Mod m) (Maybe (Mod m))))
    (let [[value modulus] (:representation modular)
          _modulus (to-int modulus)
          [vk mk gcd] (gcd+ value _modulus)
          co-prime? (i.= +1 gcd)]
      (if co-prime?
        (#.Some (mod modulus vk))
        #.None)))
  )