/*
 *~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 *
 *<LicenseText>
 *
 * CitcomS by Louis Moresi, Shijie Zhong, Lijie Han, Eh Tan,
 * Clint Conrad, Michael Gurnis, and Eun-seo Choi.
 * Copyright (C) 1994-2005, California Institute of Technology.
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *
 *</LicenseText>
 *
 *~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 */

#include <math.h>
#include "global_defs.h"
#include "parsing.h"
#include "parallel_related.h"
#include "composition_related.h"


static void allocate_composition_memory(struct All_variables *E);
static void compute_elemental_composition_ratio_method(struct All_variables *E);
static void compute_elemental_composition_absolute_method(struct All_variables *E);
static void init_bulk_composition(struct All_variables *E);
static void check_initial_composition(struct All_variables *E);
static void fill_composition_from_neighbors(struct All_variables *E);


void composition_input(struct All_variables *E)
{
    int i;
    int m = E->parallel.me;
    input_boolean("CDEPV",&(E->viscosity.CDEPV),"off",m);
    E->composition.icompositional_rheology = E->viscosity.CDEPV;

    input_boolean("chemical_buoyancy",
		  &(E->composition.ichemical_buoyancy),
		  "1,0,nomax",m);

    if (E->control.tracer && 
	(E->composition.ichemical_buoyancy || 
	 E->composition.icompositional_rheology)) {

        /* ibuoy_type=0 (absolute method) */
        /* ibuoy_type=1 (ratio method) */

        input_int("buoy_type",&(E->composition.ibuoy_type),"1,0,nomax",m);
        /* if (E->composition.ibuoy_type!=1) {
            fprintf(stderr,"Terror-Sorry, only ratio method allowed now\n");
            fflush(stderr);
            parallel_process_termination();
        } */

        if (E->composition.ibuoy_type==0)
            E->composition.ncomp = E->trace.nflavors;
        else if (E->composition.ibuoy_type==1)
            E->composition.ncomp = E->trace.nflavors - 1;
	
        E->composition.buoyancy_ratio = (double*) malloc(E->composition.ncomp
                                                         *sizeof(double));

        /* default values .... */
        for (i=0; i<E->composition.ncomp; i++)
            E->composition.buoyancy_ratio[i] = 1.0;

        input_double_vector("buoyancy_ratio", E->composition.ncomp,
                            E->composition.buoyancy_ratio,m);

    }


    /* compositional rheology */
    /* what was this about? there is a CDEPV for compositional rheology 
       
    i moved this to the top so that cdepv = on would activate tracers

    TWB  */

    /* icompositional_rheology=0 (off) */
    /* icompositional_rheology=1 (on) */
    //E->composition.icompositional_rheology = 0;


    /*
    input_int("compositional_rheology",
              &(E->composition.icompositional_rheology),"1,0,nomax",m);

    if (E->composition.icompositional_rheology==1) {
        input_double("compositional_prefactor",
                     &(E->composition.compositional_rheology_prefactor),
                     "1.0",m);
    }
    */

    return;
}



void composition_setup(struct All_variables *E)
{
    allocate_composition_memory(E);

    return;
}


void write_composition_instructions(struct All_variables *E)
{
    int k;

    if (E->composition.ichemical_buoyancy ||
        E->composition.icompositional_rheology)
        E->composition.on = 1;

    if (E->composition.on) {

        if (E->trace.nflavors < 1) {
            fprintf(E->trace.fpt, "Tracer flavors must be greater than 1 to track composition\n");
            parallel_process_termination();
        }

        if (!E->composition.ichemical_buoyancy)
	  fprintf(E->trace.fpt,"Passive Tracers\n");
	else
	  fprintf(E->trace.fpt,"Active Tracers\n");


        if (E->composition.ibuoy_type==1)
	  fprintf(E->trace.fpt,"Ratio Method\n");
        if (E->composition.ibuoy_type==0)
	  fprintf(E->trace.fpt,"Absolute Method\n");

        for(k=0; k<E->composition.ncomp; k++) {
            fprintf(E->trace.fpt,"Buoyancy Ratio: %f\n", E->composition.buoyancy_ratio[k]);
        }

        /*
        if (E->composition.icompositional_rheology==0) {
            fprintf(E->trace.fpt,"Compositional Rheology - OFF\n");
        }
        else if (E->composition.icompositional_rheology>0) {
            fprintf(E->trace.fpt,"Compositional Rheology - ON\n");
            fprintf(E->trace.fpt,"Compositional Prefactor: %f\n",
            E->composition.compositional_rheology_prefactor);
        }
        */

        fflush(E->trace.fpt);
    }

    return;
}


