import Mlang.Nickel
import Mlang.Http
import Mlang.NickelHost

namespace Mlang

open Nickel

abbrev Name := String
structure Rational where
  num : Int
  den : Nat
  deriving Repr, Inhabited, BEq

abbrev Decimal := Rational

mutual
  inductive Generator where
    | exactRatio : Rational → Generator
    | exactDecimal : Decimal → Generator
    | rationalRefiner : LazyReal → Generator
    | decimalRefiner : LazyReal → Generator
    | piChudnovsky : Generator
    | sqrt : Decimal → Generator
    | add : Generator → Generator → Generator
    | sub : Generator → Generator → Generator
    | mul : Generator → Generator → Generator
    | div : Generator → Generator → Generator
    deriving Repr, Inhabited, BEq

  inductive LazyReal where
    | exact : Decimal → LazyReal
    | ratioProcess : Rational → LazyReal
    | decimalProcess : Decimal → LazyReal
    | pi : LazyReal
    | sqrt : Decimal → LazyReal
    | add : LazyReal → LazyReal → LazyReal
    | sub : LazyReal → LazyReal → LazyReal
    | mul : LazyReal → LazyReal → LazyReal
    | div : LazyReal → LazyReal → LazyReal
    | generated : Generator → LazyReal
    deriving Repr, Inhabited, BEq
end

def normalizeDecimalParts (negative : Bool) (whole frac : String) : String :=
  let wholeChars := whole.toList.dropWhile (· = '0')
  let whole' := if wholeChars.isEmpty then "0" else String.ofList wholeChars
  let fracChars := frac.toList.reverse.dropWhile (· = '0') |>.reverse
  let frac' := if fracChars.isEmpty then "0" else String.ofList fracChars
  let sign := if negative && !(whole' = "0" && frac' = "0") then "-" else ""
  sign ++ whole' ++ "." ++ frac'

def parseDecimalParts (repr : String) : Bool × String × String :=
  let chars := repr.toList
  let (negative, chars) :=
    match chars with
    | '-' :: rest => (true, rest)
    | _ => (false, chars)
  let rec split (acc : List Char) : List Char → List Char × List Char
    | [] => (acc.reverse, [])
    | '.' :: rest => (acc.reverse, rest)
    | c :: rest => split (c :: acc) rest
  let (wholeChars, fracChars) := split [] chars
  (negative, String.ofList wholeChars, String.ofList fracChars)

def pow10 : Nat → Int
  | 0 => 1
  | n + 1 => 10 * pow10 n

def powNat (base : Nat) : Nat → Nat
  | 0 => 1
  | n + 1 => base * powNat base n

def mkRational (num : Int) (den : Nat) : Rational :=
  if den = 0 then
    { num := 0, den := 1 }
  else if num = 0 then
    { num := 0, den := 1 }
  else
    let g := Nat.gcd num.natAbs den
    { num := num / Int.ofNat g, den := den / g }

def parseDecimal (repr : String) : Decimal :=
  let (negative, whole, frac) := parseDecimalParts repr
  let digits := whole ++ frac
  let magnitude := String.toInt? digits |>.getD 0
  let signed := if negative then -magnitude else magnitude
  mkRational signed (powNat 10 frac.length)

def renderDecimalScaled (scaled : Int) (scale : Nat) : String :=
  let negative := scaled < 0
  let mag := scaled.natAbs
  let digits := toString mag
  if scale = 0 then
    normalizeDecimalParts negative digits "0"
  else if digits.length <= scale then
    let zeros := String.ofList (List.replicate (scale - digits.length + 1) '0')
    let padded := zeros ++ digits
    let wholeLen := padded.length - scale
    normalizeDecimalParts negative (padded.take wholeLen |>.toString) (padded.drop wholeLen |>.toString)
  else
    let wholeLen := digits.length - scale
    normalizeDecimalParts negative (digits.take wholeLen |>.toString) (digits.drop wholeLen |>.toString)

def decimalOfInt (n : Int) : Decimal :=
  mkRational n 1

def addDecimals (lhs rhs : Decimal) : Decimal :=
  mkRational (lhs.num * Int.ofNat rhs.den + rhs.num * Int.ofNat lhs.den) (lhs.den * rhs.den)

def subDecimals (lhs rhs : Decimal) : Decimal :=
  mkRational (lhs.num * Int.ofNat rhs.den - rhs.num * Int.ofNat lhs.den) (lhs.den * rhs.den)

def mulDecimals (lhs rhs : Decimal) : Decimal :=
  mkRational (lhs.num * rhs.num) (lhs.den * rhs.den)

def rationalEq (lhs rhs : Decimal) : Bool :=
  lhs.num == rhs.num && lhs.den == rhs.den

partial def countFactor (factor : Nat) (n : Nat) : Nat :=
  let rec loop (count current : Nat) :=
    if current = 0 then count
    else if current % factor = 0 then loop (count + 1) (current / factor)
    else count
  loop 0 n

partial def stripFactor (factor : Nat) (n : Nat) : Nat :=
  let rec loop (current : Nat) :=
    if current = 0 then 0
    else if current % factor = 0 then loop (current / factor)
    else current
  loop n

def renderDecimal (d : Decimal) : String :=
  if d.den = 1 then
    renderDecimalScaled d.num 0
  else
    let twos := countFactor 2 d.den
    let withoutTwos := stripFactor 2 d.den
    let fives := countFactor 5 withoutTwos
    let rest := stripFactor 5 withoutTwos
    if rest = 1 then
      let scale := max twos fives
      let mult2 := powNat 2 (scale - twos)
      let mult5 := powNat 5 (scale - fives)
      let scaledNum := d.num * Int.ofNat (mult2 * mult5)
      renderDecimalScaled scaledNum scale
    else
      s!"{d.num}/{d.den}"

def isTerminatingDecimal (q : Rational) : Bool :=
  if q.den = 1 then
    true
  else
    let withoutTwos := stripFactor 2 q.den
    let withoutFives := stripFactor 5 withoutTwos
    withoutFives = 1

inductive Effect where
  | error : Effect
  | io : Effect
  deriving Repr, Inhabited, BEq

abbrev Effects := List Effect

def Effects.contains (effects : Effects) (effect : Effect) : Bool :=
  effects.any (· == effect)

def Effects.insert (effects : Effects) (effect : Effect) : Effects :=
  if effects.contains effect then effects else effect :: effects

def Effects.union (lhs rhs : Effects) : Effects :=
  rhs.foldl Effects.insert lhs

def Effects.erase (effects : Effects) (target : Effect) : Effects :=
  effects.filter (· != target)

def Effects.render : Effects → String
  | [] => "{}"
  | effects =>
      let names := effects.reverse.map (fun
        | .error => "Error"
        | .io => "IO")
      "{" ++ String.intercalate ", " names ++ "}"

inductive Ty where
  | number : Ty
  | int : Ty
  | decimal : Ty
  | rational : Ty
  | lazyReal : Ty
  | stability : Ty
  | bool : Ty
  | string : Ty
  | data : Ty
  | dataInt : Ty
  | dataDecimal : Ty
  | dataRational : Ty
  | dataBool : Ty
  | dataString : Ty
  | dataNull : Ty
  | dataArray : Ty → Ty
  | dataRecord : List (String × Ty) → Ty
  | list : Ty
  | result : Ty
  | error : Ty
  | funTy : Ty → Effects → Ty → Ty
  deriving Repr, Inhabited, BEq

inductive Data where
  | null : Data
  | bool : Bool → Data
  | int : Int → Data
  | decimal : Decimal → Data
  | string : String → Data
  | array : List Data → Data
  | record : List (String × Data) → Data
  deriving Repr, Inhabited, BEq

inductive Expr where
  | int : Int → Expr
  | decimal : Decimal → Expr
  | bool : Bool → Expr
  | string : String → Expr
  | null : Expr
  | var : Name → Expr
  | arrayE : List Expr → Expr
  | recordE : List (String × Expr) → Expr
  | add : Expr → Expr → Expr
  | sub : Expr → Expr → Expr
  | mul : Expr → Expr → Expr
  | div : Expr → Expr → Expr
  | eq : Expr → Expr → Expr
  | letE : Name → Expr → Expr → Expr
  | ifE : Expr → Expr → Expr → Expr
  | tryE : Expr → Name → Expr → Expr
  | parMapE : Name → Expr → Expr → Expr
  | lam : Name → Ty → Expr → Expr
  | app : Expr → Expr → Expr
  deriving Repr, Inhabited

inductive Numeric where
  | exact : Rational → Numeric
  | composite : LazyReal → Numeric
  deriving Repr, Inhabited, BEq

