#include <iostream>
#include <cassert>
#include <vector>
#include <string.h> // memcpy()
using namespace std;
#include "CSRSparseMatrix.h"


CSRSparseMatrix::CSRSparseMatrix(void)
    : m_M(0), m_N(0), m_NNZ(0), m_A(0), m_IA(0), m_JA(0)
{
}


CSRSparseMatrix::~CSRSparseMatrix(void)
{
}


void CSRSparseMatrix::setValue(const int M, const int N, const int NNZ, 
                               int *A, int *IA, int *JA)
{
    release();
    m_M = M;
    m_N = N;
    m_NNZ = NNZ;
    m_A = A;
    m_IA = IA;
    m_JA = JA;
}


void CSRSparseMatrix::unitTest()
{
    CSRSparseMatrix A, At, B, Bt;
    A.setValue(0,0,0,0,0,0);
    At.setValue(0,0,0,0,0,0);
    B.setValue(0,0,0,0,0,0);
    Bt.setValue(0,0,0,0,0,0);

    int C0A[]={3,-1,5,1,2,2,-1},C0IA[]={0,3,5,5,6,7},C0JA[]={0,1,2,0,2,1,2};
    CSRSparseMatrix C0;
    C0.setValue(5,3,7,C0A,C0IA,C0JA);

    int C1A[]={3,2,6,3,7,5,-2},C1IA[]={0,3,6,7},C1JA[]={0,1,2,0,1,2,2};
    CSRSparseMatrix C1;
    C1.setValue(3,3,7,C1A,C1IA,C1JA);

    int C2A[]={-1,2,9,3,-13,-4,6,2,-1,1,8,10,4,-2},C2IA[]={0,2,4,8,11,14},C2JA[]={0,3,0,1,0,1,3,4,0,1,3,0,1,4};
    CSRSparseMatrix C2;
    C2.setValue(5,5,14,C2A,C2IA,C2JA);
    
    A.read();
    B.read();

    if ((A.m_M != B.m_M)||(A.m_N != B.m_N))
    {
        cout << "Matrix A and B are not of the same dimension!\n";
        return;
    }

    cout << "\n";
    cout << "Matrix A:\n";
    A.print();

    cout << "Matrix B:\n";
    B.print();

    Bt.transpose(B);
    assert(Bt.equal(C0));
    cout << "Transpose Matrix B:\n";
    Bt.print();

    At.transpose(A);

    A.multiply(Bt);
    assert(A.equal(C1));

    cout << "Matrix A * Matrix B':\n";
    A.print();

    cout << "Transpose Matrix A:\n";
    At.print();

    At.multiply(B);
    assert(At.equal(C2));

    cout << "Matrix A' * Matrix B:\n";
    At.print();

    A.release();
    At.release();
    B.release();
    Bt.release();
}


void CSRSparseMatrix::release()
{
    delete[] m_A;
    delete[] m_IA;
    delete[] m_JA;
    m_M = m_N = m_NNZ = 0;
    m_A = m_IA = m_JA = 0;
}


void CSRSparseMatrix::read()
{
    int i, row, prevRow;
    cin >> m_M >> m_N >> m_NNZ;
    m_A = new int[m_NNZ];
    m_IA = new int[m_M+1];
    m_JA = new int[m_NNZ];
    for (prevRow=-1,i=0; i<m_NNZ; i++)
    {
        cin >> row >> m_JA[i] >> m_A[i];
        if (row!=prevRow) 
        {
            m_IA[row] = i;
            while (++prevRow<row) 
                m_IA[prevRow] = m_IA[row];
        }
    }
    m_IA[m_M] = m_NNZ;
}


bool CSRSparseMatrix::equal(const CSRSparseMatrix &B) const
{
    int i;
    if (m_M!=B.m_M || m_N!=B.m_N || m_NNZ!=B.m_NNZ)
        return false;
    for (i=0; i<m_NNZ; i++)
        if (m_A[i]!=B.m_A[i] || m_JA[i]!=B.m_JA[i])
            return false;
    for (i=0; i<=m_M; i++)
        if (m_IA[i]!=B.m_IA[i])
            return false;
    return true;
}


