\import Algebra.Domain
\import Algebra.Field
\import Algebra.Group
\import Algebra.Meta
\import Algebra.Monoid
\import Algebra.Monoid.MonoidHom
\import Algebra.Monoid.Sub
\import Algebra.Pointed
\import Algebra.Pointed.Sub
\import Algebra.Ring
\import Algebra.Ring.RingHom
\import Function (IsInj, IsSurj)
\import Function.Meta
\import Logic.Meta
\import Logic.Unique
\import Logic
\import Meta
\import Order.PartialOrder
\import Paths
\import Paths.Meta
\import Relation.Equivalence
\import Set \hiding (#)
\open Monoid(Inv)
\open RingHom

\instance SubsetPoset {M : Monoid} : Poset (SubSet M)
  | <= (S S' : SubSet M) => \Pi (x : M) -> S x -> S' x
  | <=-refl x p => p
  | <=-transitive S1<=S2 S2<=S3 x p => S2<=S3 x (S1<=S2 x p)
  | <=-antisymmetric S1<=S2 S2<=S1 => path (\lam i => \new SubSet M (\lam x => propExt (S1<=S2 x) (S2<=S1 x) @ i))

\class Localization (R : CRing) (S : SubSet R) (R' : CRing) {
  | inL : RingHom R R'
  | local (x : R) : S x -> Inv (inL x)
  | local-univ {R'' : CRing} (f : RingHom R R'') (l : \Pi (x : R) -> S x -> Inv (f x)) : Contr (\Sigma (g : RingHom R' R'') (\Pi (x : R) -> g (inL x) = f x))

  \func lift {R'' : CRing} (f : RingHom R R'') (l : \Pi (x : R) -> S x -> Inv (f x)) : RingHom R' R''
    => (Contr.center {local-univ f l}).1

  \lemma lift-prop {R'' : CRing} (f : RingHom R R'') (l : \Pi (x : R) -> S x -> Inv (f x)) : lift f l  inL = f
    => equals {R} {R''} {lift f l  inL} (Contr.center {local-univ f l}).2

  \lemma lift-prop-func {R'' : CRing} (f : RingHom R R'') (l : \Pi (x : R) -> S x -> Inv (f x)) (x : R) : lift f l (inL x) = f x
    => path (\lam i => lift-prop f l i x)

  \lemma remove_inL {R'' : CRing} (f g : RingHom R' R'') (p : \Pi (x : R) -> f (inL x) = g (inL x)) : f = g
    => pmap __.1 (isContr=>isProp (local-univ {\this} {R''} (g  inL) (\lam x sx => MonoidHom.presInv g (local x sx))) (f,p) (g, \lam _ => idp))
}

-- TODO: Rewrite using LocMonoid
\instance LocRing {R : CRing} (S : SubMonoid R) : CRing \cowith
  | E => Type
  | zro => in~ (zro, ide, contains_ide)
  | + => ++
  | zro-left {in~ _} => ~-pequiv simplify
  | +-assoc {in~ _} {in~ _} {in~ _} => ~-pequiv equation.cRing
  | +-comm {in~ _} {in~ _} => ~-pequiv equation.cRing
  | ide => in~ (ide, ide, contains_ide)
  | * => **
  | ide-left {in~ _} => ~-pequiv simplify
  | *-assoc {in~ _} {in~ _} {in~ _} => ~-pequiv equation.monoid
  | ldistr {in~ _} {in~ _} {in~ _} => ~-pequiv equation.cRing
  | negative => neg
  | negative-left {in~ _} => ~-pequiv simplify
  | *-comm {in~ _} {in~ _} => path (~-equiv _ _ (pmap2 (*) *-comm *-comm))
  | natCoef n => in~ (R.natCoef n, 1, contains_ide)
  | natCoefZero => path (\lam i => in~ (R.natCoefZero i, 1, contains_ide))
  | natCoefSuc n => ~-pequiv $ simplify (R.natCoefSuc n)
  \where {
    \func SType (S : SubMonoid R) => \Sigma (x y : R) (S y)

    \func Type => Quotient {SType S} (\lam a b => a.1 * b.2 = b.1 * a.2)

    \lemma inl-surj : IsSurj (inl~ {R} {S})
      => Quotient.in-surj

    \lemma equals-lem {a b : SType S} (c : R) (s : S c) (p : a.1 * b.2 * c = b.1 * a.2 * c) : inl~ a = inl~ b
      => \let a' : SType S => (a.1 * c, a.2 * c, contains_* a.3 s)
         \in ~-pequiv equation.cMonoid *> path (~-equiv a' b $ equation.cMonoid {p})

    \lemma equals1 {a b : SType S} (p : a.1 * b.2 = b.1 * a.2) : inl~ a = inl~ b
      => equals-lem 1 S.contains_ide (simplify p)

    \lemma unequals {a b : SType S} (p : inl~ a = inl~ b) :  (c : R) (S c) (a.1 * b.2 * c = b.1 * a.2 * c)
      => Equivalence.Closure.univ (\new Equivalence (SType S) {
        | ~ x y =>  (z : R) (S z) (x.1 * y.2 * z = y.1 * x.2 * z)
        | ~-reflexive => inP (ide, contains_ide, idp)
        | ~-symmetric (inP (z,c,p)) => inP (z, c, inv p)
        | ~-transitive {x} {y} (inP (t,ct,pt)) (inP (s,cs,ps)) => inP (y.2 * (t * s), contains_* y.3 (contains_* ct cs), equation.cMonoid {pt,ps})
      }) (\lam x~y => inP (ide, contains_ide, simplify x~y))
         (Quotient.equality p)

    \lemma equality {a b : SType S} : (inl~ a = inl~ b) =  (c : R) (S c) (a.1 * b.2 * c = b.1 * a.2 * c)
      => propExt unequals (\lam (inP (c,s,p)) => equals-lem c s p)

    \lemma swap {R : CMonoid} {r1 r2 s1 s2 : R} : (r1 * r2) * (s1 * s2) = (r1 * s1) * (r2 * s2)
      => equation

    \func neg (a : Type) : Type \elim a
      | in~ (r,s,p) => in~ (negative r, s, p)
      | ~-equiv x y p => ~-pequiv (equation.cRing {p})

    \func ++ (a b : Type) : Type \elim a, b
      | in~ (r1,s1,p1), in~ (r2,s2,p2) => in~ (r1 * s2 + r2 * s1, s1 * s2, contains_* p1 p2)
      | ~-equiv x y p, in~ z => ~-pequiv (equation.cRing {p})
      | in~ x, ~-equiv y z p => ~-pequiv (equation.cRing {p})

    \func ** (a b : Type) : Type \elim a, b
      | in~ (r1,s1,p1), in~ (r2,s2,p2) => in~ (r1 * r2, s1 * s2, contains_* p1 p2)
      | ~-equiv x y p, in~ z => ~-pequiv (equation.cMonoid {p})
      | in~ x, ~-equiv y z p => ~-pequiv (equation.cMonoid {p})

    \lemma trivial {x : R} (n : R.IsNilpotent x) (s : S x) : Contr (LocRing S) \elim n
      | inP (n,p) => \new Contr {
        | center => in~ (zro, ide, contains_ide)
        | contraction => \case \elim __ \with {
          | in~ a => equals-lem 0 (transport S p $ S.contains_pow s) simplify
        }
      }
  }

\func inl~ {R : CRing} {S : SubMonoid R} (x : LocRing.SType S) : LocRing.Type => in~ x

\lemma loc_unequals_domain {D : IntegralDomain} (S : SubMonoid D) (nz : \Pi (x : D) -> S x -> x /= 0) {a b : LocRing.SType S} (p : inl~ a = inl~ b) : a.1 * b.2 = b.1 * a.2
  => \case LocRing.unequals p \with {
    | inP (c,c1,p) => Domain.nonZero-cancel-right (nz c c1) p
  }

\func locMap {R : CRing} {S : SubMonoid R} : RingHom R (LocRing S) \cowith
  | func x => in~ (x, ide, contains_ide)
  | func-ide => idp
  | func-+ => path (~-equiv _ _ (path (\lam i => ((inv ide-right @ i) + (inv ide-right @ i)) * (ide-right @ i))))
  | func-* => path (~-equiv _ _ simplify)

\lemma locMap-inj {R : CRing} {S : SubMonoid R} (p : \Pi {a : R} -> S a -> \Pi {x y : R} -> x * a = y * a -> x = y) : IsInj (locMap {R} {S})
  => \lam {a} {b} q => \case LocRing.unequals q \with {
    | inP (c,Sc,d) => p Sc $ rewrite (ide-right,ide-right) in d
  }

\func localization {R : CRing} (S : SubMonoid R) : Localization R S \cowith
  | R' => LocRing S
  | inL => locMap
  | local => elem-inv
  | local-univ {R''} (f : RingHom R R'') l => Contr.make (later (liftHom f l, liftHom-loc f l))
      \lam q => ext $ exts $ locMap-epiMap (liftHom f l) _ \lam x => liftHom-loc f l x *> inv (q.2 x)
  \where {
    \open LocRing

    \func elem-inv (s : R) (p : S s)
      => \new Inv {LocRing S} {
        | val => in~ (s, ide, contains_ide)
        | inv => in~ (ide, s, p)
        | inv-left => equals1 (*-comm *> pmap (ide *) *-comm)
        | inv-right => equals1 (*-comm *> pmap (ide *) *-comm)
      }

    \lemma locMap-epiMap {Q : CRing} (g h : RingHom (LocRing S) Q) (d : \Pi (a : R) -> g (locMap a) = h (locMap a)) (x : LocRing S) : g x = h x \elim x
      | in~ (r,s,p) =>
        \have | j => path (~-equiv {SType S} (r,s,p) (r * ide, ide * s, contains_* contains_ide p) (inv *-assoc))
              | m1 => transport (Inv {Q} __ (g (in~ (ide, s, p)))) (d s) (MonoidHom.presInv g (elem-inv s p))
              | m2 => MonoidHom.presInv h (elem-inv s p)
        \in g (in~ (r,s,p))                                       ==< pmap g j >==
            g (in~ (r * ide, ide * s, contains_* contains_ide p)) ==< g.func-* >==
            g (in~ (r, ide, contains_ide)) * g (in~ (ide, s, p))  ==< pmap2 (\lam t1 (t2 : Inv (h (in~ (s, ide, contains_ide)))) => t1 * t2.inv) (d r) (Inv.levelProp m1 m2) >==
            h (in~ (r, ide, contains_ide)) * h (in~ (ide, s, p))  ==< inv h.func-* >==
            h (in~ (r * ide, ide * s, contains_* contains_ide p)) ==< inv (pmap h j) >==
            h (in~ (r,s,p))                                       `qed

    \func lift {R' : CRing} (f : RingHom R R') (l : \Pi (x : R) -> S x -> Inv (f x)) (a : Type {R} {S}) : R' \elim a
      | in~ (r,s,p) => f r * Inv.inv {l s p}
      | ~-equiv (r,s,p) (r',s',p') rs~rs' => equation.cMonoid {
          1 inv (l s' p').inv-right,
          inv (f.func-* {r} {s'}),
          pmap f rs~rs',
          f.func-* {r'} {s},
          (l s p).inv-right
        }

    \func liftHom {R' : CRing} (f : RingHom R R') (l : \Pi (x : R) -> S x -> Inv (f x)) : RingHom (LocRing S) R' \cowith
      | func => lift f l
      | func-ide => Inv.inv-right
      | func-+ {in~ (r1,s1,p1)} {in~ (r2,s2,p2)} => equation.cRing {
          f.func-+ *> pmap2 (+) f.func-* f.func-*,
          func-inv,
          (l s1 p1).inv-right,
          (l s2 p2).inv-right
        }
      | func-* {in~ (r1,s1,p1)} {in~ (r2,s2,p2)} => equation.cMonoid {f.func-*, func-inv}
      \where
        \lemma func-inv {s1 s2 : R} {p1 : S s1} {p2 : S s2} : Inv.inv {l (s1 * s2) (contains_* p1 p2)} = Inv.inv {l s2 p2} * Inv.inv {l s1 p1}
          => pmap (\lam (t : Inv) => t.inv) $ Inv.levelProp (Inv.lmake _ $ pmap (_ *) (inv f.func-*) *> Inv.inv-left) Inv.product.func

    \lemma liftHom-loc {R' : CRing} (f : RingHom R R') (l : \Pi (x : R) -> S x -> Inv (f x)) (x : R) : liftHom f l (locMap x) = f x
      => pmap (f x *) (inv ide-right *> pmap (_ *) (inv f.func-ide) *> Inv.inv-left) *> ide-right

    \open SubMonoid(powers)
    \open Monoid

    \func pow-map (b : R) (n : Nat) : RingHom (LocRing (powers (pow b n))) (LocRing (powers b)) \cowith
      | func => \case __ \with {
        | in~ (x,y,c) => in~ (x, y, TruncP.map c (\lam p => (n Nat.* p.1, rewrite pow_* p.2)))
        | ~-equiv x y r => ~-pequiv r
      }
      | func-+ {in~ x} {in~ y} => ~-pequiv idp
      | func-ide => ~-pequiv idp
      | func-* {in~ x} {in~ y} => ~-pequiv idp

    \lemma pow-map_loc (b : R) (n : Nat) : pow-map b n  locMap = locMap
      => exts \lam x => ~-pequiv idp

    \lemma div-map_loc {a b : R} (a|b : LDiv a b) : div-map a|b  locMap = locMap
      => exts \lam x => \case (div-map.aux a|b (x, ide, contains_ide {powers a})).2 \with {
        | inP ((n,a^n=1),p) => inv p *> ~-pequiv (rewrite a^n=1 *-assoc)
      }

    \func div-map {a b : R} (a|b : LDiv a b) : RingHom (LocRing (powers a)) (LocRing (powers b)) \cowith
      | func => aux-func a|b
      | func-+ {in~ x} {in~ y} => \case (aux a|b x).2, (aux a|b y).2, (aux a|b _).2 \with {
        | inP ((n,a^n=x2),p1), inP ((m,a^m=y2),p2), inP ((k,a^k=x2y2),p3) => inv p3 *> ~-pequiv equation *> pmap2 (+) p1 p2
      }
      | func-ide => \case (aux a|b _).2 \with {
        | inP ((n,a^n=1),p) => inv p *> ~-pequiv equation
      }
      | func-* {in~ x} {in~ y} => \case (aux a|b x).2, (aux a|b y).2, (aux a|b _).2 \with {
        | inP ((n,a^n=x2),p1), inP ((m,a^m=y2),p2), inP ((k,a^k=x2y2),p3) => inv p3 *> ~-pequiv equation *> pmap2 (*) p1 p2
      }
      \where {
        \lemma aux {a b : R} (a|b : LDiv a b) (x : SType (powers a))
          => TruncP.rec-set {_} {LocRing (powers b)} x.3 (\lam (n,a^n=y) => in~ $ later (x.1 * pow a|b.inv n, pow a|b.inv n * pow a n, inP (n, rewriteI a|b.inv-right $ R.pow_*-comm *> *-comm))) (\lam (n,a^n=x2) (m,a^m=x2) => later $ path (~-equiv _ _ $ later $ rewrite (a^n=x2,a^m=x2) swap))

        \func aux-func {a b : R} (a|b : LDiv a b) (x : LocRing (powers a)) : LocRing (powers b)
          => \case \elim x \with {
            | in~ x => (aux a|b x).1
            | ~-equiv x y r => \case (aux a|b x).2, (aux a|b y).2 \with {
              | inP ((n,a^n=x2),p), inP ((m,a^m=y2),q) => inv p *> ~-pequiv equation *> q
            }
          }
      }

    \lemma map-lemma {R S : CRing} {M : SubMonoid R} (f : RingHom (LocRing M) S) {x : LocRing.SType M} (i : Inv (f (inl~ (x.2, 1, contains_ide)))) : f (inl~ x) = f (inl~ (x.1, 1, contains_ide)) * i.inv
      => \have t : inl~ x = inl~ (x.1, 1, contains_ide) * inl~ (1, x.2, x.3) => path $ ~-equiv _ _ (inv *-assoc)
         \in pmap f t *> f.func-* *> pmap (_ *) (Inv.inv-isUnique (Inv.lmake (f $ inl~ (1, x.2, x.3)) $ inv f.func-* *> pmap f (path $ ~-equiv _ _ *-assoc) *> f.func-ide) i idp)

    \lemma powers-inv {R S : CRing} (f : RingHom R S) {a : R} (p : Inv (f a)) (x : R) (c : powers a x) : Inv (f x) \elim c
      | inP (0, idp) => \new Inv {
        | inv => 1
        | inv-left => ide-left *> f.func-ide
        | inv-right => ide-right *> f.func-ide
      }
      | inP (suc n, idp) => rewrite f.func-* $ Inv.product (powers-inv f p (pow a n) (inP (n, idp))) p
  }

\lemma div_loc {R : CRing} {S : SubMonoid R} {a b : R} {Sb : S b} (a|b : Monoid.LDiv b a) : inl~ (a,b,Sb) = locMap a|b.inv
  => LocRing.equals1 $ ide-right *> inv a|b.inv-right *> *-comm

\open localization

\func loc_change-of-base1 {R S : CRing} (f : RingHom R S) {a : R} : RingHom (LocRing (SubMonoid.powers a)) (LocRing (SubMonoid.powers (f a)))
  => liftHom (locMap  f) $ powers-inv (locMap  f) $ elem-inv {_} {SubMonoid.powers (f a)} (f a) (inP (1, ide-left))

\lemma loc_change-of-base1-loc {R S : CRing} (f : RingHom R S) {a : R} : loc_change-of-base1 f {a}  locMap = locMap  f
  => exts (liftHom-loc (locMap  f) (powers-inv _ _))

\instance FieldOfQuotients (D : IntegralDomain.Dec) : DiscreteField \cowith
  | CRing => LocRing D.subMonoid
  | zro/=ide p => \case LocRing.unequals p \with {
    | inP (c,c#0,q) => D.#0-zro $ transportInv D.#0 (simplify in q) c#0
  }
  | eitherZeroOrInv => \case \elim __ \with {
    | in~ x => \case decideEq x.1 0 \with {
      | yes e => byLeft $ LocRing.equals1 $ ide-right *> e *> inv D.zro_*-left
      | no q => byRight $ Inv.lmake (inl~ (x.2, x.1, AddGroup.nonZeroApart q)) $ LocRing.equals1 $ simplify *-comm
    }
  }