/*
 * Copyright (C) 2008-2012 Delft University of Technology - All rights reserved.
 *
 * Author: Reijer Idema
 */

/*
 * PowerFlowProblem class source.
 */
#include "LogEvents.h"				// HASA LogEvents (not owned)
#include "PowerSystemExt.h"		// HASA PowerSystemExt (not owned)
#include "PowerVector.h"		// HASA PowerVector (not owned)
#include "AdmittanceMatrix.h"	// HASA AdmittanceMatrix (not owned)
#include "PowerFlowProblem.h"


POWERFLOWER_NAMESPACE_BEGIN


/////////////////////////////// PUBLIC ///////////////////////////////////////


//============================= LIFECYCLE ====================================


#undef __FUNCT__
#define __FUNCT__ "PowerFlowProblem::PowerFlowProblem"
PowerFlowProblem::PowerFlowProblem(void) :
	mErr(0), mpLogEvents(NULL), mBusCount(0), mEqCount(0), mMaxNnz(0),
	mEqMap(), mVoltageVector(), mpPowerVector(NULL), mpAdmittanceMatrix(NULL),
	mpOwnedX(PETSC_NULL), mpOwnedF(PETSC_NULL), mpOwnedJ(PETSC_NULL),
	mpX(PETSC_NULL), mpF(PETSC_NULL), mpJ(PETSC_NULL),
	mIsCalculated(false), mIsPrecalculated(false), mIsStructureSame(false)
{
	if (PetscUtils::IsInitialized())
	{
		mErr = PetscInfo(PETSC_NULL, "Creating PowerFlowProblem object.\n"); CHKERRV(mErr);
	}
} // PowerFlowProblem::PowerFlowProblem


#undef __FUNCT__
#define __FUNCT__ "PowerFlowProblem::PowerFlowProblem"
PowerFlowProblem::PowerFlowProblem(const LogEvents* pLogEvents) :
	mErr(0), mpLogEvents(pLogEvents), mBusCount(0), mEqCount(0), mMaxNnz(0),
	mEqMap(), mVoltageVector(), mpPowerVector(NULL), mpAdmittanceMatrix(NULL),
	mpOwnedX(PETSC_NULL), mpOwnedF(PETSC_NULL), mpOwnedJ(PETSC_NULL),
	mpX(PETSC_NULL), mpF(PETSC_NULL), mpJ(PETSC_NULL),
	mIsCalculated(false), mIsPrecalculated(false), mIsStructureSame(false)
{
	if (PetscUtils::IsInitialized())
	{
		mErr = PetscInfo(PETSC_NULL, "Creating PowerFlowProblem object.\n"); CHKERRV(mErr);
	}
} // PowerFlowProblem::PowerFlowProblem


#undef __FUNCT__
#define __FUNCT__ "PowerFlowProblem::~PowerFlowProblem"
PowerFlowProblem::~PowerFlowProblem(void)
{
	Clear();
	if (PetscUtils::IsInitialized())
	{
		mErr = PetscInfo(PETSC_NULL, "Destroying PowerFlowProblem object.\n"); CHKERRV(mErr);
	}
} // PowerFlowProblem::~PowerFlowProblem


//============================= OPERATIONS ===================================


#undef __FUNCT__
#define __FUNCT__ "PowerFlowProblem::Clear"
void PowerFlowProblem::Clear(void)
{
	if (PetscUtils::IsInitialized())
	{
		mErr = PetscInfo(PETSC_NULL, "Clearing PowerFlowProblem object data.\n"); CHKERRV(mErr);
	}
	// Clear owned non-PETSc data
	mErr = 0;
	mBusCount = 0;
	mEqCount = 0;
	mEqMap.Clear();
	mVoltageVector.Clear();
	// Clear pointers not owned
	mpPowerVector = NULL;
	mpAdmittanceMatrix = NULL;
	// Clear Jacobian system data
	DestroyJacobianSystem();
	return;
} // PowerFlowProblem::Clear


