static char help[] = "This example demonstrates the DMCircuit interface for a radio network flow problem.\n\
                      The available solver options are in the alohaoptions file.\n\
                      The file links.txt contains the network description.\n\
                      Run this program: mpiexec -n <n> ./aloha\n";

/*
Contributed example - Florian Meier <florian.meier@koala.de>

This computes a greatly simplified model of a radio network.
Each node in the network generates packets conforming to a poisson distribution (given rate R_gen).
These packets are routed along fixed paths to a sink. When two packets are send at the same time,
they are dropped. No channel sensing or retransmission takes place.

Each active link (i.e. each radio link that will eventually transmit data) is modeled as a vertex.
Between those vertices there exist two kinds of relations:
A relation of type REL_INTERFERENCE(affected, source), implies that whenever the source link tries
to send at the same time than the affected link, the packet will collide and is not received properly.
A relation of type REL_INFLOW(affected, source) implies that each packet that is send via the source
link and does not collide, adds up to the amount of packets send via the affected link.

The overall rate of accessing a radio link i is 
R_access,i = R_gen,i + SUM_OVER_INFLOWS( R_success,j )

The probability that no other link interfering with i accesses the channel at a given point in time is
P_free,i = PRODUCT_OVER_INTERFERERS( 1 - R_access,j ) 

Finally, the overall rate of successful transmissions over link i is
R_success,i = R_access,i * P_free,i

The input file is structured as follows:
The first line contains L, the number of links, E, the number of edges
Then L lines follow with the packet generation rate for each link.  
The subsequent E lines describe the affected and the source link connectivity.
The final E lines represent the edge relations (inflow, interference).
*/

#include <petsc.h>
#include <petscdmcircuit.h>

#include <fstream>
#include <iostream>

using namespace std;

enum {
  REL_INFLOW,
  REL_INTERFERENCE
};

enum {
  VAR_ACCESS,
  VAR_SUCCESS,
  VAR_NVARS
};

struct _p_LINKDATA{
  PetscInt      num;
  PetscScalar 	packet_generation;
};

typedef struct _p_LINKDATA *LINKDATA;

struct _p_RELATIONDATA{
  PetscInt      source;
  PetscInt      affected;
  PetscInt	type;
};

typedef struct _p_RELATIONDATA *RELATIONDATA;

typedef struct {
  LINKDATA     links;
  RELATIONDATA relations;
} PFDATA;


PetscMPIInt rank;