inductive Value where
  | number : Numeric → Value
  | stability : String → Value
  | bool : Bool → Value
  | string : String → Value
  | data : Data → Value
  | list : List Value → Value
  | resultOk : Value → Value
  | resultErr : String → Value
  | error : String → Value
  | builtinReadFile : Value
  | builtinHttpGet : Value
  | builtinReadFileNickel : Value
  | builtinParseNickel : Value
  | builtinParseJson : Value
  | builtinGetAturi : Value
  | builtinDataGet : Value
  | builtinDataGetField : Value → Value
  | builtinDataAt : Value
  | builtinDataAtIndex : Value → Value
  | builtinDataAsInt : Value
  | builtinDataAsDecimal : Value
  | builtinRationalToInt : Value
  | builtinRationalToDecimal : Value
  | builtinDataAsString : Value
  | builtinDataAsBool : Value
  | builtinDataToJson : Value
  | builtinDataToNickel : Value
  | builtinSqrt : Value
  | builtinProcess : Value
  | builtinSynthesize : Value
  | builtinFormula : Value
  | builtinApprox : Value
  | builtinApproxStep : Value → Value
  | builtinStabilize : Value
  | builtinStabilizeStep : Value → Value
  | builtinBounds : Value
  | builtinBoundsStep : Value → Value
  | builtinRationalStability : Value
  | builtinDecimalStability : Value
  | closure : Name → Expr → List (Name × Value) → Value
  deriving Repr, Inhabited

abbrev Env := List (Name × Value)

inductive Observation where
  | asRational : Observation
  | asDecimal : Observation
  | asInt : Observation
  | rationalStability : Observation
  | decimalStability : Observation
  | getField : String → Observation
  | atIndex : Nat → Observation
  | approx : Nat → Observation
  | bounds : Nat → Observation
  deriving Repr, Inhabited

inductive ObservationResult where
  | rational : Rational → ObservationResult
  | decimal : Decimal → ObservationResult
  | int : Int → ObservationResult
  | stability : String → ObservationResult
  | data : Data → ObservationResult
  | bounds : Decimal → Decimal → ObservationResult
  deriving Repr, Inhabited

inductive RuntimeError where
  | unboundVariable : Name → RuntimeError
  | typeError : String → RuntimeError
  | divisionByZero : RuntimeError
  | ioError : String → RuntimeError
  deriving Repr, Inhabited

abbrev EvalM := EIO RuntimeError

def numericOfInt (n : Int) : Numeric :=
  .exact (decimalOfInt n)

def numericOfDecimal (d : Decimal) : Numeric :=
  .exact d

def valueToDecimal : Value → EvalM Decimal
  | .number (.exact q) => pure q
  | _ => throw (.typeError "expected a numeric value")

def divDecimals (lhs rhs : Decimal) : EvalM Decimal := do
  if rhs.num = 0 then
    throw .divisionByZero
  else
    pure (mkRational (lhs.num * Int.ofNat rhs.den) (lhs.den * rhs.num.natAbs) |>
      fun q => if rhs.num < 0 then { q with num := -q.num } else q)

def valueToLazyReal : Value → EvalM LazyReal
  | .number (.composite r) => pure r
  | .number (.exact q) => pure (.exact q)
  | _ => throw (.typeError "expected a numeric or LazyReal value")

def generatorComplexity : Generator → Nat
  | .exactRatio _ => 1
  | .exactDecimal _ => 1
  | .rationalRefiner _ => 2
  | .decimalRefiner _ => 2
  | .piChudnovsky => 3
  | .sqrt _ => 2
  | .add lhs rhs => 1 + generatorComplexity lhs + generatorComplexity rhs
  | .sub lhs rhs => 1 + generatorComplexity lhs + generatorComplexity rhs
  | .mul lhs rhs => 1 + generatorComplexity lhs + generatorComplexity rhs
  | .div lhs rhs => 1 + generatorComplexity lhs + generatorComplexity rhs

def selectFirstGoodEnough [Inhabited α] (candidates : List α) (goodEnough : α → Bool) : α :=
  match candidates.find? goodEnough with
  | some candidate => candidate
  | none =>
      match candidates with
      | candidate :: _ => candidate
      | [] => default

def generatorGoodEnough : Generator → Bool
  | .exactRatio _ => true
  | .exactDecimal _ => true
  | .rationalRefiner _ => true
  | .decimalRefiner _ => true
  | .piChudnovsky => true
  | .sqrt _ => true
  | .add lhs rhs => generatorGoodEnough lhs && generatorGoodEnough rhs
  | .sub lhs rhs => generatorGoodEnough lhs && generatorGoodEnough rhs
  | .mul lhs rhs => generatorGoodEnough lhs && generatorGoodEnough rhs
  | .div lhs rhs => generatorGoodEnough lhs && generatorGoodEnough rhs

partial def synthesizedCandidates (r : LazyReal) : List Generator :=
  match r with
  | .exact q =>
      if isTerminatingDecimal q then
        [.exactDecimal q, .exactRatio q]
      else
        [.exactRatio q, .decimalRefiner (.exact q)]
  | .ratioProcess q =>
      if isTerminatingDecimal q then
        [.exactDecimal q, .exactRatio q]
      else
        [.exactRatio q, .decimalRefiner (.ratioProcess q)]
  | .decimalProcess d =>
      [.exactDecimal d, .exactRatio d]
  | .generated g =>
      [g, .decimalRefiner (.generated g), .rationalRefiner (.generated g)]
  | .pi =>
      [.decimalRefiner .pi, .rationalRefiner .pi, .piChudnovsky]
  | .sqrt d =>
      [.decimalRefiner (.sqrt d), .rationalRefiner (.sqrt d), .sqrt d]
  | .add lhs rhs =>
      let lhsGen := selectFirstGoodEnough (synthesizedCandidates lhs) generatorGoodEnough
      let rhsGen := selectFirstGoodEnough (synthesizedCandidates rhs) generatorGoodEnough
      [.add lhsGen rhsGen, .decimalRefiner (.add lhs rhs), .rationalRefiner (.add lhs rhs)]
  | .sub lhs rhs =>
      let lhsGen := selectFirstGoodEnough (synthesizedCandidates lhs) generatorGoodEnough
      let rhsGen := selectFirstGoodEnough (synthesizedCandidates rhs) generatorGoodEnough
      [.sub lhsGen rhsGen, .decimalRefiner (.sub lhs rhs), .rationalRefiner (.sub lhs rhs)]
  | .mul lhs rhs =>
      let lhsGen := selectFirstGoodEnough (synthesizedCandidates lhs) generatorGoodEnough
      let rhsGen := selectFirstGoodEnough (synthesizedCandidates rhs) generatorGoodEnough
      [.mul lhsGen rhsGen, .decimalRefiner (.mul lhs rhs), .rationalRefiner (.mul lhs rhs)]
  | .div lhs rhs =>
      let lhsGen := selectFirstGoodEnough (synthesizedCandidates lhs) generatorGoodEnough
      let rhsGen := selectFirstGoodEnough (synthesizedCandidates rhs) generatorGoodEnough
      [.div lhsGen rhsGen, .decimalRefiner (.div lhs rhs), .rationalRefiner (.div lhs rhs)]

def chooseSynthesizedGenerator (r : LazyReal) : Generator :=
  selectFirstGoodEnough (synthesizedCandidates r) generatorGoodEnough

mutual
  partial def renderGeneratorFormula : Generator → String
    | .exactRatio q => s!"{q.num}/{q.den}"
    | .exactDecimal d => renderDecimal d
    | .rationalRefiner r => s!"rationalRefine({renderLazyFormula r})"
    | .decimalRefiner r => s!"decimalRefine({renderLazyFormula r})"
    | .piChudnovsky =>
        "426880 * sqrt(10005) / sum(k => ((-1)^k * (6k)! * (13591409 + 545140134k)) / ((3k)! * (k!)^3 * 640320^(3k)), k = 0..inf)"
    | .sqrt d => s!"sqrt({renderDecimal d})"
    | .add lhs rhs => s!"({renderGeneratorFormula lhs}) + ({renderGeneratorFormula rhs})"
    | .sub lhs rhs => s!"({renderGeneratorFormula lhs}) - ({renderGeneratorFormula rhs})"
    | .mul lhs rhs => s!"({renderGeneratorFormula lhs}) * ({renderGeneratorFormula rhs})"
    | .div lhs rhs => s!"({renderGeneratorFormula lhs}) / ({renderGeneratorFormula rhs})"

  partial def renderLazyFormula : LazyReal → String
    | .exact q => renderDecimal q
    | .ratioProcess q => s!"ratio({q.num}/{q.den})"
    | .decimalProcess d => s!"decimal({renderDecimal d})"
    | .pi => "pi"
    | .sqrt d => s!"sqrt({renderDecimal d})"
    | .add lhs rhs => s!"({renderLazyFormula lhs}) + ({renderLazyFormula rhs})"
    | .sub lhs rhs => s!"({renderLazyFormula lhs}) - ({renderLazyFormula rhs})"
    | .mul lhs rhs => s!"({renderLazyFormula lhs}) * ({renderLazyFormula rhs})"
    | .div lhs rhs => s!"({renderLazyFormula lhs}) / ({renderLazyFormula rhs})"
    | .generated g => renderGeneratorFormula g
