/*  C version of setdir.c. This does not fill pits, it is intended for use AFTER
    d8.c which fills pits.  Input is a dem file name.  Output is
    slope and dinf angle file names
     
  David G Tarboton
  Utah State University   
     */

#include "lsm.h"
/* This include declares all necessary global variables  */
void incfall(int n, float **elev1, short *s1,int **spos, short *is, short *js);
void incrise(int n, float **elev1, short *s2,int **spos, short *is, 
			 short *js, short *dn);
void set2(int i,int j,float *fact,float **elev1, float **elev2);
void flatrout(int n,short *is1,short *js1,short *dn1, short *s1,short *s2,
			  int **spos,int iter,float **elev1,float **elev2, float *fact);

int setdir(char *demfile, char *angfile, char *slopefile, char *pfile)
{ 
  int err,filetype;
  float mval;
  double bndbox[4],csize;
  
/* define directions */
   d1[1]=0; d1[2]= -1; d1[3]= -1; d1[4]= -1; d1[5]=0; d1[6]=1; d1[7]=1; d1[8]=1;
   d2[1]=1; d2[2]=1; d2[3]=0; d2[4]= -1; d2[5]= -1; d2[6]= -1; d2[7]=0; d2[8]=1;


  err=gridread(demfile,(void ***)&elev,RPFLTDTYPE,&nx,&ny,&dx,&dy,
             bndbox,&csize,&mval,&filetype);
  if(err != 0)goto ERROR2;          

/*  allocate  dir and stack arrays  */
  dir = (short **) matalloc(nx, ny, RPSHRDTYPE);

  i1=0; i2=0; n1=nx; n2=ny;  /*  full grid  */ 
  err=setdfnoflood(mval);
   if(err != 0)goto ERROR1;
  err = gridwrite(pfile,(void **)dir,RPSHRDTYPE,nx,ny,dx,dy,
	  bndbox,csize,-1,filetype);

/*  allocate memory for slope and angle  */
  slope = (float **) matalloc(nx, ny, RPFLTDTYPE);
  ang = (float **) matalloc(nx, ny, RPFLTDTYPE);
  setdf2();
  err = gridwrite(slopefile,(void **)slope,RPFLTDTYPE,nx,ny,dx,dy,
             bndbox,csize,-1.,filetype);
  if (err != 0)goto ERROR3;
  err = gridwrite(angfile,(void **)ang,RPFLTDTYPE,nx,ny,dx,dy,
             bndbox,csize,-2.,filetype);
  if (err != 0)goto ERROR3;

/*  Wrapping up  */
  err=0;
ERROR3:
  free(slope[0]);
  free(slope);
  free(ang[0]);
  free(ang);
ERROR1:
  free(dir[0]);
  free(dir);
ERROR2:
  free(elev[0]);
  free(elev);
  return(err);
}

int setdfnoflood(float mval)
/*  This version is stripped of pit filling  */
{
  int i, j, k,ip, n, iter;
  float fact[9];
  short *s1, *s2, *is1, *js1, *dn1;  /*  variables for flat draining   */
  int **spos;
  float **elev2;


/*  Initialize boundaries  */
  for(i=i1; i< n1; i++)
  {
     dir[i][0]= -1;
     dir[i][n2-1]= -1;
  }
  for(i=i2; i< n2; i++)
  {
     dir[0][i]= -1;
     dir[n1-1][i]= -1;
  }
/*  iup=0; */
/*  initialize internal pointers */
  for(i=i2+1; i< n2-1; i++)for(j=i1+1; j<n1-1; j++)
  {
    if(elev[j][i] <= mval)dir[j][i]= -1;
    else dir[j][i] = 0;
  }
/*  Direction factors  */
  for(k=1; k<= 8; k++)
    fact[k]= (float) (1./sqrt(d1[k]*dy*d1[k]*dy+d2[k]*d2[k]*dx*dx));
/*  Set positive slope directions   */
   n=0;
   for(i=i2+1; i< n2-1; i++)
	  for(j=i1+1; j<n1-1; j++)
		  if(elev[j][i] > mval)
		  {
			  set(i,j,fact,mval);
			  if(dir[j][i] == 0)n++;
		  }
  if(n > 0)
  { 
/*  Now resolve flats following the Procedure of Garbrecht and Martz, Journal
   of Hydrology, 1997.  */
	  iter=1;
	  printf("Resolving %d Flats, Iteration: %d \n",n,iter);
      spos = (int **) matalloc(nx, ny, RPINTDTYPE);
      dn1 = (short *)malloc(sizeof(short) * n);
      is1 = (short *)malloc(sizeof(short) * n);
      js1 = (short *)malloc(sizeof(short) * n);
      s1 = (short *)malloc(sizeof(short) * n);
      s2 = (short *)malloc(sizeof(short) * n);
      elev2 = (float **)matalloc(nx,ny, RPFLTDTYPE);

  if(dn1 == NULL || is1 == NULL || js1 == NULL || s1 == NULL || 
	  spos == NULL || s2 == NULL || elev2 == NULL)
  {
     printf("Unable to allocate at iteration %d\n",iter);
  }
/*  Put unresolved pixels on stack  */
   ip=0;
   for(i=i2; i< n2; i++)
	  for(j=i1; j<n1; j++)
  {
    spos[j][i]=-1;   /*  Initialize stack position  */
	elev2[j][i]=0.;
    if(dir[j][i] == 0)
    {
      is1[ip]=i;
      js1[ip]=j;
	  dn1[ip]=0;
	  /*   Initialize the stage 1 and two arrays for flat routing   */
	  s1[ip] = 1;
      s2[ip] = 0;
	  spos[j][i]=ip;  /*  pointer for back tracking  */
	  ip++;
	  if(ip > n)printf("PROBLEM - Stack logic\n");
    }
  }
  flatrout(n,is1,js1,dn1,s1,s2,spos,iter,elev,elev2,fact);
  free(spos[0]); free(spos);
  free(elev2[0]); free(elev2);
  free(dn1);
  free(is1);
  free(js1);
  free(s1);
  free(s2);
  printf("Done iteration %d\nFlats resolved %d\n",iter,n);
  } /*  End if n > 0  */
   return(0);   /*  OK exit from setdir  */
}  /*  End setdfnoflood  */