#undef __FUNCT__
#define __FUNCT__ "FormFunction"
PetscErrorCode FormFunction(SNES snes,Vec X, Vec F,void *appctx)
{
  PetscErrorCode ierr;
  DM             circuitdm;
  Vec           localX,localF;
  PetscInt      e;
  PetscInt      v,vStart,vEnd,vaffected,vsource;
  const PetscScalar *xarr;
  PetscScalar   *farr;
  PetscInt      offset,offsetsource,offsetlink,offsetrel;
  DMCircuitComponentGenericDataType *arr;

  PetscFunctionBegin;
  ierr = SNESGetDM(snes,&circuitdm);CHKERRQ(ierr);
  ierr = DMGetLocalVector(circuitdm,&localX);CHKERRQ(ierr);
  ierr = DMGetLocalVector(circuitdm,&localF);CHKERRQ(ierr);
  ierr = VecSet(F,0.0);CHKERRQ(ierr);

  ierr = DMGlobalToLocalBegin(circuitdm,X,INSERT_VALUES,localX);CHKERRQ(ierr);
  ierr = DMGlobalToLocalEnd(circuitdm,X,INSERT_VALUES,localX);CHKERRQ(ierr);

  ierr = DMGlobalToLocalBegin(circuitdm,F,INSERT_VALUES,localF);CHKERRQ(ierr);
  ierr = DMGlobalToLocalEnd(circuitdm,F,INSERT_VALUES,localF);CHKERRQ(ierr);

  ierr = VecGetArrayRead(localX,&xarr);CHKERRQ(ierr);
  ierr = VecGetArray(localF,&farr);CHKERRQ(ierr);

  ierr = DMCircuitGetVertexRange(circuitdm,&vStart,&vEnd);CHKERRQ(ierr);
  ierr = DMCircuitGetComponentDataArray(circuitdm,&arr);CHKERRQ(ierr);

  for (v=vStart; v < vEnd; v++) {
    PetscBool   ghostvtex;
    ierr = DMCircuitIsGhostVertex(circuitdm,v,&ghostvtex);CHKERRQ(ierr);
    if (ghostvtex) {
      continue;
    }
    PetscInt keyv;
    ierr = DMCircuitGetComponentTypeOffset(circuitdm,v,0,&keyv,&offsetlink);CHKERRQ(ierr);

    LINKDATA link = (LINKDATA)(arr+offsetlink);

    PetscScalar inflow = 0;
    inflow = link->packet_generation;    

    PetscScalar Pfree = 1;

    PetscInt nconnedges;
    const PetscInt *connedges;
    ierr = DMCircuitGetSupportingEdges(circuitdm,v,&nconnedges,&connedges);CHKERRQ(ierr);
    for (PetscInt i = 0; i < nconnedges; i++) {
      e = connedges[i];
      PetscInt keye;
      ierr = DMCircuitGetComponentTypeOffset(circuitdm,e,0,&keye,&offsetrel);CHKERRQ(ierr);
      RELATIONDATA relation = (RELATIONDATA)(arr+offsetrel);

      const PetscInt *cone;
      ierr = DMCircuitGetConnectedNodes(circuitdm,e,&cone);CHKERRQ(ierr);
      vaffected = cone[0];
      vsource   = cone[1];

      switch (relation->type) {
      case REL_INTERFERENCE:
	PetscInt vneighbor;
	if (vaffected == v) vneighbor = vsource;
	else vneighbor = vaffected;
	ierr = DMCircuitGetVariableOffset(circuitdm,vneighbor,&offsetsource);CHKERRQ(ierr);
	Pfree *= 1 - xarr[offsetsource+VAR_ACCESS];
	break;
      case REL_INFLOW:
	if (vaffected == v) {
	  ierr = DMCircuitGetVariableOffset(circuitdm,vsource,&offsetsource);CHKERRQ(ierr);
	  inflow += xarr[offsetsource+VAR_SUCCESS];
	  break;
	}
      }
    }

    ierr = DMCircuitGetVariableOffset(circuitdm,v,&offset);CHKERRQ(ierr);
    farr[offset+VAR_ACCESS] = inflow - xarr[offset+VAR_ACCESS];
    farr[offset+VAR_SUCCESS] = xarr[offset+VAR_ACCESS]*Pfree - xarr[offset+VAR_SUCCESS];
  }

  ierr = VecRestoreArrayRead(localX,&xarr);CHKERRQ(ierr);
  ierr = VecRestoreArray(localF,&farr);CHKERRQ(ierr);
  ierr = DMRestoreLocalVector(circuitdm,&localX);CHKERRQ(ierr);

  ierr = DMLocalToGlobalBegin(circuitdm,localF,ADD_VALUES,F);CHKERRQ(ierr);
  ierr = DMLocalToGlobalEnd(circuitdm,localF,ADD_VALUES,F);CHKERRQ(ierr);
  ierr = DMRestoreLocalVector(circuitdm,&localF);CHKERRQ(ierr);
  //  ierr = VecView(F,0);
  PetscFunctionReturn(0);
}