end

def formulaOfValue (value : Value) : EvalM String := do
  match value with
  | .number (.exact q) =>
      pure (renderGeneratorFormula (chooseSynthesizedGenerator (.exact q)))
  | .number (.composite r) =>
      pure (renderGeneratorFormula (chooseSynthesizedGenerator r))
  | _ =>
      let q ← valueToDecimal value
      pure (renderGeneratorFormula (chooseSynthesizedGenerator (.exact q)))

def synthesizeNumeric (value : Value) : EvalM Numeric := do
  match value with
  | .number (.exact q) =>
      pure (.composite (.generated (chooseSynthesizedGenerator (.exact q))))
  | .number (.composite r) =>
      pure (.composite (.generated (chooseSynthesizedGenerator r)))
  | _ =>
      let q ← valueToDecimal value
      pure (.composite (.generated (chooseSynthesizedGenerator (.exact q))))

partial def natSqrtSearch (target lo hi : Nat) : Nat :=
  if lo + 1 >= hi then
    lo
  else
    let mid := lo + (hi - lo) / 2
    if mid * mid ≤ target then
      natSqrtSearch target mid hi
    else
      natSqrtSearch target lo mid

def natSqrtFloor (target : Nat) : Nat :=
  natSqrtSearch target 0 (target + 1)

def sqrtBoundsOfDecimal (d : Decimal) (steps : Nat) : EvalM (Decimal × Decimal) := do
  if d.num < 0 then
    throw (.typeError "sqrt expects a non-negative decimal")
  else
    let scale := steps
    let scaledTarget := d.num.natAbs * powNat 10 (2 * scale) / d.den
    let loNat := natSqrtFloor scaledTarget
    let lo := mkRational (Int.ofNat loNat) (powNat 10 scale)
    let step := mkRational 1 (powNat 10 scale)
    pure (lo, addDecimals lo step)

def natFactorial : Nat → Nat
  | 0 => 1
  | n + 1 => (n + 1) * natFactorial n

def chudnovskyTermMag (k : Nat) : Decimal :=
  let sixFact := natFactorial (6 * k)
  let threeFact := natFactorial (3 * k)
  let kFact := natFactorial k
  let coeff : Int := 13591409 + 545140134 * Int.ofNat k
  let num := Int.ofNat sixFact * coeff
  let den := threeFact * kFact * kFact * kFact * powNat 640320 (3 * k)
  mkRational num den

def chudnovskySignedTerm (k : Nat) : Decimal :=
  let mag := chudnovskyTermMag k
  if k % 2 = 0 then mag else { mag with num := -mag.num }

def chudnovskySum (terms : Nat) : Decimal :=
  let rec loop (k : Nat) (acc : Decimal) :=
    if k < terms then
      loop (k + 1) (addDecimals acc (chudnovskySignedTerm k))
    else
      acc
  loop 0 (decimalOfInt 0)

def piBounds (steps : Nat) : EvalM (Decimal × Decimal) := do
  let terms := steps / 14 + 1
  let sum := chudnovskySum terms
  let tail := chudnovskyTermMag terms
  let sumLo := subDecimals sum tail
  let sumHi := addDecimals sum tail
  let (sqrtLo, sqrtHi) ← sqrtBoundsOfDecimal (decimalOfInt 10005) (steps + 2)
  let coeff := decimalOfInt 426880
  let cLo := mulDecimals coeff sqrtLo
  let cHi := mulDecimals coeff sqrtHi
  let lo ← divDecimals cLo sumHi
  let hi ← divDecimals cHi sumLo
  pure (lo, hi)

def minDecimal (a b : Decimal) : Decimal :=
  if a.num * Int.ofNat b.den <= b.num * Int.ofNat a.den then a else b

def maxDecimal (a b : Decimal) : Decimal :=
  if a.num * Int.ofNat b.den >= b.num * Int.ofNat a.den then a else b

def roundRationalToScale (q : Rational) (scale : Nat) : Rational :=
  let factor := powNat 10 scale
  let scaledNum := q.num * Int.ofNat factor
  let denInt := Int.ofNat q.den
  let adjust := Int.ofNat (q.den / 2)
  let rounded :=
    if scaledNum < 0 then
      (scaledNum - adjust) / denInt
    else
      (scaledNum + adjust) / denInt
  mkRational rounded factor

def mulBounds (lo1 hi1 lo2 hi2 : Decimal) : Decimal × Decimal :=
  let p1 := mulDecimals lo1 lo2
  let p2 := mulDecimals lo1 hi2
  let p3 := mulDecimals hi1 lo2
  let p4 := mulDecimals hi1 hi2
  let lo := minDecimal (minDecimal p1 p2) (minDecimal p3 p4)
  let hi := maxDecimal (maxDecimal p1 p2) (maxDecimal p3 p4)
  (lo, hi)

mutual
  partial def boundsGenerator (g : Generator) (steps : Nat) : EvalM (Decimal × Decimal) := do
    match g with
    | .exactRatio q => pure (q, q)
    | .exactDecimal d => pure (d, d)
    | .rationalRefiner r =>
        let q ← approxLazyReal r steps
        pure (q, q)
    | .decimalRefiner r =>
        let q ← approxLazyReal r steps
        pure (roundRationalToScale q steps, roundRationalToScale q steps)
    | .piChudnovsky => piBounds steps
    | .sqrt d => sqrtBoundsOfDecimal d steps
    | .add lhs rhs =>
        let (llo, lhi) ← boundsGenerator lhs steps
        let (rlo, rhi) ← boundsGenerator rhs steps
        pure (addDecimals llo rlo, addDecimals lhi rhi)
    | .sub lhs rhs =>
        let (llo, lhi) ← boundsGenerator lhs steps
        let (rlo, rhi) ← boundsGenerator rhs steps
        pure (subDecimals llo rhi, subDecimals lhi rlo)
    | .mul lhs rhs =>
        let (llo, lhi) ← boundsGenerator lhs steps
        let (rlo, rhi) ← boundsGenerator rhs steps
        pure (mulBounds llo lhi rlo rhi)
    | .div lhs rhs =>
        let (llo, lhi) ← boundsGenerator lhs steps
        let (rlo, rhi) ← boundsGenerator rhs steps
        if rlo.num <= 0 && 0 <= rhi.num then
          throw (.typeError "generated division interval crosses zero")
        else
          let invLo ← divDecimals (decimalOfInt 1) rhi
          let invHi ← divDecimals (decimalOfInt 1) rlo
          pure (mulBounds llo lhi invLo invHi)

  partial def approxGenerator (g : Generator) (steps : Nat) : EvalM Decimal := do
    match g with
    | .exactRatio q => pure q
    | .exactDecimal d => pure d
    | .rationalRefiner r => approxLazyReal r steps
    | .decimalRefiner r =>
        let q ← approxLazyReal r steps
        pure (roundRationalToScale q steps)
    | .piChudnovsky =>
        let (lo, hi) ← piBounds steps
        divDecimals (addDecimals lo hi) (decimalOfInt 2)
    | .sqrt d =>
        let (lo, hi) ← sqrtBoundsOfDecimal d steps
        divDecimals (addDecimals lo hi) (decimalOfInt 2)
    | .add _ _ | .sub _ _ | .mul _ _ | .div _ _ =>
        let (lo, hi) ← boundsGenerator g steps
        divDecimals (addDecimals lo hi) (decimalOfInt 2)

  partial def boundsLazyReal (r : LazyReal) (steps : Nat) : EvalM (Decimal × Decimal) := do
    match r with
    | .exact d => pure (d, d)
    | .ratioProcess q => pure (q, q)
    | .decimalProcess d => pure (d, d)
    | .generated g => boundsGenerator g steps
    | .pi => piBounds steps
    | .sqrt d => sqrtBoundsOfDecimal d steps
    | .add lhs rhs =>
        let (llo, lhi) ← boundsLazyReal lhs steps
        let (rlo, rhi) ← boundsLazyReal rhs steps
        pure (addDecimals llo rlo, addDecimals lhi rhi)
    | .sub lhs rhs =>
        let (llo, lhi) ← boundsLazyReal lhs steps
        let (rlo, rhi) ← boundsLazyReal rhs steps
        pure (subDecimals llo rhi, subDecimals lhi rlo)
    | .mul lhs rhs =>
        let (llo, lhi) ← boundsLazyReal lhs steps
        let (rlo, rhi) ← boundsLazyReal rhs steps
        pure (mulBounds llo lhi rlo rhi)
    | .div lhs rhs =>
        let (llo, lhi) ← boundsLazyReal lhs steps
        let (rlo, rhi) ← boundsLazyReal rhs steps
        if rlo.num <= 0 && 0 <= rhi.num then
          throw (.typeError "lazy division interval crosses zero")
        else
          let invLo ← divDecimals (decimalOfInt 1) rhi
          let invHi ← divDecimals (decimalOfInt 1) rlo
          pure (mulBounds llo lhi invLo invHi)

  partial def approxLazyReal (r : LazyReal) (steps : Nat) : EvalM Decimal := do
    match r with
    | .exact d => pure d
    | .ratioProcess q => pure q
    | .decimalProcess d => pure d
    | .generated g => approxGenerator g steps
    | .pi =>
        let (lo, hi) ← piBounds steps
        divDecimals (addDecimals lo hi) (decimalOfInt 2)
    | .sqrt d =>
        let (lo, hi) ← sqrtBoundsOfDecimal d steps
        divDecimals (addDecimals lo hi) (decimalOfInt 2)
    | .add _ _ | .sub _ _ | .mul _ _ | .div _ _ =>
        let (lo, hi) ← boundsLazyReal r steps
        divDecimals (addDecimals lo hi) (decimalOfInt 2)
