#! /usr/bin/python

################################################################################
# Python script to generate OT Stations 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' ]
stationIDs  = [ '1', '2', '3' ]

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

def StationGeometryTemplate( buffer, indent ) :
    buffer.write( '<?xml version=%s encoding=%s?>\n\n'%( version, encoding ) )
    buffer.write( '\n' )
    buffer.write( '  <catalog name="&Station;">\n' )
    buffer.write( '    <logvolref  href="#lv&Station;" />\n' )
    buffer.write( '    <catalogref href="../../Layers/Geometry/&Station;X1@Geometry.xml#&Station;X1" />\n' )
    buffer.write( '    <catalogref href="../../Layers/Geometry/&Station;U@Geometry.xml#&Station;U" />\n' )
    buffer.write( '    <catalogref href="../../Layers/Geometry/&Station;V@Geometry.xml#&Station;V" />\n' )
    buffer.write( '    <catalogref href="../../Layers/Geometry/&Station;X2@Geometry.xml#&Station;X2" />\n' )
    buffer.write( '  </catalog>\n' )
    buffer.write( '\n' )
    buffer.write( '  <logvol name="lv&Station;" material="Air">\n' ) 
    buffer.write( '    <subtraction name="&Station;BeamPipeHole">\n' ) 
    buffer.write( '      <box name="&Station;Box"\n' )
    buffer.write( '           sizeX="StationXSize"\n' ) 
    buffer.write( '           sizeY="StationYSize"\n' )
    buffer.write( '           sizeZ="StationZSize" />\n' )
    buffer.write( '      <box name="StationSModCoupHole"\n' )
    buffer.write( '           sizeX="263*mm"\n' )
    buffer.write( '           sizeY="2*135*mm"\n' )
    buffer.write( '           sizeZ="StationZSize+Clearance" />\n' )
    buffer.write( '    </subtraction>\n' )
    buffer.write( '    <physvol name="pv&Station;X1"\n' ) 
    buffer.write( '             logvol="/dd/Geometry/AfterMagnetRegion/T/OT/&Station;/&Station;X1/lv&Station;X1">\n' )
    buffer.write( '      <transformation>\n' )
    buffer.write( '        <posXYZ />\n' )
    buffer.write( '        <rotXYZ rotX="beamAngle" />\n' )
    buffer.write( '        <posXYZ />\n' )
    buffer.write( '        <posXYZ x="0.0*mm"\n' )
    buffer.write( '                y="&X1y;"\n' )
    buffer.write( '                z="&X1z;"  />\n' )
    buffer.write( '        <rotXYZ />\n' )
    buffer.write( '        <posXYZ />\n' )
    buffer.write( '      </transformation>\n' )
    buffer.write( '    </physvol>\n' )
    buffer.write( '    <physvol name="pvXUCFrame"\n' )
    buffer.write( '             logvol="/dd/Geometry/AfterMagnetRegion/T/OT/CFrames/lvClosedCFrames">\n' )
    buffer.write( '      <transformation>\n' )
    buffer.write( '        <posXYZ />\n' )
    buffer.write( '        <rotXYZ rotX="beamAngle" />\n' )
    buffer.write( '        <posXYZ />\n' )
    buffer.write( '        <posXYZ x="0.0*mm"\n' )
    buffer.write( '                y="&XUy;"\n' )
    buffer.write( '                z="&XUz;"/>\n' )
    buffer.write( '        <rotXYZ />\n' )
    buffer.write( '        <posXYZ />\n' )
    buffer.write( '      </transformation>\n' )
    buffer.write( '    </physvol>\n' )
    buffer.write( '    <physvol name="pv&Station;U"\n' )
    buffer.write( '             logvol="/dd/Geometry/AfterMagnetRegion/T/OT/&Station;/&Station;U/lv&Station;U">\n' )
    buffer.write( '      <transformation>\n' )
    buffer.write( '        <posXYZ />\n' )
    buffer.write( '        <rotXYZ rotX="beamAngle" />\n' )
    buffer.write( '        <posXYZ />\n' )
    buffer.write( '        <posXYZ x="0.0*mm"\n' )
    buffer.write( '                y="&Uy;"\n' )
    buffer.write( '                z="&Uz;"  />\n' )
    buffer.write( '        <rotXYZ />\n' )
    buffer.write( '        <posXYZ />\n' )
    buffer.write( '      </transformation>\n' )
    buffer.write( '    </physvol>\n' )
    buffer.write( '    <physvol name="pv&Station;V"\n' )
    buffer.write( '             logvol="/dd/Geometry/AfterMagnetRegion/T/OT/&Station;/&Station;V/lv&Station;V">\n' )
    buffer.write( '      <transformation>\n' )
    buffer.write( '        <posXYZ />\n' )
    buffer.write( '        <rotXYZ rotX="beamAngle" />\n' )
    buffer.write( '        <posXYZ />\n' )
    buffer.write( '        <posXYZ x="0.0*mm"\n' )
    buffer.write( '                y="&Vy;"\n' )
    buffer.write( '                z="&Vz;"  />\n' )
    buffer.write( '        <rotXYZ />\n' )
    buffer.write( '        <posXYZ />\n' )
    buffer.write( '      </transformation>\n' )
    buffer.write( '    </physvol>\n' )
    buffer.write( '    <physvol name="pvXVCFrame"\n' )
    buffer.write( '             logvol="/dd/Geometry/AfterMagnetRegion/T/OT/CFrames/lvClosedCFrames">\n' )
    buffer.write( '      <transformation>\n' )
    buffer.write( '        <posXYZ />\n' )
    buffer.write( '        <rotXYZ rotX="beamAngle" />\n' )
    buffer.write( '        <posXYZ />\n' )
    buffer.write( '        <posXYZ x="0.0*mm"\n' )
    buffer.write( '                y="&VXy;"\n' )
    buffer.write( '                z="&VXz;"/>\n' )
    buffer.write( '        <rotXYZ />\n' )
    buffer.write( '        <posXYZ />\n' )
    buffer.write( '      </transformation>\n' )
    buffer.write( '    </physvol>\n' )
    buffer.write( '    <physvol name="pv&Station;X2"\n' )
    buffer.write( '             logvol="/dd/Geometry/AfterMagnetRegion/T/OT/&Station;/&Station;X2/lv&Station;X2">\n' )
    buffer.write( '      <transformation>\n' )
    buffer.write( '        <posXYZ />\n' )
    buffer.write( '        <rotXYZ rotX="beamAngle" />\n' )
    buffer.write( '        <posXYZ />\n' )
    buffer.write( '        <posXYZ x="0.0*mm"\n' )
    buffer.write( '                y="&X2y;"\n' )
    buffer.write( '                z="&X2z;"  />\n' )
    buffer.write( '        <rotXYZ />\n' )
    buffer.write( '        <posXYZ />\n' )
    buffer.write( '      </transformation>\n' )
    buffer.write( '    </physvol>\n' )
    buffer.write( '  </logvol>\n' )
    buffer.write( '\n' )
    