void CSRSparseMatrix::print() const
{
    int i, row;
    cout << m_M << " * " << m_N << " matrix with " << m_NNZ << " non-zero item:\n";
    for (row=i=0; i<m_NNZ; i++)
    {
        while (i>=m_IA[row+1]) row++;
        cout << "<" << row << "," << m_JA[i] << ":" << m_A[i] << ">\n";
    }
    cout << "====================\n";
}


void CSRSparseMatrix::printCSR() const
{
    int i, iRow=0;
    cout << m_M << " * " << m_N << " matrix with " << m_NNZ << " non-zero item:\n";
    cout << "A = ";
    for (i=0; i<m_NNZ; i++)
        cout << m_A[i] << " ";
    cout << "\nIA = ";
    for (i=0; i<=m_M; i++)
        cout << m_IA[i] << " ";
    cout << "\nJA = ";
    for (i=0; i<m_NNZ; i++)
        cout << m_JA[i] << " ";
    cout << "\n====================\n";
}


void CSRSparseMatrix::transpose(const CSRSparseMatrix &A)
{
    int i, iB=0, row, *index;

    release();
    m_M = A.m_N, m_N = A.m_M, m_NNZ = A.m_NNZ;
    m_A = new int[m_NNZ];
    m_IA = new int[m_M+1];
    index = new int[m_M+1];
    m_JA = new int[m_NNZ];

    for (i=0; i<=A.m_N; i++)
        m_IA[i] = 0;
    for (i=0; i<A.m_NNZ; i++)
        m_IA[A.m_JA[i]+1]++; // counting element's in B's row
    for (i=1; i<=m_M; i++)
        index[i] = m_IA[i] += m_IA[i-1]; // accumulating
    index[0] = 0; 

    for (row=0,i=0; i<A.m_NNZ; i++)
    {
        while (i>=A.m_IA[row+1]) row++;
        iB = index[A.m_JA[i]]++;
        m_A[iB] = A.m_A[i];
        m_JA[iB] = row;
    }
    delete[] index;
}


void CSRSparseMatrix::multiply(const CSRSparseMatrix &B)
{
    int i, iA, iA1, iB1, iArow, iA1row, iB1row, sum, *tmp;
    CSRSparseMatrix B1;
    B1.transpose(B);

    vector<int> CA, CIA, CJA;
    iA = iArow = 0;
    CIA.push_back(0);
    while (iA < m_NNZ)
    {
        sum = 0;
        iA1 = iA; iA1row = iArow; 
        iB1 = iB1row = 0;
        while (iB1 < B1.m_NNZ)
        {
            if (m_JA[iA1] == B1.m_JA[iB1])
            {
                sum += m_A[iA1] * B1.m_A[iB1];
                iA1++; iB1++;
            }
            else if (m_JA[iA1] > B1.m_JA[iB1])
                iB1++;
            else // if (m_JA[iA1] < B1.JA[iB1])
                iA1++;
            if (iA1 >= m_IA[iA1row+1])
                while (iB1 < B1.m_IA[iB1row+1]) iB1++;
            if (iB1 >= B1.m_IA[iB1row+1])
            {
                if (iB1 < B1.m_NNZ)
                    iA1 = iA;
                else
                    iA1 = m_IA[iA1row+1];
                if (sum!=0)
                {
                    while (iA1row >= CIA.size()) CIA.push_back(CA.size());
                    CA.push_back(sum);
                    CJA.push_back(iB1row);
                    sum = 0;
                }
                while (iA1>=m_IA[iA1row+1]) iA1row++;
                while (iB1>=B1.m_IA[iB1row+1]) iB1row++;
            }
        }
        iA = iA1, iArow = iA1row;
    }
    CIA.push_back(CA.size());

    m_N = B.m_N;
    m_NNZ = CA.size();

    delete[] m_A;
    m_A = new int[m_NNZ];
    delete[] m_JA;
    m_JA = new int[m_NNZ];
    for (i=0; i<m_NNZ; i++)
    {
        m_A[i] = CA[i];
        m_JA[i] = CJA[i];
    }
    for (i=0; i<m_M+1; i++)
        m_IA[i] = CIA[i];

    B1.release();
}