#include <iostream>
#include <iomanip>
#include "Vkey_switching.h"
#include <verilated.h>
#include <verilated_vcd_c.h>

// ciphertext
int c[] = { 0x1, 0x2, 0x3, 0x4,   // 0-3
			     0xffffffff, 0xfffffffe, 0xfffffffc, 0xfffffff8}; // 4-7

// bit-repr_ciphertext
int c_star[] = { 0x0, 0x0, 0x0, 0x0,      // 0-3
                 0x1, 0x0, 0x0, 0x0,                // 4-7
                 0x1, 0x0, 0x0, 0x0,                // 8-11
                 0x0, 0x1, 0x1, 0x0,                // 12-15
                 0x0, 0x1, 0x0, 0x0,                // 16-19
                 0x0, 0x0, 0x0, 0x0,                // 20-23
                 0xffffffff, 0x0, 0x0, 0x0,         // 24-27
                 0xffffffff, 0x0, 0x0, 0x0,         // 28-31
                 0xffffffff, 0x0, 0x0, 0x0,         // 32-35
                 0xffffffff, 0x0, 0x0, 0x0};        // 36-39

// secret key
int S[2][8] = { {0x6, 0x5, 0x0, 0x3, 0x9, 0x3, 0x3, 0x6},   // 0-7
			    {0x9, 0x7, 0x6, 0x8, 0x2, 0x0, 0x6, 0x1}}; // 8-15

// l = 4
int S_star[2][32] = { {0x30, 0x18, 0xc, 0x6,   // 0-3
                0x28, 0x14, 0xa, 0x5,               // 4-7
                0x0, 0x0, 0x0, 0x0,                 // 7-11
                0x18, 0xc, 0x6, 0x3,                // 12-15
                0x48, 0x24, 0x12, 0x9,               // 16-19
                0x18, 0xc, 0x6, 0x3,                // 20-23
                0x18, 0xc, 0x6, 0x3,                // 24-27
                0x30, 0x18, 0xc, 0x6},              // 28-31
			   {0x48, 0x24, 0x12, 0x9,               // 0-3
                0x38, 0x1c, 0xe, 0x7,               // 4-7
                0x30, 0x18, 0xc, 0x6,               // 7-11
                0x40, 0x20, 0x10, 0x8,              // 12-15
                0x10, 0x8, 0x4, 0x2,                // 16-19
                0x0, 0x0, 0x0, 0x0,                 // 20-23
                0x30, 0x18, 0xc, 0x6,               // 24-27
                0x8, 0x4, 0x2, 0x1}};               // 28-31
			     