def StationStructureTemplate( buffer, indent ) :
    buffer.write( '<?xml version=%s encoding=%s?>\n\n'%( version, encoding ) )
    buffer.write( '<detelem classID="8102" name="&Station;">\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;/lv&Station;"\n' )
    buffer.write( '              condition="/dd/Conditions/Alignment/OT/&Station;"\n' )
    buffer.write( '              support="/dd/Structure/LHCb/AfterMagnetRegion/T/OT"\n' )
    buffer.write( '              npath="pv&Station;" />\n' )
    buffer.write( '  <param name="stationID" type="int">&StationID;</param>\n' )
    buffer.write( '  <detelemref href="../../Layers/Structure/&Station;X1@Structure.xml#X1" />\n' )
    buffer.write( '  <detelemref href="../../Layers/Structure/&Station;U@Structure.xml#U"   />\n' )
    buffer.write( '  <detelemref href="../../Layers/Structure/&Station;V@Structure.xml#V"   />\n' )
    buffer.write( '  <detelemref href="../../Layers/Structure/&Station;X2@Structure.xml#X2" />\n' )
    buffer.write( '</detelem>\n' )
    buffer.write( '\n' )

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

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

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

# Nominal
zOffsets = [ ('+0.0*mm', '+0.0*mm'), ('+0.0*mm', '+0.0*mm'), ('+0.0*mm', '+0.0*mm') ] 

