import { eqType } from "../metacircular.js";
import { pretty, zip } from "../util.js";

// constructor for generic types
//   for instance, the type:
//     a -> a -> Bool
//   is created by
//     makeGeneric(a => fnType({in: a, out: fnType({in: a, out: Bool})}))
export const makeGeneric = callback => {
  // type variables to make available:
  const typeVars = ['a', 'b', 'c', 'd', 'e'].map(letter => ({
    symbol: Symbol(letter), 
    params: [],
  }));
  const type = callback(...typeVars);
  return {
    typeVars: occurring(type, new Set(typeVars)),
    type,
  };
};

// From the given set of type variables, return only those that occur in the given type.
export const occurring = (type, typeVars) => {
  if (typeVars.has(type)) {
    // type IS a type variable:
    return new Set([type]);
  }
  return new Set(type.params.flatMap(p => [...occurring(p, typeVars)]));
}

const mergeOneWay = (m1, m2) => {
  const m1copy = new Map(m1);
  const m2copy = new Map(m2);
  for (const [key1, val1] of m1copy.entries()) {
    if (m2copy.has(val1)) {
      m1copy.set(key1, m2.get(val1));
      m2copy.delete(val1);
      return [false, m1copy, m2copy, new Set([val1])];
    }
  }
  return [true, m1copy, m2copy, new Set()]; // stable
}

export const mergeTwoWay = (m1, m2) => {
  // check for conflicts:
  for (const [typeVar, actual] of m1) {
    if (m2.has(typeVar)) {
      const other = m2.get(typeVar);
      if (!eqType(actual, other)) {
        throw new Error(`conflicting substitution: ${pretty(actual)}vs. ${pretty(other)}`);
      }
    }
  }
  // actually merge
  let stable = false;
  let deletedTypeVars = new Set();
  while (!stable) {
    let d;
    // notice we swap m2 and m1, so the rewriting can happen both ways:
    [stable, m2, m1, d] = mergeOneWay(m1, m2);
    deletedTypeVars = deletedTypeVars.union(d);
  }
  return [new Map([...m1, ...m2]), deletedTypeVars];
}

// Thanks to Hans for pointing out that this algorithm exactly like "Unification" in Prolog (hence the function name):
//  https://www.dai.ed.ac.uk/groups/ssp/bookpages/quickprolog/node12.html
export const unify = (
  {typeVars: formalTypeVars, type: formalType},
  {typeVars: actualTypeVars, type: actualType},
) => {
  // console.log("unify", {formalTypeVars, formalType, actualTypeVars, actualType});
  if (formalTypeVars.has(formalType)) {
    // simplest case: formalType is a type paramater
    //   => substitute with actualType
    return {
      substitutions: new Map([[formalType, actualType]]),
      typeVars: new Set([
        ...actualTypeVars,
        ...[...formalTypeVars].filter(a => a !== formalType),
      ]),
      type: actualType,
    };
  }
  if (actualTypeVars.has(actualType)) {
    // same as above, but in the other direction
    return {
      substitutions: new Map([[actualType, formalType]]),
      typeVars: new Set([
        ...[...actualTypeVars].filter(a => a !== actualType),
        ...formalTypeVars,
      ]),
      type: formalType,
    };
  }
  // recursively unify
  if (formalType.symbol !== actualType.symbol) {
    throw new Error(`cannot unify ${pretty(formalType.symbol)} and ${pretty(actualType.symbol)}`);
  }
  else {
    const unifiedParams = zip(formalType.params, actualType.params).map(([formalParam, actualParam]) => unify({typeVars: formalTypeVars, type: formalParam}, {typeVars: actualTypeVars, type: actualParam}));
    const [unifiedSubstitusions, deleted] = unifiedParams.reduce(([substitutionsSoFar, deletedSoFar], cur) => {
      // console.log('merging', substitutionsSoFar, cur.substitutions);
      const [newSubstitutions, deleted] = mergeTwoWay(substitutionsSoFar, cur.substitutions);
      return [newSubstitutions, deletedSoFar.union(deleted)];
    }, [new Map(), new Set()]);
    const unifiedTypeVars = new Set([
      ...actualTypeVars,
      ...formalTypeVars,
    ].filter(a => !unifiedSubstitusions.has(a) && !deleted.has(a)));
    return {
      substitutions: unifiedSubstitusions,
      typeVars: unifiedTypeVars,
      type: {
        symbol: formalType.symbol,
        params: unifiedParams.map(p => p.type),
      },
    };
  }
};

export const substitute = (type, substitutions) => {
  if (substitutions.has(type)) {
    // type IS a type var to be substituted:
    return substitutions.get(type);
  }
  return {
    symbol: type.symbol,
    params: type.params.map(p => substitute(p, substitutions)),
  };
}

export const assign = (genFnType, paramType) => {
  const [inType, outType] = genFnType.type.params;
  const matchedInType = unify({
    typeVars: genFnType.typeVars,
    type: inType,
  }, paramType);
  const substitutedOutType = substitute(outType, matchedInType.substitutions);
  return {
    typeVars: matchedInType.typeVars,
    type: substitutedOutType,
  };
};