#! /usr/bin/python

################################################################################
# Python script to generate OT Layers geometry & structure
# Author: Jan Amoraal
################################################################################

# Some globals
author            = 'Jan Amoraal'
version           = '1.0'
otGeomPath        = '/dd/Geometry/AfterMagnetRegion/T/OT'
otStructPath      = '/dd/Structure/LHCb/AfterMagnetRegion/T/OT'
otChannelCondPath = '/dd/Conditions/ChannelInfo/OT'
otCalCondPath     = '/dd/Conditions/Calibration/OT'

# XML version and encoding
indent  = 3*' '
version  = '"1.0"'
encoding = '"ISO-8859-1"'
header   = '<?xml version=%s encoding=%s?>\n'%( version, encoding )
doctype  = '<!DOCTYPE DDDB SYSTEM "conddb:/DTD/geometry.dtd">\n' 
header   = header + doctype + '\n<DDDB>\n\n'
footer   = '</DDDB>\n'

# Station, Layers, and Quarters 
stations       = [ 'T1', 'T2', 'T3' ]
layers         = [ 'X1', 'U', 'V', 'X2' ]
layerIDs       = [ '0' , '1', '2', '3' ]
layerStereo    = [ '0.0*degree', '-5.0*degree', '5.0*degree', '0.0*degree' ] 
CSideCouplings = [ 'lvXCouplingsCSide' , 'lvUCouplingsCSide', 'lvVCouplingsCSide', 'lvXCouplingsCSide' ]
ASideCouplings = [ 'lvXCouplingsASide' , 'lvUCouplingsASide', 'lvVCouplingsASide', 'lvXCouplingsASide' ]

# Stereo angle -> 0, -5, 5, 0
stereo   = [ 'X1StereoAngle', 'UStereoAngle' , 'VStereoAngle', 'X2StereoAngle' ]

def LayerGeometryTemplate( buffer, indent ) :
    buffer.write( '<?xml version=%s encoding=%s?>\n\n'%( version, encoding ) )
    buffer.write( '\n' )
    buffer.write( '  <catalog name="&Station;&Layer;">\n' )
    buffer.write( '    <logvolref  href="#lv&Station;&Layer;" />\n' )
    buffer.write( '    <catalogref href="../../Quarters/QuarterGeometryCatalogs.xml#&Station;&Layer;Quarters" />\n' )
    buffer.write( '  </catalog>\n' )
    buffer.write( '\n' )
    buffer.write( '  <logvol name="lv&Station;&Layer;" material="Air">\n' )
    buffer.write( '    <subtraction name="&Station;&Layer;BeamPipeHole">\n' ) 
    buffer.write( '      <box name="&Station;&Layer;Box"\n' ) 
    buffer.write( '           sizeX="LayerXSize"\n' )
    buffer.write( '           sizeY="LayerYSize"\n' )
    buffer.write( '	      sizeZ="LayerThick" />\n' )
    buffer.write( '      <box name="&Station;&Layer;SModHole"\n' )
    buffer.write( '	      sizeX="263*mm"\n' )
    buffer.write( '	      sizeY="2*135*mm"\n' )
    buffer.write( '	      sizeZ="LayerThick+Clearance" />\n' )  
    buffer.write( '    </subtraction>\n' )
    buffer.write( '    <physvol name="pv&Station;&Layer;CSideCouplings"\n' )
    buffer.write( '	     logvol="/dd/Geometry/AfterMagnetRegion/T/OT/Couplings/&lvCSideCouplings;">\n' )
    buffer.write( '      <posXYZ x="CSideOffset" />\n' )
    buffer.write( '    </physvol>\n' )
    buffer.write( '    <physvol name="pv&Station;&Layer;ASideCouplings"\n' )
    buffer.write( '	     logvol="/dd/Geometry/AfterMagnetRegion/T/OT/Couplings/&lvASideCouplings;">\n' )
    buffer.write( '      <posXYZ x="ASideOffset" />\n' )
    buffer.write( '    </physvol>\n' )
    buffer.write( '    <physvol name="pv&Station;&Layer;Q0"\n' )
    buffer.write( '             logvol="/dd/Geometry/AfterMagnetRegion/T/OT/&Station;/&Station;&Layer;/&Station;&Layer;Quarters/lv&Station;&Layer;Q0">\n' )
    buffer.write( '      <posXYZ x="CSideOffset" />\n' )
    buffer.write( '    </physvol>\n' )
    buffer.write( '    <physvol name="pv&Station;&Layer;Q1"\n' )
    buffer.write( '             logvol="/dd/Geometry/AfterMagnetRegion/T/OT/&Station;/&Station;&Layer;/&Station;&Layer;Quarters/lv&Station;&Layer;Q1">\n' )
    buffer.write( '      <posXYZ x="ASideOffset" />\n' )
    buffer.write( '    </physvol>\n' )
    buffer.write( '    <physvol name="pv&Station;&Layer;Q2"\n' )
    buffer.write( '             logvol="/dd/Geometry/AfterMagnetRegion/T/OT/&Station;/&Station;&Layer;/&Station;&Layer;Quarters/lv&Station;&Layer;Q2">\n' )
    buffer.write( '      <posXYZ x="CSideOffset" />\n' )
    buffer.write( '    </physvol>\n' )
    buffer.write( '    <physvol name="pv&Station;&Layer;Q3"\n' )
    buffer.write( '             logvol="/dd/Geometry/AfterMagnetRegion/T/OT/&Station;/&Station;&Layer;/&Station;&Layer;Quarters/lv&Station;&Layer;Q3">\n' )
    buffer.write( '      <posXYZ x="ASideOffset" />\n' )
    buffer.write( '    </physvol>\n' )
    buffer.write( '  </logvol>\n' )
    buffer.write( '\n' )
    
