/* This file is part of Cloudy and is copyright (C) 1978-2004 by Gary J. Ferland.
 * For conditions of distribution and use, see copyright notice in license.txt */
/*atom_levelN compute an arbitrary N level atom */
#include "cddefines.h"
#include "thermal.h"
#include "physconst.h"
#include "phycon.h"
#include "trace.h"
#include "lapack.h"
#include "atoms.h"

void atom_levelN(
	/* nlev is the number of levels to compute*/ 
	long int nlev, 
	/* ABUND is total abundance of species, used for nth equation */
	float abund, 
	/* G(nlev) is stat weight of levels */
	double g[], 
	/* EX(nlev) is excitation potential of levels, deg K
	 * 0 for first one, NOT d(ENER), but energy rel to ground */
	double ex[], 
	/* populations of each level as deduced here */
	double pops[], 
	/* departure coefficient, derived below */
	double depart[],
	/* data[ilo][iho] are net A's, from up to low, A * esc prob
	 * data[ihi][ilo] is col str or collision rate from ihi to ilo */
	double ***data, 
	/* dest(ilo,ihi) is destruction rate, from up to low, A * dest prob,
	 * asserts confirm that ihi,lo is zero */
	double ***dest, 
	/* pump(lo, hi) is pumping rate from lower to upper level (s^-1), (hi,lo) must be zero  */
	double ***pump, 
	/* collision rates (s^-1), evaluated here and returned for cooling by calling function,
	 * unless following flag is true.  If true then calling function has already filled
	 * in these rates.  CollRate[i][j] is rate from i to j */
	double ***CollRate,
	/* this is an additional creation rate from continuum, normally zero, units cm-3 s-1 */
	double source[] ,
	/* this is an additional destruction rate to continuum, normally zero, units s-1 */
	double destroy[] ,
	/* flag saying whether CollRate already done, or we need to do it here,
	 * this is stored in data)[ihi][ilo] as either downward rate or collis strength*/
	int lgCollRateDone,
	/* 2D array of integer indices for line energy in continuum array, used
	 * for setting ots rates.  ignored if zero, indices are [low][high],
	 * indices are on f not c scale
	long int ***ipdest,  */
	/* total cooling and its derivative, set here but nothing done with it*/
	double *cooltl, 
	double *coolder, 
	/* string used to identify calling program in case of error */
	char *chLabel, 
	/* lgNegPop flag indicating what we have done
	 * positive if negative populations occurred
	 * zero if normal calculation done
	 * negative if too cold (for some atoms other routine will be called in this case) */
	int *lgNegPop,
	/* option to print debug information */
	int lgDeBug )
{
	long int level, 
	  ihi, 
	  ilo, 
	  i,
	  j; 

	int lgHomogeneous;

	double cool1;
	double 
	  *bvec, 
	  *work, 
	  *amat, 
	  **excit,
	  sum;

	int32 ner;
	int32 *ipiv;

#	ifdef DEBUG_FUN
	fputs( "<+>atom_levelN()\n", debug_fp );
#	endif
	/* check for zero abundance and exit if so */
	if( abund <= 0. )
	{
		*cooltl = 0.;
		*coolder = 0.;
		/* says calc was ok */
		*lgNegPop = FALSE;

		for( level=0; level < nlev; level++ )
		{
			pops[level] = 0.;
			depart[level] = 0.;
		}

		depart[0] = 1.;

		/* there are TWO abort returns in this sub,
		 * this one is for zero abundance */
#		ifdef DEBUG_FUN
		fputs( " <->atom_levelN()\n", debug_fp );
#		endif
		return;
	}

	/* >>chng 00 oct 13, MALLOC and free space for arrays */
	if( (bvec = (double *)MALLOC( sizeof(double)*(size_t)(nlev+1) )) == NULL )
		BadMalloc();
	if( (work = (double *)MALLOC( sizeof(double)*(size_t)(nlev) )) == NULL )
		BadMalloc();
	if( (ipiv = (int32 *)MALLOC( sizeof(int32)*(size_t)(nlev) )) == NULL )
		BadMalloc();
	/* amat converted to 1d array */
#	ifdef AMAT
#		undef AMAT
#	endif
#	define AMAT(I_,J_)	(*(amat+(I_)*(nlev)+(J_)))
	/* MALLOC space for the  1-d array */
	if( (amat=(double*)MALLOC(sizeof(double)*(size_t)(nlev*nlev))) == NULL )
		BadMalloc();

	/* create space for the 2D arrays */
	if( (excit = ((double **)MALLOC((size_t)(nlev+1)*sizeof(double *)))) == NULL )
		BadMalloc();
	/* create space for the 2D arrays */
	for( i=0; i<(nlev+1); ++i )
	{
		if( (excit[i] = ((double *)MALLOC((size_t)(nlev+1)*sizeof(double )))) == NULL )
			BadMalloc();
	}

	/*begin sanity checks
	 * excitation temperature of lowest level must be zero */
	ASSERT( ex[0] == 0. );

	/* both dest and pump (low, hi) are not used, should be zero */
	for( ihi=1; ihi < nlev; ihi++ )
	{
		for( ilo=0; ilo < ihi; ilo++ )
		{
			/* zero ots destruction rate */
			ASSERT( (*dest)[ihi][ilo] == 0. );
			ASSERT( (*pump)[ihi][ilo] == 0. );
			/*ASSERT( (*ipdest)[ihi][ilo] == 0 );*/
			/* these are collision strenghts and escape rates */
			ASSERT( (*data)[ihi][ilo] >= 0 );
			ASSERT( (*data)[ilo][ihi] >= 0 );
			ASSERT( (*dest)[ilo][ihi] >= 0 );
		}
	}

	/* data(level,level) should be zero */
	for( ihi=0; ihi < nlev; ihi++ )
	{
		ASSERT( (*data)[ihi][ihi] == 0 );
	}
	/*end sanity checks */

	if( lgDeBug || (trace.lgTrace && trace.lgTrLevN) )
	{
		fprintf( ioQQQ, " atom_levelN trace printout for atom=%s with tot abund %e \n", chLabel, abund);
		fprintf( ioQQQ, " dest\n" );

		fprintf( ioQQQ, "  hi  lo" );

		for( ilo=0; ilo < nlev-1; ilo++ )
		{
			fprintf( ioQQQ, "%4ld      ", ilo );
		}
		fprintf( ioQQQ, "      \n" );

		for( ihi=1; ihi < nlev; ihi++ )
		{
			fprintf( ioQQQ, "%3ld", ihi );
			for( ilo=0; ilo < ihi; ilo++ )
			{
				fprintf( ioQQQ, "%10.2e", (*dest)[ilo][ihi] );
			}
			fprintf( ioQQQ, "\n" );
		}

		fprintf( ioQQQ, " A*esc\n" );
		fprintf( ioQQQ, "  hi  lo" );
		for( ilo=0; ilo < nlev-1; ilo++ )
		{
			fprintf( ioQQQ, "%4ld      ", ilo );
		}
		fprintf( ioQQQ, "      \n" );

		for( ihi=1; ihi < nlev; ihi++ )
		{
			fprintf( ioQQQ, "%3ld", ihi );
			for( ilo=0; ilo < ihi; ilo++ )
			{
				fprintf( ioQQQ, "%10.2e", (*data)[ilo][ihi] );
			}
			fprintf( ioQQQ, "\n" );
		}

		fprintf( ioQQQ, " pump\n" );

		fprintf( ioQQQ, "  hi  lo" );
		for( ilo=0; ilo < nlev-1; ilo++ )
		{
			fprintf( ioQQQ, "%4ld      ", ilo );
		}
		fprintf( ioQQQ, "      \n" );

		for( ihi=1; ihi < nlev; ihi++ )
		{
			fprintf( ioQQQ, "%3ld", ihi );
			for( ilo=0; ilo < ihi; ilo++ )
			{
				fprintf( ioQQQ, "%10.2e", (*pump)[ilo][ihi] );
			}
			fprintf( ioQQQ, "\n" );
		}

		fprintf( ioQQQ, " coll str\n" );
		fprintf( ioQQQ, "  hi  lo" );
		for( ilo=0; ilo < nlev-1; ilo++ )
		{
			fprintf( ioQQQ, "%4ld      ", ilo );
		}
		fprintf( ioQQQ, "      \n" );

		for( ihi=1; ihi < nlev; ihi++ )
		{
			fprintf( ioQQQ, "%3ld", ihi );
			for( ilo=0; ilo < ihi; ilo++ )
			{
				fprintf( ioQQQ, "%10.2e", (*data)[ihi][ilo] );
			}
			fprintf( ioQQQ, "\n" );
		}

		fprintf( ioQQQ, " coll rate\n" );
		fprintf( ioQQQ, "  hi  lo" );
		for( ilo=0; ilo < nlev-1; ilo++ )
		{
			fprintf( ioQQQ, "%4ld      ", ilo );
		}
		fprintf( ioQQQ, "      \n" );

		for( ihi=1; ihi < nlev; ihi++ )
		{
			fprintf( ioQQQ, "%3ld", ihi );
			for( ilo=0; ilo < ihi; ilo++ )
			{
				fprintf( ioQQQ, "%10.2e", (*CollRate)[ihi][ilo] );
			}
			fprintf( ioQQQ, "\n" );
		}
	}


	for( ilo=0; ilo < (nlev - 1); ilo++ )
	{
		for( ihi=ilo + 1; ihi < nlev; ihi++ )
		{
			excit[ilo][ihi] = sexp((ex[ihi]-ex[ilo])/phycon.te);
		}
	}

	if( trace.lgTrace && trace.lgTrLevN )
	{
		fprintf( ioQQQ, " excit, te=%10.2e\n", phycon.te );
		fprintf( ioQQQ, "  hi  lo" );

		for( ilo=1; ilo <= nlev; ilo++ )
		{
			fprintf( ioQQQ, "%4ld      ", ilo );
		}
		fprintf( ioQQQ, "      \n" );

		for( ihi=0; ihi < nlev; ihi++ )
		{
			fprintf( ioQQQ, "%3ld", ihi );
			for( ilo=0; ilo < nlev; ilo++ )
			{
				fprintf( ioQQQ, "%10.2e", excit[ilo][ihi] );
			}
			fprintf( ioQQQ, "\n" );
		}
	}

	/* punt if total excitation rate is zero */
	/* >>chng 04 sep 16, add test on non-zero source */
	if( (excit[0][nlev-1] + (*pump)[0][nlev-1] < 1e-13 ) && (source[nlev-1]==0.) )
	{
		*cooltl = 0.;
		*coolder = 0.;
		/* special flag saying too cool for highest level to be computed.
		 * some routines will call another routine for lower levels in this case */
		*lgNegPop = -1;

		for( level=1; level < nlev; level++ )
		{
			pops[level] = 0.;
			/* these are off by one - lowest index is zero */
			depart[level] = 0.;
		}

		/* everything in ground */
		pops[0] = abund;
		depart[0] = 1.;

		free( bvec );
		free( work );
		free( ipiv );
		free( amat );
		for( i=0; i < (nlev+1); ++i )
		{
			free( excit[i] );
		}
		free( excit );
		
		/* there are two error exits from this routine, 
		 * previous one for zero abundance, and this one for zero excitation */
#		ifdef DEBUG_FUN
		fputs( " <->atom_levelN()\n", debug_fp );
#		endif
		return;
	}

	/* already have excitation pumping, now get deexcitation */
	for( ilo=0; ilo < (nlev - 1); ilo++ )
	{
		for( ihi=ilo + 1; ihi < nlev; ihi++ )
		{
			/* (*pump)[low][ihi] is excitation rate due to external continuum,
			 * so derive rate from upper to lower */
			(*pump)[ihi][ilo] = (*pump)[ilo][ihi]*g[ilo]/g[ihi];
		}
	}

	/* evaluate collision rates from collision strengths, but only if calling
	 * routine has not already done this */
	if( !lgCollRateDone )
	{
		/* factor is 8.629E-6 / SQRTE * EDEN */
		for( ilo=0; ilo < (nlev - 1); ilo++ )
		{
			for( ihi=ilo + 1; ihi < nlev; ihi++ )
			{
				/* this should be a collision strength */
				ASSERT( (*data)[ihi][ilo]>= 0. );
				/* this is deexcitation rate
				 * factor is 8.629E-6 / SQRTE * EDEN */
				(*CollRate)[ihi][ilo] = (*data)[ihi][ilo]/g[ihi]*phycon.cdsqte;
				/* this is excitation rate */
				(*CollRate)[ilo][ihi] = (*CollRate)[ihi][ilo]*g[ihi]/g[ilo]*
				  excit[ilo][ihi];
			}
		}
	}

	/* rate equations equal zero */
	for( level=0; level < nlev; level++ )
	{
		for( j=0; j < nlev; j++ )
		{
			AMAT(level,j) = 0.;
		}
	}

	/* following is column of vector - represents source terms from elsewhere,
	 * if this is zero then matrix is singular and must replace one row with
	 * population sum equation - if sum is non-zero then get total abundance
	 * from source and sink terms */
	sum = 0.;
	lgHomogeneous = FALSE;
	for( level=0; level < nlev; level++ )
	{
		bvec[level] = source[level];
		sum += bvec[level];
	}
	if( sum==0. )
		lgHomogeneous = TRUE;

	/* eqns for destruction of level
	 * data[ilo][iho] are A's, CollRate is coll excit in direction
	 * pump[low][high] is excitation rate from low to high */
	for( level=0; level < nlev; level++ )
	{
		AMAT(level,level) = destroy[level];

		/* leaving level to below */
		for( ilo=0; ilo < level; ilo++ )
		{
			AMAT(level,level) += (*CollRate)[level][ilo] + (*data)[ilo][level] + (*dest)[ilo][level] + 
			  (*pump)[level][ilo];
			/* >>chng 97 jul 31, added pumping down
			 * coming to level I from below */
			AMAT(ilo,level) = -(*CollRate)[ilo][level] - (*pump)[ilo][level];
		}

		/* leaving level to above */
		for( ihi=level + 1; ihi < nlev; ihi++ )
		{
			AMAT(level,level) += (*CollRate)[level][ihi] + (*pump)[level][ihi];
			/* coming to level from above */
			AMAT(ihi,level) = -(*CollRate)[ihi][level] - (*data)[level][ihi] - (*dest)[level][ihi] - 
			  (*pump)[ihi][level];
			/* >>chng 97 jul 31, added pumping down */
		}
	}

	/* singular case, all source and sink terms add up to zero, so use the population
	 * equation to replace redundant equation */
	if( lgHomogeneous )
	{
		for( level=0; level<nlev; ++level )
		{
			AMAT(level,0) = 1.0;
		}
		/* these add up to total abundance */
		bvec[0] = abund;
	}

	if( lgDeBug )
	{
		fprintf(ioQQQ," AMAT matrix follows:\n");
		for( level=0; level < nlev; level++ )
		{
			for( j=0; j < nlev; j++ )
			{
				fprintf(ioQQQ," %.4e" , AMAT(level,j));
			}
			fprintf(ioQQQ,"\n");
		}
		if( sum==0. )
		{
			fprintf(ioQQQ," Sum creation zero so pop sum used\n");
		}
		else
		{
			fprintf(ioQQQ," Sum creation non-zero (%e), vector follows:\n",sum);
			for( j=0; j < nlev; j++ )
			{
				fprintf(ioQQQ," %.4e" , bvec[j] );
			}
			fprintf(ioQQQ,"\n");
		}
	}

	/* >>chng 00 oct 13, for MALLOC'd space*/
	/*DGETRF(nlev,nlev,(double*)amat,LIMLEVELN-1,ipiv,&ner);*/
	ner = 0;
	getrf_wrapper(nlev,nlev,amat,nlev,ipiv,&ner);
	/*DGETRF(nlev,nlev,amat,nlev,ipiv,&ner);*/
	/* usage DGETRS, 'N' = no transpose
		* order of matrix,
		* number of cols in bvec, =1
		* array
		* leading dim of array */
	/*DGETRS('N',nlev,1,(double*)amat,LIMLEVELN-1,ipiv,bvec,LIMLEVELN-1,&ner);*/
	/* >>chng 00 oct 13, for MALLOC'd space*/
	getrs_wrapper('N',nlev,1,amat,nlev,ipiv,bvec,nlev,&ner);

	if( ner != 0 )
	{
		fprintf( ioQQQ, " atom_levelN: dgetrs finds singular or ill-conditioned matrix\n" );
		puts( "[Stop in leveln]" );
		cdEXIT(EXIT_FAILURE);
	}

	/* set populations */
	for( level=0; level < nlev; level++ )
	{
		/* save bvec into populations */
		pops[level] = bvec[level];
	}

	/* now find total cooling and its derivative */
	*cooltl = 0.;
	*coolder = 0.;
	for( ihi=1; ihi < nlev; ihi++ )
	{
		for( ilo=0; ilo < ihi; ilo++ )
		{
			/* this is net cooling in oK */
			cool1 = (pops[ilo]*(*CollRate)[ilo][ihi] - pops[ihi]*(*CollRate)[ihi][ilo])*
			  (ex[ihi] - ex[ilo]);
			*cooltl += cool1;

			/* simplistic way to handle line heating - better would
			 * be to go  over to method in atom_level2 and atom_level3 */
			/* >>chng 03 aug 28, use real cool1 */
			/* *coolder += (float)(MAX2(0.,cool1)*((ex[ihi] - ex[0])*thermal.tsq1 - */
			*coolder += (cool1*((ex[ihi] - ex[0])*thermal.tsq1 - 
			  thermal.halfte));
		}
	}
	/* convert to ergs */
	*cooltl *= BOLTZMANN;
	*coolder *= BOLTZMANN;

	/* fill in departure coefficients */
	if( pops[0] > 1e-20 && excit[0][nlev-1] > 1e-20 )
	{
		/* >>chng 00 aug 10, loop had been from 1 and 0 was set to total abundance */
		depart[0] = 1.;
		for( ihi=1; ihi < nlev; ihi++ )
		{
			/* these are off by one - lowest index is zero */
			depart[ihi] = (pops[ihi]/pops[0])*(g[0]/g[ihi])/
			  excit[0][ihi];
		}
	}

	else
	{
		/* >>chng 00 aug 10, loop had been from 1 and 0 was set to total abundance */
		for( ihi=0; ihi < nlev; ihi++ )
		{
			/* these are off by one - lowest index is zero */
			depart[ihi] = 0.;
		}
		depart[0] = 1.;
	}

	/* sanity check for stability of inversion */
	*lgNegPop = FALSE;
	for( level=0; level < nlev; level++ )
	{
		if( pops[level] < 0. )
		{
			*lgNegPop = TRUE;
		}
	}

	if( *lgNegPop )
	{
		fprintf( ioQQQ, "\n atom_levelN found negative population, atom=%s\n", 
		  chLabel );
		fprintf( ioQQQ, " absolute population=" );

		for( level=0; level < nlev; level++ )
		{
			fprintf( ioQQQ, "%10.2e", pops[level] );
		}

		fprintf( ioQQQ, "\n" );
		for( level=0; level < nlev; level++ )
		{
			pops[level] = (double)MAX2(0.,pops[level]);
		}
	}

	if(  lgDeBug || (trace.lgTrace && trace.lgTrLevN) )
	{
		fprintf( ioQQQ, "\n atom_leveln absolute population   " );
		for( level=0; level < nlev; level++ )
		{
			fprintf( ioQQQ, " %10.2e", pops[level] );
		}
		fprintf( ioQQQ, "\n" );

		fprintf( ioQQQ, " departure coefficients" );
		for( level=0; level < nlev; level++ )
		{
			fprintf( ioQQQ, " %10.2e", depart[level] );
		}
		fprintf( ioQQQ, "\n\n" );
	}

	free( bvec );
	free( work );
	free( ipiv );
	free( amat );
	for( i=0; i < (nlev+1); ++i )
	{
		free( excit[i] );
	}
	free( excit );

#	ifdef DEBUG_FUN
	fputs( " <->atom_levelN()\n", debug_fp );
#	endif
	return;

}
