"""
1.This Software copyright © Australian Synchrotron Research Program Inc, ("ASRP").

2.Subject to ensuring that this copyright notice and licence terms
appear on all copies and all modified versions, of PyCIFRW computer
code ("this Software"), a royalty-free non-exclusive licence is hereby
given (i) to use, copy and modify this Software including the use of
reasonable portions of it in other software and (ii) to publish,
bundle and otherwise re-distribute this Software or modified versions
of this Software to third parties, provided that this copyright notice
and terms are clearly shown as applying to all parts of software
derived from this Software on each occasion it is published, bundled
or re-distributed.  You are encouraged to communicate useful
modifications to ASRP for inclusion for future versions.

3.No part of this Software may be sold as a standalone package.

4.If any part of this Software is bundled with Software that is sold,
a free copy of the relevant version of this Software must be made
available through the same distribution channel (be that web server,
tape, CD or otherwise).

5.It is a term of exercise of any of the above royalty free licence
rights that ASRP gives no warranty, undertaking or representation
whatsoever whether express or implied by statute, common law, custom
or otherwise, in respect of this Software or any part of it.  Without
limiting the generality of the preceding sentence, ASRP will not be
liable for any injury, loss or damage (including consequential loss or
damage) or other loss, loss of profits, costs, charges or expenses
however caused which may be suffered, incurred or arise directly or
indirectly in respect of this Software.

6. This Software is not licenced for use in medical applications.
"""

from types import *
class CifFile:
    def __init__(self,datasource=None,strict=1,maxinlength=2048,maxoutlength=0):
        self.dictionary = {}
        self.maxinlength = maxinlength
        if maxoutlength == 0:
            self.maxoutlength = maxinlength
        else:
            self.maxoutlength = maxoutlength
        self.strict = strict
        if isinstance(datasource,DictType):
            for (key,value) in datasource.items():
                self.__setitem__(key,value)
        elif type(datasource) is StringType:
            self.ReadCif(datasource,strict,maxinlength)
        elif isinstance(datasource,CifFile):
            self.dictionary = datasource.dictionary.copy()
        if strict:
            self.checklengths()

    def checklengths(self):
        blocks = self.dictionary.items()
        for name,block in blocks:
            toolong = len(filter(lambda a:len(a)>75, block.keys()))
            if toolong:
                print 'Warning: block ' + name + ' has ' + `toolong` + ' overlength data names'

    def __str__(self):
        return self.WriteOut()

    def __setitem__(self,key,value):
        if isinstance(value,CifBlock):
            self.NewBlock(key,value)
        else: raise TypeError

    def __getitem__(self,key):
        return self.dictionary[key]

    def __delitem__(self,key):
        del self.dictionary[key]

    def __len__(self):
        return len(self.dictionary)

    def keys(self):
        return self.dictionary.keys()

    def has_key(self,key):
        return self.dictionary.has_key(key)

    def get(self,key,default=None):
        if self.dictionary.has_key(key):
            return self.dictionary[key]
        else: 
            return default

    def clear(self):
        self.dictionary.clear()

    def copy(self):   
        newcopy = self.dictionary.copy()
        return CifFile('',newcopy)
     
    def update(self,adict):
        for key in adict.keys():
            self.dictionary[key] = adict[key]

    def ReadCif(self,filename,strict,maxlength):
        import YappsCifParser,string
        stream = open(filename,'r')
        text = stream.read()
        stream.close()
        if not text:      # empty file, return empty block
            return
        split = string.split(text,'\n')
        if strict:
            toolong =  filter(lambda a:len(a)>maxlength, split)
            if toolong:
                pos = split.index(toolong[0])
                raise CifError, 'Line %d contains more than %d characters' % (pos+1,maxlength)
        context = {"loops":[],"latest":{}}
        try: 
            parser = YappsCifParser.CifParser(YappsCifParser.CifParserScanner(text))
            filecontents = getattr(parser,"input")()
        except YappsCifParser.SyntaxError:
            errorstring = 'Syntax error in input file: last value parsed was %s' % YappsCifParser.lastval
            errorstring = errorstring + '\nParser status: %s' % `parser._scanner`
            raise CifError, errorstring
        else:
            if not filecontents: # comments only, return empty
                return
        for block in filecontents.keys():
            self.dictionary.update({block:CifBlock(filecontents[block])})

    def NewBlock(self,blockname,blockcontents=()):
        import re
        if not blockcontents:
            blockcontents = CifBlock()
        newblockname = re.sub('\W','_',blockname)
        blocknames = self.dictionary.keys()
        i = 0
        while blocknames.count(newblockname):
            i = i + 1
            newblockname = newblockname+`i`
        if len(newblockname) > 75:
            raise CifError, 'Cif block name too long:' + newblockname
        if blockcontents.__class__.__name__[-8:] != 'CifBlock':
            raise CifError, 'Cif block initialised with non CifBlock object %s' % blockcontents.__class__.__name__[-8:]
        self.dictionary.update({newblockname:blockcontents})
        return newblockname

    def WriteOut(self,comment=''):
        import cStringIO
        if not comment:
            comment = \