end

def stabilizeNumeric (value : Value) (steps : Nat) : EvalM Rational := do
  match value with
  | .number (.exact q) => pure (roundRationalToScale q steps)
  | .number (.composite r) =>
      let q ← approxLazyReal r steps
      pure (roundRationalToScale q steps)
  | _ =>
      let q ← valueToDecimal value
      pure (roundRationalToScale q steps)

def Env.lookup (env : Env) (name : Name) : Option Value :=
  match env with
  | [] => none
  | (n, v) :: rest => if n = name then some v else rest.lookup name

def expectInt : Value → EvalM Int
  | .number (.exact q) =>
      if q.den = 1 then pure q.num else throw (.typeError "expected an integer")
  | _ => throw (.typeError "expected an integer")

def expectDecimal : Value → EvalM Decimal
  | .number (.exact q) => pure q
  | _ => throw (.typeError "expected a decimal")

def expectLazyReal : Value → EvalM LazyReal
  | .number (.composite r) => pure r
  | _ => throw (.typeError "expected a LazyReal")

def expectBool : Value → EvalM Bool
  | .bool b => pure b
  | _ => throw (.typeError "expected a boolean")

def expectString : Value → EvalM String
  | .string s => pure s
  | _ => throw (.typeError "expected a string")

def expectData : Value → EvalM Data
  | .data term => pure term
  | _ => throw (.typeError "expected a data value")

def expectWholeNat (value : Value) : EvalM Nat := do
  let q ← valueToDecimal value
  if q.den ≠ 1 then
    throw (.typeError "expected a whole rational")
  else if q.num < 0 then
    throw (.typeError "expected a non-negative whole rational")
  else
    pure q.num.toNat

def rationalStabilityOfValue (value : Value) : String :=
  match value with
  | .number (.composite (.ratioProcess _)) => "Immediate"
  | .number (.composite (.decimalProcess _)) => "Immediate"
  | .number (.composite (.generated (.exactRatio _))) => "Immediate"
  | .number (.composite (.generated (.exactDecimal _))) => "Immediate"
  | .number (.exact _) => "Immediate"
  | .number (.composite _) => "Refining"
  | _ => "Unavailable"

def decimalStabilityOfRational (q : Rational) : String :=
  if isTerminatingDecimal q then "Terminating" else "Periodic"

def decimalStabilityOfValue (value : Value) : String :=
  match value with
  | .number (.composite (.ratioProcess q)) => decimalStabilityOfRational q
  | .number (.composite (.decimalProcess d)) => decimalStabilityOfRational d
  | .number (.composite (.generated (.exactRatio q))) => decimalStabilityOfRational q
  | .number (.composite (.generated (.exactDecimal d))) => decimalStabilityOfRational d
  | .number (.exact q) => decimalStabilityOfRational q
  | .number (.composite _) => "Refining"
  | _ => "Unavailable"

def observeValue (value : Value) (obs : Observation) : EvalM ObservationResult := do
  match obs with
  | .asRational =>
      match value with
      | .number (.exact q) => pure (.rational q)
      | .number (.composite _) =>
          throw (.typeError s!"rational observation failed: number is {rationalStabilityOfValue value}, not Immediate")
      | _ =>
          pure (.rational (← valueToDecimal value))
  | .asDecimal =>
      match value with
      | .number (.exact q) =>
          if isTerminatingDecimal q then
            pure (.decimal q)
          else
            throw (.typeError s!"decimal observation failed: number is {decimalStabilityOfValue value}, not Terminating")
      | .number (.composite _) =>
          throw (.typeError s!"decimal observation failed: number is {decimalStabilityOfValue value}, not Terminating")
      | _ =>
          let q ← valueToDecimal value
          if isTerminatingDecimal q then
            pure (.decimal q)
          else
            throw (.typeError s!"decimal observation failed: number is {decimalStabilityOfValue value}, not Terminating")
  | .asInt =>
      match value with
      | .number (.exact q) =>
          if q.den = 1 then
            pure (.int q.num)
          else
            throw (.typeError "integer observation failed: number is not whole")
      | .number (.composite _) =>
          throw (.typeError s!"integer observation failed: number is {rationalStabilityOfValue value}, not Immediate")
      | _ =>
          let q ← valueToDecimal value
          if q.den = 1 then
            pure (.int q.num)
          else
            throw (.typeError "integer observation failed: number is not whole")
  | .rationalStability =>
      pure (.stability (rationalStabilityOfValue value))
  | .decimalStability =>
      pure (.stability (decimalStabilityOfValue value))
  | .getField key =>
      match value with
      | .data (.record fields) =>
          match fields.find? (fun (name, _) => name = key) with
          | some (_, field) => pure (.data field)
          | none => throw (.typeError s!"missing field: {key}")
      | _ => throw (.typeError "get expects a record")
  | .atIndex idx =>
      let rec arrayGet? : List Data → Nat → Option Data
        | [], _ => none
        | x :: _, 0 => some x
        | _ :: xs, n + 1 => arrayGet? xs n
      match value with
      | .data (.array items) =>
          match arrayGet? items idx with
          | some item => pure (.data item)
          | none => throw (RuntimeError.typeError s!"index out of bounds: {idx}")
      | _ => throw (RuntimeError.typeError "at expects an array")
  | .approx steps =>
      let r ← valueToLazyReal value
      pure (.rational (← approxLazyReal r steps))
  | .bounds steps =>
      let r ← valueToLazyReal value
      let (lo, hi) ← boundsLazyReal r steps
      pure (.bounds lo hi)

def valueToData : Value → EvalM Data
  | .data term => pure term
  | .number (.exact q) => pure (.decimal q)
  | .bool b => pure (.bool b)
  | .string s => pure (.string s)
  | _ => throw (.typeError "data literals only support Int, Decimal, Bool, String, and Data values")

partial def dataType : Data → Ty
  | .null => .dataNull
  | .bool _ => .dataBool
  | .int _ => .dataInt
  | .decimal _ => .dataRational
  | .string _ => .dataString
  | .array [] => .dataArray .data
  | .array (x :: xs) =>
      let itemTy := dataType x
      if xs.all (fun item => dataType item == itemTy) then
        .dataArray itemTy
      else
        .dataArray .data
  | .record fields =>
      .dataRecord (fields.map (fun (name, value) => (name, dataType value)))

def dataRecordField? : List (String × Ty) → String → Option Ty
  | [], _ => none
  | (name, ty) :: rest, target =>
      if name = target then some ty else dataRecordField? rest target

def liftScalarDataTy : Ty → Ty
  | .dataInt => .int
  | .dataDecimal => .decimal
  | .dataRational => .rational
  | .dataBool => .bool
  | .dataString => .string
  | ty => ty

partial def refineKnownDataTy (ty : Ty) : Ty :=
  match ty with
  | .dataInt => .int
  | .dataDecimal => .decimal
  | .dataRational => .rational
  | .dataBool => .bool
  | .dataString => .string
  | .dataNull => .data
  | .dataArray itemTy => .dataArray (refineKnownDataTy itemTy)
  | .dataRecord fields =>
      .dataRecord (fields.map (fun (name, fieldTy) => (name, refineKnownDataTy fieldTy)))
  | other => other

def dataArrayGet? : List Data → Nat → Option Data
  | [], _ => none
  | x :: _, 0 => some x
  | _ :: xs, n + 1 => dataArrayGet? xs n

def runtimeErrorToValue : RuntimeError → Value
  | .unboundVariable name => .error s!"unbound variable: {name}"
  | .typeError msg => .error s!"type error: {msg}"
  | .divisionByZero => .error "division by zero"
  | .ioError msg => .error s!"io error: {msg}"

