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 = 0 if (model, index) in zero_idx else 1
        serial[0] += 1
        handle.write(fmt % (e, serial[0], resn, resv,
            x, flag, y, flag, z, flag, c))

    with open(filename, 'w') as handle:
        cmd.iterate_state(state, selection,
                'callback(model, index, elem, resn, resv,'
                ' x, y, z, partial_charge)',
                space=locals())

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

cmd.extend('save_mopac', save_mopac)
