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

/*
 * Power System Admittance Matrix class source.
 */
#include "LogEvents.h"				// HASA LogEvents (not owned)
#include "AdmittanceMatrix.h"


#include <iostream>



POWERFLOWER_NAMESPACE_BEGIN

POWERSYSTEM_NAMESPACE_USING


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


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


#undef __FUNCT__
#define __FUNCT__ "AdmittanceMatrix::AdmittanceMatrix"
AdmittanceMatrix::AdmittanceMatrix(void)
	: mErr(0), mpLogEvents(NULL), mG(PETSC_NULL), mB(PETSC_NULL)
{
	if (PetscUtils::IsInitialized())
	{
		mErr = PetscInfo(PETSC_NULL, "Creating AdmittanceMatrix object.\n"); CHKERRV(mErr);
	}
} // AdmittanceMatrix::AdmittanceMatrix

#undef __FUNCT__
#define __FUNCT__ "AdmittanceMatrix::AdmittanceMatrix"
AdmittanceMatrix::AdmittanceMatrix(const LogEvents* pLogEvents)
	: mErr(0), mpLogEvents(pLogEvents), mG(PETSC_NULL), mB(PETSC_NULL)
{
	if (PetscUtils::IsInitialized())
	{
		mErr = PetscInfo(PETSC_NULL, "Creating AdmittanceMatrix object.\n"); CHKERRV(mErr);
	}
} // AdmittanceMatrix::AdmittanceMatrix

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


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


#undef __FUNCT__
#define __FUNCT__ "AdmittanceMatrix::Clear"
void AdmittanceMatrix::Clear(void)
{
	if (PetscUtils::IsInitialized())
	{
		mErr = PetscInfo(PETSC_NULL, "Clearing AdmittanceMatrix object data.\n"); CHKERRV(mErr);
		if (mG != PETSC_NULL) { mErr = MatDestroy(&mG); mG = PETSC_NULL; CHKERRV(mErr); }
		if (mB != PETSC_NULL) { mErr = MatDestroy(&mB); mB = PETSC_NULL; CHKERRV(mErr); }
	}
	mErr = 0;
} // AdmittanceMatrix::Clear


#undef __FUNCT__
#define __FUNCT__ "AdmittanceMatrix::Build"
bool AdmittanceMatrix::Build(const PowerSystem& rPowerSystem)
{
	if (!PetscUtils::IsInitialized()) return false;

	if (mpLogEvents) mpLogEvents->Begin(mpLogEvents->GetAMBuildEventId());
 	mErr =  PetscInfo(PETSC_NULL, "Building admittance matrix.\n"); CHKERRB(mErr);
 	bool success = Build_Private(rPowerSystem);
	if (mpLogEvents) mpLogEvents->End(mpLogEvents->GetAMBuildEventId());

	return success;
} // AdmittanceMatrix::Build


#undef __FUNCT__
#define __FUNCT__ "AdmittanceMatrix::AddBranch"
bool AdmittanceMatrix::AddBranch(const POWERSYSTEM_NAMESPACE_NAME::Branch& rBranch)
{
	if (!PetscUtils::IsInitialized()) return false;

	if (mpLogEvents) mpLogEvents->Begin(mpLogEvents->GetAMAddBranchEventId());

	mErr =  PetscInfo(PETSC_NULL, "Adding branch.\n"); CHKERRB(mErr);

	if (mG == PETSC_NULL || mB == PETSC_NULL)
	{
		mErr =  PetscInfo(PETSC_NULL, "[Failed] Admittance matrix needs to be built first.\n"); CHKERRB(mErr);
		if (mpLogEvents) mpLogEvents->End(mpLogEvents->GetAMAddBranchEventId());
		return false;
	}

 	// Get the bus indices
	PetscInt idx[2] = {rBranch.GetFromBus(), rBranch.GetToBus()};

	// Calculate the value updates
	PetscScalar valsG[4] = {0};
	PetscScalar valsB[4] = {0};
	AddValues(rBranch, valsG[0], valsG[1], valsG[2], valsG[3], valsB[0], valsB[1], valsB[2], valsB[3]);

	// Add the value updates
	MatSetValues(mG, 2, idx, 2, idx, valsG, ADD_VALUES); CHKERRB(mErr);
	mErr = MatAssemblyBegin(mG, MAT_FINAL_ASSEMBLY); CHKERRB(mErr);
	mErr = MatAssemblyEnd(mG, MAT_FINAL_ASSEMBLY); CHKERRB(mErr);
	MatSetValues(mB, 2, idx, 2, idx, valsB, ADD_VALUES); CHKERRB(mErr);
	mErr = MatAssemblyBegin(mB, MAT_FINAL_ASSEMBLY); CHKERRB(mErr);
	mErr = MatAssemblyEnd(mB, MAT_FINAL_ASSEMBLY); CHKERRB(mErr);

	if (mpLogEvents) mpLogEvents->End(mpLogEvents->GetAMAddBranchEventId());

	return true;
} // AdmittanceMatrix::AddBranch