def LayerStructureTemplate( buffer, indent ) :
    buffer.write( '<?xml version=%s encoding=%s?>\n\n'%( version, encoding ) )
    buffer.write( '<detelem classID="8103" name="&Layer;">\n' )
    buffer.write( '<author>%s</author>\n'%( author ) )
    buffer.write( '<version>%s</version>\n'%( '1.0' ) )
    buffer.write( '<geometryinfo lvname="/dd/Geometry/AfterMagnetRegion/T/OT/&Station;/&Station;&Layer;/lv&Station;&Layer;"\n' )
    buffer.write( '		 condition="/dd/Conditions/Alignment/OT/&Station;&Layer;"\n' )
    buffer.write( '		 support="/dd/Structure/LHCb/AfterMagnetRegion/T/OT/&Station;"\n' )
    buffer.write( '		 npath="pv&Station;&Layer;" />\n' )
    buffer.write( '    <param name="layerID" type="int">&LayerID;</param>\n' )
    buffer.write( '    <param name="stereoAngle">&StereoAngle;</param>\n' )
    buffer.write( '    <detelemref href="../../Quarters/Structure/&Station;&Layer;Q0@Structure.xml#Q0" />\n' )
    buffer.write( '    <detelemref href="../../Quarters/Structure/&Station;&Layer;Q1@Structure.xml#Q1" />\n' )
    buffer.write( '    <detelemref href="../../Quarters/Structure/&Station;&Layer;Q2@Structure.xml#Q2" />\n' )
    buffer.write( '    <detelemref href="../../Quarters/Structure/&Station;&Layer;Q3@Structure.xml#Q3" />\n' )
    buffer.write( '</detelem>\n' )
    buffer.write( '\n' )

# Now let's write the geometry
import cStringIO, os
# Create dirs
layerssdir  = 'Layers'
geometrydir = 'Geometry'
try :
    os.mkdir( layerssdir )
    os.chdir( layerssdir )                     
except OSError :
    print layerssdir + ' Directory exits'
    os.chdir( layerssdir )

# First the templates
# Geometry template
geometryTemplate = cStringIO.StringIO()
LayerGeometryTemplate( geometryTemplate, indent )
contents = geometryTemplate.getvalue()
geometryTemplate.close()
output = file( 'LayerGeometryTemplate.xml', 'w' )
output.write( contents )
output.close()
# Structure template
structureTemplate = cStringIO.StringIO()
LayerStructureTemplate( structureTemplate, indent )
contents = structureTemplate.getvalue()
structureTemplate.close()
output = file( 'LayerStructureTemplate.xml', 'w' )
output.write( contents )
output.close()

# Geometry
try :
    os.mkdir( geometrydir )
    os.chdir( geometrydir )                     
except OSError :
    print geometrydir + ' Directory exits'
    os.chdir( geometrydir )

for s in stations :
    for n, l in enumerate( layers ) :
        geometry = cStringIO.StringIO()
        geometry.write( '<!DOCTYPE DDDB SYSTEM "conddb:/DTD/geometry.dtd"\n' )
        geometry.write( ' [<!ENTITY Station          "%s">\n'%( s ) )
        geometry.write( '  <!ENTITY Layer            "%s">\n'%( l ) )
        geometry.write( '  <!ENTITY lvCSideCouplings "%s">\n'%( CSideCouplings[ n ] ) )
        geometry.write( '  <!ENTITY lvASideCouplings "%s">\n'%( ASideCouplings[ n ] ) )
        geometry.write( '  <!ENTITY %s%sGeometry SYSTEM "../LayerGeometryTemplate.xml">]>\n'%( s, l ) )
        geometry.write( '\n' )
        geometry.write( '<DDDB>\n' )
        geometry.write( '  &%s%sGeometry;\n'%( s, l ) )
        geometry.write( '</DDDB>\n' )
        contents = geometry.getvalue()
        geometry.close()
        fname = '%s%s@Geometry.xml'%( s, l )
        output = file( fname, 'w' )
        output.write( contents )
        output.close()

# Now generate the accompanying structure
structdir = 'Structure'
try :
    os.chdir( '..' )
    os.mkdir( structdir )
    os.chdir( structdir )
except OSError :
    print structdir + ' Directory exits'
    os.chdir( structdir )

for s in stations :
    for n, l in enumerate( layers ) :
        structure = cStringIO.StringIO()
        structure.write( '<!DOCTYPE DDDB SYSTEM "conddb:/DTD/structure.dtd"\n' )
        structure.write( ' [<!ENTITY Station         "%s">\n'%( s ) )
        structure.write( '  <!ENTITY Layer           "%s">\n'%( l ) )
        structure.write( '  <!ENTITY LayerID         "%s">\n'%( layerIDs[ n ] ) )
        structure.write( '  <!ENTITY StereoAngle     "%s">\n'%( layerStereo[ n ] ) )
        structure.write( '  <!ENTITY %s%sStructure SYSTEM "../LayerStructureTemplate.xml">]>\n'%( s, l ) )
        structure.write( '\n' )
        structure.write( '<DDDB>\n' )
        structure.write( '  &%s%sStructure;\n'%( s, l ) )
        structure.write( '</DDDB>\n' )
        contents = structure.getvalue()
        structure.close()
        fname = '%s%s@Structure.xml'%( s, l )
        output = file( fname, 'w' )
        output.write( contents )
        output.close()

