from properties import *
import numpy
from matter.Lattice import Lattice
cartesian_lattice = Lattice()
       
class AtomPropertyCurator(type):
    """meta class for Atom class to collect all properties and set up
    some class constants.
    It establishes:
      - Atom class's docstring
      - a list of names of properties
      - the list of setable attributes, "_setable", (to be used in method __setattr__ of Atom class).
    NOTE: Brandon doesn't like this third one...user-scripters should be able to set whatever 
    they want...why limit them? Alex read this note :)
    """

    def __init__(AtomClass, name, bases, dict):
        from matter import METACLASS_CTOR_NEEDS_TYPE
        if METACLASS_CTOR_NEEDS_TYPE:
            type.__init__(AtomClass, name, bases, dict)
        else:
            type.__init__(name, bases, dict)

        isotopeNumberProperties, inferredProperties, states = AtomPropertyCurator.collectProperties( AtomClass )
        properties = isotopeNumberProperties + inferredProperties + states
        AtomClass.propertyNames = [ p.name for p in properties ]
        
        propStr = '\n'.join(
            [ "  %s: %s" % (p.name, p.doc) for p in properties ] )

        #global doc
        doc = """
        Here is a list of properties:

        %s

        """ % propStr

        AtomClass.__doc__ += doc
        
        #AtomClass._setable = [ state.name for state in states ]

    #TODO: this needs to be rewritten to include diffraction's properties....or they need to be rewritten in inelastic format...either way, one should look at Paul's elements package first
    @staticmethod
    def collectProperties( klass ):
        ctorargs = []
        inferred = []
        states = []
        registry = {
            CtorArg: ctorargs,
            InferredProperty: inferred,
            State: states,
            }
        for item in klass.__dict__.values():
            if not isinstance( item, Property ):continue
            registry[ item.__class__ ].append( item )
            continue
        return ctorargs, inferred, states


class CartesianCoordinatesArray(numpy.ndarray):
    """Helper array for accessing Cartesian coordinates.
    Converts and updates related array of corresponding fractional
    coordinates.

    Data members:
        lattice -- instance of Lattice defining fractional coordinates
        xyz     -- instance of numpy.array storing fractional coordinates
    """

    def __new__(self, lattice, xyz):
        return numpy.zeros(3, dtype=float).view(self)

    def __init__(self, lattice, xyz):
        self.lattice = lattice
        self.xyz = xyz
        self[:] = self.lattice.cartesian(self.xyz)
        pass

    def __setitem__(self, idx, value):
        """Set idx-th coordinate and update linked self.xyz

        idx     -- index in xyz array
        value   -- new value of x, y or z
        """
        numpy.ndarray.__setitem__(self, idx, value)
        self.xyz[:] = self.lattice.fractional(self)
        return
        

# conversion constants
_BtoU = 1.0/(8 * numpy.pi**2)
_UtoB = 1.0/_BtoU
tol_anisotropy = 1.0e-6
class Atom(object):
    """Create an atom of a specified type at given lattice coordinates:
    >>> Fe = Atom( 'Fe' )
    >>> Fe57 = Atom( 'Fe', [0,0,0], mass=57)

    Atom(a) creates a copy of Atom instance a.

    atype       -- element symbol string or Atom instance
    xyz         -- fractional coordinates
    label        -- atom label
    occupancy   -- fractional occupancy
    anisotropy  -- flag for anisotropic thermal displacements
    U           -- anisotropic thermal displacement tensor, property
    Uisoequiv   -- isotropic thermal displacement or equivalent value,
                   property
    lattice     -- coordinate system for fractional coordinates

    You can obtain its property (for example, scattering length) in a similar
    way one accesses a property of a normal python object:

    TODO: We need to list all atomic properties here, and have a method to show
    what properties are available
    
    >>> print Fe.scattering_length
    
    """


    def __init__(self, atype='H', xyz=[0,0,0], mass=None, label='', 
                 occupancy=1.0, charge=0.0, anisotropy=None, U=None, Uisoequiv=None, lattice=None):
        object.__init__(self)
        # declare non-singleton data members
        self.xyz = numpy.zeros(3, dtype=float)
        self.label = label
        self.charge = charge
        self.occupancy = occupancy
        self._anisotropy = None
        self._U = numpy.zeros((3,3), dtype=float)
        self._Uisoequiv = 0.0
        self._Usynced = True
        self.lattice = None

        # assign them as needed
        if isinstance(atype, Atom):
            atype_dup = atype.__copy__()
            self.__dict__.update(atype_dup.__dict__)
        #elif isinstance(atype, str):
        else:
            self.initializeProperties(atype, mass)
        #else:
            # just leave symbol unspecified for now
            #pass
        #   self.__dict__['symbol'] = None
            
        # take care of remaining arguments
        if xyz is not None:         self.xyz[:] = xyz
        if anisotropy is not None:  self._anisotropy = bool(anisotropy)
        if U is not None:           self._U = U
        if Uisoequiv is not None:   self._Uisoequiv = Uisoequiv
        if lattice is not None:     self.lattice = lattice
        
        return


