import fs from "fs";
import path from "path";
import { Scalar, ZqField, utils } from "ffjavascript";
const { unstringifyBigInts } = utils;

// Version to write in  hexadecimal
function stringifyBigInts(o) {
  if (typeof o == "bigint" || o.eq !== undefined) {
    return "0x" + o.toString(16);
  } else if (o instanceof Uint8Array) {
    return Scalar.fromRprLE(o, 0);
  } else if (Array.isArray(o)) {
    return o.map(stringifyBigInts);
  } else if (typeof o == "object") {
    const res = {};
    const keys = Object.keys(o);
    keys.forEach((k) => {
      res[k] = stringifyBigInts(o[k]);
    });
    return res;
  } else {
    return o;
  }
}

const { C, M } = unstringifyBigInts(
  JSON.parse(
    fs.readFileSync(path.join("src", "poseidon_constants.json"), "utf8")
  )
);

const N_ROUNDS_F = 8;
const N_ROUNDS_P = [
  56, 57, 56, 60, 60, 63, 64, 63, 60, 66, 60, 65, 70, 60, 64, 68,
];

function matrix_inverse(Fr, A) {
  const m = A.length;
  const B = [];
  for (let i = 0; i < m; i++) {
    if (A[i].length != m) throw new Error("Matrix is not square");
    B[i] = [];
    for (let j = 0; j < m; j++) B[i][j] = A[i][j];
    for (let j = 0; j < m; j++) B[i][m + j] = i == j ? Fr.one : Fr.zero;
  }
  const n = 2 * m;
  let h = 0; // curRow
  let k = 0; // curColumn

  // Triangularize
  while (h < m && k < n) {
    // Find the pivot
    let i_max;
    for (i_max = h; i_max < m && Fr.isZero(B[i_max][k]); i_max++);
    if (i_max == m) throw new Error("Matrix is not inveritible");

    [B[h], B[i_max]] = [B[i_max], B[h]];

    const inv_pivot = Fr.inv(B[h][k]);

    for (let i = h + 1; i < m; i++) {
      const f = Fr.mul(B[i][k], inv_pivot);
      B[i][k] = Fr.zero;
      for (let j = k + 1; j < n; j++)
        B[i][j] = Fr.sub(B[i][j], Fr.mul(B[h][j], f));
    }

    B[h][k] = Fr.one;
    for (let j = k + 1; j < n; j++) B[h][j] = Fr.mul(B[h][j], inv_pivot);
    h++;
    k++;
  }

  for (let i = m - 2; i >= 0; i--) {
    for (let i2 = i + 1; i2 < m; i2++) {
      const f = B[i][i2];
      for (let k = m; k < n; k++)
        B[i][k] = Fr.sub(B[i][k], Fr.mul(f, B[i2][k]));
      B[i][i2] = Fr.zero;
    }
  }

  for (let i = 0; i < m; i++) {
    B[i] = B[i].slice(-m);
  }
  return B;
}

function matrix_print(Fr, A) {
  for (let i = 0; i < A.length; i++) {
    let S = "";
    for (let j = 0; j < A[i].length; j++) {
      if (j > 0) S = S + ", ";
      S = S + Fr.toString(A[i][j]);
    }
    console.log(S);
  }
}

function vec_print(Fr, v) {
  for (let i = 0; i < v.length; i++) {
    console.log(Fr.toString(v[i]));
  }
}

function vec_mul_matrix(Fr, v, A) {
  checkSquare(A);
  if (v.length != A.length) throw new Error("Invalid vec mul sizes");
  const res = [];
  for (let j = 0; j < v.length; j++) {
    res[j] = Fr.zero;
    for (let i = 0; i < A.length; i++) {
      res[j] = Fr.add(res[j], Fr.mul(v[i], A[i][j]));
    }
  }
  return res;
}

function matrix_mul_vec(Fr, A, v) {
  checkSquare(A);
  if (v.length != A.length) throw new Error("Invalid vec mul sizes");
  const res = [];
  for (let i = 0; i < v.length; i++) {
    res[i] = Fr.zero;
    for (let j = 0; j < A.length; j++) {
      res[i] = Fr.add(res[i], Fr.mul(A[i][j], v[j]));
    }
  }
  return res;
}

function matrix_mul(Fr, A, B) {
  checkSquare(A);
  if (A.length != B.length) throw new Error("multiplication different sizes");
  const res = [];
  for (let i = 0; i < A.length; i++) {
    res[i] = [];
    for (let j = 0; j < A.length; j++) {
      res[i][j] = Fr.zero;
    }
  }
  for (let i = 0; i < A.length; i++) {
    for (let j = 0; j < A.length; j++) {
      for (let k = 0; k < A.length; k++) {
        res[i][j] = Fr.add(res[i][j], Fr.mul(A[i][k], B[k][j]));
      }
    }
  }

  return res;
}

function matrix_traspose(A) {
  checkSquare(A);
  const res = [];
  for (let i = 0; i < A.length; i++) {
    res[i] = [];
    for (let j = 0; j < A.length; j++) {
      res[i][j] = A[j][i];
    }
  }
  return res;
}

