#!/usr/bin/python
import sys
from math import sin, cos, radians
from subprocess import call
from collections import Mapping
from tempfile import NamedTemporaryFile

output_format = '{n:6d}{h:4d}{k:4d}{l:4d}'\
    '{i1:10.2f}{sigi1:10.2f}'\
    '{th1:8.2f}{p1:8.2f}{chi:8.2f}{phi:8.2f}'

if len(sys.argv) > 0 and sys.argv[0] == '-fit':
    call_fit = True
    sys.argv.pop(0)
else:
    call_fit = False
    
if len(sys.argv) > 0:
    datfile = sys.argv[0]
    sys.argv.pop(0)
else:
    datfile = 'fitlor_data.txt'

if len(sys.argv) > 0:
    outfile = sys.argv[0]
    sys.argv.pop(0)
else:
    outfile = 'fitlor_out.txt'

if call_fit:
    with NamedTemporaryFile('w') as cmdfil:
        cmdfil.write('k h,k,l,i1,p1,intexp,two_theta,chi,phi\n')
        cmdfil.write('open %s\n' % outfile)
        call(('fit', '-F', cmdfil.name))
    
tilt_geometry = False
lorentz_correction = True

class Row(object):
    def __init__(self, keys, values):
        if keys is None:
            raise ValueError('no keys')
        for k,v in zip(keys, values):
            v = float(v)
            if v == int(v):
                v = int(v)
            setattr(self, k.lower(), v)

keys = None
rows = []
with open(datfile) as inp:
    for line in inp:
        values = line.split()
        try:
            row = Row(keys, values)
        except ValueError:
            keys = values
            continue
        rows.append(row)
        row.n = len(rows)
        row.th1=row.two_theta * 0.2
        if lorentz_correction:
            lor_cor = sin(abs(radians(row.two_theta)))
            if tilt_geometry:
                lor_cor *= cos(radians(row.chi))
        row.intcor = row.i1 * lor_cor
        row.sigcor = row.sigi1 * lor_cor
    
with open(outfile, 'w') as out:
    for row in rows:
        out.write(output_format.format(row.__dict__) + '\n')