#  Copyright (C) 2021 The authors of Py-ChemShell
#
#  This file is part of Py-ChemShell.
#
#  Py-ChemShell is free software: you can redistribute it and/or modify
#  it under the terms of the GNU Lesser General Public License as
#  published by the Free Software Foundation, either version 3 of the
#  License, or (at your option) any later version.
#
#  Py-ChemShell is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU Lesser General Public License for more details.
#
#  You should have received a copy of the GNU Lesser General Public
#  License along with Py-ChemShell.  If not, see
#  <http://www.gnu.org/licenses/>.


__author__ = 'You Lu <you.lu@stfc.ukri.org>'


'''Tools module'''

__all__ = 'plot'


# shortcuts of direct importable names for the user's convenience
from .md import takeSnapshots


def getUnits(units):
    '''Return the unit names and conversion factors'''

    from ..data  import unitconvert

    # default units and factors for length and time
    lunit   = 'a.u.'
    tunit   = 'fs'
    lfactor = unitconvert.au2unit[lunit]
    tfactor = unitconvert.fs2unit[tunit]

    # dict
    if type(units) is dict:
        for key, val in units.items():
            val_lower = val.strip().lower()
            if val_lower in unitconvert.au2unit.keys():
                lunit = val
                lfactor = unitconvert.au2unit[val_lower]
            if val_lower in unitconvert.fs2unit.keys():
                tunit = val
                tfactor = unitconvert.fs2unit[val_lower]

    # str
    elif type(units) is str:
        units_lower = units.strip().lower()
        if units_lower in unitconvert.au2unit.keys():
            lunit = units
            lfactor = unitconvert.au2unit[units_lower]
        if units_lower in unitconvert.fs2unit.keys():
            tunit = units
            tfactor = unitconvert.au2unit[units_lower]

    # other iterables such as tuple/list
    elif type(units) is not str and hasattr(units, '__iter__'):
        for u in units:
            u_lower = u.strip().lower()
            if u_lower in unitconvert.au2unit.keys():
                lunit = u
                lfactor = unitconvert.au2unit[u_lower]
            if u_lower in unitconvert.fs2unit.keys():
                tunit = u
                tfactor = unitconvert.fs2unit[u_lower]

    else:
        print('\n WARNING: the keyword argument units should be type dict, str, or list/tuple')

    return lunit, tunit, lfactor, tfactor 