void flatrout(int n,short *is1,short *js1,short *dn1, short *s1,short *s2,
			  int **spos,int iter,float **elev1,float **elev2, float *fact)
{
	int ip,nu,i,j;
  short *s12, *s22, *is2, *js2, *dn2;  /*  variables for flat draining   */
  int **spos2;
  float **elev3;

  incfall(n,elev1,s1,spos,is1,js1);
  incrise(n,elev1,s2,spos,is1,js1,dn1);
  for(ip=0; ip < n; ip++)
  {
	  elev2[js1[ip]][is1[ip]]=(float)(s1[ip]+s2[ip]);
  }
  nu=0;
  for(ip=0; ip < n; ip++)
  {
	  set2(is1[ip],js1[ip],fact,elev1,elev2);
	  if(dir[js1[ip]][is1[ip]] == 0)nu++;
  }
  if(nu > 0)
  { 
/*  Iterate Recursively   */
/*  Now resolve flats following the Procedure of Garbrecht and Martz, Journal
   of Hydrology, 1997.  */
	  iter=iter+1;
	  printf("Resolving %d Flats, Iteration: %d \n",nu,iter);
      spos2 = (int **) matalloc(nx, ny, RPINTDTYPE);
      dn2 = (short *)malloc(sizeof(short) * nu);
      is2 = (short *)malloc(sizeof(short) * nu);
      js2 = (short *)malloc(sizeof(short) * nu);
      s12 = (short *)malloc(sizeof(short) * nu);
      s22 = (short *)malloc(sizeof(short) * nu);
      elev3 = (float **)matalloc(nx,ny, RPFLTDTYPE);

  if(dn2 == NULL || is2 == NULL || js2 == NULL || s12 == NULL || 
	  spos2 == NULL || s22 == NULL || elev3 == NULL)
  {
     printf("Unable to allocate at iteration %d\n",iter);
  }
/*  Put unresolved pixels on stack  */
	ip=0;
	for(i=i2; i< n2; i++)
	  for(j=i1; j<n1; j++)
  {
    spos2[j][i]= -1;   /*  Initialize stack position  */
	elev3[j][i]=0.;
    if(dir[j][i] == 0)
    {
      is2[ip]=i;
      js2[ip]=j;
	  dn2[ip]=0;
	  /*   Initialize the stage 1 and two arrays for flat routing   */
	  s12[ip] = 1;
      s22[ip] = 0;
	  spos2[j][i]=ip;  /*  pointer for back tracking  */
	  ip++;
	  if(ip > nu)printf("PROBLEM - Stack logic\n");
   }
  }
  flatrout(nu,is2,js2,dn2,s12,s22,spos2,iter,elev2,elev3,fact);
  free(spos2[0]); free(spos2);
  free(elev3[0]); free(elev3);
  free(dn2);
  free(is2);
  free(js2);
  free(s12);
  free(s22);
  printf("Done iteration %d\nFlats resolved %d\n",iter,n);
  } /*  end if nu > 0  */
}   /*  End flatrout  */


