//////////////////////////////////////////////////////////////////////////////
// This file was originally written by Bradley S. Meyer.
//
// This 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 software 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.
//
////////////////////////////////////////////////////////////////////////////////

////////////////////////////////////////////////////////////////////////////////
//!
//! \file
//! \brief Code for user-defined screening functions.
//!
////////////////////////////////////////////////////////////////////////////////

#include "my_screen.h"

namespace my_user
{

//##############################################################################
// set_screening_function().
//##############################################################################

void
set_screening_function( nnt::Zone& zone )
{

  Libnucnet__Zone__setScreeningFunction(
    zone.getNucnetZone(),
    (Libnucnet__Zone__screeningFunction) screening_function,
    NULL
  );

  zone.updateFunction(
    nnt::s_SCREENING_DATA_FUNCTION,
    static_cast<boost::function<boost::any()> >(
      boost::bind( get_screening_data, boost::ref( zone ) )
    )
  );

}

//##############################################################################
// get_screening_data.
//##############################################################################

screening_data_t
get_screening_data( nnt::Zone& zone )
{

  screening_data_t screening_data( zone );
  return screening_data;

}

//##############################################################################
// screening_function().
//##############################################################################

void
screening_function(
  Libnucnet__Zone * p_zone,
  Libnucnet__Reaction * p_reaction,
  double d_t9,
  double d_rho,
  double d_Ye,
  double * p_forward,
  double * p_reverse
)
{

  double d_screen_f, d_screen_r, d_nse_corr;

  reaction_screening_function(
    Libnucnet__Zone__getNet( p_zone ),
    p_reaction,
    d_t9,
    d_rho,
    d_Ye,
    Libnucnet__Zone__getScreeningData( p_zone ),
    &d_screen_f,
    &d_screen_r
  );

  d_nse_corr =
    Libnucnet__Net__computeReverseRatioCorrectionFactorForReaction(
      Libnucnet__Zone__getNet( p_zone ),
      p_reaction,
      d_t9,
      d_rho,
      d_Ye,
      Libnucnet__Zone__getNseCorrectionFactorFunction( p_zone ),
      Libnucnet__Zone__getNseCorrectionFactorData( p_zone )
    );

  *p_forward *= d_screen_f;
  *p_reverse *= d_screen_f * d_nse_corr;

}

//##############################################################################
// reaction_screening_function().
//##############################################################################

void
reaction_screening_function(
  Libnucnet__Net * p_net,
  Libnucnet__Reaction * p_reaction,
  double d_t9,
  double d_rho,
  double d_Ye,
  void *p_user_data,
  double *p_screen_f,
  double *p_screen_r
)
{

  unsigned int i_z1, i_a1;

  //============================================================================
  // Compute the screening in the forward direction.
  //============================================================================

  nnt::reaction_element_list_t reactant_list =
    nnt::make_reaction_nuclide_reactant_list( p_reaction );

  i_z1 = 0;
  i_a1 = 0;
  *p_screen_f = 1;

  BOOST_FOREACH( nnt::ReactionElement reactant, reactant_list )
  {

    unsigned int i_z2 =
      Libnucnet__Species__getZ(
        Libnucnet__Nuc__getSpeciesByName(
          Libnucnet__Net__getNuc( p_net ),
          Libnucnet__Reaction__Element__getName(
            reactant.getNucnetReactionElement()
          )
        )
      ); 
    
    unsigned int i_a2 =
      Libnucnet__Species__getA(
        Libnucnet__Nuc__getSpeciesByName(
          Libnucnet__Net__getNuc( p_net ),
          Libnucnet__Reaction__Element__getName(
            reactant.getNucnetReactionElement()
          )
        )
      ); 
    
    if( i_z1 != 0 && i_a1 != 0 )
    {
      *p_screen_f *=
        pair_screening_function(
          d_t9,
          d_rho,
          d_Ye,
          i_z1,
          i_a1,
          i_z2,
          i_a2,
          p_user_data
        );
    }

    i_z1 += i_z2;
    i_a1 += i_a2;

  }

  //============================================================================
  // Compute the screening in the reverse direction.
  //============================================================================

  nnt::reaction_element_list_t product_list =
    nnt::make_reaction_nuclide_product_list( p_reaction );

  i_z1 = 0;
  i_a1 = 0;
  *p_screen_r = 1;

  BOOST_FOREACH( nnt::ReactionElement product, product_list )
  {

    unsigned int i_z2 =
      Libnucnet__Species__getZ(
        Libnucnet__Nuc__getSpeciesByName(
          Libnucnet__Net__getNuc( p_net ),
          Libnucnet__Reaction__Element__getName(
            product.getNucnetReactionElement()
          )
        )
      ); 
    
    unsigned int i_a2 =
      Libnucnet__Species__getA(
        Libnucnet__Nuc__getSpeciesByName(
          Libnucnet__Net__getNuc( p_net ),
          Libnucnet__Reaction__Element__getName(
            product.getNucnetReactionElement()
          )
        )
      ); 
    
    if( i_z1 != 0 && i_a1 != 0 )
    {
      *p_screen_r *=
        pair_screening_function(
          d_t9,
          d_rho,
          d_Ye,
          i_z1,
          i_a1,
          i_z2,
          i_a2,
          p_user_data
        );
    }

    i_z1 += i_z2;
    i_a1 += i_a2;

  }

}

//##############################################################################
// Base screening function.
//##############################################################################

double
pair_screening_function(
  double d_T9,
  double d_rho,
  double d_Ye,
  unsigned int i_Z1,
  unsigned int i_A1,
  unsigned int i_Z2,
  unsigned int i_A2,
  void * p_data 
)
{

  double d_L12, d_my_H120, d_my_H120A;

  if( !d_T9 || !d_rho || !d_Ye || !i_A1 || !i_A2 || !p_data )
  {
    std::cerr << "Screening function data invalid." << std::endl;
    exit( EXIT_FAILURE );
  }

  screening_data_t sd =
    boost::any_cast<screening_data_t>( *(boost::any *) p_data );

  d_L12 = sd.dL0 * static_cast<double>( i_Z1 * i_Z2 ) * sd.dZschl;

  if( d_L12 < 0.1 )
  {
    return exp( d_L12 );
  }

  d_my_H120 =
    sd.dH120M *
    (
      pow( static_cast<double>( i_Z1 + i_Z2 ), 1.86 ) -
      pow( static_cast<double>( i_Z1 ), 1.86 ) -
      pow( static_cast<double>( i_Z2 ), 1.86 )
    );

  if( d_L12 >= 0.1 && d_L12 <= 2.0 )
  {
    return exp( d_my_H120 );
  }

  d_my_H120A =
    sd.dH120A *
    (
      pow( static_cast<double>( i_Z1 + i_Z2 ), sd.F5 ) -
      pow( static_cast<double>( i_Z1 ), sd.F5 ) -
      pow( static_cast<double>( i_Z2 ), sd.F5 )
    )
    + 
    (
      0.316 * pow( ( sd.dYe / sd.dYtot ), sd.F1 ) *
      (
        pow( static_cast<double>( i_Z1 + i_Z2 ), sd.F4 ) -
        pow( static_cast<double>( i_Z1 ), sd.F4 ) -
        pow( static_cast<double>( i_Z2 ), sd.F4 )
      )
    )
    +
    (
      ( 0.737 /
        (
          sd.dYe *
          pow( sd.dYe / sd.dYtot, sd.F2 )
        )
      ) *
      (
        pow( static_cast<double>( i_Z1 + i_Z2 ), sd.F2 ) -
        pow( static_cast<double>( i_Z1 ), sd.F2 ) -
        pow( static_cast<double>( i_Z2 ), sd.F2 )
      )
    );
            
  if( d_L12 > 2.0 && d_L12 < 5.0 )
  {
    return exp( GSL_MIN( d_my_H120, d_my_H120A ) );
  }
  else
  {
    return exp( d_my_H120A );
  }
         
}

} // namespace my_user