import sys
import string


class FdbTemplateNodes:
    """FdbTemplateNodes class looks for TEMPLATENODE IDs 
       and the SUBMODEL_NODE_MAP in fdb file
    """

    def __init__(self,file):
        """Read the fdb file:

           file - input fdb file 
        """
        
        #TEMPLATENODES
        #{
        #VERSION: 2
        #NUM_FRONTS: 1
        #FRONT: 0 0 97 0 1
        #NODES: 5215
        # 451379 479144 477032 472787 470012 458933 458931 480714 480310 480309

        #SUBMODEL_NODE_MAP
        #(
        # 14137 4962 15832 4963 17527 4965 5085 4964 6780 4966 8475 4967

        self.num_tn = 0
        self.tn = {}
        self.map = {}

        buff = self.GetLine(file)
        while buff != None:
            vals = buff.split()
            if len(vals) < 1: continue 

            if vals[0] == "TEMPLATENODES":
                # read template node IDs
                buff = self.GetLine(file) # {
                buff = self.GetLine(file) # VERSION
                buff = self.GetLine(file) # NUM_FRONTS:
                buff = self.GetLine(file) # FRONT: 
                buff = self.GetLine(file) # NODES: num_tn
                vals = buff.split()
                self.num_tn = int(vals[1])
                
                cnt = 0
                while buff != None:
                    buff = self.GetLine(file) 
                    vals = buff.split()
                    for j in range(len(vals)):
                        n = int(vals[j])
                        self.tn[n] = n
                        cnt+=1
                    if cnt >= self.num_tn: break

            if vals[0] == "SUBMODEL_NODE_MAP":
                # read submodel node ID map
                buff = self.GetLine(file) # (
                
                cnt = 0
                while buff != None:
                    buff = self.GetLine(file) 
                    vals = buff.split()
                    if vals[0] == ")" : break
                    for j in range(int(len(vals)/2)):
                        i = j*2 ;
                        lid = int(vals[i])
                        if lid in self.tn:
                            fid = int(vals[i+1])
                            self.map[fid] = lid 

            buff = self.GetLine(file)

    def GetLine(self,file):
        """helper function that scans the input file"""
        buff = file.readline()
        if len(buff) == 0: return None
        while buff[0] == '#':
            buff = file.readline()
            if len(buff) == 0: return None
        return buff


class PchDisplacement:
    """PchDisplacement class reads pch file"""

    def __init__(self,file,nids):
        """Read the pch file:

           file - input pch file 
           nids - template node IDs
        """

        # $TITLE   = MODEL FROM FRANC3D                                                  1
        # $SUBTITLE=                                                                     2
        # $LABEL   = ENDTIME=101075                                                      3
        # $DISPLACEMENTS                                                                 4
        # $REAL OUTPUT                                                                   5
        # $SUBCASE ID =           1                                                      6
        # $TIME =   0.1010750E+06                                                        7
        #     212853       G      4.361379E-03      1.806554E-03      1.013106E-03       8
        # -CONT-                  0.000000E+00      0.000000E+00      0.000000E+00       9


        self.num_ls = 0
        self.ls_dsp = {}

        buff = self.GetLine(file)
        
        while buff != None:
            vals = buff.split()
            if vals[0] == "$SUBCASE":
                if len(vals) > 3: ls = int(vals[3])

                buff = self.GetLine(file)
                vals = buff.split()
                if vals[0] == "$TIME":
                    buff = self.GetLine(file)
                    vals = buff.split()

                dsp = {}
                while buff != None:
                    if vals[1] == "G":
                        nid = int(vals[0])
                        if nid in nids:
                            vals = buff.split()
                            ndsp = []
                            for j in range(3): ndsp.append(float(vals[j+2]))
                            dsp[nid] = ndsp
                        buff = self.GetLine(file) # -CONT- skip 3_rotations

                        buff = self.GetLine(file)
                        if buff != None:
                            vals = buff.split()
                        else:
                            break
                    else:
                        break

                self.num_ls = ls
                self.ls_dsp[ls] = dsp

            buff = self.GetLine(file)

    def GetLine(self,file):
        """helper function that scans the input file"""
        buff = file.readline()
        if len(buff) == 0: return None
        while buff[0] == '#':
            buff = file.readline()
            if len(buff) == 0: return None
        if len(buff) > 73: buff = buff[:72]
        return buff


class DtpTemplate:
    """DtpTemplate class writes dtp file"""

    def __init__(self,fo,nls,ls_dsp):
        """Write the dtp file:

           fo - output dtp file 
           nls - number of load steps
           dsp - template node displacements 
        """
        # NUM STEP n
        # LOADSTEP i
        # DISPLACEMENT
        # 245 -0.00232298974879086 0.008318539708852768 -0.003217903897166252

        print('NUM STEP', nls, file=fo)
        for ls in range(nls):
            print('LOADSTEP', ls+1, file=fo)
            print('DISPLACEMENT', file=fo)
            dsp = ls_dsp[ls+1]
            for n in dsp:
                print(n, dsp[n][0], dsp[n][1], dsp[n][2], file=fo)
            

# -----------------------------------------------------------------

if __name__ == "__main__":

    if len(sys.argv) < 2:
        print("%s -fdb fn.fdb -pch fn.pch" % (sys.argv[0]))
        raise SystemExit

    if sys.argv[1] == "-fdb":
        fd = open(sys.argv[2], 'r')
        fdb_tn = FdbTemplateNodes(fd)
        fd.close()

    if sys.argv[3] == "-pch":
        fp = open(sys.argv[4], 'r')
        pch_dsp = PchDisplacement(fp,fdb_tn.map)
        fp.close()

    if sys.argv[5] == "-dtp" and pch_dsp.num_ls > 0:
        fo = open(sys.argv[6], 'w')
        DtpTemplate(fo,pch_dsp.num_ls,pch_dsp.ls_dsp)
        fo.close()