// CreateMaterialInputFile.C
// Author: Andrew Edmonds
// Created: 06.09.2013
// This script create the MATER.INP file needed for MARS by reading directly from the material list in the ROOT file
// NB Requires ROOT v5.34 or later to be able to use the GetNelements() and GetElementProp() functions
//
// The format of the MATER.INP file is:
// ID# 'material name' density n_elements
// atomic_mass1 atomic_charge1 weight_fraction1
// atomic_mass2 atomic_charge2 weight_fraction2
// ...

void CreateMaterialInputFile(std::string geometry_filename) {

        TFile* file = new TFile(geometry_filename.c_str(), "READ"); // open the file

  // MARS doesn't seem to have data for all atomic masses
  const int n_unknown_atomic_masses = 3;
  int unknown_atomic_masses[n_unknown_atomic_masses] = {13, 18, 180};
  
  TIter nextkey(file->GetListOfKeys()); // get the list of keys in the file
     TKey *key;
     
     while (key = (TKey*)nextkey()) {
       // Check that the class name is TGeoManager
       if (strcmp(key->ReadObj()->ClassName(), "TGeoManager") == 0) {
      gGeoManager = (TGeoManager*) key->ReadObj();
      break;
    }
  }
  
  TList* listOfMedia = gGeoManager->GetListOfMedia();
  TIterator* mediumIter = listOfMedia->MakeIterator();
  
  std::ofstream out_file("MATER.INP", std::ofstream::out);
  out_file << "MATER.INP generated from tgeo_init.cc" << std::endl << std::endl;
  out_file << "1 'VAC'" << std::endl; // need to put in a dummy material because ROOT starts at IM=2
  
  TGeoMedium* medium = (TGeoMedium*) mediumIter->Next();	
  
  while (medium != NULL) {
    int n_elements = medium->GetMaterial()->GetNelements();

    // Check for any unknown atomic masses
    int n_missing_elements = 0;
    for (int i = 0; i < n_elements; ++i) {
      double a, z, w;
      medium->GetMaterial()->GetElementProp(a, z, w, i);

      for (int j = 0; j < n_unknown_atomic_masses; ++j) {
        if (floor(a+0.5) == unknown_atomic_masses[j]) // add 0.5 so that it rounds to the nearest value
          ++n_missing_elements;
      }
    }
    // make sure the number of elements is consistent after taking into account the missing ones
    out_file << medium->GetId() << " \'" << medium->GetName() << "\' " << medium->GetMaterial()->GetDensity() << " " << n_elements - n_missing_elements << std::endl;
    
    for (int i = 0; i < n_elements; i++) {
      double a, z, w;
      medium->GetMaterial()->GetElementProp(a, z, w, i);

      bool missing_element = false;
      for (int j = 0; j < n_unknown_atomic_masses; ++j) {
        if (floor(a+0.5) == unknown_atomic_masses[j]) // add 0.5 so that it rounds to the nearest value
          missing_element = true;
      }
      if (missing_element == false)
        out_file << "\t" << a << " " << z << " " << w << std::endl;
    }
    out_file << std::endl;
    
    medium = (TGeoMedium*) mediumIter->Next();
  }
  out_file << "STOP" << std::endl;
}