#include <mbedtls/bignum.h>
#include <mbedtls/version.h>
#include "poseidon_constants.h"
// #include "bignum.h"
// #include "esp_bignum.c"


void setup() {
  // put your setup code here, to run once:
  Serial.begin(115200);
  Serial.println("Hello, ESP32-C3!");


  // Print mbedTLS version
  Serial.print("mbedTLS version: ");
  Serial.println(MBEDTLS_VERSION_STRING);
  
  
  mbedtls_mpi inputs[2], result;
  char result_str[256];
  size_t olen;

  mbedtls_mpi_init(&inputs[0]);
  mbedtls_mpi_init(&inputs[1]);
  mbedtls_mpi_init(&result);

  mbedtls_mpi_lset(&inputs[0], 11);
  mbedtls_mpi_lset(&inputs[1], 0);

  unsigned int start = millis();
  // 2136ms - original poseidon
  poseidon(&result, inputs, 2);
  Serial.println("Time: ");
  Serial.println(millis() - start);

  // mbedtls_mpi_write_string(&result, 10, result_str, sizeof(result_str), &olen);
  
  // Serial.println("XYW");
  // Serial.println(esp_mpi_hardware_words(2));

  // Print the result
  Serial.println("Poseidon Hash:");
  Serial.println(result_str);


  mbedtls_mpi_free(&inputs[0]);
  mbedtls_mpi_free(&inputs[1]);
  mbedtls_mpi_free(&result);

}


void poseidon_fast(mbedtls_mpi *result, const mbedtls_mpi *inputs, size_t num_inputs) {
    mbedtls_mpi P, FIVE_BIGNUM, temp, RR;
    mbedtls_mpi state[MAX_INPUTS+1], s2[MAX_INPUTS+1];
    size_t t = num_inputs + 1;
    size_t coff = POSEIDON_C_OFF[t-2];
    size_t moff = (2*t*t*t - 3*t*t + t - 6) / 6;

    mbedtls_mpi_init(&P);
    mbedtls_mpi_init(&FIVE_BIGNUM);
    mbedtls_mpi_init(&temp);
    mbedtls_mpi_init(&RR);

    for (size_t i = 0; i < t; i++) {
        mbedtls_mpi_init(&state[i]);
        mbedtls_mpi_init(&s2[i]);
    }

    // Initialize constants
    mbedtls_mpi_read_string(&P, 10, "21888242871839275222246405745257275088548364400416034343698204186575808495617");
    mbedtls_mpi_lset(&FIVE_BIGNUM, 5);

    // Initialize state
    mbedtls_mpi_lset(&state[0], 0);
    for (size_t i = 1; i < t; i++) {
        mbedtls_mpi_copy(&state[i], &inputs[i-1]);
    }

    for (size_t r = 0; r < 8 + N_ROUNDS_P[t-2]; r++) {
        for (size_t i = 0; i < t; i++) {
            mbedtls_mpi_read_binary(&temp, POSEIDON_C[coff + r*t + i], 32);
            mbedtls_mpi_add_mpi(&state[i], &state[i], &temp);
            mbedtls_mpi_mod_mpi(&state[i], &state[i], &P);

            if (i == 0 || r < 4 || r >= 4 + N_ROUNDS_P[t-2]) {
                mbedtls_mpi_exp_mod(&state[i], &state[i], &FIVE_BIGNUM, &P, &RR);
            }
        }

        for (size_t i = 0; i < t; i++) {
            mbedtls_mpi_lset(&s2[i], 0);
            for (size_t j = 0; j < t; j++) {
                mbedtls_mpi_read_binary(&temp, POSEIDON_M[moff + t*i + j], 32);
                esp_mpi_mul_mpi_mod(&temp, &temp, &state[j], &P);
                mbedtls_mpi_add_mpi(&s2[i], &s2[i], &temp);
                mbedtls_mpi_mod_mpi(&s2[i], &s2[i], &P);
            }
        }
        for (size_t i = 0; i < t; i++) {
            mbedtls_mpi_copy(&state[i], &s2[i]);
        }
    }

    mbedtls_mpi_copy(result, &state[0]);
    mbedtls_mpi_mod_mpi(result, result, &P);

    // Free allocated memory
    mbedtls_mpi_free(&P);
    mbedtls_mpi_free(&FIVE_BIGNUM);
    mbedtls_mpi_free(&temp);
    mbedtls_mpi_free(&RR);
    for (size_t i = 0; i < t; i++) {
        mbedtls_mpi_free(&state[i]);
        mbedtls_mpi_free(&s2[i]);
    }
}

