#!/usr/bin/python
# above line tells us which program to handle this script
#
# M.Zolliker June 2081
#
usage='''
Usage:

> fitlor -fit <output-file> <input-file>

   fit is called first, you should treat your peaks:
 
   fit> data 5182       get data file 5182
   fit> data [34]       get peak number 34 from above datafile
   fit> fun 0           as an example: select one gaussian
   fit> fit             you might play with parameters and reapeat
   fit> k               saving the data on the intermediate <input-file>
                        k means: keep data of this peak, might also be used
                        in the plot window
   fit> exit            exit and run the conversion program

when starting fit from fitlor, the commands 'open' and 'k' are used to
configure the outputfile for the needs of single crystal analysis

> fitlor <output-file> <input-file>

   do the conversion of already existing input files

if file names are omitted, default values are used
'''

import sys
from math import sin, cos, radians
from subprocess import call
from collections import Mapping
from tempfile import NamedTemporaryFile

# each entry is contained witinh curly brackets {} and
# contains the variable name and the format separated by ':'
# <n>d      n characters wide integer
# <n>.<m>d  n characterws wide fixed point with m decimal digits

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}'

# treating the arguments
if len(sys.argv) > 1 and sys.argv[1] == '-fit':
    call_fit = True
    sys.argv.pop(1)
else:
    call_fit = False
    
if len(sys.argv) > 1:
    outfile = sys.argv[1]
    sys.argv.pop(0)
else:
    outfile = 'fitlor_out.txt'
    print('output-file: %s' % outfile)

if len(sys.argv) > 1:
    datafile = sys.argv[1]
    sys.argv.pop(1)
else:
    datafile = 'fitlor_data.txt'
    print('input-file: %s' % datafile)

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' % datafile)
        cmdfil.flush()
        call(('fit', '-F', cmdfil.name))
    
tilt_geometry = False
lorentz_correction = True

class Row(object):
    def __init__(self, keys, values, fmt):
        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)
	self.fmt = fmt

    def write(self, fil):
	fil.write(self.fmt.format(**self.__dict__) + '\n')

keys = None
rows = []
try:
  with open(datafile) as inp:
    for line in inp:
        values = line.split()
        try:
            row = Row(keys, values, output_format)
        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
except IOError:
    print(usage)
  
with open(outfile, 'w') as out:
    for row in rows:
        row.write(out)