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

// encrypted data
int c_1[16] = {0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 0x8,   // 0-7
			    0xffffffff, 0xfffffffe, 0xfffffffd, 0xfffffffc, 0xfffffffb, 0xfffffffa, 0xfffffff9, 0xfffffff8}; // 8-15

int c_2[16] = {0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 0x8, 0x9,   // 0-7
			    0x1, 0x2, 0x3, 0x4, 0xffffffff, 0xfffffffe, 0xfffffffd, 0xfffffffc}; // 8-15

// expected result
int c_out[16] = {0x3, 0x5, 0x7, 0x9, 0xb, 0xd, 0xf, 0x11,   // 0-7
			    0x0, 0x0, 0x0, 0x0, 0xfffffffa, 0xfffffff8, 0xfffffff6, 0xfffffff4}; // 8-15

// encrypted data
int c[8] = {0x1, 0x2, 0x3, 0x4, 0xffffffff, 0xfffffffe, 0xfffffffd, 0xfffffffc};   // 0-7

// encrypted linear operator
int M[4][8] = { {0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 0x8},   // 0-7
			    {0xfffffffe, 0xfffffffd, 0xfffffffc, 0xfffffffb, 0xfffffffa, 0xfffffff9, 0xfffffff8, 0xfffffff7},   // 8-15
                {0x3, 0x6, 0x9, 0xc, 0xfffffffd, 0xfffffffa, 0xfffffff7, 0xfffffff4},   // 16-23
                {0xfffffffe, 0xfffffffd, 0xfffffffc, 0xfffffffb, 0x6, 0x7, 0x8, 0x9}   // 24-31
               };

// result
int a[4] = {0xffffffd8, 0x28, 0xb4, 0xffffff88};    // 0-4

// encrypted data
int c_1_out[2] = {0x1, 0xfffffffe};   // 0-1

int c_2_out[2] = {0x5, 0xfffffffd};   // 0-1

// encrypted linear operator
int M_wi[4][4] = { {0x1, 0x2, 0x3, 0x4},   // 0-7
			    {0xfffffffe, 0xfffffffd, 0xfffffffc, 0xfffffffb},   // 8-15
                {0x3, 0x6, 0x9, 0xc},   // 16-23
                {0xfffffffc, 0xfffffffb, 0x9, 0xc}   // 24-31
               };

// reminder: vectorized result
int x[4] = {0x5, 0xfffffffd, 0xfffffff6, 0x6};    // 0-4

int w = 1;

