import argparse
import io
import os
import sys

from iotbx.cif import CifParserError, reader


def read_cif(f):
    """opens cif and returns cctbx data object"""
    try:
        if isinstance(f, io.IOBase):
            structures = reader(file_object=f).build_crystal_structures()
        elif isinstance(f, str):
            structures = reader(file_path=f).build_crystal_structures()
        else:
            raise TypeError(f'read_cif: Can not deal with type {type(f)}')
    except CifParserError as e:
        print(e)
        print("Error parsing cif file, check if the data tag does not contain any spaces.")
        sys.exit()
    for key, val in list(structures.items()):
        print("\nstructure:", key)
        val.show_summary().show_scatterers()
    return structures

def main():
    usage = """cif2patterson structure.cif"""

    description = """Notes: Takes any cif file and generated patterson map
    """

    parser = argparse.ArgumentParser(  # usage=usage,
        description=description,
        formatter_class=argparse.RawDescriptionHelpFormatter)


    parser.add_argument("args",
                        type=str, metavar="FILE",
                        help="Path to input cif")


    parser.set_defaults(
        spgr="P1"
    )

    options = parser.parse_args()

    cif = options.args
    s = list(read_cif(cif).values())[0]
    s = s.expand_to_p1()
    print(f"Expanded to P1 => {s.scatterers().size()} atoms")
    print()

    root, ext = os.path.splitext(cif)

    uc = s.unit_cell()

    scatterers = s.scatterers()

    fout1 = open("patterson_dists.txt", 'w')
    fout2 = open("patterson_full.txt", 'w')

    distances_all = []

    verbose = False
    for atom1 in scatterers:
        x1, y1, z1 = atom1.site

        if verbose:
            print()
            atom1.show()
        distances = []
        for atom2 in scatterers:
            if verbose:
                atom2.show()

            x2, y2, z2 = atom2.site

            dx, dy, dz = x1-x2, y1-y2, z1-z2

            dx = dx % 1
            dy = dy % 1
            dz = dz % 1

            length = uc.length((dx, dy, dz))

            distances.append((atom2.label, dx, dy, dz, length))

            if verbose:
                print(f' --> {atom2.label:>4s} {dx:9.5f} {dy:9.5f} {dz:9.5f}  {length:9.5f}')

            # print atom1.label, '-->', atom2.label, '=', uc.length((dx,dy,dz))

        distances_all.extend(distances)

        atom1.show(fout2)
        for label, dx, dy, dz, distance in sorted(distances, key=lambda x: x[-1]):
            print(' --> {:>4s} {:9.5f} {:9.5f} {:9.5f}  {:9.5f}'.format(
                label, dx, dy, dz, distance), file=fout2)
        print(file=fout2)

    print('Wrote file', fout2.name)

    for label, dx, dy, dz, distance in sorted(distances_all, key=lambda x: x[-1]):
        print(f'{distance:9.5f}', file=fout1)

    print('Wrote file', fout1.name)

    fout1.close()
    fout2.close()
