#include <iostream>
#include <vector>
#include <cmath>
#include <cassert>
#include <chrono>
#include <random> 
using namespace std;

double sigmoid(double x) {
    return 1.0 / (1.0 + exp(-x));
}

class DenseMatrix {
private:
    int n, m;
    vector<double> data; 

public:
    // Constructor
    DenseMatrix(int n_, int m_) : n(n_), m(m_), data(n_ * m_) { //initalizes n,m, allocates data with n*m elements = 0
        int nm;
        if (n > m)
            nm = n;
        else
            nm = m;

        for (int i = 0; i < n; ++i) { //for each row index compute xi
            double xi = 10.0 * i / (nm - 1) - 5.0; 
            for (int j = 0; j < m; ++j) { //for each column index compute xj
                double xj = 10.0 * j / (nm - 1) - 5.0;
                data[i * m + j] = sigmoid(xi) * sigmoid(xj); //compute every element
            }
        }
    }

    vector<double> Mult(const vector<double> &u) const {
        assert((int)u.size() == m); //checks if the dimensions match 
        vector<double> result(n, 0.0); //prepares the empty vector, size n
        for (int i = 0; i < n; ++i) {
            for (int j = 0; j < m; ++j) {
                result[i] += data[i * m + j] * u[j]; //row wise acces
            }
        }
        return result;
    }

    vector<double> MultT(const vector<double> &v) const { //same as above, but the other way 
        assert((int)v.size() == n); 
        vector<double> result(m, 0.0);
        for (int j = 0; j < m; ++j) {
            for (int i = 0; i < n; ++i) {
                result[j] += data[i * m + j] * v[i];
            }
        }
        return result;
    }
};

class TensorMatrix { //rank one matrix M=u*v^T
private:
    int n;
    vector<double> u, v;

public:
    TensorMatrix(int n_) : n(n_), u(n_), v(n_) {
        for (int k = 0; k < n; ++k) {
            double xk = 10.0 * k / (n - 1) - 5.0;
            u[k] = sigmoid(xk);
            v[k] = sigmoid(xk); //ensures simmetry 
        }
    }

    vector<double> Mult(const vector<double> &x) const {
        assert((int)x.size() == n); //does u*(v^T*x) so the parenthesis first 
        vector<double> result(n, 0.0);
        double dot_vx = 0.0;
        for (int k = 0; k < n; ++k)
            dot_vx += v[k] * x[k];
        for (int i = 0; i < n; ++i)
            result[i] = u[i] * dot_vx;
        return result;
    }

    vector<double> MultT(const vector<double> &y) const { //same as above 
        assert((int)y.size() == n);
        vector<double> result(n, 0.0);
        double dot_uy = 0.0;
        for (int k = 0; k < n; ++k)
            dot_uy += u[k] * y[k];
        for (int j = 0; j < n; ++j)
            result[j] = v[j] * dot_uy;
        return result;
    }
};


int main() {
    DenseMatrix M(5, 3); //example with a small dense matrix 
    vector<double> u = {1, 2, 3};
    vector<double> f1 = M.Mult(u);
    vector<double> v = {-1, 2, -3, 4, -5};
    vector<double> f2 = M.MultT(v);

    cout << "Result of M.Mult(u): ";
    for (double x : f1) cout << x << " ";
    cout << endl;
    cout << "Result of M.MultT(v): ";
    for (double x : f2) cout << x << " ";
    cout << endl;

    int n = 1000; //example with a big matrix 
    
    vector<double> w(n); //random non trivial test vector w 
    std::mt19937 gen(42);
    std::uniform_real_distribution<double> dist(-1.0, 1.0); //values uniformly distributed in [-1,1]
    for (int i = 0; i < n; ++i)
        w[i] = dist(gen);

    int NLOOPS = 50; //how many repetitions 
    
    DenseMatrix M2(n, n); //Matrix 
    vector<double> res1, res2;

    auto t0 = chrono::system_clock::now();
    for (int k = 0; k < NLOOPS; ++k)
        res1 = M2.Mult(w);
    auto t1 = chrono::system_clock::now();
    chrono::duration<double> timeMult = (t1 - t0) / NLOOPS;

    t0 = chrono::system_clock::now();
    for (int k = 0; k < NLOOPS; ++k)
        res2 = M2.MultT(w);
    t1 = chrono::system_clock::now();
    chrono::duration<double> timeMultT = (t1 - t0) / NLOOPS;

    cout << "\nAverage time per Mult:  " << timeMult.count() << " sec\n";
    cout << "Average time per MultT: " << timeMultT.count() << " sec\n";

    double diff = 0.0;
    for (int i = 0; i < n; ++i)
        diff += fabs(res1[i] - res2[i]);
    diff /= n;
    cout << "Average difference between Mult and MultT results: " << diff << endl;

    TensorMatrix T(n); //Tensor 
    vector<double> res1T, res2T;

    auto t0T = chrono::system_clock::now();
    for (int k = 0; k < NLOOPS; ++k)
        res1T = T.Mult(w);
    auto t1T = chrono::system_clock::now();
    chrono::duration<double> timeMult2 = (t1T - t0T) / NLOOPS;

    t0T = chrono::system_clock::now();
    for (int k = 0; k < NLOOPS; ++k)
        res2T = T.MultT(w);
    t1T = chrono::system_clock::now();
    chrono::duration<double> timeMultT2 = (t1T - t0T) / NLOOPS;

    cout << "\nTensorMatrix Mult time:  " << timeMult2.count() << " sec\n";
    cout << "TensorMatrix MultT time: " << timeMultT2.count() << " sec\n";

    double diffT = 0.0;
    for (int i = 0; i < n; ++i)
        diffT += fabs(res1T[i] - res2T[i]);
    diffT /= n;
    cout << "Average difference between Tensor Mult and MultT results: " << diffT << endl;

    return 0;
}