// expected result
int a_out[4] = {0xfffffff9, 0x9, 0xffffffeb, 0xffffffe9};    // 0-4


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;
  int width_l, length_l;
  int width_w, length_w;
  if (argc > 1 && argv[1][0] != '+') {
    width_v = atoi(argv[1]);
    width_l = atoi(argv[2]);
    length_l = atoi(argv[3]);
    width_w = atoi(argv[4]);
  } else {
    width_v = 16; // Default
    width_l = 8;
    length_l = 4;
    width_w = 2;
    length_w = 4;
  }

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

  dut->clk = 0;
  
  int time;
  std::cout << "width_v: " << width_v;
  std::cout << std::endl;
  std::cout << "width_l: " << width_l;
  std::cout << std::endl;
  std::cout << "length_l: " << length_l;
  std::cout << std::endl;
  std::cout << "width_w: " << width_w;
  std::cout << std::endl;

  // vector_addition operation test
  int i = 0;
  int j = 0;
  int err = 0;
  int load_stamp = 100;
  int start_stamp = load_stamp + 100 + 2 * 20 * width_v;
  int read_stamp = load_stamp + 120 + 2 * 20 * width_v;

  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
    // load operation type
    if (time >= load_stamp && time <= load_stamp + 20 && dut->clk == 1) {
        dut->address = 0;
        // vector_addition
        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 = 2;
        dut->writedata = width_v;
        std::cout <<"width received: " << (int) dut->writedata;
        std::cout << std::endl;
    }

    // load data_type c_1
    if (time >= load_stamp + 40 && time <= load_stamp + 60 && dut->clk == 1) {
        dut->address = 1;
        dut->writedata = 0; // load c_1
        std::cout << "start writing to input type: " << (int) dut->writedata;
        std::cout << std::endl;
    }

    // load c_1
    if (time >= load_stamp + 60 && time <= load_stamp + 60 + 20 * width_v && dut->clk == 1) {
        dut->address = 4;
        dut->writedata = c_1[i]; // load c_1
        std::cout << "input: " << (int) dut->writedata;
        std::cout << std::endl;
        i += 1;
        if (i == width_v) i = 0;
    }

    // load data_type c_2
    if (time >= load_stamp + 60 + 20 * width_v && time <= load_stamp + 80 + 20 * width_v && dut->clk == 1) {
        dut->address = 1;
        dut->writedata = 1; // load c_2
        std::cout << "start writing to input type: " << (int) dut->writedata;
        std::cout << std::endl;
    }

    // load c_2
    if (time >= load_stamp + 80 + 20 * width_v && time <= load_stamp + 80 + 2 * 20 * width_v && dut->clk == 1) {
        dut->address = 4;
        dut->writedata = c_2[i]; // load c_2
        std::cout << "input: " << (int) dut->writedata;
        std::cout << std::endl;
        i += 1;
    }

    // everything loaded and ready
    if (time >= load_stamp + 80 + 2 * 20 * width_v && time <= load_stamp + 100 + 2 * 20 * width_v && dut->clk == 1) {
        dut->all_loaded = 1;
        i = 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 && dut->clk == 1) {
        std::cout << "vector output received: " << (int) dut->DATA_OUT_0;
        std::cout << std::endl;
        if (dut->DATA_OUT_0 == c_out[0])
            std::cout << " OK";
        else {
            std::cout << " INCORRECT expected " << std::setfill('0') << (int) c_out[0];
            err += 1;
        }
        std::cout << std::endl;

        std::cout << "vector output received: " << (int) dut->DATA_OUT_1;
        std::cout << std::endl;
        if (dut->DATA_OUT_1 == c_out[1])
            std::cout << " OK";
        else {
            std::cout << " INCORRECT expected " << std::setfill('0') << (int) c_out[1];
            err += 1;
        }
        std::cout << std::endl;

        std::cout << "vector output received: " << (int) dut->DATA_OUT_2;
        std::cout << std::endl;
        if (dut->DATA_OUT_2 == c_out[2])
            std::cout << " OK";
        else {
            std::cout << " INCORRECT expected " << std::setfill('0') << (int) c_out[2];
            err += 1;
        }
        std::cout << std::endl;

        std::cout << "vector output received: " << (int) dut->DATA_OUT_3;
        std::cout << std::endl;
        if (dut->DATA_OUT_3 == c_out[3])
            std::cout << " OK";
        else {
            std::cout << " INCORRECT expected " << std::setfill('0') << (int) c_out[3];
            err += 1;
        }
        std::cout << std::endl;

        std::cout << "vector output received: " << (int) dut->DATA_OUT_4;
        std::cout << std::endl;
        if (dut->DATA_OUT_4 == c_out[4])
            std::cout << " OK";
        else {
            std::cout << " INCORRECT expected " << std::setfill('0') << (int) c_out[4];
            err += 1;
        }
        std::cout << std::endl;

        std::cout << "vector output received: " << (int) dut->DATA_OUT_5;
        std::cout << std::endl;
        if (dut->DATA_OUT_5 == c_out[5])
            std::cout << " OK";
        else {
            std::cout << " INCORRECT expected " << std::setfill('0') << (int) c_out[5];
            err += 1;
        }
        std::cout << std::endl;

        std::cout << "vector output received: " << (int) dut->DATA_OUT_6;
        std::cout << std::endl;
        if (dut->DATA_OUT_6 == c_out[6])
            std::cout << " OK";
        else {
            std::cout << " INCORRECT expected " << std::setfill('0') << (int) c_out[6];
            err += 1;
        }
        std::cout << std::endl;

        std::cout << "vector output received: " << (int) dut->DATA_OUT_7;
        std::cout << std::endl;
        if (dut->DATA_OUT_7 == c_out[7])
            std::cout << " OK";
        else {
            std::cout << " INCORRECT expected " << std::setfill('0') << (int) c_out[7];
            err += 1;
        }
        std::cout << std::endl;

        std::cout << "vector output received: " << (int) dut->DATA_OUT_8;
        std::cout << std::endl;
        if (dut->DATA_OUT_8 == c_out[8])
            std::cout << " OK";
        else {
            std::cout << " INCORRECT expected " << std::setfill('0') << (int) c_out[8];
            err += 1;
        }
        std::cout << std::endl;

        std::cout << "vector output received: " << (int) dut->DATA_OUT_9;
        std::cout << std::endl;
        if (dut->DATA_OUT_9 == c_out[9])
            std::cout << " OK";
        else {
            std::cout << " INCORRECT expected " << std::setfill('0') << (int) c_out[9];
            err += 1;
        }
        std::cout << std::endl;

        std::cout << "vector output received: " << (int) dut->DATA_OUT_10;
        std::cout << std::endl;
        if (dut->DATA_OUT_10 == c_out[10])
            std::cout << " OK";
        else {
            std::cout << " INCORRECT expected " << std::setfill('0') << (int) c_out[10];
            err += 1;
        }
        std::cout << std::endl;

        std::cout << "vector output received: " << (int) dut->DATA_OUT_11;
        std::cout << std::endl;
        if (dut->DATA_OUT_11 == c_out[11])
            std::cout << " OK";
        else {
            std::cout << " INCORRECT expected " << std::setfill('0') << (int) c_out[11];
            err += 1;
        }
        std::cout << std::endl;

        std::cout << "vector output received: " << (int) dut->DATA_OUT_12;
        std::cout << std::endl;
        if (dut->DATA_OUT_12 == c_out[12])
            std::cout << " OK";
        else {
            std::cout << " INCORRECT expected " << std::setfill('0') << (int) c_out[12];
            err += 1;
        }
        std::cout << std::endl;

        std::cout << "vector output received: " << (int) dut->DATA_OUT_13;
        std::cout << std::endl;
        if (dut->DATA_OUT_13 == c_out[13])
            std::cout << " OK";
        else {
            std::cout << " INCORRECT expected " << std::setfill('0') << (int) c_out[13];
            err += 1;
        }
        std::cout << std::endl;

        std::cout << "vector output received: " << (int) dut->DATA_OUT_14;
        std::cout << std::endl;
        if (dut->DATA_OUT_14 == c_out[14])
            std::cout << " OK";
        else {
            std::cout << " INCORRECT expected " << std::setfill('0') << (int) c_out[14];
            err += 1;
        }
        std::cout << std::endl;

        std::cout << "vector output received: " << (int) dut->DATA_OUT_15;
        std::cout << std::endl;
        if (dut->DATA_OUT_15 == c_out[15])
            std::cout << " OK";
        else {
            std::cout << " INCORRECT expected " << std::setfill('0') << (int) c_out[15];
            err += 1;
        }
        std::cout << std::endl;
    }
    
    if (time >= read_stamp + 20 && dut->clk == 1) {
        dut->all_loaded = 0;
        break;
    }
  }

  int run1_time = time + 10;

  // linear transform operation test
  int r = 0;
  i = 0;
  err = 0;
  int comp_flag = 0;
  load_stamp = run1_time + 100;
  // after c been loaded
  int M_stamp = load_stamp + 80 + 20 * 16;
  // after one round of M been loaded
  int eval_stamp = load_stamp + 80 + 20 * 16 + 20 * 16;

  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
    if (time >= load_stamp && time <= load_stamp + 20 && dut->clk == 1) {
        dut->address = 0;
        // linear_transform, activate module
        dut->writedata = 1;
        std::cout <<"operation type received: " << (int) dut->writedata;
        std::cout << std::endl;
    }


    // load data_type c
    if (time >= load_stamp + 40 && time <= load_stamp + 60 && dut->clk == 1) {
        dut->address = 1;
        dut->writedata = 3; // load c
        std::cout << "start writing to input type: " << (int) dut->writedata;
        std::cout << std::endl;
    }

    // load c
    // 16 entires supported
    if (time >= load_stamp + 60 && time <= load_stamp + 60 + 20 * 16 && dut->clk == 1) {
        dut->address = 4;
        if (i >= width_l) {
            dut->writedata = 0;
            std::cout << "input: " << (int) 0;
            std::cout << std::endl;
        } else {
            dut->writedata = c[i]; // load c
            std::cout << "input: " << (int) dut->writedata;
            std::cout << std::endl;
        }

        i += 1;
        if (i == 16) i = 0;
    }

    // load data_type M
    if (time >= load_stamp + 60 + 20 * 16 && time <= load_stamp + 80 + 20 * 16 && dut->clk == 1) {
        dut->address = 1;
        dut->writedata = 2; // load M
        std::cout << "start writing to input type: " << (int) dut->writedata;
        std::cout << std::endl;
    }

    // load M
    if (time >= M_stamp && r < 16 && dut->clk == 1) {
        dut->all_loaded = 0;

        dut->address = 4;
        if (i >= width_l || r >= length_l) {
            dut->writedata = 0;
            std::cout << "input: " << (int) 0;
            std::cout << std::endl;
        } else {
            dut->writedata = M[r][i]; // load M
            std::cout << "input: " << (int) dut->writedata;
            std::cout << std::endl;
        }
        i += 1;
        if (i == 16) {
            // let it compute
            i = 0;
            r += 1;
            dut->all_loaded = 1;
            comp_flag = 1;
            M_stamp += 16 * 20 + 5 * 20;
        }
    }

    // start computation once 'all_loaded' activated

    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 >= eval_stamp && dut->DONE && comp_flag == 1 && dut->clk == 1) {
        // whenever it yields an output
        std::cout << "vector output received: " << (int) dut->DATA_OUT_0;
        std::cout << std::endl;
        if (j >= length_l) {
            if (dut->DATA_OUT_0 == 0)
            std::cout << " OK";
            else {
                std::cout << " INCORRECT expected " << std::setfill('0') << (int) 0;
                err += 1;
            }
        } else {
            if (dut->DATA_OUT_0 == a[j])
            std::cout << " OK";
            else {
                std::cout << " INCORRECT expected " << std::setfill('0') << (int) a[j];
                err += 1;
            }
        }
        std::cout << std::endl;
        comp_flag = 0;
        j += 1;
    }

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

  }


  // test weighted_inner_product module functionality

  r = 0;
  i = 0;
  j = 0;
  err = 0;
  comp_flag = 0;

  int run2_time = time + 10;

  // linear transform operation test
  r = 0;
  i = 0;
  err = 0;
  load_stamp = run2_time + 100;
  // after c been loaded
  start_stamp = load_stamp + 140 + 2* 20 * 4;
  M_stamp = start_stamp + (2 * 4 * 4 + 9) * 20;
  eval_stamp = M_stamp + 16 * 20;

  for (time = run2_time ; time < 100000 ; time += 10) {
    dut->clk = ((time % 20) >= 10) ? 1 : 0; // Simulate a 50 MHz clock
    if (time == run2_time + 20) dut->reset = 1; // Handle "reset" on for four cycles
    if (time == run2_time + 100) {
      dut->reset = 0;
      dut->write = 1;           // Put "write" on
      dut->chipselect = 1;
    }

    // load operation type
    if (time >= load_stamp && time <= load_stamp + 20 && dut->clk == 1) {
        dut->address = 0;
        // weighted_inner_product, activate module
        dut->writedata = 2;
        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 = 2;
        // load full length
        dut->writedata = width_w;
        std::cout <<"width received: " << (int) dut->writedata;
        std::cout << std::endl;
    }

    // load data_type c_1_out
    if (time >= load_stamp + 40 && time <= load_stamp + 60 && dut->clk == 1) {
        dut->address = 1;
        dut->writedata = 4; // load c_1_out
        std::cout << "start writing to input type: " << (int) dut->writedata;
        std::cout << std::endl;
    }

    // load c_1_out
    // 4 entires supported
    if (time >= load_stamp + 60 && time <= load_stamp + 60 + 20 * 4 && dut->clk == 1) {
        dut->address = 4;
        if (i >= width_w) {
            dut->writedata = 0;
            std::cout << "input: " << (int) 0;
            std::cout << std::endl;
        } else {
            dut->writedata = c_1_out[i]; // load c
            std::cout << "input: " << (int) dut->writedata;
            std::cout << std::endl;
        }

        i += 1;
        if (i == 4) i = 0;
    }

    // load data_type c_2_out
    if (time >= load_stamp + 60 + 20 * 4 && time <= load_stamp + 80 + 20 * 4 && dut->clk == 1) {
        dut->address = 1;
        dut->writedata = 5; // load c_2_out
        std::cout << "start writing to input type: " << (int) dut->writedata;
        std::cout << std::endl;
    }

    // load c_2_out
    // 4 entires supported
    if (time >= load_stamp + 80 + 20 * 4 && time <= load_stamp + 80 + 2* 20 * 4 && dut->clk == 1) {
        dut->address = 4;
        if (i >= width_w) {
            dut->writedata = 0;
            std::cout << "input: " << (int) 0;
            std::cout << std::endl;
        } else {
            dut->writedata = c_2_out[i]; // load c
            std::cout << "input: " << (int) dut->writedata;
            std::cout << std::endl;
        }

        i += 1;
        if (i == 4) i = 0;
    }

    // load data_type w
    if (time >= load_stamp + 80 + 2* 20 * 4 && time <= load_stamp + 100 + 2* 20 * 4 && dut->clk == 1) {
        dut->address = 1;
        dut->writedata = 6; // load w
        std::cout << "start writing to input type: " << (int) dut->writedata;
        std::cout << std::endl;
    }

    // load w, write 1
    if (time >= load_stamp + 100 + 2* 20 * 4 && time <= load_stamp + 120 + 2* 20 * 4 && dut->clk == 1) {
        dut->address = 4;
        dut->writedata = 1; // load w
        std::cout << "input: " << (int) dut->writedata;
        std::cout << std::endl;
    }

    // load data_type M
    if (time >= load_stamp + 120 + 2 * 20 * 4 && time <= load_stamp + 140 + 2* 20 * 4 && dut->clk == 1) {
        dut->address = 1;
        dut->writedata = 2; // load M
        std::cout << "start writing to input type: " << (int) dut->writedata;
        std::cout << std::endl;
    }


    // after which the rest of inputs will be M
    // load M_wi
    if (time >= M_stamp && r < 16 && dut->clk == 1) {
        dut->all_loaded = 0;

        dut->address = 4;
        if (i >= length_w || r >= length_w) {
            dut->writedata = 0;
            std::cout << "input: " << (int) 0;
            std::cout << std::endl;
        } else {
            dut->writedata = M_wi[r][i]; // load M
            std::cout << "input: " << (int) dut->writedata;
            std::cout << std::endl;
        }
        i += 1;
        if (i == 16) {
            // let it compute
            i = 0;
            r += 1;
            dut->all_loaded = 1;
            comp_flag = 1;
            M_stamp += 16 * 20 + (2 * 4 * 4 + 9) * 20;
        }
    }


    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 >= eval_stamp && dut->DONE && comp_flag == 1 && dut->clk == 1) {
        // whenever it yields an output
        std::cout << "vector output received: " << (int) dut->DATA_OUT_0;
        std::cout << std::endl;
        if (j >= length_w) {
            if (dut->DATA_OUT_0 == 0)
            std::cout << " OK";
            else {
                std::cout << " INCORRECT expected " << std::setfill('0') << (int) 0;
                err += 1;
            }
        } else {
            if (dut->DATA_OUT_0 == a_out[j])
            std::cout << " OK";
            else {
                std::cout << " INCORRECT expected " << std::setfill('0') << (int) a_out[j];
                err += 1;
            }
        }
        std::cout << std::endl;
        comp_flag = 0;
        j += 1;
    }

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

  }




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

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

  return 0;
}