function checkSquare(A) {
  const n = A.length;
  for (let i = 0; i < A.length; i++) {
    if (A[i].length != n) throw new Error("Matrix is not square");
  }
}

function convertConstants(Fr, t, C, M) {
  let res = [];
  const nRoundsF = N_ROUNDS_F;
  const nRoundsP = N_ROUNDS_P[t - 2];

  const Minv = matrix_inverse(Fr, M);

  for (let k = 0; k < t; k++) res.push(C[k]);

  for (let r = 0; r < nRoundsF / 2 - 1; r++) {
    const cr = C.slice((r + 1) * t, (r + 1) * t + t);
    const crt = vec_mul_matrix(Fr, cr, Minv);
    for (let k = 0; k < t; k++) res.push(crt[k]);
  }

  const partialConst = [];
  let acc = [];
  for (let k = 0; k < t; k++)
    acc[k] = C[(nRoundsF / 2 + nRoundsP - 1 + 1) * t + k];

  for (let r = nRoundsF / 2 + nRoundsP - 1; r >= nRoundsF / 2; r--) {
    const accp = vec_mul_matrix(Fr, acc, Minv);
    partialConst.push(accp[0]);
    accp[0] = Fr.zero;
    for (let k = 0; k < t; k++)
      acc[k] = Fr.add(accp[k], C[(r - 1 + 1) * t + k]);
  }

  const accp = vec_mul_matrix(Fr, acc, Minv);
  for (let k = 0; k < t; k++) res.push(accp[k]);

  for (let i = 0; i < partialConst.length; i++)
    res.push(partialConst[partialConst.length - 1 - i]);

  for (let r = nRoundsF / 2 + nRoundsP; r < nRoundsF + nRoundsP - 1; r++) {
    const cr = C.slice((r + 1) * t, (r + 1) * t + t);
    const crt = vec_mul_matrix(Fr, cr, Minv);
    for (let k = 0; k < t; k++) res.push(crt[k]);
  }

  return res;
}
/*
                m_{0,0}     m_{0,1}     m_{0,2}   ...........   m_{0,t-1}
                w_0         1           0         ...........   0
                w_1         0           1         ...........   0
                .           .           .            .          .
                .           .           .                .      .
                w_{t-2}     0           0         ...........   1
*/
// Returns sparse matrix with format: [m_{0,0}, m_0, ..... m_{t-2}, m_{0,1}, ... m_{0, t-1}]

function sparseFactorize(Fr, m) {
  const m_hat = [];
  const mp = [];
  const w = [];
  for (let i = 0; i < m.length; i++) {
    mp[i] = [];
    if (i < m.length - 1) m_hat[i] = [];
    for (let j = 0; j < m.length; j++) {
      if (i > 0 && j > 0) {
        m_hat[i - 1][j - 1] = m[i][j];
        mp[i][j] = m[i][j];
      } else {
        mp[i][j] = i == 0 && j == 0 ? Fr.one : Fr.zero;
      }
      if (i > 0 && j == 0) w[i - 1] = m[i][j];
    }
  }
  const m_hat_inv = matrix_inverse(Fr, m_hat);
  const wp = matrix_mul_vec(Fr, m_hat_inv, w);

  const S = [];

  S.push(m[0][0]);
  for (let k = 0; k < wp.length; k++) S.push(wp[k]);
  for (let k = 1; k < m.length; k++) S.push(m[0][k]);

  return [mp, S];
}

function calculatePS(Fr, t, M) {
  const nRoundsP = N_ROUNDS_P[t - 2];
  const sparse = [];
  let m = M;

  for (let i = 0; i < nRoundsP; i++) {
    const [mp, mpp] = sparseFactorize(Fr, m);
    sparse.push(mpp);
    m = matrix_mul(Fr, M, mp);
  }
  const P = m;
  const S = [];

  for (let i = 0; i < sparse.length; i++) {
    for (let k = 0; k < sparse[sparse.length - 1 - i].length; k++) {
      S.push(sparse[sparse.length - 1 - i][k]);
    }
  }

  return [P, S];
}

async function run() {
  // Prime bn128
  const Fr = new ZqField(
    Scalar.fromString(
      "21888242871839275222246405745257275088548364400416034343698204186575808495617"
    )
  );

  let opt = {
    C: [],
    M: [],
    P: [],
    S: [],
  };
  for (let i = 1; i <= 16; i++) {
    const oM = matrix_traspose(M[i - 1]);
    const oC = convertConstants(Fr, i + 1, C[i - 1], oM);
    const [oP, oS] = calculatePS(Fr, i + 1, oM);
    opt.C.push(oC);
    opt.M.push(oM);
    opt.P.push(oP);
    opt.S.push(oS);
  }

  opt = stringifyBigInts(opt);

  fs.writeFileSync(
    path.join("src", "poseidon_constants_opt.js"),
    "export default " + JSON.stringify(opt, null, 1),
    "utf8"
  );
}

run().then(
  () => {
    process.exit(0);
  },
  (err) => {
    //    console.log(err);
    console.log(err.stack);
    process.exit(1);
  }
);