import numpy as np
from .symm_util import bond_angle

write_str = "{:2} {:11.7f} {:11.7f} {:11.7f}\n"


def write_coordinates(symbols, coordinates, f):

    f.write("elem \t x (Å) \t  y (Å) \t  z (Å)\n")
    for elem, coord in zip(symbols, coordinates):
        f.write(write_str.format(elem, *coord))
    f.write("\n")



def write_xyz(symbols, coordinates, fname):
    with open(fname, 'w') as f:
        f.write('{}\n'.format(len(coordinates)))
        f.write('Generated by symmetry_analysis.\n')
        for elem, coord in zip(symbols, coordinates):
            f.write(write_str.format(elem, *coord))


def write_bonding_environment(bond_dict, f):

    f.write('Bonding environments\n')
    f.write('--------------------\n')

    for atom, neighbours in bond_dict.items():

        if len(neighbours) > 1:

            f.write('Atom: {:4} \t Index: {} \t Position: {:11.7f} {:11.7f} {:11.7f} (Å)\n'.format(atom.symbol, atom.idx, *atom.coordinate))
            f.write('Bonded atoms:\n')
            for n in neighbours:
                f.write('\tBonded atom: {:4} \t Index: {} \t Position: {:11.7f} {:11.7f} {:11.7f} (Å) \t Bond length: {:11.7f} (Å)\n'.format(n.symbol, n.idx, *n.coordinate, np.linalg.norm(atom.coordinate-n.coordinate)))
            f.write('--\n')


def write_symmetry_axes(Cn_dict, f):
    f.write('Reporting all symmetry axes found\n')
    f.write('---------------------------------\n')

    if len(Cn_dict) == 0:
        f.write("Could not locate any symmetry axes\n--\n")
        return

    for rotation_order, axes in Cn_dict.items():

        f.write("C{}\n".format(rotation_order))

        for axis in axes:
            f.write(' \t {:11.7f} {:11.7f} {:11.7f}\n'.format(*axis))

        f.write('--\n')


def write_angles(bond_dict, f):
    f.write('\nBond angles\n')
    f.write('-----------\n')
    for atom1 in bond_dict:
        for atom2 in bond_dict[atom1]:
            for atom3 in bond_dict[atom2]:
                if not atom1 == atom3:
                    c1 = atom1.coordinate
                    s1 = atom1.symbol
                    i1 = atom1.idx
                    c2 = atom2.coordinate
                    s2 = atom2.symbol
                    i2 = atom2.idx
                    c3 = atom3.coordinate
                    s3 = atom3.symbol
                    i3 = atom3.idx
                    v1 = c2 - c1
                    v2 = c2 - c3
                    theta = bond_angle(v1, v2)
                    f.write("{:35} \t {:25}\n".format('Bond: {:>2}({:>3}) - {:>2}({:>3}) - {:>2}({:>3})'.format(s1,i1,s2,i2,s3,i3), 'Angle: {:11.7f} deg'.format(theta * 180 / np.pi)))