#undef __FUNCT__
#define __FUNCT__ "PowerFlowProblem::SetPowerSystem"
bool PowerFlowProblem::SetPowerSystem(const PowerSystemExt& rPowerSystem)
{
	// :TODO: Check if set power system data is valid?

	// Get bus count
	mBusCount = rPowerSystem.GetBusCount();
	if (mBusCount <= 0) return false;

	// Build equation map
	mEqMap.Clear();		// Clear map in case this is a PowerFlowProblem object being re-used.
	mEqMap.Initialize(mBusCount);
	const PowerSystemExt::BusType* pBus;
	for (PetscInt i=0; i<mBusCount; ++i)
	{
		pBus = rPowerSystem.GetBus(i);
		if (pBus == NULL) return false;
		if (pBus->IsSlackType()) mEqMap.SetSlackBus(i);
		else if (pBus->IsGeneratorType()) mEqMap.SetGeneratorBus(i);
		else if (pBus->IsLoadType()) mEqMap.SetLoadBus(i);
	}
	mEqMap.Finalize();
	if (!mEqMap.IsValid()) return false;
	mEqCount = mEqMap.GetEquationCount();

	// Get pointer to admittance matrix
	mpAdmittanceMatrix = rPowerSystem.GetAdmittanceMatrix();
	if (mpAdmittanceMatrix == NULL) return false;

	// Get pointer to power vector
	mpPowerVector = rPowerSystem.GetPowerVector();
	if (mpPowerVector == NULL) return false;

	// Build flat start solution vector (:TODO: make option for type of starting solution?)
	if (!mVoltageVector.BuildFlat(rPowerSystem)) return false;

	// Clear Jacobian system data
	if (!DestroyJacobianSystem()) return false;

	return true;
} // PowerFlowProblem::SetPowerSystem


#undef __FUNCT__
#define __FUNCT__ "PowerFlowProblem::UpdatePowerSystem"
bool PowerFlowProblem::UpdatePowerSystem(const PowerSystemExt& rPowerSystem)
{
	// Check if a power system is currently set
	if (mBusCount <= 0) return false;

	// :TODO: Is it possibly to do a cheap check whether the structure is the same?

	// Bus count remains the same
	// Equation map remains the same
	// Equation count remains the same

	// Get pointer to admittance matrix (just in case it changed)
	mpAdmittanceMatrix = rPowerSystem.GetAdmittanceMatrix();
	if (mpAdmittanceMatrix == NULL) return false;

	// Get pointer to power vector (just in case it changed)
	mpPowerVector = rPowerSystem.GetPowerVector();
	if (mpPowerVector == NULL) return false;

	// Build flat start solution vector (:TODO: make option for type of starting solution?)
	if (!mVoltageVector.BuildFlat(rPowerSystem)) return false;

	// Do not clear Jacobian system data! The whole point of this function is to keep the nonzero structure.
	mIsCalculated = false;
	mIsPrecalculated = false;

	return true;
} // PowerFlowProblem::SetPowerSystem


#undef __FUNCT__
#define __FUNCT__ "PowerFlowProblem::SetInitialSolution"
bool PowerFlowProblem::SetInitialSolution(const VoltageVector& rVoltageVector)
{
	// Check if rVoltageVector has the correct size
	PetscInt n;
	mErr = VecGetSize(rVoltageVector.mVa, &n);
	if (n != mBusCount) return false;

	// Create mVoltageVector PETSc vectors if they were not created yet
	if (mVoltageVector.mVa == PETSC_NULL)
	{
		mErr = VecDuplicate(rVoltageVector.mVa, &mVoltageVector.mVa); CHKERRB(mErr);
	}
	if (mVoltageVector.mVm == PETSC_NULL)
	{
		mErr = VecDuplicate(rVoltageVector.mVm, &mVoltageVector.mVm); CHKERRB(mErr);
	}

	// Copy vector contents
	mErr = VecCopy(rVoltageVector.mVa, mVoltageVector.mVa); CHKERRB(mErr);
	mErr = VecCopy(rVoltageVector.mVm, mVoltageVector.mVm); CHKERRB(mErr);

	// Jacobian system structure still valid, but values are not.
	mIsCalculated = false;
	mIsPrecalculated = false;

	return true;
} // PowerFlowProblem::SetInitialSolution


#undef __FUNCT__
#define __FUNCT__ "PowerFlowProblem::CalculateFunctionVector"
PetscErrorCode PowerFlowProblem::CalculateFunctionVector(SNES snes, Vec pX, Vec pF)
{
	// Check if function was already calculated for this pX.
	// This is always the case in the first Newton iteration.
	// :WARNING: Cannot simply check if (pX == mpX && pF == mpF) because PETSc sometimes changes
	// the vector without the pointer changing (probably when line search is invoked)
	// :WARNING: Cannot simply check if Newton iteration count is 0 because PETSc recalculates
	//  the function and Jacobian at the end of iteration 0 to check for convergence.
	if (mIsPrecalculated)
	{
		mIsPrecalculated = false;
		return 0;
	}

	mpF = pF;	// Make sure that the member function vector points to the correct data
	mpX = pX;	// Make sure that the member solution vector points to the correct data

	if (mpLogEvents) mpLogEvents->Begin(mpLogEvents->GetPFPCalculateJacobianSystemEventId());
	// Update voltage vector
	if (!CalculateVoltage()) SETERRQ(PETSC_COMM_SELF, PETSC_ERR_GENERAL, "Error calculating voltage vector");
	// Update Jacobian matrix and power mismatch vector
	if (!CalculateMismatchAndJacobian()) SETERRQ(PETSC_COMM_SELF, PETSC_ERR_GENERAL, "Error calculating mismatch vector and Jacobian matrix");
	if (mpLogEvents) mpLogEvents->End(mpLogEvents->GetPFPCalculateJacobianSystemEventId());

	mIsCalculated = true;
	return 0;
} // PowerFlowProblem::CalculateFunctionVector