/************ FILL COMPOSITION ************************/
void fill_composition(struct All_variables *E)
{

    /* ratio method */

    if (E->composition.ibuoy_type==1) {
        compute_elemental_composition_ratio_method(E);
    }

    /* absolute method */
    if (E->composition.ibuoy_type==0) {
        compute_elemental_composition_absolute_method(E);
    }

    /* if (E->composition.ibuoy_type!=1) {
        fprintf(E->trace.fpt,"Error(compute...)-only ratio method now\n");
        fflush(E->trace.fpt);
        exit(10);
    } */

    /* Map elemental composition to nodal points */

    map_composition_to_nodes(E);

    return;
}



static void allocate_composition_memory(struct All_variables *E)
{
    int i, j;

    for (i=0; i<E->composition.ncomp; i++) {
        E->composition.bulk_composition = (double*) malloc(E->composition.ncomp*sizeof(double));
        E->composition.initial_bulk_composition = (double*) malloc(E->composition.ncomp*sizeof(double));
        E->composition.error_fraction = (double*) malloc(E->composition.ncomp*sizeof(double));
    }


    /* for horizontal average */
    E->Have.C = (float **)malloc((E->composition.ncomp+1)*sizeof(float*));
    for (i=0; i<E->composition.ncomp; i++) {
        E->Have.C[i] = (float *)malloc((E->lmesh.noz+2)*sizeof(float));
    }


    /* allocat memory for composition fields at the nodes and elements */
    if ((E->composition.comp_el=(double **)malloc((E->composition.ncomp)*sizeof(double*)))==NULL) {
        fprintf(E->trace.fpt,"AKM(allocate_composition_memory)-no memory 8987y\n");
        fflush(E->trace.fpt);
        exit(10);
    }
    if ((E->composition.comp_node=(double **)malloc((E->composition.ncomp)*sizeof(double*)))==NULL) {
        fprintf(E->trace.fpt,"AKM(allocate_composition_memory)-no memory 8988y\n");
        fflush(E->trace.fpt);
        exit(10);
    }

    for (i=0; i<E->composition.ncomp; i++) {
        if ((E->composition.comp_el[i]=(double *)malloc((E->lmesh.nel+1)*sizeof(double)))==NULL) {
            fprintf(E->trace.fpt,"AKM(allocate_composition_memory)-no memory 8989y\n");
            fflush(E->trace.fpt);
            exit(10);
        }

        if ((E->composition.comp_node[i]=(double *)malloc((E->lmesh.nno+1)*sizeof(double)))==NULL) {
            fprintf(E->trace.fpt,"AKM(allocate_composition_memory)-no memory 983rk\n");
            fflush(E->trace.fpt);
            exit(10);
        }
    }
}

void init_composition(struct All_variables *E)
{
    /* ratio method */
    if (E->composition.ibuoy_type==1) {
        compute_elemental_composition_ratio_method(E);
    }

    /* absolute method */
    if (E->composition.ibuoy_type==0) {
        compute_elemental_composition_absolute_method(E);
    }

    /* if (E->composition.ibuoy_type!=1) {
        fprintf(E->trace.fpt,"Error(compute...)-only ratio method now\n");
        fflush(E->trace.fpt);
        exit(10);
    } */

    /* for empty elements */
    check_initial_composition(E);

    /* Map elemental composition to nodal points */
    map_composition_to_nodes(E);

    init_bulk_composition(E);

    return;
}


static void check_initial_composition(struct All_variables *E)
{
    /* check empty element if using ratio method */
    if (E->composition.ibuoy_type == 1) {
        if (E->trace.istat_iempty) {
            /* using the composition of neighboring elements to determine
               the initial composition of empty elements. */
            fill_composition_from_neighbors(E);
        }
    }

    return;
}



