# pymol

from pymol import cmd
from chempy import io
from chempy import feedback
from chempy.bmin.state import State

import threading
import traceback
import os

state = None
model = None

def assign(sele,preserve=0):

    from molobj import MolObj
    from typer import Typer,Rules
    import rules

    result = 1
    
    state = State()

    model = cmd.get_model(sele)
    # now assign atom types

    ruleSet = Rules()

    ruleSet.fromList(rules.mmff_types)
    ruleSet.mappingFromList(rules.mmff_mapping)
    
    mobj = MolObj()
    mobj.fromChemPyModel(model)

    typed = Typer(molObj = mobj)
    
    print " realtime: assigning atom types"
    typed.applyRules(ruleSet)

    c = 0
    for a in typed.getTypes():
        at = model.atom[c]
        if (at.text_type == '??') or (not preserve):
            if a==-99:
                print " warning: unable to assign atom type to atom %d"%c
                result = 0
            else:
                cmd.alter("((%s) and (index %s))" % (sele,at.index),
                             "numeric_type ='%s'" % a)
                if feedback['tinker']:
                    print " "+str(__name__)+': %s is a %s' % (at.name,a)
                at.numeric_type = a
        c = c + 1

    sm = 0
    for a in model.atom:
        a.resi = str(a.resi_number)
        sm = sm + a.partial_charge

    return result

def setup(sele,preserve=0):
    
    global state
    global model
    
    state = State()
    model = cmd.get_model(sele)

    sm = 0
    for a in model.atom:
        a.resi = str(a.resi_number)
        sm = sm + a.partial_charge

    state.load_model(model)
    return 1

def check(obj='check'):
    global state
    global model
    
    if not state:
        if not model:
            print " realtime.reload: please run setup first."
        else:
            cmd.load_model(model,obj,1)
    else:
        model = state.model
        cmd.load_model(model,obj,1)

def mini(total_steps=500,
            gradient=0.001,
            interval=100,
            object='rt',
            fix_flag=None,
            rest_flag=None,
            solvation=None,
            finish=None):

    global state
    if not state:
        print " realtime.mini: please run setup first..."
    else:
        model = state.model
        print " realtime.mini: %d atoms total\n" % model.nAtom
        try:
            while total_steps>0:
                total_steps = total_steps - interval
                state.minimize(max_iter=interval,
                                    fix_flag=fix_flag,
                                    rest_flag=rest_flag,
                                    solvation=solvation)
                cmd.delete(object)
                cmd.load_model(state.model,object,1)
                cmd.refresh()
                if finish!=None:
                    apply(finish[0],finish[1],finish[2])
        except:
            cmd.load_model(state.model,'ref')
            traceback.print_exc()
        print " realtime.mini: complete."

def mini_threaded(*args,**kwargs):
    t = threading.Thread(target=mini,
                                args=args,
                                kwargs=kwargs)
    t.setDaemon(1)
    t.start()