#undef __FUNCT__
#define __FUNCT__ "PowerFlowProblem::CalculateJacobianMatrix"
PetscErrorCode PowerFlowProblem::CalculateJacobianMatrix(SNES snes, Vec pX, Mat* ppJ)
{
	// Assumes that the desired Jacobian has already been calculated in CalculateFunctionVector()
	// Check if the SNES Jacobian matrix points to the correct data
	// If not it can probably be fixed with a simple: *ppJ = mpJ;
	if (*ppJ != mpJ) SETERRQ(PETSC_COMM_SELF, PETSC_ERR_GENERAL, "SNES Jacobian matrix no longer points to the correct data");
	mpX = pX;	// Make sure that the member solution vector points to the correct data
	return 0;
} // PowerFlowProblem::CalculateJacobianMatrix


//============================= ACCESS     ===================================


Vec PowerFlowProblem::GetSolutionVector(void)
{
	if (!mIsCalculated && !UpdateJacobianSystem()) return PETSC_NULL;
	return mpOwnedX;
}

Vec PowerFlowProblem::GetFunctionVector(void)
{
	if (!mIsCalculated && !UpdateJacobianSystem()) return PETSC_NULL;
	return mpOwnedF;
}

Mat PowerFlowProblem::GetJacobianMatrix(void)
{
	if (!mIsCalculated && !UpdateJacobianSystem()) return PETSC_NULL;
	return mpOwnedJ;
}

bool PowerFlowProblem::IsJacobianMatrixStructureSame(void) const { return mIsStructureSame; }
const VoltageVector* PowerFlowProblem::GetVoltageVector(void) const { return &mVoltageVector; }


/////////////////////////////// PRIVATE    ///////////////////////////////////


#undef __FUNCT__
#define __FUNCT__ "PowerFlowProblem::CreateJacobianSystem"
bool PowerFlowProblem::CreateJacobianSystem(void)
{
	if (mpLogEvents) mpLogEvents->Begin(mpLogEvents->GetPFPCreateJacobianSystemEventId());

	mErr = PetscInfo(PETSC_NULL, "Creating Jacobian system data.\n"); CHKERRB(mErr);
	// Build solution vector data structure
	if (mpOwnedX != PETSC_NULL) { mErr = VecDestroy(&mpOwnedX); mpOwnedX = PETSC_NULL; CHKERRB(mErr); }
	if (!BuildSolutionVectorStructure()) return false;
	mpX = mpOwnedX;
	// Build mismatch vector data structure
	if (mpOwnedF != PETSC_NULL) { mErr = VecDestroy(&mpOwnedF); mpOwnedF = PETSC_NULL; CHKERRB(mErr); }
	if (!BuildMismatchVectorStructure()) return false;
	mpF = mpOwnedF;
	// Build Jacobian matrix data structure
	if (mpOwnedJ != PETSC_NULL) { mErr = MatDestroy(&mpOwnedJ); mpOwnedJ = PETSC_NULL; CHKERRB(mErr); }
	if (!BuildJacobianMatrixStructure()) return false;
	mpJ = mpOwnedJ;

	if (mpLogEvents) mpLogEvents->End(mpLogEvents->GetPFPCreateJacobianSystemEventId());
	return true;
} // PowerFlowProblem::CreateJacobianSystem


#undef __FUNCT__
#define __FUNCT__ "PowerFlowProblem::UpdateJacobianSystem"
bool PowerFlowProblem::UpdateJacobianSystem(void)
{
	if (mpOwnedX == PETSC_NULL || mpOwnedF == PETSC_NULL || mpOwnedJ == PETSC_NULL)
	{
		if (!CreateJacobianSystem()) return false;
	}

	// Reset pointers to point to owned objects
	mpX = mpOwnedX;
	mpF = mpOwnedF;
	mpJ = mpOwnedJ;

	// Calculate PETSc solution vector, mismatch vector, and Jacobian matrix
	if (mpLogEvents) mpLogEvents->Begin(mpLogEvents->GetPFPCalculateJacobianSystemEventId());
	if (!CalculateSolution()) return false;
	if (!CalculateMismatchAndJacobian()) return false;
	if (mpLogEvents) mpLogEvents->End(mpLogEvents->GetPFPCalculateJacobianSystemEventId());

	mIsCalculated = true;
	mIsPrecalculated = true;

	return true;
} // PowerFlowProblem::UpdateJacobianSystem