def runtimeErrorToMessage : RuntimeError → String
  | .unboundVariable name => s!"unbound variable: {name}"
  | .typeError msg => s!"type error: {msg}"
  | .divisionByZero => "division by zero"
  | .ioError msg => s!"io error: {msg}"

partial def Data.ofNickel : Nickel.Term → Data
  | .null => .null
  | .bool b => .bool b
  | .int n => .int n
  | .decimal d => .decimal (parseDecimal d)
  | .string s => .string s
  | .array xs => .array (xs.map Data.ofNickel)
  | .record fields => .record (fields.map (fun (k, v) => (k, Data.ofNickel v)))

def parseRenderedData (rendered : String) : EvalM Data := do
  match Nickel.parse rendered with
  | .ok term => pure (Data.ofNickel term)
  | .error err => throw (.typeError s!"host Nickel render parse error: {err}")

def encodeMapResult : Except RuntimeError Value → Value
  | .ok value => .resultOk value
  | .error err => .resultErr (runtimeErrorToMessage err)

def escapeJsonChar : Char → String
  | '"' => "\\\""
  | '\\' => "\\\\"
  | '\n' => "\\n"
  | '\r' => "\\r"
  | '\t' => "\\t"
  | c => c.toString

def renderJsonString (s : String) : String :=
  "\"" ++ String.join (s.toList.map escapeJsonChar) ++ "\""

def isNickelKeyStart (c : Char) : Bool :=
  c.isAlpha || c = '_'

def isNickelKeyContinue (c : Char) : Bool :=
  c.isAlpha || c.isDigit || c = '_' || c = '-'

def renderNickelKey (s : String) : String :=
  match s.toList with
  | [] => renderJsonString s
  | c :: cs =>
      if isNickelKeyStart c && cs.all isNickelKeyContinue then
        s
      else
        renderJsonString s

partial def renderNickelData : Data → String
  | .null => "null"
  | .bool b => toString b
  | .int n => toString n
  | .decimal d => renderDecimal d
  | .string s => renderJsonString s
  | .array xs => "[" ++ String.intercalate ", " (xs.map renderNickelData) ++ "]"
  | .record fields =>
      let rendered := fields.map (fun (k, v) => s!"{renderNickelKey k} = {renderNickelData v}")
      "{ " ++ String.intercalate ", " rendered ++ " }"

partial def renderJsonData : Data → String
  | .null => "null"
  | .bool b => if b then "true" else "false"
  | .int n => toString n
  | .decimal d => renderDecimal d
  | .string s => renderJsonString s
  | .array xs => "[" ++ String.intercalate ", " (xs.map renderJsonData) ++ "]"
  | .record fields =>
      let rendered := fields.map (fun (k, v) => s!"{renderJsonString k}: {renderJsonData v}")
      "{" ++ String.intercalate ", " rendered ++ "}"

def builtinType? (name : Name) : Option Ty :=
  match name with
  | "readFile" => some (.funTy .string [.error, .io] .string)
  | "httpGet" => some (.funTy .string [.error, .io] .string)
  | "readFileNickel" => some (.funTy .string [.error, .io] .data)
  | "parseNickel" => some (.funTy .string [.error] .data)
  | "parseJson" => some (.funTy .string [.error] .data)
  | "getAturi" => some (.funTy .string [.error, .io] .data)
  | "get" => some (.funTy .data [] (.funTy .string [.error] .data))
  | "at" => some (.funTy .data [] (.funTy .number [.error] .data))
  | "asInt" => some (.funTy .data [.error] .int)
  | "asDecimal" => some (.funTy .data [.error] .decimal)
  | "toInt" => some (.funTy .number [.error] .int)
  | "toDecimal" => some (.funTy .number [.error] .decimal)
  | "asString" => some (.funTy .data [.error] .string)
  | "asBool" => some (.funTy .data [.error] .bool)
  | "toJson" => some (.funTy .data [] .string)
  | "toNickel" => some (.funTy .data [] .string)
  | "rationalStability" => some (.funTy .number [] .stability)
  | "decimalStability" => some (.funTy .number [] .stability)
  | "pi" => some .number
  | "sqrt" => some (.funTy .number [.error] .number)
  | "process" => some (.funTy .number [] .number)
  | "synthesize" => some (.funTy .number [.error] .number)
  | "formula" => some (.funTy .number [.error] .string)
  | "lazyRationalStability" => some (.funTy .number [] .stability)
  | "lazyDecimalStability" => some (.funTy .number [] .stability)
  | "stabilize" => some (.funTy .number [] (.funTy .number [.error] .number))
  | "approx" => some (.funTy .number [] (.funTy .number [.error] .number))
  | "bounds" => some (.funTy .number [] (.funTy .number [.error] .data))
  | _ => none

def builtinValue? (name : Name) : Option Value :=
  match name with
  | "readFile" => some .builtinReadFile
  | "httpGet" => some .builtinHttpGet
  | "readFileNickel" => some .builtinReadFileNickel
  | "parseNickel" => some .builtinParseNickel
  | "parseJson" => some .builtinParseJson
  | "getAturi" => some .builtinGetAturi
  | "get" => some .builtinDataGet
  | "at" => some .builtinDataAt
  | "asInt" => some .builtinDataAsInt
  | "asDecimal" => some .builtinDataAsDecimal
  | "toInt" => some .builtinRationalToInt
  | "toDecimal" => some .builtinRationalToDecimal
  | "asString" => some .builtinDataAsString
  | "asBool" => some .builtinDataAsBool
  | "toJson" => some .builtinDataToJson
  | "toNickel" => some .builtinDataToNickel
  | "rationalStability" => some .builtinRationalStability
  | "decimalStability" => some .builtinDecimalStability
  | "pi" => some (.number (.composite .pi))
  | "sqrt" => some .builtinSqrt
  | "process" => some .builtinProcess
  | "synthesize" => some .builtinSynthesize
  | "formula" => some .builtinFormula
  | "lazyRationalStability" => some .builtinRationalStability
  | "lazyDecimalStability" => some .builtinDecimalStability
  | "stabilize" => some .builtinStabilize
  | "approx" => some .builtinApprox
  | "bounds" => some .builtinBounds
  | _ => none