void incfall(int n, float **elev1, short *s1,int **spos, short *is, short *js)
{
	/* This routine implements drainage towards lower areas - stage 1 */
	int done=0,donothing,k,ip,ninc;
	short st=1;
	float ed;
	while(done < 1)
	{
		done=1;
		ninc=0;
		for(ip=0; ip<n; ip++)
		{
/*			if 	adjacent to same level or lower that drains or 
				adjacent to pixel with s1 < st and dir not set
				do nothing  */
			donothing=0;
			for(k=1; k<=8; k++)
			{
				ed=elev1[js[ip]][is[ip]]-elev1[js[ip]+d2[k]][is[ip]+d1[k]];
				if(ed >= 0. && dir[js[ip]+d2[k]][is[ip]+d1[k]] != 0)donothing = 1;
				if(spos[js[ip]+d2[k]][is[ip]+d1[k]] >= 0)
					if(s1[spos[js[ip]+d2[k]][is[ip]+d1[k]]] < st 
					&& dir[js[ip]+d2[k]][is[ip]+d1[k]]  == 0)donothing = 1;
			}
			if(donothing == 0)
			{
				s1[ip]++;
				ninc++;
				done=0;
			}
		}
		st=st+1;
		printf("%d %d \n",ninc,n);
	}
}


void incrise(int n, float **elev1, short *s2,int **spos, short *is, 
			 short *js, short *dn)
{
	/*  This routine implements stage 2 drainage away from higher ground 
	dn is used to flag pixels still being incremented */
	int done=0,ip,k,ninc,nincold;
	float ed;
	nincold=0;
	while(done < 1)
	{
		done=1;
		ninc=0; 
		for(ip=0; ip<n; ip++)
		{
			for(k=1; k<=8; k++)
			{
				ed=elev1[js[ip]][is[ip]]-elev1[js[ip]+d2[k]][is[ip]+d1[k]];
				if(ed < 0.)dn[ip]=1;
				if(spos[js[ip]+d2[k]][is[ip]+d1[k]]>=0)
					if(s2[spos[js[ip]+d2[k]][is[ip]+d1[k]]] > 0)dn[ip] = 1;
			}
		}
		for(ip=0; ip<n; ip++)
		{
			s2[ip]=s2[ip]+dn[ip];
			ninc=ninc+dn[ip];
			if(dn[ip] == 0)done=0;  /*  if still some not being incremented continue 
									looping  */
		}
		printf("incfall %d %d\n",ninc,n);   
		if(ninc == nincold)done=1;   /*  If there are no new cells incremented  
									 stop - this is the case when a flat has
									 no higher ground around it.  */
/*		{
			int err,i,j;
			double bndbox[4],csize;
			for(i=0; i<ny; i++)for(j=0;j<nx;j++)elev1[j][i]=0.;
			for(ip=0; ip<n; ip++)
				elev1[js[ip]][is[ip]]=s2[ip];
			err=gridwrite("s2.asc",(void **)elev1, RPFLTDTYPE,nx,ny,dx,dy,bndbox,
				csize,-2.,0);
			for(ip=0; ip<n; ip++)
				elev1[js[ip]][is[ip]]=1;
			err=gridwrite("pit.asc",(void **)elev1, RPFLTDTYPE,nx,ny,dx,dy,bndbox,
				csize,-2.,0);
		}    */
		nincold=ninc;
	}
}

void set2(int i,int j,float *fact,float **elev1, float **elev2)
{
/*  This function sets directions based upon secondary elevations for
  assignment of flow directions across flats according to Garbrecht and Martz
  scheme   */
  float slope,slope2,smax;
  int k;
  smax=0.;
  for(k=1; k<=8; k++)
  {
    slope=fact[k]*(elev1[j][i]-elev1[j+d2[k]][i+d1[k]]);
	slope2=fact[k]*(elev2[j][i]-elev2[j+d2[k]][i+d1[k]]);
    if(slope2 > smax && slope >= 0.)  /*  Only if actual slope is positive or flat  */
    {
       smax=slope2;
       dir[j][i]=k;
    }
  }
}