#undef __FUNCT__
#define __FUNCT__ "PowerFlowProblem::DestroyJacobianSystem"
bool PowerFlowProblem::DestroyJacobianSystem(void)
{
	// Reset all member data related to the Jacobian system
	mMaxNnz = 0;
	mIsCalculated = false;
	mIsPrecalculated = false;
	mIsStructureSame = false;
	if (mpOwnedX != PETSC_NULL) { mErr = VecDestroy(&mpOwnedX); mpOwnedX = PETSC_NULL; CHKERRB(mErr); }
	mpX = PETSC_NULL;
	if (mpOwnedF != PETSC_NULL) { mErr = VecDestroy(&mpOwnedF); mpOwnedF = PETSC_NULL; CHKERRB(mErr); }
	mpF = PETSC_NULL;
	if (mpOwnedJ != PETSC_NULL) { mErr = MatDestroy(&mpOwnedJ); mpOwnedJ = PETSC_NULL; CHKERRB(mErr); }
	mpJ = PETSC_NULL;
	return true;
} // PowerFlowProblem::DestroyJacobianSystem


#undef __FUNCT__
#define __FUNCT__ "PowerFlowProblem::BuildSolutionVectorStructure"
bool PowerFlowProblem::BuildSolutionVectorStructure(void)
{
	mErr = VecCreate(PETSC_COMM_WORLD, &mpOwnedX); CHKERRB(mErr);
	mErr = VecSetSizes(mpOwnedX, mEqCount, PETSC_DETERMINE); CHKERRB(mErr);
	mErr = VecSetType(mpOwnedX, VECSEQ); CHKERRB(mErr);
	return true;
} // PowerFlowProblem::BuildSolutionVectorStructure


#undef __FUNCT__
#define __FUNCT__ "PowerFlowProblem::BuildMismatchVectorStructure"
bool PowerFlowProblem::BuildMismatchVectorStructure(void)
{
	mErr = VecCreate(PETSC_COMM_WORLD, &mpOwnedF); CHKERRB(mErr);
	mErr = VecSetSizes(mpOwnedF, mEqCount, PETSC_DETERMINE); CHKERRB(mErr);
	mErr = VecSetType(mpOwnedF, VECSEQ); CHKERRB(mErr);
	return true;
} // PowerFlowProblem::BuildMismatchVectorStructure


