'''
http://pymolwiki.org/index.php/Save_Mopac

(c) 2012 Thomas Holder

License: BSD-2-Clause
'''

from __future__ import print_function

from pymol import cmd, CmdException


def save_mopac(filename, selection='all', zero='none', state=-1, quiet=1):
    '''
DESCRIPTION

    Save to MOPAC format

ARGUMENTS

    filename = string: file path to be written

    selection = string: atoms to save {default: all}

    zero = string: atoms to save with zero flag {default: none}

    state = integer: state to save {default: -1 (current state)}
    '''
    state, quiet = int(state), int(quiet)

    fmt = '%5s(%6i %3s%4i) %12.8f +%i %12.8f +%i %12.8f +%i %26.4f\n'

    zero_idx = set()
    cmd.iterate(zero, 'zero_idx.add((model,index))', space=locals())

    serial = [0]

    def callback(model, index, e, resn, resv, x, y, z, c):
        flag = (model, index) not in zero_idx and 1 or 0
        serial[0] += 1
        handle.write(fmt % (e, serial[0], resn, resv,
                            x, flag, y, flag, z, flag, c))

    handle = open(filename, 'w')
    handle.write('PM6\n\nGenerated by PyMOL\n')

    cmd.iterate_state(state, selection,
                      'callback(model, index, elem, resn, resv, x, y, z, partial_charge)',
                      space=locals())

    handle.close()

    if not quiet:
        print(' Save-MOPAC: Wrote %i atoms to file' % (serial[0]))


cmd.extend('save_mopac', save_mopac)