import argparse

parser = argparse.ArgumentParser(description="Translate a dx file to cube format")
parser.add_argument("dx_input", help="Name of the dx_input file (required arguement)")
parser.add_argument("pqr_input", help="Name of the pqr_input file (required arguement)")
parser.add_argument("output", help="Name of the output file (required arguement)")

args = parser.parse_args()
counter = 0

# DX STUFF
if args.dx_input.endswith(".dx"):
    print "Success"
else:
    print "Error converting file"

try:
    with open(args.dx_input, "r") as in_f, open(args.output, "w") as out_f, open(
        args.pqr_input, "r"
    ) as in_pqr:
        out_f.write(
            "CPMD CUBE FILE.\n" "OUTER LOOP: X, MIDDLE LOOP: Y, INNER LOOP: Z\n"
        )

        # Discard comments at top of file.
        line = in_f.readline()
        newline = in_pqr.readline()
        while line.startswith("#"):
            line = in_f.readline()
        while newline.startswith("REMARK"):
            newline = in_pqr.readline()

        split_line = line.split()
        grid_sizes = [int(x) * -1 for x in split_line[-3:]]

        split_line = in_f.readline().split()

        origin = [float(x) for x in split_line[-3:]]

        parameter_fmt = "{:>4} {:>11.6f} {:>11.6f} {:>11.6f}\n"
        try:
            while newline.startswith("ATOM"):
                pqr_line = in_pqr.readline()
                new_split_line = pqr_line.split()
                atom_num = new_split_line[1]
        except IndexError:
            pass
        in_pqr.seek(0)
        newline = in_pqr.readline()
        while newline.startswith("REMARK"):
            newline = in_pqr.readline()

        origin_line = parameter_fmt.format(atom_num, *origin)
        out_f.write(origin_line)

        for x in xrange(3):
            split_line = in_f.readline().split()
            grid_dims = [float(item) for item in split_line[-3:]]

            dim_lin = parameter_fmt.format(grid_sizes[x], *grid_dims)
            out_f.write(dim_lin)

        atoms_parameter_fmt = "{:>4} {:>11.6f} {:>11.6f} {:>11.6f} {:>11.6f}\n"
        a = True
        xreal_center = []
        yreal_center = []
        zreal_center = []
        try:
            while a == True:
                if not newline.startswith("TER"):
                    new_split_line = newline.split()
                    radius = new_split_line[-1]
                    xyz = new_split_line[-5:-2]
                    line_atom_num = new_split_line[1]
                    atom_radius = new_split_line[-1]
                    pqr_lin = atoms_parameter_fmt.format(
                        int(line_atom_num),
                        float(new_split_line[-2]),
                        float(xyz[0]),
                        float(xyz[1]),
                        float(xyz[2]),
                    )
                    out_f.write(pqr_lin)
                    newline = in_pqr.readline()
                    xreal_center.append(float(xyz[0]))
                    yreal_center.append(float(xyz[1]))
                    zreal_center.append(float(xyz[2]))
                else:
                    a = False
        except IndexError:
            pass

        x_avg = sum(xreal_center) / float(atom_num)
        y_avg = sum(yreal_center) / float(atom_num)
        z_avg = sum(zreal_center) / float(atom_num)
        print x_avg, y_avg, z_avg

        # print origin
        # new_origin = []
        # for item in origin:
        #    newitem = item/0.529177
        # new_new = item/2 + newitem/2
        #    new_origin.append(newitem)
        # print new_origin

        # Consume unneeded object lines.
        in_f.readline()
        in_f.readline()

        ##TODO: put atoms here

        value_format = ["{:< 13.5E}"]
        value_format = " ".join(value_format * 6) + "\n"
        print value_format
        group = []
        line = in_f.readline()
        while not line.startswith("attribute"):
            values = [float(item) for item in line.split()]
            group.extend(values)

            if len(group) >= 6:
                out_f.write(value_format.format(*group))
                group = []

            line = in_f.readline()

        if group:
            group_strs = ["{:< 13.5E}".format(item) for item in group]
            out_f.write(" ".join(group_strs))


except IOError:
    print "file doesn't exist"
"""
for line in dx_data:
    dx_data_2 = dx_data
    data_point = line.strip("\n")
    dx_data_2.pop(0)
    dx_data._2append(data_point)
    

for line in dx_data_2:
    if line.startswith('#'):
        dx_data_2.remove(line)
        continue
    elif "gridpositions" in line:
        new_line = line.split()
        grid_values = new_line[-3:]
        #grid values for cube file, will need grid_valies[0]
        break
    elif line.startswith('delta'):
        line.insert('0')
        line.remove(delta)



#PQR STUFF
if args.pqr_input.endswith('.pqr'):
    print "yay"
else:
    print "nay"

try:
    with open(args.pqr_input, 'r') as f:
        pqr_data = f.readlines()
        
except IOError:
    print "file doesn't exist"
    
for line in pqr_data:
    data_point = line.strip("\n")
    pqr_data.pop(0)
    pqr_data.append(data_point)

for line in pqr_data:
    if line.startswith('REMARK'):
        pqr_data.remove(line)
        continue
    counter +=1
    #the number next to the origin in cube format
"""