"""#\\#CIF1.1
##########################################################################
#               Crystallographic Information Format file 
#               Produced by PyCifRW module
# 
#  This is a CIF file.  CIF has been adopted by the International
#  Union of Crystallography as the standard for data archiving and 
#  transmission.
#
#  For information on this file format, follow the CIF links at
#  http://www.iucr.org
##########################################################################
"""
        outstring = cStringIO.StringIO()
        outstring.write(comment)
        for datablock in self.dictionary.keys():
            outstring.write('\ndata_'+datablock+'\n')
            outstring.write(str(self.dictionary[datablock]))
        returnstring =  outstring.getvalue()
        outstring.close()
        return returnstring


class CifBlock:
    def __init__(self,data = (),maxoutlength=2048,wraplength=80):
        self.block = {"loops":[]}
        self.maxoutlength = maxoutlength
        self.wraplength = wraplength
        if type(data) is DictType:     #direct placement
            self.block = data
            if not self.block.has_key("loops"):
                self.block.update({"loops":[]})
        elif type(data) is TupleType:
            for item in data:
                self.AddCifItem(item)
        else: raise TypeError

    def __str__(self):
        return self.printsection()

    def __setitem__(self,key,value):
        self.AddCifItem((key,value))

    def __getitem__(self,key):
        return self.GetCifItem(key)

    def __delitem__(self,key):
        self.RemoveCifItem(key)

    def __len__(self):
        blen = len(self.block) - 1   #non-looped items
        for aloop in self.block["loops"]:
            blen = blen + len(aloop.keys())
        return blen    

    def __nonzero__(self):
        if len(self.block) == 1 and len(self.block["loops"]) == 0:
            return 0
        return 1

    def keys(self):
        thesekeys = self.block.keys()
        for aloop in self.block["loops"]:
            thesekeys.extend(aloop.keys())
        return thesekeys

    def has_key(self,key):
        if self.block.has_key(key):
            return 1
        for aloop in self.block["loops"]:
            if aloop.has_key(key):
                return 1
        return 0

    def get(self,key,default=None):
        if self.has_key(key):
            retval = self.GetCifItem(key)
        else:
            retval = default
        return retval

    def clear(self):
        self.block = self.NewBlock()

    def copy(self):
        newcopy = self.block.copy()
        newcopy["loops"] = []
        for aloop in self.block["loops"]:  # do a deeper copy
            newcopy["loops"].append(aloop.copy())
        return CifBlock(newcopy)
     
    def update(self,adict):
        loopdone = []
        if not isinstance(adict,CifBlock):
            raise TypeError
        for key in adict.block.keys():
            if key!="loops":
                self.AddCifItem((key,adict[key]))
            else:
                for aloop in adict.block["loops"]:
                    self.AddCifItem((aloop.keys(),aloop.values()))

    def GetCifItem(self,itemname):
        if self.block.has_key(itemname):
            return self.block[itemname]
        else:
            for aloop in self.block["loops"]:
                if aloop.has_key(itemname):
                    return aloop[itemname]
        raise KeyError, 'Item not in Cif block'

    def RemoveCifItem(self,itemname):
        if self.block.has_key(itemname):
            del self.block[itemname]
            return
        for aloop in self.block["loops"]:
            if aloop.has_key(itemname):
                del aloop[itemname]
        self.block["loops"] = filter(None, self.block["loops"])

    def AddCifItem(self,data):
        # we accept only tuples, strings and lists!!
        if not (isinstance(data[0],StringType) or isinstance(data[0],TupleType)
              or isinstance(data[0],ListType)):
                  raise TypeError, 'Cif datanames are either a string, tuple or list'
        # now put into the dictionary properly...
        if isinstance(data[0],StringType):   # a single name
            self.check_data_name(data[0])    # make sure no nasty characters   
            # now make sure the data is OK
            regval = self.regularise_data(data[1])
            self.check_item_value(regval)
            self.block.update({data[0]:regval})  # trust the data is OK
            for aloop in self.block["loops"]:
                if aloop.has_key(data[0]):
                    del aloop[data[0]]
            self.block["loops"] = filter(len,self.block["loops"])
        else:                                # we loop
           if(len(data[0])!=len(data[1])):
               raise TypeError, 'Length mismatch between itemnames and values'
           dellist = []
           map (self.check_data_name,data[0])
           for itemname in data[0]:
               self.block["loops"] = filter(lambda a,b=itemname:b not in a.keys(),self.block["loops"])
           newdict = {}
           proper_vals = map(self.regularise_data,data[1])
           map(self.check_item_value,proper_vals)
           map(lambda a,b,c=newdict:c.update({a:b}),data[0],proper_vals)
           self.block["loops"].append(newdict)
        return

    def check_data_name(self,dataname): 
        if len(dataname) > 75:
            raise CifError, 'Dataname ' + dataname + ' too long.'
        if dataname[0]!='_':
            raise CifError, 'Dataname ' + dataname + ' does not begin with _'
        if len (filter (lambda a: ord(a) < 33 or ord(a) > 126, dataname)) > 0:
            raise CifError, 'Dataname ' + dataname + ' contains forbidden characters'
     
    def check_item_value(self,item):
        test_item = item
        if type(item) != TupleType and type(item) != ListType:
           test_item = [item]         #single item list
        def check_one (it):
            if type(it) == StringType:
                if len (filter (lambda a:ord(a) != 9 and ord(a) != 10 and ord(a) != 13
                        and (ord(a) < 32 or ord(a) > 126), it)) > 0:
                    raise CifError, 'Data item ' + it + '... contains forbidden characters'
        map(check_one,test_item)

    def regularise_data(self,dataitem):
        alrighttypes = [IntType, LongType, 
                        FloatType, StringType]
        okmappingtypes = [TupleType, ListType]
        thistype = type(dataitem)
        if thistype in alrighttypes or thistype in okmappingtypes:
            return dataitem
        # so try to make into a list
        try:
            regval = list(dataitem)
        except TypeError, value:
            raise CifError, str(dataitem) + ' is wrong type for data value\n' 
        return regval
        
    def GetLoop(self,itemname):
        for aloop in self.block["loops"]:
            if aloop.has_key(itemname):
                return aloop.items()
        # not a looped item
        if self.block.has_key(itemname):
            raise TypeError, 'Non-looped item'
        raise KeyError, 'Item not in loop'

    def AddToLoop(self,dataname,loopdata):
        found = 0
        for aloop in self.block["loops"]:
            if aloop.has_key(dataname):
                found = 1
                for itemname in loopdata.keys():
                    if len(loopdata[itemname])!= len(aloop[dataname]):
                        raise CifError, 'Datalength mismatch adding to loop: %s and %s' % (itemname, dataname)
                    aloop.update({itemname:loopdata[itemname]})
        if not found:
            raise KeyError, 'No such looped item name: %s' % dataname

    def printsection(self,order=[]):
        import cStringIO
        import string
        # first make an ordering
        if not order:
            order = self.block.keys()
            order.sort()
        # now prune that ordering...
        order = filter(lambda a,b=self.block:b.has_key(a),order)
        order.remove('loops')
        # now do it...
        outstring = cStringIO.StringIO()       # the returned string
        for itemname in order:
            itemvalue = self.block[itemname]
            if isinstance(itemvalue,StringType):
                  thisstring = self._formatstring(itemvalue)
                  if len(thisstring) + len(itemname) < (self.wraplength-2):
                          outstring.write('%s %s\n' % (itemname,thisstring))
                  else:
                          outstring.write('%s\n %s\n' % (itemname, thisstring))
            else: 
                      if len(str(itemvalue)) + len(itemname) < (self.wraplength-2):
                          outstring.write('%s %s\n' % (itemname, itemvalue))
                      else:
                          outstring.write('%s\n %s\n' % (itemname, itemvalue))
            continue
        #do the loops
        for aloop in self.block["loops"]:
               outstring.write('\n loop_\n')
               loopnames = aloop.keys()
               loopnames.sort()
               numdata = len(aloop[loopnames[0]])
               for name in loopnames: 
                   outstring.write('   %-75s\n' % name)
                   if len(aloop[name]) != numdata:
                       raise CifError,'Loop data mismatch for ' + name + ':output aborted'
               curstring = ''      
               # when adding to outstring, make sure to add spaces
               for position in range(numdata):
                   for name in loopnames:
                       # at each point, get the next data value
                       datapoint = aloop[name][position]
                       if isinstance(datapoint,StringType):
                           thisstring = '%s' % (self._formatstring(datapoint)) #no spaces yet
                           if '\n' in thisstring:
                               # we try to wrap if the first <eol> is bigger than our wrap length
                               if len(curstring) + string.find(thisstring,'\n') > (self.wraplength):
                                   outstring.write(' ' + curstring + '\n' + thisstring)
                               else:
                                   outstring.write(' ' + curstring + ' ' + thisstring) #a space
                               curstring = ''
                               continue
                       else: 
                           thisstring = ' %s ' % datapoint
                       if len(curstring) + len(thisstring)> self.wraplength-2: #past end of line with space
                           outstring.write(' ' + curstring+'\n') #add the space
                           curstring = ''
                       curstring = curstring + ' ' + thisstring
                   outstring.write(' ' + curstring + '\n')    #last time through
                   curstring = ''
        returnstring = outstring.getvalue()
        outstring.close()
        return returnstring

    def _formatstring(self,instring):
        import re, string
        if len(instring)< (self.maxoutlength-2) and '\n' not in instring and not ('"' in instring and '\'' in instring):
            if not ' ' in instring and not '\t' in instring and not '\v' \
              in instring:                  # no blanks
                return ' %s ' % (instring)
            if not "'" in instring:                                       #use apostrophes
                return "'%s'" % (instring)
            elif not "\"" in instring:
                return '"%s"' % (instring)
        # is a long one or one that needs semicolons due to carriage returns
        outstring = "\n;\n"
        # if there are returns in the string, try to work with them
        while 1:
            retin = string.find(instring,'\n')+1
            if retin < self.maxoutlength and retin > 0:      # honour this break
                outstring = outstring + instring[:retin]
                instring = instring[retin:]
            elif len(instring)<self.maxoutlength:            # finished
                outstring = outstring + instring + '\n;\n'
                break
            else:                             # find a space
                for letter in range(self.maxoutlength-1,40,-1): 
                    if instring[letter] in ' \t\f': break
                outstring = outstring + instring[:letter+1]
                outstring = outstring + '\n'
                instring = instring[letter+1:]            
        return outstring


class CifError(Exception):
    def __init__(self,value):
        self.value = value
    def __str__(self):
        print '\nCif Format error: '+ self.value