/*

Converted from FORTRAN July 05, 1997  K. Tarbet

C
C---SUBROUTINE TO SET VARIABLE DIRECTIONS.
C    DIRECTIONS MEASURED ANTICLOCKWISE FROM EAST IN RADIANS.
C
*/
void setdf2(void )
{
float FANG[9];
float dxx[3];
int i,j;
float dd;

/*       INITIALISE ANGLES and slopes */

	for (i=0; i<nx; i++)
        {
         ang[i][0]=-2;
         ang[i][ny-1]=-2;
         slope[i][0]=-1;
         slope[i][ny-1]=-1;
         }

       for(i=0; i<ny; i++)
        {
         ang[0][i]=-2;
         ang[nx-1][i]=-2;
         slope[0][i]=-1;
         slope[nx-1][i]=-1;
        }

         FANG[1]=0.;
         FANG[2]=(float)atan2(dy,dx);
         FANG[3]=(float) PI2; 
         FANG[4]=2*FANG[3]-FANG[2];
         FANG[5]=2*FANG[3];
         FANG[6]=2*FANG[3]+FANG[2];
         FANG[7]=3*FANG[3];
         FANG[8]=4*FANG[3]-FANG[2];

/* --INITIALISE INTERNAL POINTERS (-VE POINTER INDICATES OUTSIDE DOMAIN) */

	for(i=1; i<ny-1; i++)
           for(j=1; j<nx-1; j++)
              {
                if(dir[j][i] <  0)   
                                         
                 {
                  ang[j][i]=-2.;
		  slope[j][i]=-1.; /*  !  -1 slope indicates no data */
                 }
                else if(dir[j][i] == 0)
                 {
                  ang[j][i]=-1.;  /*  DGT 9/2/97 since without   */
                  slope[j][i]=0.;  /*  pit filling dir = 0 is possible  */
/*  ang = -1 designates unfilled pit, ang = -2 designates no data.  */                  
                 }
                else
                 ang[j][i]=0.;
               }



/*    TEST ALL INTERNAL ELEVATIONS AND SET SLOPE AND ANGLE */

       dxx[1]=dx;
       dxx[2]=dy;
       dd=(float)sqrt(dx*dx+dy*dy);
	for(i=1; i<ny-1; i++)
    {
    SetWorkingStatus(); 
    for(j=1; j<nx-1; j++)
              {
               if(dir[j][i]>0 )
                   {
                   SET2(i,j,dxx,dd);
                   if(ang[j][i] < -0.5)
                        ang[j][i]=FANG[dir[j][i]];
                   }
              }
    }
 }

/*     SUBROUTINE TO RETURN SLOPE AND ANGLE OF STEEPEST DECENT IF POSITIVE */

void   SET2(int I, int J,float *DXX,float DD)
{

float SK[9];
float ANGLE[9];
float SMAX;
int K;
int KD;

/*    int ID1[]= {NULL,1,2,2,1,1,2,2,1 };
      int ID2[]= {NULL,2,1,1,2,2,1,1,2};
      int I1[] = {NULL,0,-1,-1,0,0,1,1,0 };
      int I2[] = {NULL,-1,-1,-1,-1,1,1,1,1};
      int J1[] = {NULL,1,0,0,-1,-1,0,0,1};
      int J2[] = {NULL,1,1,-1,-1,-1,-1,1,1};
   float  ANGC[]={NULL,0.,1.,1.,2.,2.,3.,3.,4.};
   float  ANGF[]={NULL,1.,-1.,1.,-1.,1.,-1.,1.,-1.}; */
      int ID1[]= {0,1,2,2,1,1,2,2,1 };
/*		int *ID1;  */
      int ID2[]= {0,2,1,1,2,2,1,1,2};
      int I1[] = {0,0,-1,-1,0,0,1,1,0 };
      int I2[] = {0,-1,-1,-1,-1,1,1,1,1};
      int J1[] = {0,1,0,0,-1,-1,0,0,1};
      int J2[] = {0,1,1,-1,-1,-1,-1,1,1};
   float  ANGC[]={0,0.,1.,1.,2.,2.,3.,3.,4.};
   float  ANGF[]={0,1.,-1.,1.,-1.,1.,-1.,1.,-1.};
 /*  ID1=ID1m-1;  */

       for(K=1; K<=8; K++)
       VSLOPE(elev[J][I],
              elev[J+J1[K]][I+I1[K]],
              elev[J+J2[K]][I+I2[K]],
              DXX[ID1[K]],DXX[ID2[K]],DD,&SK[K],&ANGLE[K]);

SMAX=0.;
KD=0;
      ang[J][I]=-1.;  /* USE -1 TO INDICATE DIRECTION NOT YET SET */
      for(K=1; K<=8; K++)
        {
        if(SK[K] >  SMAX)
           {
            SMAX=SK[K];
            KD=K;
           }
         }

      if(KD  > 0)
         ang[J][I]= (float) (ANGC[KD]*PI2+ANGF[KD]*ANGLE[KD]);
      slope[J][I]=SMAX;
}



void   VSLOPE(float E0,float E1, float E2,
              float D1,float D2,float DD,
              float *S,float *A)
{
/*    SUBROUTINE TO RETURN THE SLOPE AND ANGLE ASSOCIATED WITH A DEM PANEL */
float S1,S2,AD;


if(D1!=0)
      S1=(E0-E1)/D1;
if(D2!=0)
      S2=(E1-E2)/D2;

if(S2==0 && S1==0) *A=0;
      else
      *A= (float) atan2(S2,S1);
      AD= (float) atan2(D2,D1);
      if(*A  <   0.)
       {
       *A=0.;
       *S=S1;
       }
      else if(*A > AD)
       {
        *A=AD;
        *S=(E0-E2)/DD;
       }
      else
        *S= (float) sqrt(S1*S1+S2*S2);
 }