// TODO: This setup is more confusing than it should be, there's more work to do in cleanup here.
module Inputs = {
  module SamplingInputs = {
    type t = {
      sampleCount: option<int>,
      outputXYPoints: option<int>,
      kernelWidth: option<float>,
      pointDistLength: option<int>,
    }
  }
  let defaultRecommendedLength = 100
  let defaultShouldDownsample = true

  type inputs = {
    squiggleString: string,
    samplingInputs: SamplingInputs.t,
    environment: ASTTypes.environment,
  }

  let empty: SamplingInputs.t = {
    sampleCount: None,
    outputXYPoints: None,
    kernelWidth: None,
    pointDistLength: None,
  }

  let make = (
    ~samplingInputs=empty,
    ~squiggleString,
    ~environment=ASTTypes.Environment.empty,
    (),
  ): inputs => {
    samplingInputs: samplingInputs,
    squiggleString: squiggleString,
    environment: environment,
  }
}

type exportDistribution = [
  | #DistPlus(DistPlus.t)
  | #Float(float)
  | #Function(float => Belt.Result.t<DistPlus.t, string>)
]

type exportEnv = array<(string, ASTTypes.node)>

type exportType = {
  environment: exportEnv,
  exports: array<exportDistribution>,
}

module Internals = {
  let addVariable = (
    {samplingInputs, squiggleString, environment}: Inputs.inputs,
    str,
    node,
  ): Inputs.inputs => {
    samplingInputs: samplingInputs,
    squiggleString: squiggleString,
    environment: ASTTypes.Environment.update(environment, str, _ => Some(node)),
  }

  type outputs = {
    graph: ASTTypes.node,
    pointSetDist: PointSetTypes.pointSetDist,
  }
  let makeOutputs = (graph, shape): outputs => {graph: graph, pointSetDist: shape}

  let makeInputs = (inputs: Inputs.inputs): SamplingInputs.samplingInputs => {
    sampleCount: inputs.samplingInputs.sampleCount |> E.O.default(10000),
    outputXYPoints: inputs.samplingInputs.outputXYPoints |> E.O.default(10000),
    kernelWidth: inputs.samplingInputs.kernelWidth,
    pointSetDistLength: inputs.samplingInputs.pointDistLength |> E.O.default(10000),
  }

  let runNode = (inputs, node) => AST.toLeaf(makeInputs(inputs), inputs.environment, node)

  let renderIfNeeded = (inputs: Inputs.inputs, node: ASTTypes.node): result<
    ASTTypes.node,
    string,
  > =>
    node |> (
      x =>
        switch x {
        | #Normalize(_) as n
        | #SymbolicDist(_) as n =>
          #Render(n)
          |> runNode(inputs)
          |> (
            x =>
              switch x {
              | Ok(#RenderedDist(_)) as r => r
              | Error(r) => Error(r)
              | _ => Error("Didn't render, but intended to")
              }
          )

        | n => Ok(n)
        }
    )

  let outputToDistPlus = (inputs: Inputs.inputs, pointSetDist: PointSetTypes.pointSetDist) =>
    DistPlus.make(~pointSetDist, ~squiggleString=Some(inputs.squiggleString), ())

  let rec returnDist = (
    functionInfo: (array<string>, ASTTypes.node),
    inputs: Inputs.inputs,
    env: ASTTypes.environment,
  ) => {
    (input: float) => {
      let foo: Inputs.inputs = {...inputs, environment: env}
      evaluateFunction(foo, functionInfo, [#SymbolicDist(#Float(input))]) |> E.R.bind(_, a =>
        switch a {
        | #DistPlus(d) => Ok(DistPlus.T.normalize(d))
        | n =>
          Js.log2("Error here", n)
          Error("wrong type")
        }
      )
    }
  }
  // TODO: Consider using ExpressionTypes.ExpressionTree.getFloat or similar in this function
  and coersionToExportedTypes = (inputs, env: ASTTypes.environment, ex: ASTTypes.node): result<
    exportDistribution,
    string,
  > =>
    ex
    |> renderIfNeeded(inputs)
    |> E.R.bind(_, x =>
      switch x {
      | #RenderedDist(Discrete({xyShape: {xs: [x], ys: [1.0]}})) => Ok(#Float(x))
      | #SymbolicDist(#Float(x)) => Ok(#Float(x))
      | #RenderedDist(n) => Ok(#DistPlus(outputToDistPlus(inputs, n)))
      | #Function(n) => Ok(#Function(returnDist(n, inputs, env)))
      | n => Error("Didn't output a rendered distribution. Format:" ++ AST.toString(n))
      }
    )

  and evaluateFunction = (inputs: Inputs.inputs, fn: (array<string>, ASTTypes.node), fnInputs) => {
    let output = AST.runFunction(makeInputs(inputs), inputs.environment, fnInputs, fn)
    output |> E.R.bind(_, coersionToExportedTypes(inputs, inputs.environment))
  }

  let runProgram = (inputs: Inputs.inputs, p: ASTTypes.program) => {
    let ins = ref(inputs)
    p
    |> E.A.fmap(x =>
      switch x {
      | #Assignment(name, node) =>
        ins := addVariable(ins.contents, name, node)
        None
      | #Expression(node) => Some(runNode(ins.contents, node))
      }
    )
    |> E.A.O.concatSomes
    |> E.A.R.firstErrorOrOpen
    |> E.R.bind(_, d =>
      d
      |> E.A.fmap(x => coersionToExportedTypes(inputs, ins.contents.environment, x))
      |> E.A.R.firstErrorOrOpen
    )
    |> E.R.fmap(ex => {
      environment: Belt.Map.String.toArray(ins.contents.environment),
      exports: ex,
    })
  }

  let inputsToLeaf = (inputs: Inputs.inputs) =>
    Parser.fromString(inputs.squiggleString) |> E.R.bind(_, g => runProgram(inputs, g))
}

@genType
let runAll: (string, Inputs.SamplingInputs.t, exportEnv) => result<exportType, string> = (
  squiggleString,
  samplingInputs,
  environment,
) => {
  let inputs = Inputs.make(
    ~samplingInputs,
    ~squiggleString,
    ~environment=Belt.Map.String.fromArray(environment),
    (),
  )
  Internals.inputsToLeaf(inputs)
}