#undef __FUNCT__
#define __FUNCT__ "FormJacobian"
PetscErrorCode FormJacobian(SNES snes,Vec X, Mat *J,Mat *Jpre,MatStructure *flg,void *appctx)
{
  PetscErrorCode ierr;
  DM            circuitdm;
  Vec           localX;
  PetscInt      e;
  PetscInt      v,vStart,vEnd,vaffected,vsource;
  const PetscScalar *xarr;
  PetscInt      offsetrel;
  DMCircuitComponentGenericDataType *arr;
  PetscInt      row[1],col[1];
  PetscScalar   values[1];

  PetscFunctionBegin;
  *flg = SAME_NONZERO_PATTERN;
  ierr = MatZeroEntries(*J);CHKERRQ(ierr);

  ierr = SNESGetDM(snes,&circuitdm);CHKERRQ(ierr);
  ierr = DMGetLocalVector(circuitdm,&localX);CHKERRQ(ierr);

  ierr = DMGlobalToLocalBegin(circuitdm,X,INSERT_VALUES,localX);CHKERRQ(ierr);
  ierr = DMGlobalToLocalEnd(circuitdm,X,INSERT_VALUES,localX);CHKERRQ(ierr);

  ierr = VecGetArrayRead(localX,&xarr);CHKERRQ(ierr);

  ierr = DMCircuitGetVertexRange(circuitdm,&vStart,&vEnd);CHKERRQ(ierr);
  ierr = DMCircuitGetComponentDataArray(circuitdm,&arr);CHKERRQ(ierr);

  for (v=vStart; v < vEnd; v++) {
    PetscInt    offset,goffset,offsetsource,goffsetsource;
    PetscBool   ghostvtex;

    ierr = DMCircuitIsGhostVertex(circuitdm,v,&ghostvtex);CHKERRQ(ierr);
    if (ghostvtex) {
      continue;
    }

    ierr = DMCircuitGetVariableGlobalOffset(circuitdm,v,&goffset);CHKERRQ(ierr);

    row[0] = goffset+VAR_ACCESS;
    col[0] = goffset+VAR_ACCESS;
    values[0] = -1.0;
    ierr = MatSetValues(*J,1,row,1,col,values,ADD_VALUES);CHKERRQ(ierr);
    row[0] = goffset+VAR_SUCCESS;
    col[0] = goffset+VAR_SUCCESS;
    values[0] = -1.0;
    ierr = MatSetValues(*J,1,row,1,col,values,ADD_VALUES);CHKERRQ(ierr);

    // Calculate Pfree
    // TODO Pfree was already calculated in FromFunction, can I access that?
    PetscScalar Pfree = 1;

    PetscInt nconnedges;
    const PetscInt *connedges;
    ierr = DMCircuitGetSupportingEdges(circuitdm,v,&nconnedges,&connedges);CHKERRQ(ierr);
    for (PetscInt i = 0; i < nconnedges; i++) {
      e = connedges[i];

      const PetscInt *cone;
      ierr = DMCircuitGetConnectedNodes(circuitdm,e,&cone);CHKERRQ(ierr);
      vaffected = cone[0];
      vsource   = cone[1];

      if (vaffected == v) {
	PetscInt keye;
        ierr = DMCircuitGetComponentTypeOffset(circuitdm,e,0,&keye,&offsetrel);CHKERRQ(ierr);
        RELATIONDATA relation = (RELATIONDATA)(arr+offsetrel);

        ierr = DMCircuitGetVariableOffset(circuitdm,vsource,&offsetsource);CHKERRQ(ierr);

        if (relation->type == REL_INTERFERENCE) {
          Pfree *= 1 - xarr[offsetsource+VAR_ACCESS];
        }
      }
    }

    row[0] = goffset+VAR_SUCCESS;
    col[0] = goffset+VAR_ACCESS;
    values[0] = Pfree;
    ierr = MatSetValues(*J,1,row,1,col,values,ADD_VALUES);CHKERRQ(ierr);

    // Set the derivatives
    ierr = DMCircuitGetVariableOffset(circuitdm,v,&offset);CHKERRQ(ierr);
    PetscScalar inflow = xarr[offset+VAR_ACCESS];

    for (PetscInt i = 0; i < nconnedges; i++) {
      e = connedges[i];

      const PetscInt *cone;
      ierr = DMCircuitGetConnectedNodes(circuitdm,e,&cone);CHKERRQ(ierr);
      vaffected = cone[0];
      vsource   = cone[1];

      if (vaffected == v) {
	PetscInt keye;
        ierr = DMCircuitGetComponentTypeOffset(circuitdm,e,0,&keye,&offsetrel);CHKERRQ(ierr);
        RELATIONDATA relation = (RELATIONDATA)(arr+offsetrel);

        ierr = DMCircuitGetVariableOffset(circuitdm,vsource,&offsetsource);CHKERRQ(ierr);
        ierr = DMCircuitGetVariableGlobalOffset(circuitdm,vsource,&goffsetsource);CHKERRQ(ierr);

        switch (relation->type) {
          case REL_INTERFERENCE:
            // TODO Ok to set the entries one by one or is there a better way?
            row[0] = goffset+VAR_SUCCESS;
            col[0] = goffsetsource+VAR_ACCESS;
            values[0] = -inflow*(Pfree/(1-xarr[offsetsource+VAR_ACCESS]));
            ierr = MatSetValues(*J,1,row,1,col,values,ADD_VALUES);CHKERRQ(ierr);
            break;
          case REL_INFLOW:
            row[0] = goffset+VAR_ACCESS;
            col[0] = goffsetsource+VAR_SUCCESS;
            values[0] = 1.0;
            ierr = MatSetValues(*J,1,row,1,col,values,ADD_VALUES);CHKERRQ(ierr);
            break;
        }
      }
    }
  }

  ierr = VecRestoreArrayRead(localX,&xarr);CHKERRQ(ierr);
  ierr = DMRestoreLocalVector(circuitdm,&localX);CHKERRQ(ierr);

  ierr = MatAssemblyBegin(*J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
  ierr = MatAssemblyEnd(*J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);

  PetscFunctionReturn(0);
}



#undef __FUNCT__
#define __FUNCT__ "SetInitialValues"
PetscErrorCode SetInitialValues(DM circuitdm,Vec X,void *appctx)
{
  PetscErrorCode ierr;
  PetscInt       v, vStart, vEnd;
  Vec            localX;
  PetscScalar    *xarr;
  DMCircuitComponentGenericDataType *arr;
  
  PetscFunctionBegin;
  ierr = DMCircuitGetVertexRange(circuitdm,&vStart,&vEnd);CHKERRQ(ierr);

  ierr = DMGetLocalVector(circuitdm,&localX);CHKERRQ(ierr);

  ierr = VecSet(X,0.0);CHKERRQ(ierr);
  ierr = DMGlobalToLocalBegin(circuitdm,X,INSERT_VALUES,localX);CHKERRQ(ierr);
  ierr = DMGlobalToLocalEnd(circuitdm,X,INSERT_VALUES,localX);CHKERRQ(ierr);

  ierr = DMCircuitGetComponentDataArray(circuitdm,&arr);CHKERRQ(ierr);
  ierr = VecGetArray(localX,&xarr);CHKERRQ(ierr);
  for (v = vStart; v < vEnd; v++) {
    PetscBool   ghostvtex;
    ierr = DMCircuitIsGhostVertex(circuitdm,v,&ghostvtex);CHKERRQ(ierr);
    if (ghostvtex) {
      continue;
    }

    PetscInt offsetlink, offset;
    PetscInt keyv;
    ierr = DMCircuitGetVariableOffset(circuitdm,v,&offset);CHKERRQ(ierr);
    ierr = DMCircuitGetComponentTypeOffset(circuitdm,v,0,&keyv,&offsetlink);CHKERRQ(ierr);

    LINKDATA link = (LINKDATA)(arr+offsetlink);

    xarr[offset+VAR_ACCESS] = link->packet_generation;
    xarr[offset+VAR_SUCCESS] = link->packet_generation;
  }

  ierr = VecRestoreArray(localX,&xarr);CHKERRQ(ierr);
  ierr = DMLocalToGlobalBegin(circuitdm,localX,ADD_VALUES,X);CHKERRQ(ierr);
  ierr = DMLocalToGlobalEnd(circuitdm,localX,ADD_VALUES,X);CHKERRQ(ierr);
  ierr = DMRestoreLocalVector(circuitdm,&localX);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}


#undef __FUNCT__
#define __FUNCT__ "main"
int main(int argc,char ** argv)
{
  string	       inputFile = "links1.txt";
  PetscErrorCode       ierr;
  PFDATA               pfdata;
  PetscInt             numEdges=0,numVertices=0;
  int                  *edges = NULL;
  PetscInt             i;  
  DM                   circuitdm;
  PetscInt             componentkey[2];
  PetscLogStage        stage1,stage2;
  PetscInt             size;

  PetscInitialize(&argc,&argv,"alohaoptions",help);

  ierr = MPI_Comm_rank(PETSC_COMM_WORLD,&rank);CHKERRQ(ierr);

  /* Create an empty circuit object */
  ierr = DMCircuitCreate(PETSC_COMM_WORLD,&circuitdm);CHKERRQ(ierr);

  /* Register the components in the circuit */
  ierr = DMCircuitRegisterComponent(circuitdm,"linkstruct",sizeof(struct _p_LINKDATA),&componentkey[0]);CHKERRQ(ierr);
  ierr = DMCircuitRegisterComponent(circuitdm,"relationstruct",sizeof(struct _p_RELATIONDATA),&componentkey[1]);CHKERRQ(ierr);

  ierr = PetscLogStageRegister("Read Data",&stage1);CHKERRQ(ierr);
  PetscLogStagePush(stage1);

  /* READ THE DATA */
  if (!rank) {
    /* Only rank 0 reads the data */
    ifstream linksF(inputFile.c_str());

    linksF >> numVertices >> numEdges;

    ierr = PetscMalloc(numVertices*sizeof(struct _p_LINKDATA),&pfdata.links);CHKERRQ(ierr);
    ierr = PetscMalloc(numEdges*sizeof(struct _p_RELATIONDATA),&pfdata.relations);CHKERRQ(ierr);

    for(int i = 0; i < numVertices; i++) {
      linksF >> pfdata.links[i].packet_generation;
      pfdata.links[i].num = i;
    }

    ierr = PetscMalloc(2*numEdges*sizeof(PetscInt),&edges);CHKERRQ(ierr);

    for(int i = 0; i < numEdges; i++) {
      linksF >> pfdata.relations[i].affected >> pfdata.relations[i].source;

      edges[2*i] = pfdata.relations[i].affected;
      edges[2*i+1] = pfdata.relations[i].source;
    }

    for(int i = 0; i < numEdges; i++) {
      linksF >> pfdata.relations[i].type;
    }

    linksF.close();
  }

  PetscLogStagePop();
  MPI_Barrier(PETSC_COMM_WORLD);

  ierr = PetscLogStageRegister("Create circuit",&stage2);CHKERRQ(ierr);
  PetscLogStagePush(stage2);
  /* Set number of nodes/edges */
  ierr = DMCircuitSetSizes(circuitdm,numVertices,numEdges,PETSC_DETERMINE,PETSC_DETERMINE);CHKERRQ(ierr);
  /* Add edge connectivity */
  ierr = DMCircuitSetEdgeList(circuitdm,edges);CHKERRQ(ierr);
  /* Set up the circuit layout */
  ierr = DMCircuitLayoutSetUp(circuitdm);CHKERRQ(ierr);

  if (!rank) {
    ierr = PetscFree(edges);CHKERRQ(ierr);
  }

  /* Add circuit components */
  PetscInt eStart, eEnd, vStart, vEnd;

  ierr = DMCircuitGetVertexRange(circuitdm,&vStart,&vEnd);CHKERRQ(ierr);
  for (i = vStart; i < vEnd; i++) {
    ierr = DMCircuitAddComponent(circuitdm,i,componentkey[0],&pfdata.links[i-vStart]);CHKERRQ(ierr);

    /* Add number of variables */
    ierr = DMCircuitAddNumVariables(circuitdm,i,VAR_NVARS);CHKERRQ(ierr);
  }

  ierr = DMCircuitGetEdgeRange(circuitdm,&eStart,&eEnd);CHKERRQ(ierr);
  for (i = eStart; i < eEnd; i++) {
    ierr = DMCircuitAddComponent(circuitdm,i,componentkey[1],&pfdata.relations[i-eStart]);CHKERRQ(ierr);
  }

  /* Set up DM for use */
  ierr = DMSetUp(circuitdm);CHKERRQ(ierr);

  if (!rank) {
    ierr = PetscFree(pfdata.links);CHKERRQ(ierr);
    ierr = PetscFree(pfdata.relations);CHKERRQ(ierr);
  }


  ierr = MPI_Comm_size(PETSC_COMM_WORLD,&size);CHKERRQ(ierr);
  if (size > 1) {
    DM distcircuitdm;
    /* Circuit partitioning and distribution of data */
    ierr = DMCircuitDistribute(circuitdm,"metis",1,&distcircuitdm);CHKERRQ(ierr);
    ierr = DMDestroy(&circuitdm);CHKERRQ(ierr);
    circuitdm = distcircuitdm;
  }

  DMCircuitComponentGenericDataType *arr;
  ierr = DMCircuitGetComponentDataArray(circuitdm,&arr);CHKERRQ(ierr);

  ierr = DMCircuitGetVertexRange(circuitdm,&vStart,&vEnd);CHKERRQ(ierr);
  for (i=vStart; i < vEnd; i++) {
    PetscInt keyv,offsetlink;
    ierr = DMCircuitGetComponentTypeOffset(circuitdm,i,0,&keyv,&offsetlink);CHKERRQ(ierr);
    LINKDATA link = (LINKDATA)(arr+offsetlink);
    PetscBool   ghostvtex;
    ierr = DMCircuitIsGhostVertex(circuitdm,i,&ghostvtex);CHKERRQ(ierr);

    ierr = PetscPrintf(PETSC_COMM_SELF,"Rank[%d]: Vertex %d, Ghostvertex = %d\n",rank,link->num,ghostvtex);CHKERRQ(ierr);

  }
  MPI_Barrier(PETSC_COMM_WORLD);
  ierr = DMCircuitGetEdgeRange(circuitdm,&eStart,&eEnd);CHKERRQ(ierr);
  for (i=eStart; i < eEnd; i++) {
    PetscInt keye,offsetrel;
    ierr = DMCircuitGetComponentTypeOffset(circuitdm,i,0,&keye,&offsetrel);CHKERRQ(ierr);
    RELATIONDATA relation = (RELATIONDATA)(arr+offsetrel);
    ierr = PetscPrintf(PETSC_COMM_SELF,"Rank[%d]: %d -- %d\n",rank,relation->affected,relation->source);CHKERRQ(ierr);
  }

  PetscLogStagePop();

  Vec X,F;
  ierr = DMCreateGlobalVector(circuitdm,&X);CHKERRQ(ierr);
  ierr = VecDuplicate(X,&F);CHKERRQ(ierr);

  Mat J;
  ierr = DMCreateMatrix(circuitdm,&J);CHKERRQ(ierr);
  ierr = MatSetOption(J,MAT_NEW_NONZERO_ALLOCATION_ERR,PETSC_FALSE);CHKERRQ(ierr);

  ierr = SetInitialValues(circuitdm,X,NULL);CHKERRQ(ierr);

  SNES snes;
  /* HOOK UP SOLVER */
  ierr = SNESCreate(PETSC_COMM_WORLD,&snes);CHKERRQ(ierr);
  ierr = SNESSetDM(snes,circuitdm);CHKERRQ(ierr);
  ierr = SNESSetFunction(snes,F,FormFunction,NULL);CHKERRQ(ierr);
  ierr = SNESSetJacobian(snes,J,J,FormJacobian,NULL);CHKERRQ(ierr);
  ierr = SNESSetFromOptions(snes);CHKERRQ(ierr);

  ierr = SNESSolve(snes,NULL,X);CHKERRQ(ierr);
  ierr = VecView(X,NULL);CHKERRQ(ierr);

  ierr = VecDestroy(&X);CHKERRQ(ierr);
  ierr = VecDestroy(&F);CHKERRQ(ierr);
  ierr = MatDestroy(&J);CHKERRQ(ierr);

  ierr = SNESDestroy(&snes);CHKERRQ(ierr);
  ierr = DMDestroy(&circuitdm);CHKERRQ(ierr);

  PetscFinalize();
  return 0;
}