# Copyright 2016 Autodesk Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import moldesign as mdt
from moldesign.parameters import Parameter, WhenParam
from moldesign.utils import exports

from . import ForceField


@exports
class GAFF(ForceField):
    """ Model the energy using the GAFF forcefield

    This is implemented as a special case of the ForceField energy model; it automates small
    parameterization process
    """
    # TODO: mechanism to store partial charges so they don't need to be constantly recomputed

    PARAMETERS = [Parameter('partial_charges',
                            'Partial charge model',
                            type=str,
                            default='am1-bcc',
                            choices=['am1-bcc', 'gasteiger', 'esp']),
                  Parameter('gaff_version',
                            'GAFF version',
                            type=str,
                            choices='gaff gaff2'.split(),
                            default='gaff2')
                  ] + ForceField.PARAMETERS

    def prep(self, force=False):
        self._parameterize()
        return super(GAFF, self).prep()

    def calculate(self, requests=None):
        if not self._prepped:
            self._parameterize()
        return super(GAFF, self).calculate(requests=requests)

    def _parameterize(self):
        if not self.mol.ff:
            mdt.parameterize(self.mol,
                             charges=self.params.partial_charges,
                             ffname=self.params.gaff_version)


@@ -729,19 +729,32 @@
 %0Aclass G
-AFF
+affSmallMolecule
 (ForceFi
@@ -1640,35 +1640,48 @@
   return super(G
-AFF
+affSmallMolecule
 , self).prep()%0A%0A
@@ -1805,19 +1805,32 @@
  super(G
-AFF
+affSmallMolecule
 , self).
@@ -1932,12 +1932,31 @@
     
-mdt.
+params = mdt.create_ff_
 para
@@ -1956,27 +1956,25 @@
 ff_parameter
-ize
+s
 (self.mol,%0A 
@@ -1972,16 +1972,33 @@
 lf.mol,%0A
+                 
         
@@ -2088,14 +2088,31 @@
     
-ffname
+                 baseff
 =sel
@@ -2133,10 +2133,46 @@
 version)
+%0A            params.assign(self.mol)
 %0A%0A

Rename GAFF energy model to distinguish it from the forcefield