open Component

module Term = HOTerm
module Judgment = HOTerm
module JudgmentView = TermViewAsJudgmentView.Make(Term, Judgment, HOTermView)
module Rule = Rule.Make(Term, Judgment)
module RuleView = RuleView.Make(Term, Judgment, JudgmentView)
module Ports = Ports(Term, Judgment)
type state = {rules: dict<Rule.t>}
type props = {
  content: state,
  imports: Ports.t,
  onChange: (state, ~exports: Ports.t=?) => unit,
  reset: unit => unit,
}

type predicateGroup = {
  name: string,
  arity: int,
  rules: array<(string, Rule.t)>,
}

let makeKey = (name, arity) => name ++ "§" ++ Int.toString(arity)
let letters = [
  "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m",
  "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z",
]

let nameForIndex = i => {
  let letter = letters[mod(i, 26)]->Option.getOr("v")
  let cycle = i / 26
  cycle == 0 ? letter : letter ++ Int.toString(cycle)
}
let upperLetters = [
  "A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L", "M",
  "N", "O", "P", "Q", "R", "S", "T", "U", "V", "W", "X", "Y", "Z",
]

let predicateNameForIndex = i => {
  let shifted = i + 15 // start at "P" (index 15)
  let letter = upperLetters[mod(shifted, 26)]->Option.getOr("P")
  let cycle = shifted / 26
  cycle == 0 ? letter : letter ++ Int.toString(cycle)
}

let makeVarNames = (arity: int) =>
  Array.fromInitializer(~length=arity, nameForIndex)
let makeVarArgs = (arity: int) => Array.fromInitializer(~length=arity, i => Term.Var({idx: i}))

let extractPredicateSignature = (rule: Rule.t): option<(string, int)> => {
  let (head, args) = Term.strip(rule.conclusion)
  switch head {
  | Symbol({name: name}) => Some((name, Array.length(args)))
  | _ => None
  }
}

let groupByPredicate = (rules: dict<Rule.t>): array<predicateGroup> =>
  rules
  ->Dict.toArray
  ->Array.filterMap(((name, rule)) =>
    extractPredicateSignature(rule)->Option.map(sig => (name, rule, sig))
  )
  ->Array.reduce(Dict.make(), (acc, (name, rule, (cname, arity))) => {
    let key = makeKey(cname, arity)
    Dict.set(acc, key, Array.concat(Dict.get(acc, key)->Option.getOr([]), [(name, rule)]))
    acc
  })
  ->Dict.valuesToArray
  ->Array.map(predicates => {
    let (_, firstRule) = predicates[0]->Option.getExn
    let (name, arity) = extractPredicateSignature(firstRule)->Option.getExn
    {name, arity, rules: predicates}
  })
  let generateInductionRule = (group: predicateGroup, allGroups: array<predicateGroup>): Rule.t => {
    let {name: str, arity: i} = group
    let numFormers = Array.length(allGroups)
    let groupIndex = Util.mustFindIndex(allGroups, g => g.name == str && g.arity == i)

    let findFormerIndex = (name, arity) =>
      Util.mustFindIndex(allGroups, g => g.name == name && g.arity == arity)

    let outerVars = Array.concat(
      makeVarNames(i),
      Array.fromInitializer(~length=numFormers, predicateNameForIndex),
    )

    let generateInductiveHypothesis = (
      premise: Rule.t,
      offset: int,
      scope: array<string>,
    ): option<Rule.t> => {
      let (head, args) = Term.strip(premise.conclusion)
      switch head {
      | Symbol({name: name}) =>
        let formerIndex = findFormerIndex(name, Array.length(args))
        let freshVars = Term.freshenMetas(~existing=scope, ~incoming=premise.vars)
        Some({
          Rule.vars: freshVars,
          premises: premise.premises,
          conclusion: Term.unstrip(
            Term.Var({idx: offset + Array.length(premise.vars) + i + formerIndex}),
            args,
          ),
        })
      | _ => None
      }
    }

    let caseSubgoal = (constructorRule: Rule.t): Rule.t => {
      let offset = Array.length(constructorRule.vars)
      let freshConstructorVars = Term.freshenMetas(~existing=outerVars, ~incoming=constructorRule.vars)
      let initialScope = Array.concat(outerVars, freshConstructorVars)

      let (inductiveHypotheses, _finalScope) = constructorRule.premises->Array.reduce(
        ([], initialScope),
        ((acc, scope), premise) =>
          switch generateInductiveHypothesis(premise, offset, scope) {
          | Some(ihRule) => (Array.concat(acc, [ihRule]), Array.concat(scope, ihRule.vars))
          | None => (acc, scope)
          },
      )

      let (conclusionHead, conclusionArgs) = Term.strip(constructorRule.conclusion)
      let typeIndex = switch conclusionHead {
      | Symbol({name: name}) =>
        findFormerIndex(name, Array.length(conclusionArgs))
      | _ => throw(Util.Unreachable("Constructor conclusion must have a Symbol head"))
      }

      {
        Rule.vars: freshConstructorVars,
        premises: Array.concat(constructorRule.premises, inductiveHypotheses),
        conclusion: Term.unstrip(Term.Var({idx: offset + i + typeIndex}), conclusionArgs),
      }
    }

    let allConstructors = Array.flatMap(allGroups, g => g.rules)
    let subgoals = Array.map(allConstructors, ((_, rule)) => caseSubgoal(rule))

    {
      Rule.vars: outerVars,
      premises: [
        {
          Rule.vars: [],
          premises: [],
          conclusion: Term.unstrip(
            Term.Symbol({name: str, constructor: false}),
            makeVarArgs(i),
          ),
        },
        ...subgoals,
      ],
      conclusion: Term.unstrip(Term.Var({idx: i + groupIndex}), makeVarArgs(i)),
    }
  }
  