mutual
  partial def applyValue (fnVal argVal : Value) : EvalM Value := do
    match fnVal with
    | .closure param body closureEnv =>
        eval ((param, argVal) :: closureEnv) body
    | .builtinReadFile => do
        let path ← expectString argVal
        let contents ← IO.toEIO (fun err => .ioError (toString err)) (IO.FS.readFile path)
        pure (.string contents)
    | .builtinHttpGet => do
        let url ← expectString argVal
        let body ← IO.toEIO (fun err => .ioError (toString err)) (Http.httpGet url)
        pure (.string body)
    | .builtinReadFileNickel => do
        let path ← expectString argVal
        let rendered ← IO.toEIO (fun err => .ioError (toString err)) (NickelHost.evalFile path)
        pure (.data (← parseRenderedData rendered))
    | .builtinParseNickel => do
        let source ← expectString argVal
        let rendered ← IO.toEIO (fun err => .ioError (toString err)) (NickelHost.evalString source)
        pure (.data (← parseRenderedData rendered))
    | .builtinParseJson => do
        let source ← expectString argVal
        let rendered ← IO.toEIO (fun err => .ioError (toString err)) (NickelHost.evalJsonString source)
        pure (.data (← parseRenderedData rendered))
    | .builtinGetAturi => do
        let source ← expectString argVal
        let rendered ← IO.toEIO (fun err => .ioError (toString err)) (Http.parseAturi source)
        pure (.data (← parseRenderedData rendered))
    | .builtinDataGet => pure (.builtinDataGetField argVal)
    | .builtinDataGetField target => do
        let key ← expectString argVal
        match (← observeValue target (.getField key)) with
        | .data value => pure (.data value)
        | _ => throw (.typeError "get observation returned non-data")
    | .builtinDataAt => pure (.builtinDataAtIndex argVal)
    | .builtinDataAtIndex target => do
        let idx ← expectWholeNat argVal
        match (← observeValue target (.atIndex idx)) with
        | .data value => pure (.data value)
        | _ => throw (.typeError "at observation returned non-data")
    | .builtinDataAsInt => do
        let term ← expectData argVal
        match term with
        | .int n => pure (.number (numericOfInt n))
        | _ => throw (.typeError "asInt expects an integer")
    | .builtinDataAsDecimal => do
        let term ← expectData argVal
        match term with
        | .decimal d => pure (.number (numericOfDecimal d))
        | _ => throw (.typeError "asDecimal expects a decimal")
    | .builtinRationalToInt => do
        match (← observeValue argVal .asInt) with
        | .int n => pure (.number (numericOfInt n))
        | _ => throw (.typeError "toInt observation returned non-int")
    | .builtinRationalToDecimal => do
        match (← observeValue argVal .asDecimal) with
        | .decimal d => pure (.number (numericOfDecimal d))
        | _ => throw (.typeError "toDecimal observation returned non-decimal")
    | .builtinDataAsString => do
        let term ← expectData argVal
        match term with
        | .string s => pure (.string s)
        | _ => throw (.typeError "asString expects a string")
    | .builtinDataAsBool => do
        let term ← expectData argVal
        match term with
        | .bool b => pure (.bool b)
        | _ => throw (.typeError "asBool expects a boolean")
    | .builtinDataToJson => do
        let term ← expectData argVal
        pure (.string (renderJsonData term))
    | .builtinDataToNickel => do
        let term ← expectData argVal
        pure (.string (renderNickelData term))
    | .builtinRationalStability => do
        match (← observeValue argVal .rationalStability) with
        | .stability label => pure (.stability label)
        | _ => throw (.typeError "rationalStability observation returned non-stability")
    | .builtinDecimalStability => do
        match (← observeValue argVal .decimalStability) with
        | .stability label => pure (.stability label)
        | _ => throw (.typeError "decimalStability observation returned non-stability")
    | .builtinSqrt => do
        let d ← valueToDecimal argVal
        if d.num < 0 then
          throw (.typeError "sqrt expects a non-negative decimal")
        else
          pure (.number (.composite (.sqrt d)))
    | .builtinProcess => do
        match argVal with
        | .number (.exact q) => pure (.number (.composite (.exact q)))
        | .number (.composite r) => pure (.number (.composite r))
        | _ => throw (.typeError "process expects a number")
    | .builtinSynthesize => do
        pure (.number (← synthesizeNumeric argVal))
    | .builtinFormula => do
        pure (.string (← formulaOfValue argVal))
    | .builtinStabilize => pure (.builtinStabilizeStep argVal)
    | .builtinStabilizeStep target => do
        let steps ← expectWholeNat argVal
        pure (.number (.exact (← stabilizeNumeric target steps)))
    | .builtinApprox => pure (.builtinApproxStep argVal)
    | .builtinApproxStep target => do
        let steps ← expectWholeNat argVal
        match (← observeValue target (.approx steps)) with
        | .rational q => pure (.number (.exact q))
        | _ => throw (.typeError "approx observation returned non-rational")
    | .builtinBounds => pure (.builtinBoundsStep argVal)
    | .builtinBoundsStep target => do
        let steps ← expectWholeNat argVal
        match (← observeValue target (.bounds steps)) with
        | .bounds lo hi => pure (.data (.record [("lo", .decimal lo), ("hi", .decimal hi)]))
        | _ => throw (.typeError "bounds observation returned non-bounds")
    | .data (.record fields) =>
        let key ← expectString argVal
        match fields.find? (fun (name, _) => name = key) with
        | some (_, value) => pure (.data value)
        | none => throw (.typeError s!"missing field: {key}")
    | .data (.array items) =>
        let idx ← expectWholeNat argVal
        match dataArrayGet? items idx with
        | some value => pure (.data value)
        | none => throw (RuntimeError.typeError s!"index out of bounds: {idx}")
    | .data _ =>
        throw (.typeError "data value is not callable with this argument")
    | _ =>
        throw (.typeError "expected a function")

  partial def eval (env : Env) : Expr → EvalM Value
  | .int n => pure (.number (numericOfInt n))
  | .decimal d => pure (.number (numericOfDecimal d))
  | .bool b => pure (.bool b)
  | .string s => pure (.string s)
  | .null => pure (.data .null)
  | .var name =>
      match env.lookup name with
      | some v => pure v
      | none =>
          match builtinValue? name with
          | some v => pure v
          | none => throw (.unboundVariable name)
  | .arrayE items => do
      let values ← items.mapM (eval env)
      pure (.data (.array (← values.mapM valueToData)))
  | .recordE fields => do
      let fields' ← fields.mapM (fun (name, value) => do
        let value' ← eval env value
        pure (name, (← valueToData value')))
      pure (.data (.record fields'))
  | .add lhs rhs => do
      let lVal ← eval env lhs
      let rVal ← eval env rhs
      match lVal, rVal with
      | .number (.composite _), _ | _, .number (.composite _) =>
          pure (.number (.composite (.add (← valueToLazyReal lVal) (← valueToLazyReal rVal))))
      | _, _ =>
          pure (.number (.exact (addDecimals (← valueToDecimal lVal) (← valueToDecimal rVal))))
  | .sub lhs rhs => do
      let lVal ← eval env lhs
      let rVal ← eval env rhs
      match lVal, rVal with
      | .number (.composite _), _ | _, .number (.composite _) =>
          pure (.number (.composite (.sub (← valueToLazyReal lVal) (← valueToLazyReal rVal))))
      | _, _ =>
          pure (.number (.exact (subDecimals (← valueToDecimal lVal) (← valueToDecimal rVal))))
  | .mul lhs rhs => do
      let lVal ← eval env lhs
      let rVal ← eval env rhs
      match lVal, rVal with
      | .number (.composite _), _ | _, .number (.composite _) =>
          pure (.number (.composite (.mul (← valueToLazyReal lVal) (← valueToLazyReal rVal))))
      | _, _ =>
          pure (.number (.exact (mulDecimals (← valueToDecimal lVal) (← valueToDecimal rVal))))
  | .div lhs rhs => do
      let lVal ← eval env lhs
      let rVal ← eval env rhs
      match lVal, rVal with
      | .number (.composite _), _ | _, .number (.composite _) =>
          pure (.number (.composite (.div (← valueToLazyReal lVal) (← valueToLazyReal rVal))))
      | _, _ =>
          pure (.number (.exact (← divDecimals (← valueToDecimal lVal) (← valueToDecimal rVal))))
  | .eq lhs rhs => do
      let l ← eval env lhs
      let r ← eval env rhs
      match l, r with
      | .number _, .number _ => pure (.bool (rationalEq (← valueToDecimal l) (← valueToDecimal r)))
      | _, _ => throw (.typeError "equality expects matching Int or Decimal operands")
  | .letE name value body => do
      let value' ← eval env value
      eval ((name, value') :: env) body
  | .ifE cond thenBranch elseBranch => do
      let c ← expectBool (← eval env cond)
      if c then
        eval env thenBranch
      else
        eval env elseBranch
  | .tryE body errName handler => do
      try
        eval env body
      catch err =>
        eval ((errName, runtimeErrorToValue err) :: env) handler
  | .parMapE itemName collection body => do
      let source ← expectData (← eval env collection)
      match source with
      | .array items =>
          let tasks ← items.mapM (fun item =>
            EIO.asTask (prio := .dedicated) (eval ((itemName, .data item) :: env) body))
          let results : List (Except RuntimeError Value) := List.map Task.get tasks
          pure (.list (results.map encodeMapResult))
      | _ =>
          throw (.typeError "pmap expects an array value")
  | .lam param _ body =>
      pure (.closure param body env)
  | .app fn arg => do
      let fnVal ← eval env fn
      let argVal ← eval env arg
      applyValue fnVal argVal
end

def renderData : Data → String := renderNickelData

def renderObservedScalar (value : Value) : Option String :=
  match value with
  | .number (.exact q) => some s!"{q.num}/{q.den}"
  | _ => none

def renderValue : Value → String
  | value@(.number (.exact _)) =>
      match renderObservedScalar value with
      | some rendered => rendered
      | none => "<number>"
  | .stability label => label
  | .number (.composite (.ratioProcess q)) => s!"<process:ratio({q.num}/{q.den})>"
  | .number (.composite (.decimalProcess d)) => s!"<process:decimal({renderDecimal d})>"
  | .number (.composite (.generated (.exactRatio q))) => s!"<generator:ratio({q.num}/{q.den})>"
  | .number (.composite (.generated (.exactDecimal d))) => s!"<generator:decimal({renderDecimal d})>"
  | .number (.composite (.generated (.rationalRefiner _))) => "<generator:rational-refiner>"
  | .number (.composite (.generated (.decimalRefiner _))) => "<generator:decimal-refiner>"
  | .number (.composite (.generated .piChudnovsky)) => "<generator:pi[chudnovsky]>"
  | .number (.composite (.generated (.sqrt d))) => s!"<generator:sqrt({renderDecimal d})>"
  | .number (.composite (.generated (.add _ _))) => "<generator:add>"
  | .number (.composite (.generated (.sub _ _))) => "<generator:sub>"
  | .number (.composite (.generated (.mul _ _))) => "<generator:mul>"
  | .number (.composite (.generated (.div _ _))) => "<generator:div>"
  | .number (.composite (.exact d)) => s!"<lazy:{renderDecimal d}>"
  | .number (.composite .pi) => "<lazy:pi>"
  | .number (.composite (.sqrt d)) => s!"<lazy:sqrt({renderDecimal d})>"
  | .number (.composite (.add _ _)) => "<lazy:add>"
  | .number (.composite (.sub _ _)) => "<lazy:sub>"
  | .number (.composite (.mul _ _)) => "<lazy:mul>"
  | .number (.composite (.div _ _)) => "<lazy:div>"
  | .bool b => toString b
  | .string s => s!"\"{s}\""
  | .data data => renderData data
  | .list items => "[" ++ String.intercalate ", " (items.map renderValue) ++ "]"
  | .resultOk value => s!"ok({renderValue value})"
  | .resultErr msg => s!"err(\"{msg}\")"
  | .error msg => s!"<error:{msg}>"
  | .builtinReadFile => "<builtin:readFile>"
  | .builtinHttpGet => "<builtin:httpGet>"
  | .builtinReadFileNickel => "<builtin:readFileNickel>"
  | .builtinParseNickel => "<builtin:parseNickel>"
  | .builtinParseJson => "<builtin:parseJson>"
  | .builtinGetAturi => "<builtin:getAturi>"
  | .builtinDataGet => "<builtin:get>"
  | .builtinDataGetField _ => "<builtin:getField>"
  | .builtinDataAt => "<builtin:at>"
  | .builtinDataAtIndex _ => "<builtin:atIndex>"
  | .builtinDataAsInt => "<builtin:asInt>"
  | .builtinDataAsDecimal => "<builtin:asDecimal>"
  | .builtinRationalToInt => "<builtin:toInt>"
  | .builtinRationalToDecimal => "<builtin:toDecimal>"
  | .builtinDataAsString => "<builtin:asString>"
  | .builtinDataAsBool => "<builtin:asBool>"
  | .builtinDataToJson => "<builtin:toJson>"
  | .builtinDataToNickel => "<builtin:toNickel>"
  | .builtinRationalStability => "<builtin:rationalStability>"
  | .builtinDecimalStability => "<builtin:decimalStability>"
  | .builtinSqrt => "<builtin:sqrt>"
  | .builtinProcess => "<builtin:process>"
  | .builtinSynthesize => "<builtin:synthesize>"
  | .builtinFormula => "<builtin:formula>"
  | .builtinStabilize => "<builtin:stabilize>"
  | .builtinStabilizeStep _ => "<builtin:stabilizeStep>"
  | .builtinApprox => "<builtin:approx>"
  | .builtinApproxStep _ => "<builtin:approxStep>"
  | .builtinBounds => "<builtin:bounds>"
  | .builtinBoundsStep _ => "<builtin:boundsStep>"
  | .closure _ _ _ => "<closure>"

abbrev TyEnv := List (Name × Ty)

structure Judgment where
  ty : Ty
  effects : Effects
  deriving Repr, Inhabited

inductive TypeError where
  | unboundVariable : Name → TypeError
  | mismatch : Ty → Ty → TypeError
  | expectedFunction : Ty → TypeError
  deriving Repr, Inhabited

abbrev CheckM := Except TypeError

def TyEnv.lookup (env : TyEnv) (name : Name) : Option Ty :=
  match env with
  | [] => none
  | (n, ty) :: rest => if n = name then some ty else rest.lookup name

def ensureType (expected actual : Ty) : CheckM Unit :=
  if expected == actual then
    pure ()
  else
    throw (.mismatch expected actual)

def isNumericTy : Ty → Bool
  | .number | .int | .decimal | .rational | .lazyReal => true
  | _ => false

def pureJudgment (ty : Ty) : Judgment :=
  { ty := ty, effects := [] }

def eraseDataRefinement : Ty → Ty
  | .dataInt | .dataDecimal | .dataRational | .dataBool | .dataString | .dataNull | .dataArray _ | .dataRecord _ => .data
  | ty => ty

def canFlowTo (actual expected : Ty) : Bool :=
  if actual == expected then
    true
  else
    match actual, expected with
    | .number, .number => true
    | .int, .number => true
    | .decimal, .number => true
    | .rational, .number => true
    | .lazyReal, .number => true
    | .number, .decimal => true
    | .number, .rational => true
    | .int, .decimal => true
    | .int, .rational => true
    | .decimal, .rational => true
    | .rational, .decimal => true
    | .dataInt, .data => true
    | .dataDecimal, .data => true
    | .dataRational, .data => true
    | .dataBool, .data => true
    | .dataString, .data => true
    | .dataNull, .data => true
    | .dataArray _, .data => true
    | .dataRecord _, .data => true
    | .dataArray _, .dataArray .data => true
    | _, _ => false

partial def constString? : Expr → Option String
  | .string s => some s
  | _ => none

unsafe def evalConstData? (env : List (Name × Data)) : Expr → Option Data
  | .null => some .null
  | .string s => some (.string s)
  | .int n => some (.int n)
  | .decimal d => some (.decimal d)
  | .bool b => some (.bool b)
  | .var name => env.lookup name
  | .arrayE items => do
      some (.array (← items.mapM (evalConstData? env)))
  | .recordE fields => do
      some (.record (← fields.mapM (fun (name, value) => do
        let value' ← evalConstData? env value
        pure (name, value'))))
  | .letE name value body => do
      let value' ← evalConstData? env value
      evalConstData? ((name, value') :: env) body
  | .ifE cond thenBranch elseBranch => do
      match (← evalConstData? env cond) with
      | .bool true => evalConstData? env thenBranch
      | .bool false => evalConstData? env elseBranch
      | _ => none
  | .app fn arg => do
      match fn with
      | .var "parseNickel" =>
          let source ← evalConstData? env arg
          match source with
          | .string s =>
              match unsafeIO (NickelHost.evalString s) with
              | .ok rendered =>
                  match Nickel.parse rendered with
                  | .ok term => some (Data.ofNickel term)
                  | .error _ => none
              | .error _ => none
          | _ => none
      | .var "parseJson" =>
          let source ← evalConstData? env arg
          match source with
          | .string s =>
              match unsafeIO (NickelHost.evalJsonString s) with
              | .ok rendered =>
                  match Nickel.parse rendered with
                  | .ok term => some (Data.ofNickel term)
                  | .error _ => none
              | .error _ => none
          | _ => none
      | .var "httpGet" =>
          let url ← evalConstData? env arg
          match url with
          | .string s =>
              match unsafeIO (Http.httpGet s) with
              | .ok body => some (.string body)
              | .error _ => none
          | _ => none
      | .var "getAturi" =>
          let source ← evalConstData? env arg
          match source with
          | .string s =>
              match unsafeIO (Http.parseAturi s) with
              | .ok rendered =>
                  match Nickel.parse rendered with
                  | .ok term => some (Data.ofNickel term)
                  | .error _ => none
              | .error _ => none
          | _ => none
      | .var "readFileNickel" =>
          let path ← evalConstData? env arg
          match path with
          | .string p =>
              match unsafeIO (NickelHost.evalFile p) with
              | .ok rendered =>
                  match Nickel.parse rendered with
                  | .ok term => some (Data.ofNickel term)
                  | .error _ => none
              | .error _ => none
          | _ => none
      | .app (.var "get") base => do
          let container ← evalConstData? env base
          let key ← evalConstData? env arg
          match container, key with
          | .record fields, .string field =>
              match fields.find? (fun (name, _) => name = field) with
              | some (_, value) => some value
              | none => none
          | _, _ => none
      | .app (.var "at") base => do
          let container ← evalConstData? env base
          let index ← evalConstData? env arg
          match container, index with
          | .array items, .int idx =>
              if idx < 0 then none else dataArrayGet? items idx.toNat
          | _, _ => none
      | _ => none
  | _ => none

unsafe def inferType (env : TyEnv) : Expr → CheckM Judgment
  | .int _ => pure (pureJudgment .number)
  | .decimal _ => pure (pureJudgment .number)
  | .bool _ => pure (pureJudgment .bool)
  | .string _ => pure (pureJudgment .string)
  | .null => pure (pureJudgment .dataNull)
  | .var name =>
      match env.lookup name with
      | some ty => pure (pureJudgment ty)
      | none =>
          match builtinType? name with
          | some ty => pure (pureJudgment ty)
          | none => throw (.unboundVariable name)
  | .arrayE items => do
      let itemJs ← items.mapM (inferType env)
      let dataTys ← itemJs.mapM (fun j =>
        match j.ty with
        | .number => pure .dataRational
        | .int => pure .dataInt
        | .decimal => pure .dataDecimal
        | .rational => pure .dataRational
        | .bool => pure .dataBool
        | .string => pure .dataString
        | .data => pure .data
        | .dataInt => pure .dataInt
        | .dataDecimal => pure .dataDecimal
        | .dataRational => pure .dataRational
        | .dataBool => pure .dataBool
        | .dataString => pure .dataString
        | .dataNull => pure .dataNull
        | .dataArray ty => pure (.dataArray ty)
        | .dataRecord fields => pure (.dataRecord fields)
        | ty => throw (.mismatch .data ty))
      let itemTy :=
        match dataTys with
        | [] => .data
        | first :: rest => if rest.all (· == first) then first else .data
      pure {
        ty := .dataArray itemTy
        effects := itemJs.foldl (fun acc j => acc.union j.effects) []
      }
  | .recordE fields => do
      let fieldJs ← fields.mapM (fun (name, value) => do pure (name, ← inferType env value))
      let fieldTys ← fieldJs.mapM (fun (name, j) => do
        let ty ← match j.ty with
          | .number => pure .dataRational
          | .int => pure .dataInt
          | .decimal => pure .dataDecimal
          | .rational => pure .dataRational
          | .bool => pure .dataBool
          | .string => pure .dataString
          | .data => pure .data
          | .dataInt => pure .dataInt
          | .dataDecimal => pure .dataDecimal
          | .dataRational => pure .dataRational
          | .dataBool => pure .dataBool
          | .dataString => pure .dataString
          | .dataNull => pure .dataNull
          | .dataArray ty => pure (.dataArray ty)
          | .dataRecord fields => pure (.dataRecord fields)
          | ty => throw (.mismatch .data ty)
        pure (name, ty, j.effects))
      pure {
        ty := .dataRecord (fieldTys.map (fun (name, ty, _) => (name, ty)))
        effects := fieldTys.foldl (fun acc (_, _, effects) => acc.union effects) []
      }
  | .add lhs rhs
  | .sub lhs rhs
  | .mul lhs rhs => do
      let lhsJ ← inferType env lhs
      let rhsJ ← inferType env rhs
      if isNumericTy lhsJ.ty && isNumericTy rhsJ.ty then
        pure {
          ty := .number
          effects := lhsJ.effects.union rhsJ.effects
        }
      else
        throw (.mismatch lhsJ.ty rhsJ.ty)
  | .div lhs rhs => do
      let lhsJ ← inferType env lhs
      let rhsJ ← inferType env rhs
      if isNumericTy lhsJ.ty && isNumericTy rhsJ.ty then
        pure {
          ty := .number
          effects := (lhsJ.effects.union rhsJ.effects).insert .error
        }
      else
        throw (.mismatch lhsJ.ty rhsJ.ty)
  | .eq lhs rhs => do
      let lhsJ ← inferType env lhs
      let rhsJ ← inferType env rhs
      if isNumericTy lhsJ.ty && isNumericTy rhsJ.ty then
        pure { ty := .bool, effects := lhsJ.effects.union rhsJ.effects }
      else
        throw (.mismatch lhsJ.ty rhsJ.ty)
  | .letE name value body => do
      let valueJ ← inferType env value
      let bodyJ ← inferType ((name, valueJ.ty) :: env) body
      pure {
        ty := bodyJ.ty
        effects := valueJ.effects.union bodyJ.effects
      }
  | .ifE cond thenBranch elseBranch => do
      let condJ ← inferType env cond
      ensureType .bool condJ.ty
      let thenJ ← inferType env thenBranch
      let elseJ ← inferType env elseBranch
      ensureType thenJ.ty elseJ.ty
      pure {
        ty := thenJ.ty
        effects := condJ.effects.union (thenJ.effects.union elseJ.effects)
      }
  | .tryE body errName handler => do
      let bodyJ ← inferType env body
      let handlerJ ← inferType ((errName, .error) :: env) handler
      ensureType bodyJ.ty handlerJ.ty
      pure {
        ty := bodyJ.ty
        effects := (bodyJ.effects.erase .error).union handlerJ.effects
      }
  | .parMapE itemName collection body => do
      let collectionJ ← inferType env collection
      match collectionJ.ty with
      | .data
      | .dataArray _ =>
          let itemTy :=
            match collectionJ.ty with
            | .dataArray ty => ty
            | _ => .data
          let bodyJ ← inferType ((itemName, itemTy) :: env) body
          pure {
            ty := .list
            effects := collectionJ.effects.union bodyJ.effects
          }
      | _ =>
          throw (.mismatch .data collectionJ.ty)
  | .lam param paramTy body => do
      let bodyJ ← inferType ((param, paramTy) :: env) body
      pure (pureJudgment (.funTy paramTy bodyJ.effects bodyJ.ty))
  | .app fn arg => do
      let fnJ ← inferType env fn
      let argJ ← inferType env arg
      match fnJ.ty with
      | .funTy paramTy latentEffects resultTy =>
          if canFlowTo argJ.ty paramTy then
            let refinedResultTy :=
              match fn, arg, resultTy with
              | .app (.var "get") base, .string key, .data =>
                  match inferType env base with
                  | .ok baseJ =>
                      match baseJ.ty with
                      | .dataRecord fields =>
                          match dataRecordField? fields key with
                          | some ty => liftScalarDataTy ty
                          | none => .data
                      | _ =>
                          match evalConstData? [] (.app fn arg) with
                          | some data => liftScalarDataTy (dataType data)
                          | none => .data
                  | .error _ => .data
              | .app (.var "at") base, _, .data =>
                  match evalConstData? [] (.app fn arg) with
                  | some data => liftScalarDataTy (dataType data)
                  | none =>
                      match inferType env base with
                      | .ok baseJ =>
                          match baseJ.ty with
                          | .dataArray itemTy => liftScalarDataTy itemTy
                          | _ => .data
                      | .error _ => .data
              | .var "parseNickel", _, .data =>
                  match evalConstData? [] (.app fn arg) with
                  | some data => refineKnownDataTy (dataType data)
                  | none => .data
              | .var "parseJson", _, .data =>
                  match evalConstData? [] (.app fn arg) with
                  | some data => refineKnownDataTy (dataType data)
                  | none => .data
              | .var "readFileNickel", _, .data =>
                  match evalConstData? [] (.app fn arg) with
                  | some data => refineKnownDataTy (dataType data)
                  | none => .data
              | .var "asInt", _, .int =>
                  match argJ.ty with
                  | .dataInt => .int
                  | _ => .int
              | .var "asDecimal", _, .decimal =>
                  match argJ.ty with
                  | .dataDecimal => .decimal
                  | _ => .decimal
              | .var "toInt", _, .int =>
                  .int
              | .var "toDecimal", _, .decimal =>
                  .decimal
              | .var "asString", _, .string =>
                  match argJ.ty with
                  | .dataString => .string
                  | _ => .string
              | .var "asBool", _, .bool =>
                  match argJ.ty with
                  | .dataBool => .bool
                  | _ => .bool
              | .app (.var "approx") _, _, .number =>
                  .number
              | .app (.var "stabilize") _, _, .number =>
                  .number
              | .app (.var "bounds") _, _, .data =>
                  .data
              | _, _, _ => resultTy
            pure {
            ty := refinedResultTy
            effects := fnJ.effects.union (argJ.effects.union latentEffects)
          }
          else
            throw (.mismatch paramTy argJ.ty)
      | _ =>
          throw (.expectedFunction fnJ.ty)

partial def renderType : Ty → String
  | .number => "Number"
  | .int => "Int"
  | .decimal => "Decimal"
  | .rational => "Number"
  | .lazyReal => "Number"
  | .stability => "Stability"
  | .bool => "Bool"
  | .string => "String"
  | .data => "Data"
  | .dataInt => "Data:Int"
  | .dataDecimal => "Data:Decimal"
  | .dataRational => "Data:Rational"
  | .dataBool => "Data:Bool"
  | .dataString => "Data:String"
  | .dataNull => "Data:Null"
  | .dataArray itemTy => s!"Data:[{renderType itemTy}]"
  | .dataRecord fields =>
      let rendered := fields.map (fun (name, ty) => s!"{name}: {renderType ty}")
      "Data:{" ++ String.intercalate ", " rendered ++ "}"
  | .list => "List"
  | .result => "Result"
  | .error => "Error"
  | .funTy lhs effects rhs => s!"({renderType lhs} -> {renderType rhs} ! {effects.render})"

def renderTypeError : TypeError → String
  | .unboundVariable name => s!"unbound variable {name}"
  | .mismatch expected actual => s!"expected {renderType expected}, got {renderType actual}"
  | .expectedFunction actual => s!"expected function, got {renderType actual}"

def renderJudgment (judgment : Judgment) : String :=
  s!"{renderType judgment.ty} ! {judgment.effects.render}"

def sampleProgram : Expr :=
  .letE "inc" (.lam "x" .int (.add (.var "x") (.int 1)))
    (.app (.var "inc") (.int 41))

def lexicalScopeProgram : Expr :=
  .letE "x" (.int 10)
    (.letE "f" (.lam "y" .int (.add (.var "x") (.var "y")))
      (.letE "x" (.int 100)
        (.app (.var "f") (.int 5))))

def conditionalProgram : Expr :=
  .ifE (.eq (.mul (.int 6) (.int 7)) (.int 42))
    (.int 1)
    (.int 0)

end Mlang