#undef __FUNCT__
#define __FUNCT__ "AdmittanceMatrix::RemoveBranch"
bool AdmittanceMatrix::RemoveBranch(const POWERSYSTEM_NAMESPACE_NAME::Branch& rBranch)
{
	if (!PetscUtils::IsInitialized()) return false;

	if (mpLogEvents) mpLogEvents->Begin(mpLogEvents->GetAMRemBranchEventId());

	mErr =  PetscInfo(PETSC_NULL, "Removing branch.\n"); CHKERRB(mErr);

	if (mG == PETSC_NULL || mB == PETSC_NULL)
	{
		if (mpLogEvents) mpLogEvents->End(mpLogEvents->GetAMRemBranchEventId());
		mErr =  PetscInfo(PETSC_NULL, "[Failed] Admittance matrix needs to be built first.\n"); CHKERRB(mErr);
		return false;
	}

	// Get the bus indices
	PetscInt idx[2] = {rBranch.GetFromBus(), rBranch.GetToBus()};

	// Calculate the value updates
	PetscScalar valsG[4] = {0};
	PetscScalar valsB[4] = {0};
	AddValues(rBranch, valsG[0], valsG[1], valsG[2], valsG[3], valsB[0], valsB[1], valsB[2], valsB[3]);

	// Negate value updates
	valsG[0] = -valsG[0]; valsG[1] = -valsG[1]; valsG[2] = -valsG[2]; valsG[3] = -valsG[3];
	valsB[0] = -valsB[0]; valsB[1] = -valsB[1]; valsB[2] = -valsB[2]; valsB[3] = -valsB[3];

	// Add the value updates
	MatSetValues(mG, 2, idx, 2, idx, valsG, ADD_VALUES); CHKERRB(mErr);
	mErr = MatAssemblyBegin(mG, MAT_FINAL_ASSEMBLY); CHKERRB(mErr);
	mErr = MatAssemblyEnd(mG, MAT_FINAL_ASSEMBLY); CHKERRB(mErr);
	MatSetValues(mB, 2, idx, 2, idx, valsB, ADD_VALUES); CHKERRB(mErr);
	mErr = MatAssemblyBegin(mB, MAT_FINAL_ASSEMBLY); CHKERRB(mErr);
	mErr = MatAssemblyEnd(mB, MAT_FINAL_ASSEMBLY); CHKERRB(mErr);

	if (mpLogEvents) mpLogEvents->End(mpLogEvents->GetAMRemBranchEventId());

	return true;
} // AdmittanceMatrix::RemoveBranch


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


//Mat AdmittanceMatrix::GetReal(void) { return mG; }
//Mat AdmittanceMatrix::GetImag(void) { return mB; }

PetscUtils::ConstMat AdmittanceMatrix::GetReal(void) const { return mG; }
PetscUtils::ConstMat AdmittanceMatrix::GetImag(void) const { return mB; }


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


