#include <G4Material.hh>
#include <G4Element.hh>
#include <G4Isotope.hh>
#include <G4NistManager.hh>

#include <RAT/Materials.hh>
#include <RAT/Log.hh>
#include <RAT/DB.hh>
using namespace RAT;

#include <vector>
#include <CLHEP/Units/SystemOfUnits.h>
#include <CLHEP/Units/PhysicalConstants.h>
#include <string>
using namespace CLHEP;
using namespace std;

void
Materials::ConstructElements()
{
  DBLinkGroup elements = DB::Get()->GetLinkGroup( "ELEMENT" );
  for( DBLinkGroup::iterator iElement = elements.begin(); iElement != elements.end(); ++iElement )
    {
      const string name = iElement->first;
      DBLinkPtr dbTable = iElement->second;
      // First load the elemental symbol and Z
      int z;
      string symbol;
      try
        {
          z = dbTable->GetI( "z" );
          symbol = dbTable->GetS( "symbol" );
        }
      catch( DBNotFoundError& e ) { Log::Die( "Materials::ConstructElements: cannot load " + name + " " + e.field + " not found." ); }
      // Now assume it is a single isotope element
      try
        {
          int a = dbTable->GetI( "a" );
          new G4Element( name, symbol, z, a * g / mole );
// Element loaded
          continue;
        }
      catch( DBNotFoundError& e ) { /* Not a concern if multiple isotopes exist.*/ }
      // Now assume it has multiple isotopes
      try
        {
          const vector<int>& isotopes = dbTable->GetIArray("isotopes");
          const vector<double>& isotopeFractions = dbTable->GetDArray("isotopes_frac");
          Log::Assert( isotopes.size() == isotopeFractions.size(), "Materials::ConstructElements: Number of isotopes is different to number of isotopes_frac for " + name );

          G4Element* element = new G4Element( name, symbol, isotopes.size() );
          vector<int>::const_iterator iIsotope;
          vector<double>::const_iterator iFraction;
          for( iIsotope = isotopes.begin(), iFraction = isotopeFractions.begin();
               iIsotope != isotopes.end() && iFraction != isotopeFractions.end();
               ++iIsotope, ++iFraction )
            {
              // Leave out last field of constructor so mass/mole of isotope is pulled from NIST database.
              Log::Assert( *iFraction != 0, "Materials::ConstructElements: isotope_frac for " + name + " should not be set to zero." );
              G4Isotope* isotope = new G4Isotope( name, z, *iIsotope );
              element->AddIsotope( isotope, *iFraction );
            }
// Element including all isotopes loaded
          continue;
        }
      catch( DBNotFoundError& e ) { Log::Die( "Materials::ConstructElements: cannot load " + name + " a or isotope list not found." ); }
    }
}

void
Materials::ConstructMaterials()
{
  // First the elements
  ConstructElements();
  DBLinkGroup materials = DB::Get()->GetLinkGroup( "MATERIAL" );
  while( !materials.empty() ) // Loop until ALL materials are constructed
    {
      DBLinkGroup::iterator iMaterial;
      for( iMaterial = materials.begin(); iMaterial != materials.end(); ++iMaterial )
        {
          const string name = iMaterial->first;
          DBLinkPtr dbTable = iMaterial->second;
          // First check if this material has any unbuilt dependent materials
          vector<string> subMaterials;
          try
            {
              subMaterials = dbTable->GetSArray( "materials" );
              bool buildable = true;
              for( vector<string>::const_iterator iSubMaterial = subMaterials.begin(); iSubMaterial != subMaterials.end(); ++iSubMaterial )
                if( G4Material::GetMaterial( *iSubMaterial, false ) == NULL && // Note the first term is guaranteed to be evaluated first
                    G4NistManager::Instance()->FindOrBuildMaterial( *iSubMaterial ) == NULL )
                  {
                    buildable = false;
                    break;
                  }
              if( !buildable )
// Can't build material, skip to next one
                continue;
            }
          catch( DBNotFoundError& e ) { /* This material has no sub materials => can always be built assuming subElements exist.*/ }
          ConstructMaterial( name, dbTable );
// Material built
          materials.erase( iMaterial );
          break;
        }
      // Check for circular dependencies
      if( iMaterial == materials.end() )
        {
          warn << "Materials::ConstructMaterials: The following materials cannot be constructed." << newline;
          for( iMaterial = materials.begin(); iMaterial != materials.end(); ++iMaterial )
            warn << iMaterial->first << newline;
          Log::Die( "Materials::ConstructMaterials: Circular dependency in materials database." );
        }
    }
}