/*********** COMPUTE ELEMENTAL COMPOSITION RATIO METHOD ***/
/*                                                        */
/* This function computes the composition per element.    */
/* The concentration of material i in an element is       */
/* defined as:                                            */
/*   (# of tracers of flavor i) / (# of all tracers)      */

static void compute_elemental_composition_ratio_method(struct All_variables *E)
{
    int i, e, flavor, numtracers;
    int iempty = 0;

    for (e=1; e<=E->lmesh.nel; e++) {
      numtracers = 0;
      for (flavor=0; flavor<E->trace.nflavors; flavor++)
        numtracers += E->trace.ntracer_flavor[flavor][e];

      /* Check for empty entries and compute ratio.  */
      /* If no tracers are in an element, skip this element, */
      /* use previous composition. */
      if (numtracers == 0) {
        iempty++;
        /* fprintf(E->trace.fpt, "No tracer in element %d!\n", e); */
        continue;
      }
      for(i=0;i<E->composition.ncomp;i++) {
        flavor = i + 1;
        E->composition.comp_el[i][e] = 
          E->trace.ntracer_flavor[flavor][e] / (double)numtracers;
      }
    }


    if (iempty) {

      if ((1.0*iempty/E->lmesh.nel)>0.80) {
          fprintf(E->trace.fpt,"WARNING(compute_elemental...)-number of tracers is REALLY LOW\n");
          fflush(E->trace.fpt);
          if (E->trace.itracer_warnings) 
            exit(10);
      }
    }

    E->trace.istat_iempty += iempty;
}

/*********** COMPUTE ELEMENTAL ABSOLUTE METHOD *************************/
/*                                                                     */
/* This function computes the composition per element.                 */
/* The concentration of material i in an element is                    */
/* defined as:                                                         */
/*   tracer i mass * (# of tracers of flavor i) / ( volume of element) */
/*                                                                     */ 
/* where tracer i mass = tot tracer volume / tot num of tracers        */
/*                                                                     */
/* Added by DJB 01/14/10.  Known caveats: (will address at some point) */
/*     1) XXX: needs better error handling                             */
/*     2) XXX: only tested with one flavor of tracer (flavor 0)        */
/*     3) XXX: only tested by reading in a file of tracer locations    */
/*     4) XXX: the initial volume of the tracer domain is HARD-CODED!  */

static void compute_elemental_composition_absolute_method(struct All_variables *E)
{
  int i, j, e, flavor, numtracers;
  double domain_volume;
  float comp;
  float one = 1.0;

  /* This needs to be `manually' changed for the total volume */
  /*  occupied by your tracers */
  domain_volume = 1e-2;

  for (e=1; e<=E->lmesh.nel; e++) {
    numtracers = 0;
    for (flavor=0; flavor<E->trace.nflavors; flavor++)
        numtracers += E->trace.ntracer_flavor[flavor][e];

    /* Check for empty entries */
    /* If no tracers are in an element, comp = 0.0 (i.e. is ambient) */
    if (numtracers == 0) {
      for(i=0;i<E->composition.ncomp;i++) {
          E->composition.comp_el[i][e] = 0.0;
      }
      continue;
    }

    /* Composition is proportional to (local) tracer density */
    for(i=0;i<E->composition.ncomp;i++) {
      flavor = i;
      comp = E->trace.ntracer_flavor[flavor][e] / E->eco[e].area
              * domain_volume / E->trace.number_of_tracers;

      /* truncate composition at 1.0 */
      /* This violates mass conservation but prevents unphysical C */
      /* XXX: make truncation a switch for the user to specify */
      E->composition.comp_el[i][e] = min(comp,one);
    }
  }
}

/********** MAP COMPOSITION TO NODES ****************/
void map_composition_to_nodes(struct All_variables *E)
{
    double *tmp;
    int i, n, kk;
    int nelem, nodenum;


    /* first, initialize node array */
    for(i=0;i<E->composition.ncomp;i++) {
      for (kk=1;kk<=E->lmesh.nno;kk++)
        E->composition.comp_node[i][kk]=0.0;
    }

    /* Loop through all elements */
    for (nelem=1;nelem<=E->lmesh.nel;nelem++) {

      /* for each element, loop through element nodes */

      /* weight composition */
      for (nodenum=1;nodenum<=8;nodenum++) {
        n = E->ien[nelem].node[nodenum];
        for(i=0;i<E->composition.ncomp;i++) {
          E->composition.comp_node[i][n] +=
            E->composition.comp_el[i][nelem]*
            E->TWW[E->mesh.levmax][nelem].node[nodenum];
        }
      }
    } /* end nelem */

    for(i=0;i<E->composition.ncomp;i++) {
      tmp = E->composition.comp_node[i];
      (E->exchange_node_d)(E,tmp,E->mesh.levmax);
    }

    /* Divide by nodal volume */
    for(i=0;i<E->composition.ncomp;i++)
      for (kk=1;kk<=E->lmesh.nno;kk++)
        E->composition.comp_node[i][kk] *= E->MASS[E->mesh.levmax][kk];
}