#undef __FUNCT__
#define __FUNCT__ "PowerFlowProblem::BuildJacobianMatrixStructure"
bool PowerFlowProblem::BuildJacobianMatrixStructure(void)
{
	PetscInt i, j, k, l;

	// Non-zero structure will (most likely) change
	mIsStructureSame = false;

	// Get the structure matrix
	PetscUtils::ConstMat pStructure = mpAdmittanceMatrix->GetReal();

	// Get equation map
	const EquationMap::IndexType* pMapY2J = mEqMap.GetArray();

	// Allocate memory for the non-zero count
	PetscInt* pNnz;
	mErr = PetscMalloc(mEqCount * sizeof(PetscInt), &pNnz); CHKERRB(mErr);
	mErr = PetscMemzero(pNnz, mEqCount * sizeof(PetscInt)); CHKERRB(mErr);

    // Count number of non-zeros for each row
	PetscInt rowExp, colExp;			// Number of rows in J for the current row in Y (0,1,2)
	PetscInt colCountY;					// Number of non-zeros in the current row of Y
	const PetscInt* pColsY;				// Column indices of the non-zeros in the current row of Y
	l = 0;								// Index of the current row in J
	mMaxNnz = 0;						// Maximum number of non-zeros on a row in J
	for (i=0; i<mBusCount; ++i)			// Run through the rows of Y
	{
		rowExp = pMapY2J[i+1] - pMapY2J[i];		// Number of rows in J for row i of Y
		if (rowExp == 0) continue;				// Skip slack bus row

		// Process row i of Y
		mErr = PetscUtils::MatSeqAijGetRow(pStructure, i, &colCountY, &pColsY, PETSC_NULL); CHKERRB(mErr);
		for (j=0; j<colCountY; ++j)				// Run through the non-zeros of the current row in Y
		{
			colExp = pMapY2J[pColsY[j]+1] - pMapY2J[pColsY[j]];	// Number of columns in J for column i of Y
			pNnz[l] += colExp;					// Update non-zero count in row l of J
		}
		mErr = PetscUtils::MatSeqAijRestoreRow(pStructure, i, &colCountY, &pColsY, PETSC_NULL); CHKERRB(mErr);
		if (rowExp>1) pNnz[l+1] = pNnz[l];		// Update non-zero count in row l+1 of J if row is a load bus

		mMaxNnz = PetscMax(mMaxNnz, pNnz[l]);	// Update max numbers of non-zeros found in a row of J
		l += rowExp;							// Update the current row index in J
	} // for (i=0; i<mEqCount; ++i)

	// Initialize Jacobian matrix J
	if (mpOwnedJ != PETSC_NULL) { mErr = MatDestroy(&mpOwnedJ); mpOwnedJ = PETSC_NULL; CHKERRB(mErr); }
	mErr = MatCreate(PETSC_COMM_WORLD, &mpOwnedJ); CHKERRB(mErr);
	mErr = MatSetSizes(mpOwnedJ, mEqCount, mEqCount, PETSC_DECIDE, PETSC_DECIDE); CHKERRB(mErr);
	mErr = MatSetType(mpOwnedJ, MATSEQAIJ); CHKERRB(mErr);
	mErr = MatSeqAIJSetPreallocation(mpOwnedJ, 0, pNnz); CHKERRB(mErr);
	mErr = MatSetOption(mpOwnedJ, MAT_IGNORE_ZERO_ENTRIES, PETSC_FALSE); CHKERRB(mErr);	// Make sure zero entries are stored

	// Free the non-zero count array
	mErr = PetscFree(pNnz); CHKERRB(mErr);

	// Allocate arrays to hold the column indices and matrix values for the current row of Y
	PetscInt* pColsJ;
	mErr = PetscMalloc(mMaxNnz * sizeof(PetscInt), &pColsJ); CHKERRB(mErr);
	PetscScalar* pValsJ;
	mErr = PetscMalloc(mMaxNnz * sizeof(PetscScalar), &pValsJ); CHKERRB(mErr);
	mErr = PetscMemzero(pValsJ, mMaxNnz * sizeof(PetscScalar)); CHKERRB(mErr);

	// Build arrays with column indices and (dummy) matrix values for the current row of Y
	PetscInt rowIndexJ, colIndexJ;		// Current row and column index in J
	for (i=0; i<mBusCount; ++i)			// Run through the rows of Y
	{
		l = 0;		// Non-zero index in the current row of J, of a non-zero in the current row of Y
		rowIndexJ = pMapY2J[i];
		rowExp = pMapY2J[i+1] - rowIndexJ;		// Number of rows in J for row i of Y
		if (rowExp == 0) continue;		// Skip slack bus row

		mErr = PetscUtils::MatSeqAijGetRow(pStructure, i, &colCountY, &pColsY, PETSC_NULL); CHKERRB(mErr);
		for (k=0; k<colCountY; ++k)		// Run through the non-zeros of the current row in Y
		{
			j = pColsY[k];
			colIndexJ = pMapY2J[j];
			colExp = pMapY2J[j+1] - colIndexJ;	// Number of columns in J for column i of Y
			if (colExp == 0) continue;	// Skip slack bus column

			pColsJ[l] = colIndexJ;		// Set column index of the non-zero in J
			++l;

			if (colExp>1)				// Column corresponds to a load bus
			{
				pColsJ[l] = colIndexJ+1;// Set global column index of the non-zero in J
				++l;
			}
		} // for (k=0; k<colCountY; ++k)
		mErr = PetscUtils::MatSeqAijRestoreRow(pStructure, i, &colCountY, &pColsY, PETSC_NULL); CHKERRB(mErr);

		// Set the non-zeros of the current row into J, and the next row if applicable
		mErr = MatSetValues(mpOwnedJ, 1, &rowIndexJ, l, pColsJ, pValsJ, INSERT_VALUES); CHKERRB(mErr);
		if (rowExp>1)	// Load bus row
		{
			++rowIndexJ;
			mErr = MatSetValues(mpOwnedJ, 1, &rowIndexJ, l, pColsJ, pValsJ, INSERT_VALUES); CHKERRB(mErr);
		}
	} // for (i=0; i<mEqCount; ++i)

	// Free column index and value arrays
	mErr = PetscFree(pColsJ); CHKERRB(mErr);
	mErr = PetscFree(pValsJ); CHKERRB(mErr);

	// Assemble Jacobian matrix
	mErr = MatAssemblyBegin(mpOwnedJ, MAT_FINAL_ASSEMBLY); CHKERRB(mErr);
	mErr = MatAssemblyEnd(mpOwnedJ, MAT_FINAL_ASSEMBLY); CHKERRB(mErr);

	// Set to produce an error if the code tries to add new nonzeros to the Jacobian matrix
	mErr = MatSetOption(mpOwnedJ, MAT_NEW_NONZERO_LOCATION_ERR, PETSC_TRUE); CHKERRB(mErr);

	return true;
} // PowerFlowProblem::BuildJacobianMatrixStructure