module StringCmp = Belt.Id.MakeComparable({
  type t = string
  let cmp = Pervasives.compare
})

let extractInductiveType = (premise: Rule.t): option<(string, int)> => {
  let (head, args) = Term.strip(premise.conclusion)
  switch head {
  | Symbol({name: name, constructor: false}) =>
    Some((name, Array.length(args)))
  | _ => None
  }
}

let isSelfReference = (group: predicateGroup, (name, arity)): bool =>
  name == group.name && arity == group.arity

let findDependencies = (group: predicateGroup): array<(string, int)> =>
  group.rules
  ->Array.flatMap(((_name, rule)) => rule.premises->Array.filterMap(extractInductiveType))
  ->Array.filter(dep => !isSelfReference(group, dep))

let rec collectReachable = (
  toVisit: array<(string, int)>,
  visited: Belt.Set.t<string, StringCmp.identity>,
  allGroups: array<predicateGroup>,
): Belt.Set.t<string, StringCmp.identity> =>
  switch toVisit {
  | [] => visited
  | _ =>
    let (name, arity) = toVisit[0]->Option.getExn
    let rest = Array.sliceToEnd(toVisit, ~start=1)
    let key = makeKey(name, arity)

    if Belt.Set.has(visited, key) {
      collectReachable(rest, visited, allGroups)
    } else {
      let visited = Belt.Set.add(visited, key)
      let newDeps =
        allGroups
        ->Array.find(g => g.name == name && g.arity == arity)
        ->Option.map(findDependencies)
        ->Option.getOr([])
      collectReachable(Array.concat(rest, newDeps), visited, allGroups)
    }
  }

let findMutuallyInductiveComponent = (
  targetGroup: predicateGroup,
  allGroups: array<predicateGroup>,
): array<predicateGroup> => {
  let reachableKeys = collectReachable(
    [(targetGroup.name, targetGroup.arity)],
    Belt.Set.make(~id=module(StringCmp)),
    allGroups,
  )
  allGroups->Array.filter(g => Belt.Set.has(reachableKeys, makeKey(g.name, g.arity)))
}

