#!@PYTHON@

# Boxmaker 1.0
# Creates tremolo-datafiles with arbitrary size and specific density from a single input molecule,
# supporting numerous pre- and postprocessing features such as unit conversion.
# Gregor Bollerhey - bollerhe@ins.uni-bonn.de


import re, os, os.path, sys, operator
import pyMoleCuilder as mol

avogadro =  6.022143e23

class c_opt():
    basename = None
    tremofiledir = './'
    potentialsfiledir = './'
    outfilename = 'out'

    source = None
    molarmass = None
    density = None
    temp = None

    number = '1000'
    
    cubicdomain = 'on'
    cubiccell = 'off'
    autorotate = 'off'
    autodim = 'on'
    postprocess = 'on'
    automass = 'on'

    def update(self, name, value):
        shortcuts = {'tf': 'temofiledir', 'pf': 'potentialsfiledir', 'o': 'outfilename',
        'i': 'source', 'm': 'molarmass', 'rho': 'density',
        't': 'temp', 'n': 'number', 'cd': 'cubicdomain',
        'cc': 'cubiccell', 'ar': 'autorotate', 'ad': 'autodim',
        'pp': 'postprocess', 'am': 'automass'}
        
        if name in shortcuts:
            name = shortcuts[name]
        
        if name in dir(self):
            exec('self.%s = "%s"' % (name, value))
        else:
            print 'Warning: Unknown option:', name


def ReadSettings(opt):
    # Obtain basename
    if len(sys.argv) >= 2:
        opt.basename = sys.argv[1]
    else:
        print 'Usage: boxmaker.py <basename> [options]'
        exit()

    # Read settings file
    try:
        with open('boxmaker.' + opt.basename) as f:
            for line in f:
                if len(line) > 0 and line[0] != '#':
                    L, S, R = line.partition('=')
                    opt.update(L.strip(), R.strip())
    except IOError:
        print 'Warning: Configuration file not readable, CLI only'

    # Parse parameters
    i = 2
    while i < len(sys.argv):
        L = sys.argv[i]

        if L[0] in '+-':
            LN = L[1:]

            if L[0] == '+':
                R = 'on'
            else:
                R = 'off'
        else:
            LN = L
            i += 1
            R = sys.argv[i]
        
        opt.update(LN, R)
        i += 1


def ReadUnits(opt):
    lines = [] # The file needs to be processed twice, so we save the lines in the first run

    with open(opt.tremofiledir + opt.basename + '.tremolo') as f:
        for line in f:
            if len(line) > 0 and line[0] != '#':
                line = line.strip()
                lines.append(line)

                if 'systemofunits' in line:
                    L, S, SOU = line.partition('=')
                    SOU = SOU.strip()[:-1] # Remove semicolon

    if SOU == 'custom':
        units = {}
        quantities = ['length', 'mass', 'temperature']

        for quantity in quantities:
            units[quantity] = [None, None] # Init with scaling factor and unit 'None'.

        for line in lines:
            for quantity in quantities:
                if quantity in line:
                    L, S, R = line.partition('=')
                    R = R.strip()[:-1] # Remove semicolon

                    if 'scalingfactor' in line:
                        units[quantity][0] = float(R)
                    else:
                        units[quantity][1] = R

    elif SOU == 'kcalpermole':
        units = {'length': [1.0, 'angstrom'], 'mass': [1.0, 'u'], 'temperature': [503.556, 'K']}

    elif SOU == 'evolt':
        units = {'length': [1.0, 'angstrom'], 'mass': [1.0, 'u'], 'temperature': [11604.0, 'K']}

    else: # SI
        units = {'length': [1.0, 'm'], 'mass': [1.0, 'kg'], 'temperature': [1.0, 'K']}

    return units


def ConvertUnits(have, want):
    if have[0] == '!':
        return float(have[1:])
    
    # Redo with pipes?
    #sys.stdout.write("units '%s' '%s'\n" % (have, want))
    ret = os.system("units '%s' '%s' > temp_units_output" % (have, want))

    if ret == 0:
        with open('temp_units_output') as f:
            line = f.readline()

        os.system('rm temp_units_output')

        return float(line[3:-1])
    else:
        raise NameError('UnitError')
        
        
def GetSourceMolareMass(opt, units):
    mol.SelectionAllAtoms()
    mol.wait()
    mass_sum = mol.getSelectedMolarMass()
    have = ("%f atomicmassunit" % mass_sum)
    want = ("%f*%s" % tuple(units['mass']))
    return ConvertUnits(have, want)*avogadro