#undef __FUNCT__
#define __FUNCT__ "PowerFlowProblem::CalculateVoltage"
bool PowerFlowProblem::CalculateVoltage(void)
{
	// Assumes valid mEqMap, mVoltageVector of size mBusCount, and mpX of size mEqCount

	// Get pointers to PETSc voltage vectors
	Vec pVa = mVoltageVector.GetAngPrivate();
	if (pVa == PETSC_NULL) return false;
	Vec pVm = mVoltageVector.GetMagPrivate();
	if (pVm == PETSC_NULL) return false;

	const PetscScalar* pValsX;
	PetscScalar *pValsVa, *pValsVm;
	mErr = PetscUtils::VecSeqGetArray(mpX, &pValsX); CHKERRB(mErr);
	mErr = VecGetArray(pVa, &pValsVa); CHKERRB(mErr);
	mErr = VecGetArray(pVm, &pValsVm); CHKERRB(mErr);
	PetscInt rowExp, l = 0;
	for (PetscInt i=0; i<mBusCount; ++i)
	{
		rowExp = mEqMap.GetEquationCount(i);
		if (rowExp == 0) continue;	// Skip slack bus
		pValsVa[i] = pValsX[l];		// Update voltage angle value
		++l;
		if (rowExp>1)				// Load bus
		{
			pValsVm[i] = pValsX[l];	// Update voltage magnitude value
			++l;
		}
	}
	mErr = PetscUtils::VecSeqRestoreArray(mpX, &pValsX); CHKERRB(mErr);
	mErr = VecRestoreArray(pVa, &pValsVa); CHKERRB(mErr);
	mErr = VecRestoreArray(pVm, &pValsVm); CHKERRB(mErr);
	return true;
} // PowerFlowProblem::CalculateVoltage


#undef __FUNCT__
#define __FUNCT__ "PowerFlowProblem::CalculateSolution"
bool PowerFlowProblem::CalculateSolution(void)
{
	// Assumes valid mEqMap, mVoltageVector of size mBusCount, and mpX of size mEqCount

	// Get pointers to PETSc voltage vectors
	PetscUtils::ConstVec pVa = mVoltageVector.GetAng();
	if (pVa == PETSC_NULL) return false;
	PetscUtils::ConstVec pVm = mVoltageVector.GetMag();
	if (pVm == PETSC_NULL) return false;

	// Set the appropriate values of pVa and pVm into mpX
	const PetscScalar *pValsVa, *pValsVm;
	PetscScalar* pValsX;
	mErr = PetscUtils::VecSeqGetArray(pVa, &pValsVa); CHKERRB(mErr);
	mErr = PetscUtils::VecSeqGetArray(pVm, &pValsVm); CHKERRB(mErr);
	mErr = VecGetArray(mpX, &pValsX); CHKERRB(mErr);
	PetscInt rowExp, l = 0;
	for (PetscInt i=0; i<mBusCount; ++i)
	{
		rowExp = mEqMap.GetEquationCount(i);
		if (rowExp == 0) continue;	// Skip slack bus
		pValsX[l] = pValsVa[i];
		++l;
		if (rowExp>1)				// Load bus
		{
			pValsX[l] = pValsVm[i];
			++l;
		}
	}
	mErr = PetscUtils::VecSeqRestoreArray(pVa, &pValsVa); CHKERRB(mErr);
	mErr = PetscUtils::VecSeqRestoreArray(pVm, &pValsVm); CHKERRB(mErr);
	mErr = VecRestoreArray(mpX, &pValsX); CHKERRB(mErr);
	return true;
} // PowerFlowProblem::CalculateSolution