#undef __FUNCT__
#define __FUNCT__ "AdmittanceMatrix::Build_Private"
bool AdmittanceMatrix::Build_Private(const PowerSystem& rPowerSystem)
{
 	mErr =  PetscInfo(PETSC_NULL, "Creating and preallocating admittance matrix.\n"); CHKERRB(mErr);

 	// Get the number of rows in the admittance matrix
	PetscInt n = rPowerSystem.GetBusCount();
	if (n == 0) return false;

	// Row and column indices
	PetscInt i, j;

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

	// Count number of non-zeros in each row by iterating over all branches
	for (i = 0; i < n; ++i) pNnz[i] = 1;		// Set the diagonal non-zero for each row in advance
	PowerSystem::ConstBranchIterator it = rPowerSystem.GetBranchBegin();
	PowerSystem::ConstBranchIterator itEnd = rPowerSystem.GetBranchEnd();
	while (it != itEnd)
	{
		++pNnz[it->GetFromBus()];	// Add one non-zero for the from bus
		++pNnz[it->GetToBus()];		// Add one non-zero for the to bus
		// Skip all branches between the same two busses
		it = rPowerSystem.GetBranchNextUnique(it);
	} // while (it != itEnd)

	// Create the real part of the admittance matrix (G)
	if (mG != PETSC_NULL) { mErr = MatDestroy(&mG);  mG = PETSC_NULL; CHKERRB(mErr); }
	mErr = MatCreate(PETSC_COMM_WORLD, &mG); CHKERRB(mErr);
	mErr = MatSetSizes(mG, n, n, PETSC_DECIDE, PETSC_DECIDE); CHKERRB(mErr);
	mErr = MatSetType(mG, MATSEQAIJ); CHKERRB(mErr);
	mErr = MatSeqAIJSetPreallocation(mG, 0, pNnz); CHKERRB(mErr);
	mErr = MatSetOption(mG, MAT_IGNORE_ZERO_ENTRIES, PETSC_FALSE); CHKERRB(mErr);	// Make sure zero entries are stored

	// Create the imaginary part of the admittance matrix (B)
	if (mB != PETSC_NULL) { mErr = MatDestroy(&mB); mB = PETSC_NULL; CHKERRB(mErr); }
	mErr = MatCreate(PETSC_COMM_WORLD, &mB); CHKERRB(mErr);
	mErr = MatSetSizes(mB, n, n, PETSC_DECIDE, PETSC_DECIDE); CHKERRB(mErr);
	mErr = MatSetType(mB, MATSEQAIJ); CHKERRB(mErr);
	mErr = MatSeqAIJSetPreallocation(mB, 0, pNnz); CHKERRB(mErr);
	mErr = MatSetOption(mB, MAT_IGNORE_ZERO_ENTRIES, PETSC_FALSE); CHKERRB(mErr);	// Make sure zero entries are stored

 	mErr =  PetscInfo(PETSC_NULL, "Calculating admittance matrix values.\n"); CHKERRB(mErr);

	// Allocate arrays of pointers to column index and value arrays
	PetscInt** ppCols;
	PetscScalar** ppValsG;
	PetscScalar** ppValsB;
	mErr = PetscMalloc(n * sizeof(PetscInt*), &ppCols); CHKERRB(mErr);
	mErr = PetscMalloc(n * sizeof(PetscScalar*), &ppValsG); CHKERRB(mErr);
	mErr = PetscMalloc(n * sizeof(PetscScalar*), &ppValsB); CHKERRB(mErr);

	// Allocate memory and set initial data of the diagonal entries
	for (i = 0; i < n; ++i)
	{
		// Allocate the column index and value arrays and set values to 0 so we can work with additions
		mErr = PetscMalloc(pNnz[i] * sizeof(PetscInt), &ppCols[i]); CHKERRB(mErr);
		mErr = PetscMalloc(pNnz[i] * sizeof(PetscScalar), &ppValsG[i]); CHKERRB(mErr);
		mErr = PetscMemzero(ppValsG[i], pNnz[i] * sizeof(PetscScalar)); CHKERRB(mErr);
		mErr = PetscMalloc(pNnz[i] * sizeof(PetscScalar), &ppValsB[i]); CHKERRB(mErr);
		mErr = PetscMemzero(ppValsB[i], pNnz[i] * sizeof(PetscScalar)); CHKERRB(mErr);

		// Reserve the first element in each row for the diagonal non-zero
		pNnz[i] = 0;		// Reset array with number of non-zeros per row to act as non-zero index
		ppCols[i][0] = i;	// First element is diagonal non-zero
		// Initialize diagonal non-zero values with the bus shunt values
		ppValsG[i][0] = rPowerSystem.GetBus(i)->GetShuntConductance() / rPowerSystem.GetBasePower();
		ppValsB[i][0] = rPowerSystem.GetBus(i)->GetShuntSusceptance() / rPowerSystem.GetBasePower();
	} // for (i = 0; i < n; ++i)

	// Set the non-zero values
	PowerSystem::ConstBranchIterator itNextUnique;
	it = rPowerSystem.GetBranchBegin();
	while (it != itEnd)
	{
		// Add values of first branch between a bus pair
		i = it->GetFromBus();
		j = it->GetToBus();
		++pNnz[i];				// Update column index for row i
		++pNnz[j];				// Update column index for row j
		ppCols[i][pNnz[i]] = j;	// Set column index for the new non-zeros in row i
		ppCols[j][pNnz[j]] = i;	// Set column index for the new non-zeros in row j
		AddValues(*it, ppValsG[i][0], ppValsG[i][pNnz[i]], ppValsG[j][pNnz[j]], ppValsG[j][0], ppValsB[i][0], ppValsB[i][pNnz[i]], ppValsB[j][pNnz[j]], ppValsB[j][0]);
		itNextUnique = rPowerSystem.GetBranchNextUnique(it);
		++it;
		// Add values of the rest of the branches between the same bus pair
		while (it != itNextUnique)
		{
			i = it->GetFromBus();
			j = it->GetToBus();
			AddValues(*it, ppValsG[i][0], ppValsG[i][pNnz[i]], ppValsG[j][pNnz[j]], ppValsG[j][0], ppValsB[i][0], ppValsB[i][pNnz[i]], ppValsB[j][pNnz[j]], ppValsB[j][0]);
			++it;
		} // while (it != itNextUnique)
	} // while (it != itEnd)
	for (i = 0; i < n; ++i) ++pNnz[i];	// Update from non-zero index into non-zero count

 	mErr =  PetscInfo(PETSC_NULL, "Setting admittance matrix values.\n"); CHKERRB(mErr);

	// Set the calculated values into the matrices and free the column index and value arrays
	for (PetscInt i = 0; i < n; ++i)
	{
		mErr = MatSetValues(mG, 1, &i, pNnz[i], ppCols[i], ppValsG[i], INSERT_VALUES); CHKERRB(mErr);
		mErr = PetscFree(ppValsG[i]); CHKERRB(mErr);
		mErr = MatSetValues(mB, 1, &i, pNnz[i], ppCols[i], ppValsB[i], INSERT_VALUES); CHKERRB(mErr);
		mErr = PetscFree(ppValsB[i]); CHKERRB(mErr);
		mErr = PetscFree(ppCols[i]); CHKERRB(mErr);
	}

	// Free arrays with pointers to column index and value arrays
	mErr = PetscFree(ppValsB); CHKERRB(mErr);
	mErr = PetscFree(ppValsG); CHKERRB(mErr);
	mErr = PetscFree(ppCols); CHKERRB(mErr);

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

 	mErr =  PetscInfo(PETSC_NULL, "Assembling admittance matrix.\n"); CHKERRB(mErr);

 	// Assemble the matrices
	mErr = MatAssemblyBegin(mG, MAT_FINAL_ASSEMBLY); CHKERRB(mErr);
	mErr = MatAssemblyEnd(mG, MAT_FINAL_ASSEMBLY); CHKERRB(mErr);
	mErr = MatAssemblyBegin(mB, MAT_FINAL_ASSEMBLY); CHKERRB(mErr);
	mErr = MatAssemblyEnd(mB, MAT_FINAL_ASSEMBLY); CHKERRB(mErr);

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

	return true;
} // AdmittanceMatrix::Build_Private