let generateCasesRule = (group: predicateGroup): Rule.t => {
  let {name: str, arity} = group
  let outerVars = Array.concat(makeVarNames(arity), ["P"])

  let caseSubgoal = ((_constructorName: string, predicateRule: Rule.t)): Rule.t => {
    let offset = Array.length(predicateRule.vars)
    let freshVars = Term.freshenMetas(~existing=outerVars, ~incoming=predicateRule.vars)

    // Extract the argument from the predicate conclusion
    // e.g., from (Nat 0) extract 0, from (Nat (S n)) extract (S n)
    let (_head, args) = Term.strip(predicateRule.conclusion)
    assert(Array.length(args) == arity)

    let equalityPremises = args->Array.mapWithIndex((arg, idx) => {
      {
        Rule.vars: [],
        premises: [],
        conclusion: Term.mkEquation(Term.Var({idx: offset + idx}), arg),
      }
    })

    {
      Rule.vars: freshVars,
      premises: Array.concat(equalityPremises, predicateRule.premises),
      conclusion: Term.Var({idx: offset + arity}),
    }
  }

  let subgoals = Array.map(group.rules, ((name, rule)) => caseSubgoal((name, rule)))

  {
    Rule.vars: outerVars,
    premises: [
      {
        Rule.vars: [],
        premises: [],
        conclusion: Term.unstrip(
          Term.Symbol({name: str, constructor: false}),
          makeVarArgs(arity),
        ),
      },
      ...subgoals,
    ],
    conclusion: Term.Var({idx: arity}),
  }
}

let derived = (state: dict<Rule.t>): dict<Rule.t> =>
  state
  ->groupByPredicate
  ->Array.flatMap(group => {
    let mutualComponent = findMutuallyInductiveComponent(group, groupByPredicate(state))
    let inductionRule = generateInductionRule(group, mutualComponent)
    let casesRule = generateCasesRule(group)
    [("§induction-" ++ group.name, inductionRule), ("§cases-" ++ group.name, casesRule)]
  })
  ->Dict.fromArray
let serialise = (state: state, ~imports: Ports.t) =>
  state.rules
  ->Dict.toArray
  ->Array.map(((k, r)) => r->Rule.prettyPrintTopLevel(~name=k, ~grammar=imports.grammar))
  ->Array.join("\n")
let deserialise = (str: string, ~imports: Ports.t) => {
  let cur = ref(str)
  let go = ref(true)
  let results = Dict.make()
  let ret = ref(Error("impossible"))
  while go.contents {
    switch Rule.parseTopLevel(cur.contents, ~grammar=imports.grammar, ~scope=[]) {
    | Ok((t, n), rest) =>
      if n->String.trim == "" {
        go := false
        ret := Error("Rule given with no name")
      } else {
        Dict.set(results, n, t)
        if rest->String.trim == "" {
          go := false
          ret := Ok(results)
        } else {
          cur := rest
        }
      }
    | Error(e) => {
        go := false
        ret := Error(e)
      }
    }
  }
  ret.contents->Result.map(state => (
    {rules:state},
    {Ports.facts: state->Dict.copy->Dict.assign(derived(state)), ruleStyle: None, grammar: Term.emptyGrammar},
  ))
}

let make = props => {
  <div
    className={"axiom-set axiom-set-"->String.concat(
      String.make(props.imports.ruleStyle->Option.getOr(Hybrid)),
    )}
  >
    {Dict.toArray(props.content.rules)
    ->Array.mapWithIndex(((n, r), i) =>
      <RuleView
        rule={r}
        scope={[]}
        key={String.make(i)}
        grammar={props.imports.grammar}
        style={props.imports.ruleStyle->Option.getOr(Hybrid)}
      >
        <span className="rule-rulename-global"><IdentifierView identifier=n /></span>
      </RuleView>
    )
    ->React.array}
    <section className="block">
      <details>
        <summary>
          <header>
            <h1> {React.string("Derived Rules")} </h1>
          </header>
        </summary>
        {Dict.toArray(derived(props.content.rules))
        ->Array.mapWithIndex(((n, r), i) =>
          <RuleView
            rule={r}
            scope={[]}
            grammar={props.imports.grammar}
            key={String.make(i)}
            style={props.imports.ruleStyle->Option.getOr(Hybrid)}
          >
            <span className="rule-rulename-global"><IdentifierView identifier=n /></span>
          </RuleView>
        )
        ->React.array}
      </details>
    </section>
  </div>
}