void poseidon(mbedtls_mpi *result, const mbedtls_mpi *inputs, size_t num_inputs) {
    mbedtls_mpi P, FIVE_BIGNUM, temp, temp2, RR;
    mbedtls_mpi state[MAX_INPUTS+1], s2[MAX_INPUTS+1];
    size_t t = num_inputs + 1;
    size_t coff = POSEIDON_C_OFF[t-2];
    size_t moff = (2*t*t*t - 3*t*t + t - 6) / 6;

    mbedtls_mpi_init(&P);
    mbedtls_mpi_init(&FIVE_BIGNUM);
    mbedtls_mpi_init(&temp);
    mbedtls_mpi_init(&temp2);
    mbedtls_mpi_init(&RR);

    for (size_t i = 0; i < t; i++) {
        mbedtls_mpi_init(&state[i]);
        mbedtls_mpi_init(&s2[i]);
    }

    // Initialize constants
    mbedtls_mpi_read_string(&P, 10, "21888242871839275222246405745257275088548364400416034343698204186575808495617");
    mbedtls_mpi_lset(&FIVE_BIGNUM, 5);

    // Initialize state
    mbedtls_mpi_lset(&state[0], 0);
    for (size_t i = 1; i < t; i++) {
        mbedtls_mpi_copy(&state[i], &inputs[i-1]);
    }

    for (size_t r = 0; r < 8 + N_ROUNDS_P[t-2]; r++) {
        for (size_t i = 0; i < t; i++) {
            mbedtls_mpi_read_binary(&temp, POSEIDON_C[coff + r*t + i], 32);
            mbedtls_mpi_add_mpi(&state[i], &state[i], &temp);
            mbedtls_mpi_mod_mpi(&state[i], &state[i], &P);

            if (i == 0 || r < 4 || r >= 4 + N_ROUNDS_P[t-2]) {
                mbedtls_mpi_exp_mod(&state[i], &state[i], &FIVE_BIGNUM, &P, &RR);
            }
        }

        for (size_t i = 0; i < t; i++) {
            mbedtls_mpi_lset(&s2[i], 0);
            for (size_t j = 0; j < t; j++) {
                mbedtls_mpi_read_binary(&temp, POSEIDON_M[moff + t*i + j], 32);
                mbedtls_mpi_mul_mpi(&temp2, &temp, &state[j]);
                mbedtls_mpi_mod_mpi(&temp2, &temp2, &P);
                mbedtls_mpi_add_mpi(&s2[i], &s2[i], &temp2);
                mbedtls_mpi_mod_mpi(&s2[i], &s2[i], &P);
            }
        }
        for (size_t i = 0; i < t; i++) {
            mbedtls_mpi_copy(&state[i], &s2[i]);
        }
    }

    mbedtls_mpi_copy(result, &state[0]);
    mbedtls_mpi_mod_mpi(result, result, &P);

    // Free allocated memory
    mbedtls_mpi_free(&P);
    mbedtls_mpi_free(&FIVE_BIGNUM);
    mbedtls_mpi_free(&temp);
    mbedtls_mpi_free(&temp2);
    mbedtls_mpi_free(&RR);
    for (size_t i = 0; i < t; i++) {
        mbedtls_mpi_free(&state[i]);
        mbedtls_mpi_free(&s2[i]);
    }
}


void loop() {
  // put your main code here, to run repeatedly:
  delay(10); // this speeds up the simulation
}