/****************************************************************/

static void fill_composition_from_neighbors(struct All_variables *E)
{
    int i, j, k, e, ee, n, flavor, numtracers, count;
    double *sum;
    const int n_nghbrs = 4;
    int nghbrs[n_nghbrs];
    int *is_empty;

    fprintf(E->trace.fpt,"WARNING(check_initial_composition)-number of tracers is low, %d elements contain no tracer initially\n", E->trace.istat_iempty);

    fprintf(E->trace.fpt,"Using neighboring elements for initial composition...\n");

    /* index shift for neighboring elements in horizontal direction */
    nghbrs[0] = E->lmesh.elz;
    nghbrs[1] = -E->lmesh.elz;
    nghbrs[2] = E->lmesh.elz * E->lmesh.elx;
    nghbrs[3] = -E->lmesh.elz * E->lmesh.elx;

    is_empty = (int *)calloc(E->lmesh.nel+1, sizeof(int));
    sum = (double *)malloc(E->composition.ncomp * sizeof(double));

    /* which element is empty? */
    for (e=1; e<=E->lmesh.nel; e++) {
      numtracers = 0;
      for (flavor=0; flavor<E->trace.nflavors; flavor++)
        numtracers += E->trace.ntracer_flavor[flavor][e];

      if (numtracers == 0)
        is_empty[e] = 1;
    }

    /* using the average comp_el from neighboring elements */
    for (e=1; e<=E->lmesh.nel; e++) {
      if(is_empty[e]) {
        count = 0;
        for (i=0; i<E->composition.ncomp; i++)
          sum[i] = 0.0;

        for(n=0; n<n_nghbrs; n++) {
          ee = e + nghbrs[n];
          /* is ee a valid element number and the elemnt is not empty? */
          if((ee>0) && (ee<=E->lmesh.nel) && (!is_empty[ee])) {
            count++;
              for (i=0; i<E->composition.ncomp; i++)
                sum[i] += E->composition.comp_el[i][ee];
          }
        }

        if(count == 0) {
          fprintf(E->trace.fpt,"Error(fill_composition_from_neighbors)-all neighboring elements are empty\n");
          fflush(E->trace.fpt);
          exit(10);
        }

        for (i=0; i<E->composition.ncomp; i++)
          E->composition.comp_el[i][e] = sum[i] / count;
      }
    }

    free(is_empty);
    free(sum);

    fprintf(E->trace.fpt,"Done.\n");
    fflush(E->trace.fpt);
}


/*********** GET BULK COMPOSITION *******************************/

static void init_bulk_composition(struct All_variables *E)
{

    double return_bulk_value_d();
    double volume;
    double *tmp;
    int i, m;
    const int ival=0;


    for (i=0; i<E->composition.ncomp; i++) {

      tmp = E->composition.comp_node[i];

      /* ival=0 returns integral not average */
      volume = return_bulk_value_d(E,tmp,ival);

      E->composition.bulk_composition[i] = volume;
      E->composition.initial_bulk_composition[i] = volume;
    }
}


void get_bulk_composition(struct All_variables *E)
{

    double return_bulk_value_d();
    double volume;
    double *tmp;
    int i, m;
    const int ival = 0;

    for (i=0; i<E->composition.ncomp; i++) {

      tmp = E->composition.comp_node[i];

      /* ival=0 returns integral not average */
      volume = return_bulk_value_d(E,tmp,ival);

      E->composition.bulk_composition[i] = volume;

      E->composition.error_fraction[i] = 
        (volume - E->composition.initial_bulk_composition[i]) / 
        E->composition.initial_bulk_composition[i];
    }
}