int main(int argc, const char ** argv, const char ** env) {
  Verilated::commandArgs(argc, argv);

  // Treat the argument on the command-line as the place to start
  int width_v, ell_v;
  int width_m, length_m, ell_m;
  if (argc > 1 && argv[1][0] != '+') {
    width_v = atoi(argv[1]);
    ell_v = atoi(argv[2]);
    width_m = atoi(argv[3]);
    length_m = atoi(argv[4]);
    ell_m = atoi(argv[5]);
  } else {
    width_v = 8; // Default
    ell_v = 5;  // Default
    width_m = 8; // Default
    length_m = 2; // Default
    ell_m = 4;  // Default
  }

  Vkey_switching * dut = new Vkey_switching;  // Instantiate a new module
  
  // Enable dumping a VCD file
  Verilated::traceEverOn(true);
  VerilatedVcdC * tfp = new VerilatedVcdC;
  dut->trace(tfp, 99);
  tfp->open("key_switching.vcd");

  dut->clk = 0;
  
  int time;
  std::cout << "width_v: " << width_v;
  std::cout << std::endl;
  std::cout << "ell_v: " << ell_v;
  std::cout << std::endl;
  std::cout << "width_m: " << width_m;
  std::cout << std::endl;
  std::cout << "length_m: " << length_m;
  std::cout << std::endl;
  std::cout << "ell_m: " << ell_m;
  std::cout << std::endl;

  // bit_repr_vector operation test
  int i = 0;
  int j = 0;
  int err = 0;
  int load_stamp = 100;
  int start_stamp = 160;
  // let run (width * ell) cycles for computations
  // extra 11 cycles also needed
  int read_stamp = start_stamp + width_v * ell_v * 20 + 17 * 20;
  for (time = 0 ; time < 10000 ; time += 10) {
    dut->clk = ((time % 20) >= 10) ? 1 : 0; // Simulate a 50 MHz clock
    if (time == 20) dut->reset = 1; // Handle "reset" on for four cycles
    if (time == 100) {
      dut->reset = 0;
      dut->write = 1;           // Put "write" on
      dut->chipselect = 1;
    }

    // load operation type, width, ell
    // load operation type
    if (time >= load_stamp && time <= load_stamp + 20 && dut->clk == 1) {
        dut->address = 0;
        // bit_repr_vector
        dut->writedata = 0;
        std::cout <<"operation type received: " << (int) dut->writedata;
        std::cout << std::endl;
    }

    // load width
    if (time >= load_stamp + 20 && time <= load_stamp + 40 && dut->clk == 1) {
        dut->address = 1;
        dut->writedata = width_v;
        std::cout <<"width received: " << (int) dut->writedata;
        std::cout << std::endl;
    }

    // load ell
    // all loaded
    if (time >= load_stamp + 40 && time <= load_stamp + 60 && dut->clk == 1) {
        dut->address = 3;
        dut->all_loaded = 1;
        dut->writedata = ell_v;
        std::cout << "ell received: " << (int) dut->writedata;
        std::cout << std::endl;
    }
    
    // taking inputs for (width) cycles
    if (time >= start_stamp && time <= start_stamp + width_v * 20 && dut->clk == 1) {
        dut->address = 4;
        dut->writedata = c[i];
        std::cout << i << "-th input received: " << (int) c[i];
        std::cout << std::endl;
        i += 1;
    }

    dut->eval();     // Run the simulation for a cycle
    tfp->dump(time); // Write the VCD file for this cycle

    // compare outputs for (width * ell) cycles
    if (time >= read_stamp && time <= read_stamp + width_v * ell_v * 20 && dut->clk == 1) {
        std::cout << ' ' <<  std::setfill('0') << std::setw(0) << (int) dut->DATA_OUT;
        if (dut->DATA_OUT == c_star[j])
            std::cout << " OK";
        else {
            std::cout << " INCORRECT expected " << std::setfill('0') << (int) c_star[j];
            err += 1;
        }
        std::cout << std::endl;
        j += 1;
    }

    if (dut->DONE) {
        dut->all_loaded = 0;
        break;
    }
  }

  int run1_time = time + 10;

  // bit_repr_matrix operation test
  i = 0;
  int r = 0;
  err = 0;
  load_stamp = run1_time + 100;
  start_stamp = run1_time + 180;
  // let run ((width * ell) * length) cycles for computations
  // extra 16 cycles are needed in addition for computations
  read_stamp = start_stamp + length_m * width_m * ell_m * 20 + 18 * 20;
  
  for (time = run1_time; time < 10000 ; time += 10) {
    dut->clk = ((time % 20) >= 10) ? 1 : 0; // Simulate a 50 MHz clock
    if (time == run1_time + 20) dut->reset = 1; // Handle "reset" on for four cycles
    if (time == run1_time + 100) {
      dut->reset = 0;
      dut->write = 1;           // Put "write" on
      dut->chipselect = 1;
    }

    // load operation type, width, length, ell
    // load operation type
    if (time >= load_stamp && time <= load_stamp + 20 && dut->clk == 1) {
        dut->address = 0;
        // bit_repr_matrix
        dut->writedata = 1;
        std::cout <<"operation type received: " << (int) dut->writedata;
        std::cout << std::endl;
    }

    // load width
    if (time >= load_stamp + 20 && time <= load_stamp + 40 && dut->clk == 1) {
        dut->address = 1;
        dut->writedata = width_m;
        std::cout <<"width received: " << (int) dut->writedata;
        std::cout << std::endl;
    }

    // load length
    if (time >= load_stamp + 40 && time <= load_stamp + 60 && dut->clk == 1) {
        dut->address = 2;
        dut->writedata = length_m;
        std::cout <<"length received: " << (int) dut->writedata;
        std::cout << std::endl;
    }

    // load ell
    if (time >= load_stamp + 60 && time <= load_stamp + 80 && dut->clk == 1) {
        dut->address = 3;
        dut->writedata = ell_m;
        std::cout << "ell received: " << (int) dut->writedata;
        std::cout << std::endl;
    }

    // take inputs for (length) rows
    // for each row, take inputs for (width) cycles
    if (time >= start_stamp && time <= start_stamp + width_m * length_m * 20 && dut->clk == 1) {
        dut->address = 4;
        dut->all_loaded = 1;
        dut->writedata = S[r][i];
        std::cout << "input received: " << (int) S[r][i];
        std::cout << std::endl;
        i += 1;
        if (i == width_m) {
            i = 0;
            r += 1;
            if (r == length_m) {
                r = 0;
            }
        }
    }

    dut->eval();     // Run the simulation for a cycle
    tfp->dump(time); // Write the VCD file for this cycle

    // compare outputs for (length) rows
    // for each row, compare outputs for (width * ell) cycles
    if (time >= read_stamp && time <= read_stamp + width_m * length_m * ell_m * 20 && dut->clk == 1) {
        std::cout << ' ' <<  std::setfill('0') << std::setw(0) << (int) dut->DATA_OUT;
        if (dut->DATA_OUT == S_star[r][i])
            std::cout << " OK";
        else {
            std::cout << " INCORRECT expected " << std::setfill('0') << (int) S_star[r][i];
            err += 1;
        }
        std::cout << std::endl;
        i += 1;
        if (i == width_m * ell_m) {
            i = 0;
            r += 1;
            if (r == length_m) {
                r = 0;
            }
        }
    }

    if (time >= read_stamp && dut->DONE) {
        dut->all_loaded = 0;
        break;
    }
  }

  tfp->close(); // Stop dumping the VCD file
  delete tfp;

  dut->final(); // Stop the simulation
  delete dut;

  return 0;
}