void
Materials::ConstructMaterial( const std::string& name,
                              DBLinkPtr dbTable )
{
  vector<string> materials;
  vector<double> materialProportions;
  try
    {
      materials = dbTable->GetSArray( "materials" );
      materialProportions = dbTable->GetDArray( "matprop" );
    }
  catch (DBNotFoundError &e) { /* Not an error, just made out of materials. */ }
  Log::Assert( materials.size() == materialProportions.size(), "Materials::ConstructMaterial: materials and matprop different sizes." );
  vector<string> elements;
  vector<double> elementProportions;
  try
    {
      elements = dbTable->GetSArray( "elements" );
      elementProportions = dbTable->GetDArray( "elemprop" );
    }
  catch (DBNotFoundError &e) { /* Not an error, just made out of materials. */ }
  Log::Assert( elements.size() == elementProportions.size(), "Materials::ConstructMaterial: elements and elemprop different sizes." );
  Log::Assert( !elements.empty() || !materials.empty(), "Materials::ConstructMaterial: Material " + name + " must be made of elements or other materials." );
  // Can now go ahead and build this material, first the required...
  double density = 0.0;
  try { density = dbTable->GetD( "density" ) * g / cm3; }
  catch( DBNotFoundError& e ) { Log::Die( "Materials::ConstructMaterial: Material must have a density." ); }
  // ... and now the optional definitions
  G4State state;
  try
    {
      const string dbState = dbTable->GetS("state");
      if( dbState == "gas")
        state = kStateGas;
      else if( dbState == "solid" )
        state = kStateSolid;
      else if( dbState == "liquid" )
        state = kStateLiquid;
      else
        Log::Die( "Materials::ConstructMaterial: Unknown material state " + dbState );
    }
  catch( DBNotFoundError& e ) { state = kStateUndefined; } // Default to undefined state
  double temperature;
  try { temperature = dbTable->GetD( "temperature" ); }
  catch( DBNotFoundError& e ) { temperature = STP_Temperature; } // Default to STP Temperature
  double pressure;
  try { pressure = dbTable->GetD( "pressure" ) * STP_Pressure; }
  catch( DBNotFoundError& e ) { pressure = STP_Pressure; } // Default to STP pressure
  G4Material* material = new G4Material( name, density, elements.size() + materials.size(), state, temperature, pressure );
  material->SetMaterialPropertiesTable( new G4MaterialPropertiesTable() ); // Ensure all materials have a properties table
  // Now optionally set the chemical name & molar volume
  try
    {
      material->SetChemicalFormula( dbTable->GetS( "formula" ) ); // This seems to be just a label - PGJ
      material->GetMaterialPropertiesTable()->AddConstProperty( "MOL", dbTable->GetD( "mol" ) / g );
    }
  catch( DBNotFoundError& e ) { /* Not required. */ }
  // Now build the material from elements...
  vector<string>::const_iterator iElement;
  vector<double>::const_iterator iProportion;
  for( iElement = elements.begin(), iProportion = elementProportions.begin();
       iElement != elements.end() && iProportion != elementProportions.end();
       ++iElement, ++iProportion )
    {
      Log::Assert( *iProportion != 0, "Materials::ConstructMaterial: elemprop for " + name + " should not be set to zero." );
      material->AddElement( G4Element::GetElement( *iElement ), *iProportion );
    }
  // ... and/or materials...
  vector<string>::const_iterator iMaterial;
  for( iMaterial = materials.begin(), iProportion = materialProportions.begin();
       iMaterial != materials.end() && iProportion != materialProportions.end();
       ++iMaterial, ++iProportion )
    {
      Log::Assert( *iProportion != 0, "Materials::ConstructMaterial: matprop for " + name + " should not be set to zero." );
      G4Material* subMaterial = G4Material::GetMaterial( *iMaterial );
      if( subMaterial == NULL )
        subMaterial = G4NistManager::Instance()->FindOrBuildMaterial( *iMaterial );
      material->AddMaterial( subMaterial, *iProportion );
    }
  // ...geant4 then checks the proportions sum up to 1 etc...
}