#undef __FUNCT__
#define __FUNCT__ "PowerFlowProblem::CalculateJacobianAndMismatch"
bool PowerFlowProblem::CalculateMismatchAndJacobian(void)
{
	// Get equation map
	const EquationMap::IndexType* pMapY2J = mEqMap.GetArray();

	// Allocate arrays to hold the P and Q row sums
	PetscScalar *pRowSumPij, *pRowSumQij;
	mErr = PetscMalloc(mBusCount * sizeof(PetscScalar), &pRowSumPij); CHKERRB(mErr);
	mErr = PetscMalloc(mBusCount * sizeof(PetscScalar), &pRowSumQij); CHKERRB(mErr);

	/*** CALCULATE THE JACOBIAN MATRIX J ***/

	// Allocate arrays to hold the column indices and matrix values for the current row
	PetscInt* pColsJ;					// Column indices
	mErr = PetscMalloc(mMaxNnz * sizeof(PetscInt), &pColsJ); CHKERRB(mErr);
	PetscScalar* pValsJ1;				// Non-zero values of first row (dP/d* row)
	PetscScalar* pValsJ2;				// Non-zero values of second row (dQ/d* row)
	mErr = PetscMalloc(mMaxNnz * sizeof(PetscScalar), &pValsJ1); CHKERRB(mErr);
	mErr = PetscMalloc(mMaxNnz * sizeof(PetscScalar), &pValsJ2); CHKERRB(mErr);

	// Variable declarations
	PetscInt i, j, k, l, ldiag = 0;
	PetscInt rowIndexJ, colIndexJ, colCount;
	PetscInt rowExp, colExp;
	PetscScalar va_i, vm_i, vm_j, g_ij, b_ij, Pij, Qij;
	PetscScalar vm_ij, va_ij, sin_ij, cos_ij;
	const PetscInt* pCols;
	const PetscScalar *pValsG, *pValsB;

	// Get the voltage angle and voltage magnitude arrays
	PetscUtils::ConstVec pVa = mVoltageVector.GetAng();
	if (pVa == PETSC_NULL) return false;
	PetscUtils::ConstVec pVm = mVoltageVector.GetMag();
	if (pVm == PETSC_NULL) return false;
	const PetscScalar *pValsVa, *pValsVm;
	mErr = PetscUtils::VecSeqGetArray(pVa, &pValsVa); CHKERRB(mErr);
	mErr = PetscUtils::VecSeqGetArray(pVm, &pValsVm); CHKERRB(mErr);

	// Get pointers to PETSc admittance matrices
	if (mpAdmittanceMatrix == NULL) return false;
	PetscUtils::ConstMat pG = mpAdmittanceMatrix->GetReal();
	if (pG == PETSC_NULL) return false;
	PetscUtils::ConstMat pB = mpAdmittanceMatrix->GetImag();
	if (pB == PETSC_NULL) return false;

	for (i=0; i<mBusCount; ++i)			// Run through the rows of Y
	{
		pRowSumPij[i] = 0;				// Initialize P row sum
		pRowSumQij[i] = 0;				// Initialize Q row sum

		rowIndexJ = pMapY2J[i];				// First row index in J for row i of Y
		rowExp = pMapY2J[i+1] - rowIndexJ;	// Number of rows in J for row i of Y
		if (rowExp == 0) continue;			// Skip the slack bus row

		l = 0;							// Non-zero index in the current row of J
		va_i = pValsVa[i];				// Get the voltage angle for row i of Y
		vm_i = pValsVm[i];				// Get the voltage magnitude for row i of Y

		// Process the non-zeros of row i of Y
		mErr = PetscUtils::MatSeqAijGetRow(pG, i, &colCount, &pCols, &pValsG); CHKERRB(mErr);		// Get row i of G = real(Y)
		mErr = PetscUtils::MatSeqAijGetRow(pB, i, PETSC_NULL, PETSC_NULL, &pValsB); CHKERRB(mErr);	// Get row i of B = imag(Y)
		for (k=0; k<colCount; ++k)		// Run through the non-zeros in the current row of Y
		{
			j = pCols[k];				// Column index
			if (i==j) ldiag = l;		// Non-zero index in J for the diagonal element of row i of Y

			// Get the admittance for column j of Y
			g_ij = pValsG[k];			// Conductance
			b_ij = pValsB[k];			// Susceptance

			// Precalculate values used for the calculation of Pij and Qij
			vm_j = pValsVm[j];			// Voltage magnitude
			vm_ij = vm_i * vm_j;		// Voltage magnitude product
			va_ij = va_i - pValsVa[j];	// Voltage angle difference
			sin_ij = sin(va_ij);
			cos_ij = cos(va_ij);

			// Evaluate Pij and Qij using the precalculated values
			Pij = vm_ij * (g_ij*cos_ij + b_ij*sin_ij);
			Qij = vm_ij * (g_ij*sin_ij - b_ij*cos_ij);
			pRowSumPij[i] += Pij;	// :NOTE: currently not calculated (nor used) for slack bus rows
			pRowSumQij[i] += Qij;	// :NOTE: currently not calculated (nor used) for slack bus rows

			colIndexJ = pMapY2J[j];				// First column index in J for column j of Y
			colExp = pMapY2J[j+1] - colIndexJ;	// Number of columns in J for column j of Y
			if (colExp == 0) continue;	// :WARNING: update pRowSumPij and pRowSumQij BEFORE skipping the slack bus column

			// Set first value of J for element (i,j) of Y
			pColsJ[l] = colIndexJ;			// Set column index in J
			pValsJ1[l] = -Qij;				// Set value in the first row of J
			pValsJ2[l] = Pij;				// Set value in the second row of J (:NOTE: can be skipped if rowExp<2)
			++l;

			if (colExp>1)
			{	// Set second value of J for element (i,j) of Y
				pColsJ[l] = colIndexJ+1;	// Set column index in J
				pValsJ1[l] = -Pij/vm_j;		// Set value in the first row of J
				pValsJ2[l] = -Qij/vm_j;		// Set value in the second row of J (:NOTE: can be skipped if rowExp<2)
				++l;
			}
		} // for (k=0; k<colCount; ++k)
		mErr = PetscUtils::MatSeqAijRestoreRow(pG, i, &colCount, &pCols, &pValsG); CHKERRB(mErr);		// Restore row i of G
		mErr = PetscUtils::MatSeqAijRestoreRow(pB, i, &colCount, PETSC_NULL, &pValsB); CHKERRB(mErr);	// Restore row i of B

		// Fix the non-zeros in J for the diagonal element of row i of Y
		pValsJ1[ldiag] += pRowSumQij[i];	// Fix first value of first row
		pValsJ2[ldiag] -= pRowSumPij[i];	// Fix first value of second row (:NOTE: can be skipped if rowExp<2)
		if (rowExp>1)	// :NOTE: rowExp = colExp at the diagonal
		{
			++ldiag;
			pValsJ1[ldiag] -= pRowSumPij[i]/vm_i;	// Fix second value of first row (:NOTE: vm_i = vm_j at the diagonal)
			pValsJ2[ldiag] -= pRowSumQij[i]/vm_i;	// Fix second value of second row (:NOTE: cannot be skipped since rowExp>1)
		}

		// Set the non-zeros of J for row i of Y into the Jacobian matrix
		mErr = MatSetValuesRow(mpJ, rowIndexJ, pValsJ1); CHKERRB(mErr);
//		mErr = MatSetValues(mpJ, 1, &rowIndexJ, l, pColsJ, pValsJ1, INSERT_VALUES); CHKERRB(mErr);
		if (rowExp>1)
		{
			++rowIndexJ;
			mErr = MatSetValuesRow(mpJ, rowIndexJ, pValsJ2); CHKERRB(mErr);
//			mErr = MatSetValues(mpJ, 1, &rowIndexJ, l, pColsJ, pValsJ2, INSERT_VALUES); CHKERRB(mErr);
		}
	} // for (i=0; i<mBusCount; ++i)

	// Restore voltage angle and voltage magnitude arrays
	mErr = PetscUtils::VecSeqRestoreArray(pVa, &pValsVa); CHKERRB(mErr);
	mErr = PetscUtils::VecSeqRestoreArray(pVm, &pValsVm); CHKERRB(mErr);

	// Free the column index and value arrays
	mErr = PetscFree(pColsJ); CHKERRB(mErr);
	mErr = PetscFree(pValsJ1); CHKERRB(mErr);
	mErr = PetscFree(pValsJ2); CHKERRB(mErr);

	// Assemble the Jacobian matrix
	mErr = MatAssemblyBegin(mpJ, MAT_FINAL_ASSEMBLY); CHKERRB(mErr);	// :NOTE: does nothing for MATSEQAIJ matrices
	mErr = MatAssemblyEnd(mpJ, MAT_FINAL_ASSEMBLY); CHKERRB(mErr);		// :NOTE: calls MatAssemblyEnd_SeqAIJ() for MATSEQAIJ matrices

	/*** CALCULATE THE POWER MISMATCH VECTOR F ***/

	PetscUtils::ConstVec pP = mpPowerVector->GetReal();
	if (pP == PETSC_NULL) return false;
	PetscUtils::ConstVec pQ = mpPowerVector->GetImag();
	if (pQ == PETSC_NULL) return false;
	const PetscScalar *pValsP, *pValsQ;
	mErr = PetscUtils::VecSeqGetArray(pP, &pValsP); CHKERRB(mErr);	// Get active power array
	mErr = PetscUtils::VecSeqGetArray(pQ, &pValsQ); CHKERRB(mErr);	// Get reactive power array
	PetscScalar* pValsF;
	mErr = VecGetArray(mpF, &pValsF); CHKERRB(mErr);				// Get power mismatch array
	l = 0;							// Index in F for the current index in P&Q
	for (i=0; i<mBusCount; ++i)		// Run through the elements of P&Q
	{
		rowExp = pMapY2J[i+1] - pMapY2J[i];
		if (rowExp == 0) continue;	// Skip slap bus row

		pValsF[l] = pValsP[i] - pRowSumPij[i];			// Set first (P) value in F
		++l;
		if (rowExp>1)
		{	// Index in P&Q corresponds to a load bus
			pValsF[l] = pValsQ[i] - pRowSumQij[i];		// Set second (Q) value in F
			++l;
		}
	}
	mErr = PetscUtils::VecSeqRestoreArray(pP, &pValsP); CHKERRB(mErr);	// Restore active power array
	mErr = PetscUtils::VecSeqRestoreArray(pQ, &pValsQ); CHKERRB(mErr);	// Restore reactive power array
	mErr = VecRestoreArray(mpF, &pValsF); CHKERRB(mErr);				// Restore power mismatch vector

	// Free the memory allocated for the P and Q row sums
	mErr = PetscFree(pRowSumPij); CHKERRB(mErr);
	mErr = PetscFree(pRowSumQij); CHKERRB(mErr);

	// Non-zero structure was not changed
	mIsStructureSame = true;

	return true;
} // PowerFlowProblem::CalculateJacobianAndMismatch


POWERFLOWER_NAMESPACE_END