# Stations are closer to magnet compared to nominal
#zOffsets = [ ('-0.0*mm', '-0.5*mm'), ('-2.5*mm', '-3.5*mm'), ('-0.0*mm', '-2.5*mm') ] 

for n, s in enumerate( stations ) :
    geometry = cStringIO.StringIO()
    geometry.write( '<!DOCTYPE DDDB SYSTEM "conddb:/DTD/geometry.dtd"\n' )
    geometry.write( ' [<!ENTITY Station         "%s">\n'%( s ) )
    geometry.write( '  <!ENTITY X1y             "%s">\n'%( s + 'X1Y' ) )
    geometry.write( '  <!ENTITY X1z             "%s">\n'%( s + 'X1Z' + zOffsets[n][0] ) )
    geometry.write( '  <!ENTITY XUy             "%s">\n'%( s + 'XUY' ) )
    geometry.write( '  <!ENTITY XUz             "%s">\n'%( s + 'XUZ' + zOffsets[n][0] ) )
    geometry.write( '  <!ENTITY Uy              "%s">\n'%( s + 'UY'  ) )
    geometry.write( '  <!ENTITY Uz              "%s">\n'%( s + 'UZ'  + zOffsets[n][0] ) )
    geometry.write( '  <!ENTITY Vy              "%s">\n'%( s + 'VY'  ) )
    geometry.write( '  <!ENTITY Vz              "%s">\n'%( s + 'VZ'  + zOffsets[n][1] ) )
    geometry.write( '  <!ENTITY VXy             "%s">\n'%( s + 'VXY' ) )
    geometry.write( '  <!ENTITY VXz             "%s">\n'%( s + 'VXZ' + zOffsets[n][1] ) )
    geometry.write( '  <!ENTITY X2y             "%s">\n'%( s + 'X2Y' ) )
    geometry.write( '  <!ENTITY X2z             "%s">\n'%( s + 'X2Z' + zOffsets[n][1] ) )
    geometry.write( '  <!ENTITY %sGeometry SYSTEM "../StationGeometryTemplate.xml">]>\n'%( s ) )
    geometry.write( '\n' )
    geometry.write( '<DDDB>\n' )
    geometry.write( '  &%sGeometry;\n'%( s ) )
    geometry.write( '</DDDB>\n' )
    contents = geometry.getvalue()
    geometry.close()
    fname = '%s@Geometry.xml'%( s )
    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 n, s in enumerate( stations ) :
    structure = cStringIO.StringIO()
    structure.write( '<!DOCTYPE DDDB SYSTEM "conddb:/DTD/structure.dtd"\n' )
    structure.write( ' [<!ENTITY Station         "%s">\n'%( s ) )
    structure.write( '  <!ENTITY StationID       "%s">\n'%( stationIDs[ n ] ) )
    structure.write( '  <!ENTITY %sStructure SYSTEM "../StationStructureTemplate.xml">]>\n'%( s ) )
    structure.write( '\n' )
    structure.write( '<DDDB>\n' )
    structure.write( '  &%sStructure;\n'%( s ) )
    structure.write( '</DDDB>\n' )
    contents = structure.getvalue()
    structure.close()
    fname = '%s@Structure.xml'%( s )
    output = file( fname, 'w' )
    output.write( contents )
    output.close()