def UpdateSettingsAndSource(opt):
    # Map boolean values
    boolmap = {'on': True, 'off': False}
    
    for name in ['cubicdomain', 'cubiccell', 'autorotate', 'autodim', 'postprocess', 'automass']:
        value = eval('opt.' + name)
        
        if value in boolmap:
            value = boolmap[value]
        else:
            print 'Not a boolean value:', value
            exit()
            
        exec('opt.' + name + '= value')

    # Convert dimensions
    if opt.autodim:
        units = ReadUnits(opt)

        if not opt.automass:
            have = opt.molarmass
            want = '%f*%s / mol' % tuple(units['mass'])
            opt.molarmass = ConvertUnits(have, want)

        have = opt.density
        want = '(%f*%s) ' % tuple(units['mass']) + '/ (%f*%s)**3' % tuple(units['length'])
        opt.density = ConvertUnits(have, want)
        
        if opt.temp:
            have = opt.temp
            want = '%f*%s' % tuple(units['temperature'])
            opt.temp = ConvertUnits(have, want)
    else:
        if not opt.automass:
            opt.molarmass = float(opt.molarmass)
        
        opt.density = float(opt.density)
        
        if opt.temp:
            opt.temp = float(opt.temp)

    # Number might be an integer or a 3-vector
    nvec = opt.number.split()
    if len(nvec) == 3:
        opt.number = [0]*3

        for i in range(3):
            opt.number[i] = int(nvec[i])
    else:
        opt.number = int(opt.number)
        
    InitialiseSource(opt)
        
    # Automatic source mass
    if opt.automass:
        opt.molarmass = GetSourceMolareMass(opt, units)
        print '======== MOLAR MASS:', opt.molarmass


def FindBestCube(opt):
    newroot = int( round(opt.number**(1./3)) )
    newnumber = newroot**3

    if newnumber != opt.number:
        print 'Warning: Number changed to %d.' % newnumber

    return [newroot] * 3


def FindBestCuboid(opt):
    n = opt.number

    # Prime factors of n
    factors = []

    for i in [2, 3]:
        while n % i == 0:
            factors.append(i)
            n /= 2

    t = 5
    diff = 2

    while t*t <= n:
        while n % t == 0:
            factors.append(t)
            n /= t

        t = t + diff
        diff = 6 - diff

    if n > 1:
        factors.append(n)

    # Even distribution of current biggest prime to each vector -> similar sizes
    if len(factors) < 3:
        print 'Warning: Not enough prime factors - falling back to cubic placement'
        return FindBestCube(opt)

    factors.sort()
    distri = [[],[],[]]
    current = 0

    for factor in factors:
        distri[current].append(factor)
        current += 1
        if current == 3:
            current = 0

    result = [0]*3

    print '======== CUBOID USED:',

    for i in range(3):
        result[i] = int( reduce(operator.mul, distri[i]) )

    print result
    return result


def GetSourceBBabs(opt):
    mol.WorldCenterOnEdge()
    mol.wait()
    bb = mol.getBoundingBox()
    bbmin = [bb[0], bb[2], bb[4]]
    bbmax = [bb[1], bb[3], bb[5]]

    bb = [0.0]*3

    for i in range(3):
        bb[i] = abs(bbmax[i] - bbmin[i])

    return bb
    
    
def InitialiseSource(opt):
    potfilepath = opt.potentialsfiledir + opt.basename + '.potentials'
    mol.ParserParseTremoloPotentials(potfilepath)
    mol.MoleculeLoad(opt.source)
    
    if opt.autorotate:
        mol.SelectAllAtoms()
        mol.RotateToPrincipalAxisSystem("0 1 0")
    mol.wait()

# Global options with sensible default parameters
opt = c_opt()

ReadSettings(opt)
UpdateSettingsAndSource(opt)

if type(opt.number) == type([]):
    # Number is a vector - use it without any modification
    nbox = opt.number
else:
    if opt.cubicdomain:
        nbox = FindBestCube(opt)
    else:
        nbox = FindBestCuboid(opt)

VolumePerMolecule = opt.molarmass / (avogadro * opt.density)
cell = [VolumePerMolecule**(1./3)] * 3

if not opt.cubiccell:
    try:
        bb = GetSourceBBabs(opt)
        print '======== BBOX:', bb
        # Scaling factor - the molecules bounding box is scaled to fit the volume suiting the density
        s = (VolumePerMolecule / (bb[0]*bb[1]*bb[2])) ** (1./3)

        if s < 1:
            print 'Warning: Molecular cells will overlap.'

        for i in range(3):
            cell[i] = bb[i]*s
    except ZeroDivisionError:
        print 'Warning:  Singularity in bounding box, falling back to cubic cell.'
        

print '======== CELL: ', cell

mol.CommandVerbose('0')
mol.WorldCenterInBox('%f 0 0 %f 0 %f' % tuple(cell))
mol.WorldRepeatBox('%d %d %d' % tuple(nbox))
mol.WorldOutputAs(opt.outfilename + '.data')
mol.WorldOutputAs(opt.outfilename + '.xyz')
mol.wait()

domain = [0.0]*3

for i in range(3):
    domain[i] = cell[i]*nbox[i]

print  '======== DOMAIN: ', domain

# Postprocessing

if opt.postprocess:
    with open(opt.outfilename + '.data') as f:
        ofile = f.read()
    
    with open(opt.outfilename + '.data', 'w') as f:
        f.write('# INPUTCONV shift center\n')
        
        if opt.temp:
            f.write('# INPUTCONV temp %.4f\n' % opt.temp)
            
        f.write(ofile)
        
os.system('rm temp_source.data temp_source.xyz')