from tempfile import NamedTemporaryFile

class Pose:
    def assign(self, other):
        pass

def compute_similarity(pose, filespec):
    pose_copy = Pose()
    pose_copy.assign(pose)

    #Setup the 9-mer fragment (9-mer is better than 3-mer for this analysis)
    fragset = ConstantLengthFragSet(9)
    fragset.read_fragment_file(filespec)

    movemap = MoveMap()
    movemap.set_bb(True)

    #Setup and apply the fragment inserting mover
    mover = ClassicFragmentMover(fragset, movemap)
    mover.apply(pose_copy, count)

    #Measure the RMSD difference between the original pose and the new changed pose (the copy)
    difference = rosetta.core.scoring.CA_rmsd(pose , pose_copy)
    return difference

def get_frag_suffix(position, fragment):
    """Return a string matching the end of fragment lines for this pos/frag"""
    return 'P{:3d} F{:3d}'.format(position, fragment)

def process_fragment(pose, tempfile, position, outfile):
    """Process the fragment data in tempfile, set up for next fragment"""
    tempfile.close()
    diff = compute_similarity(pose, tempfile.name)
    print(diff, position, sep='\t', file=outfile)

def RMSD(pose):
    # Isolate fragments
        rmsd = NamedTemporaryFile(mode='w', prefix='RMSD')
        temp_fragfile = None

        with open('aat000_09_05.200_v1_3' , 'r') as frag:
            for line in frag:
                if line.startswith(' position: '):
                    # Start of a new position block

                    if temp_fragfile is not None:
                        # Process end of prior position
                        process_fragment(pose, temp_fragfile, position, rmsd)

                    pos_line = line
                    position = int(line.split()[1])
                    fragment = 1
                    frag_suffix = get_frag_suffix(position, fragment)

                elif line.strip() == '':
                    continue

                elif line.endswith(frag_suffix):
                    print(line, file=temp_fragfile)

                else:
                    # End of previous fragment, start of a new one.
                    process_fragment(pose, temp_fragfile, position, rmsd)


    count = 0
    for x in range (int(size)):
        count +=1
        fragnum = 0
        for x in range(200):
            frag = open('aat000_09_05.200_v1_3' , 'r')
            temp = open('temp' , 'w')
            for line in frag:
                if line.startswith(' position:{:13d}'.format(count)):
                    temp.write(line)
                    #print(line)
            temp.write('\n')
            fragnum += 1
            frag = open('aat000_09_05.200_v1_3' , 'r')
            for line in frag:
                if line.endswith('P{:3d} F{:3d}'.format(count , fragnum) + '\n'):
                    temp.write(line)
                    #print(line)
                    sys.stdout.flush()
            temp.write('\n')
            temp.close()
            frag.close()
                        rmsd_diff = compute_similarity(pose, temp_fragfile.name)
                        #print(RMSD , '\t' , count , '\t' , fragnum)
                        rmsd.write(str(RMSD) + '\t' + str(count) + '\n')
            #Reset the copy pose to original pose
            pose_copy.assign(pose)
            os.remove('temp')
    rmsd.close()
    #Analyse the RMSD file to get the lowest RMSD for each position
    data = open('RMSDvsPosition.dat' , 'w')
    lowest = {}                         #Mapping group number -> lowest value found
    for line in open('RMSD.temp'):
        parts = line.split()
        if len(parts) != 2:             #Only lines with two items on it
            continue
        first = float(parts[0])
        second = int(parts[1])
        if first == 0:                  #Skip line with 0.0 RMSD (this is an error from the 9-mer fragment file). I don't know why it happens
            continue
        if second not in lowest:
            lowest[second] = first
        else:
            if first < lowest[second]:
                lowest[second] = first
    for position, rmsd in lowest.items():
        print(str(rmsd) + '\t' + str(position))
        data.write(str(position) + '\t' + str(rmsd) + '\n')
    data.close()
    os.remove('RMSD.temp')
    gnuplot = open('gnuplot_sets' , 'w')
    gnuplot.write("""set terminal postscript
set output './plot.pdf'
set encoding iso_8859_1
set term post eps enh color
set xlabel 'Position'
set ylabel 'RMSD (\\305)'
set yrange [0:]
set xrange [0:]
set xtics 1
set title 'Fragment Quality'
set key off
set boxwidth 0.5
set style fill solid
plot 'RMSDvsPosition.dat' with boxes
exit
""")
    gnuplot.close()
    os.system('gnuplot < gnuplot_sets')
    os.remove('gnuplot_sets')
#----------------------------------------------------------------------
pose = pose_from_pdb('structure.pdb')
RMSD(pose)