#    def __setattr__(self, name, value):
#        if name not in Atom._setable:
#            raise AttributeError, "Unknown attribute %s" % name
#        return object.__setattr__(self, name, value)
        
    def initializeProperties(self, atype, mass):
        # if symbol is specified we should check that the symbol passed is a valid chemical element symbol
        self.__dict__['symbol'] = atype
        try:
            Z = self.atomic_number
            self.__dict__['Z'] = Z
            if mass is None: mass = self.average_mass
            self.__dict__['mass'] = mass
        except KeyError:
            #raise AttributeError, 'Invalid chemical element symbol.'
            pass

#    def __str__(self):
#        l = []
#        for prop in self.propertyNames:
#
#            value = self.__dict__.get( prop )
#            if value is None: continue
#
#            l.append( ( prop, value ) )
#            continue
#
#        rt = ','.join( ['%s=%s' % (name, value) for name,value in l ] )
#        return "Atom " + rt
    
#    def __str__(self):
#        """simple string representation"""
#        xyz = self.xyz
#        s = "%-4s %8.6f %8.6f %8.6f %6.4f" % \
#                (self.symbol, xyz[0], xyz[1], xyz[2], self.occupancy)
#        return s
      
    def __repr__(self):
        """simple string representation"""
        xyz = self.xyz
        s = "%-4s %8.6f %8.6f %8.6f %6.4f" % \
                (self.symbol, xyz[0], xyz[1], xyz[2], self.occupancy)
        return s
    
    def __copy__(self):
        """Return a copy of this instance.
        """
        adup = Atom(self.symbol)
        adup.__dict__.update(self.__dict__)
        # create copies for what should be copied
        adup.xyz = numpy.array(self.xyz)
        return adup
    
    ####################################################################
    # property handlers
    ####################################################################

    # xyz_cartn

    def _get_xyz_cartn(self):
        if not self.lattice:
            rv = self.xyz
        else:
            rv = CartesianCoordinatesArray(self.lattice, self.xyz)
        return rv

    def _set_xyz_cartn(self, value):
        if not self.lattice:
            self.xyz[:] = value
        else:
            self.xyz = self.lattice.fractional(value)
        return

    xyz_cartn = property(_get_xyz_cartn, _set_xyz_cartn, doc =
        """absolute Cartesian coordinates of an atom
        """ )
    
    # anisotropy

    def _get_anisotropy(self):
        # determine when unknown
        if self._anisotropy is None:
            Uisoequiv = self._get_Uisoequiv()
            # calculate isotropic tensor Uisoij
            lat = self.lattice or cartesian_lattice
            Tu = lat.recnormbase
            Uisoij = numpy.dot(numpy.transpose(Tu), Uisoequiv*Tu)
            # compare with new value
            maxUdiff = numpy.max(numpy.fabs(self._U - Uisoij))
            self._anisotropy = maxUdiff > tol_anisotropy
            self._Uijsynced = False
        return self._anisotropy

    def _set_anisotropy(self, value):
        if bool(value) == self._anisotropy: return
        # convert from isotropic to anisotropic
        if value:
            self._U = self._get_U()
        # otherwise convert from anisotropic to isotropic
        else:
            self._Uisoequiv = self._get_Uisoequiv()
            self._Uijsynced = False
        self._anisotropy = bool(value)
        return

    anisotropy = property(_get_anisotropy, _set_anisotropy, doc =
        """flag for anisotropic thermal displacements.
        """ )

    # U

    def _get_U(self):
        # for isotropic non-synced case we need to
        # calculate _U from _Uisoequiv
        if self._anisotropy is False and not self._Uijsynced:
            lat = self.lattice or cartesian_lattice
            Tu = lat.recnormbase
            self._U = numpy.dot(numpy.transpose(Tu), self._Uisoequiv*Tu)
            self._Uijsynced = True
        # handle can be changed by the caller
        self._anisotropy = None
        return self._U

    def _set_U(self, value):
        self._anisotropy = None
        self._U = numpy.array(value, dtype=float)
        return

    U = property(_get_U, _set_U, doc =
        "anisotropic thermal displacement tensor.")

    # Uij elements

    def _get_Uij(self, i, j):
        Uij = self._get_U()
        return Uij[i,j]

    def _set_Uij(self, i, j, value):
        self._anisotropy = None
        self._U[i,j] = value
        self._U[j,i] = value

    U11 = property(lambda self: self._get_Uij(0, 0),
            lambda self, value: self._set_Uij(0, 0, value), doc =
            "U11 element of anisotropic displacement tensor")
    U22 = property(lambda self: self._get_Uij(1, 1),
            lambda self, value: self._set_Uij(1, 1, value), doc =
            "U22 element of anisotropic displacement tensor")
    U33 = property(lambda self: self._get_Uij(2, 2),
            lambda self, value: self._set_Uij(2, 2, value), doc =
            "U33 element of anisotropic displacement tensor")
    U12 = property(lambda self: self._get_Uij(0, 1),
            lambda self, value: self._set_Uij(0, 1, value), doc =
            "U12 element of anisotropic displacement tensor")
    U13 = property(lambda self: self._get_Uij(0, 2),
            lambda self, value: self._set_Uij(0, 2, value), doc =
            "U13 element of anisotropic displacement tensor")
    U23 = property(lambda self: self._get_Uij(1, 2),
            lambda self, value: self._set_Uij(1, 2, value), doc =
            "U23 element of anisotropic displacement tensor")

    # Uisoequiv

    def _get_Uisoequiv(self):
        if self._anisotropy is None or self._anisotropy is True:
            lat = self.lattice or cartesian_lattice
            Uequiv = (
                    self._U[0,0]*lat.ar*lat.ar*lat.a*lat.a +
                    self._U[1,1]*lat.br*lat.br*lat.b*lat.b +
                    self._U[2,2]*lat.cr*lat.cr*lat.c*lat.c +
                    2*self._U[0,1]*lat.ar*lat.br*lat.a*lat.b*lat.cg +
                    2*self._U[0,2]*lat.ar*lat.cr*lat.a*lat.c*lat.cb +
                    2*self._U[1,2]*lat.br*lat.cr*lat.b*lat.c*lat.ca ) / 3.0
            self._Uisoequiv = Uequiv
        else:
            self._Uisoequiv = self._U[0,0]
        return self._Uisoequiv

    def _set_Uisoequiv(self, value):
        double_eps = (1.0 + numpy.sqrt(2.0**-52)) - 1.0
        self._Uisoequiv = float(value)
        self._Uijsynced = False
        if self._get_anisotropy():
            Uequiv = self._get_Uisoequiv()
            # scale if Uequiv is not zero
            if numpy.fabs(Uequiv) > double_eps:
                self._U *= value/Uequiv
        # otherwise just convert from Uiso value
        else:
            lat = self.lattice or cartesian_lattice
            Tu = lat.recnormbase
            self._U = numpy.dot(numpy.transpose(Tu), value*Tu)
            self._Uijsynced = True
        return

    Uisoequiv = property(_get_Uisoequiv, _set_Uisoequiv, doc =
            "isotropic thermal displacement or equivalent value")

    # Bij elements

    B11 = property(lambda self: _UtoB*self._get_Uij(0, 0),
            lambda self, value: self._set_Uij(0, 0, _BtoU*value), doc =
            "B11 element of Debye-Waler displacement tensor")
    B22 = property(lambda self: _UtoB*self._get_Uij(1, 1),
            lambda self, value: self._set_Uij(1, 1, _BtoU*value), doc =
            "B22 element of Debye-Waler displacement tensor")
    B33 = property(lambda self: _UtoB*self._get_Uij(2, 2),
            lambda self, value: self._set_Uij(2, 2, _BtoU*value), doc =
            "B33 element of Debye-Waler displacement tensor")
    B12 = property(lambda self: _UtoB*self._get_Uij(0, 1),
            lambda self, value: self._set_Uij(0, 1, _BtoU*value), doc =
            "B12 element of Debye-Waler displacement tensor")
    B13 = property(lambda self: _UtoB*self._get_Uij(0, 2),
            lambda self, value: self._set_Uij(0, 2, _BtoU*value), doc =
            "B13 element of Debye-Waler displacement tensor")
    B23 = property(lambda self: _UtoB*self._get_Uij(1, 2),
            lambda self, value: self._set_Uij(1, 2, _BtoU*value), doc =
            "B23 element of Debye-Waler displacement tensor")

    # Bisoequiv

    def _get_Bisoequiv(self):
        return _UtoB * self._get_Uisoequiv()

    def _set_Bisoequiv(self, value):
        self._set_Uisoequiv(_BtoU*value)

    Bisoequiv = property(_get_Bisoequiv, _set_Bisoequiv, doc =
            "Debye-Waler isotropic thermal displacement or equivalent value")
    
    
    # inelastic group properties 
    
    # Z and mass
    Z = CtorArg( 'Z', 'atomic number' )
    symbol = CtorArg( 'symbol', 'chemical symbol' )
    mass = CtorArg( 'mass', 'atomic mass number' )


    # read-only, inferred properties
    import atomic_properties
    from utils import getModules
    modules = getModules( atomic_properties )
    del getModules, atomic_properties

    for module in modules:
        name = module.__name__.split( '.' )[-1]
        doc = module.__doc__
        lookup = module.lookup
        #print name, lookup
        cmd = "%s=InferredProperty( name, doc, lookup )" % name
        exec cmd
        del name, doc, lookup, module, cmd
        continue
    del modules


    # states
    velocity = State('velocity', 'velocity of the atom')
    displacement = State('displacement', 'displacement vector of atom')
    force = State('force', 'force on atom (vector)')
    ### should add all possible states here including but
    ### not limit to:
    ###   position( displacement ), pseudopotential, force


    __metaclass__ = AtomPropertyCurator