#undef __FUNCT__
#define __FUNCT__ "AdmittanceMatrix::AddValues"
void AdmittanceMatrix::AddValues(const Branch& rBranch,
		PetscScalar& rGii, PetscScalar& rGij, PetscScalar& rGji, PetscScalar& rGjj,
		PetscScalar& rBii, PetscScalar& rBij, PetscScalar& rBji, PetscScalar& rBjj)
{
	// y = (R - jX) / (R^2 + X^2)
	// T = ratio * e^(j*angle)
	// Y(i,i) = y / |T|^2 + jB/2
	// Y(j,j) = y + jB/2
	// Y(i,j) = -y * T / |T|^2
	// Y(j,i) = -y * conj(T) / |T|^2

	// Do not add values for branches that are not in service
	if (!rBranch.IsActive()) return;

	// Precalc some values
	PetscScalar temp = rBranch.GetResistance() * rBranch.GetResistance() + rBranch.GetReactance() * rBranch.GetReactance();
	PetscScalar admReal = rBranch.GetResistance() / temp;		// Real part of the line admittance [Re(y)]
	PetscScalar admImag = -rBranch.GetReactance() / temp;		// Imag part of the line admittance [Im(y)]
	PetscScalar susImagHalf = rBranch.GetSusceptance() / 2;	// Half of the susceptance [B/2]

	// Y(j,j) = y + jB/2
	rGjj += admReal;
	rBjj += admImag + susImagHalf;

	// Distinguish different cases based on tap ratio and shift angle (speed optimized)
	PetscScalar angle = rBranch.GetAngle();
	PetscScalar ratio = rBranch.GetRatio();
	if (angle == 0) // No shift: T = ratio
	{
		if (ratio == 0) // No tap or shift: T = 1
		{
			// Y(i,i) = y + jB/2
			rGii += admReal;
			rBii += admImag + susImagHalf;
			// Y(i,j) = Y(j,i) = -y
			rGij -= admReal;
			rGji -= admReal;
			rBij -= admImag;
			rBji -= admImag;
		}
		else			// Tap but no shift: T = ratio
		{
			// Y(i,i) = y / |T|^2 + jB/2
			temp = ratio * ratio;
			rGii += admReal / temp;
			rBii += admImag / temp + susImagHalf;
			// Y(i,j) = Y(j,i) = -y / T
			temp = admReal / ratio;
			rGij -= temp;
			rGji -= temp;
			temp = admImag / ratio;
			rBij -= temp;
			rBji -= temp;
		}
	}
	else // Shift (angle != 0)
	{
		PetscScalar tapReal = cos(angle);		// Real part of the tap with shift [T]
		PetscScalar tapImag = sin(angle);		// Imag part of the tap with shift [T]
		if (ratio == 0) // Shift but no tap: T = e^(j*angle*pi/180), |T|^2 = 1
		{
			// Y(i,i) = y + jB/2
			rGii += admReal;
			rBii += admImag + susImagHalf;
		}
		else // Tap and shift
		{
			// Scale tap with shift to T/|T|^2
			tapReal /= ratio;
			tapImag /= ratio;
			// Y(i,i) = y / |T|^2 + jB/2
			temp = ratio * ratio;
			rGii += admReal / temp;
			rBii += admImag / temp + susImagHalf;
		}
		// Precalc products
		PetscScalar prodRR = admReal * tapReal;
		PetscScalar prodRI = admReal * tapImag;
		PetscScalar prodIR = admImag * tapReal;
		PetscScalar prodII = admImag * tapImag;
		// Y(i,j) = -y * (T / |T|^2)
		rGij -= prodRR - prodII;
		rBij -= prodIR + prodRI;
		// Y(j,i) = -y * (conj(T) / |T|^2)
		rGji -= prodRR + prodII;
		rBji -= prodIR - prodRI;
	}
} // AdmittanceMatrix::AddValues


POWERFLOWER_NAMESPACE_END
