from numpy       import float16, float32, zeros
from sys         import stdout
from .           import charmm
from .atomicdata import nspecies
from ..dl_py2f   import utils

names2znums = {
'H'   :1,
'HA'  :1,
'HB'  :1,
'HE'  :1,
'HG'  :1,
'HH'  :1,
'HN'  :1,
'HZ'  :1,
'HG11':1,
'HG12':1,
'HG21':1,
'HG22':1,
'HS'  :1,
'HW'  :1,
'HY'  :1,
'QQH' :1,
'C'   :6,
'C1C' :6,
'C2C' :6,
'C3C' :6,
'C4C' :6,
'CA'  :6,
'CAA' :6,
'CAB' :6,
'CAC' :6,
'CAD' :6,
'CB'  :6,
'CBA' :6,
'CBB' :6,
'CBC' :6,
'CBD' :6,
'CD'  :6,
'CE'  :6,
'CG'  :6,
'CGA' :6,
'CGD' :6,
'CH'  :6,
'CHA' :6,
'CHB' :6,
'CHC' :6,
'CHD' :6,
'CMA' :6,
'CMB' :6,
'CMC' :6,
'CMD' :6,
'CY'  :6,
'CG1' :6,
'CS'  :6,
'CW'  :6,
'CZ'  :6,
'N'   :7,
'NA'  :7,
'NB'  :7,
'NC'  :7,
'ND'  :7,
'NE'  :7,
'NG1' :7,
'NG2' :7,
'NH1' :7,
'NH2' :7,
'NZ'  :7,
'O'   :8,
'OPP' :8,
'OXT' :8,
'OY'  :8,
'OH'  :8,
'OD'  :8,
'OG'  :8,
'OE'  :8,
'OA'  :8,
'OT'  :8,
'F'   :9,
'SOD' :11,
'Na+_cation':11,
'FE'  :26,
'CU'  :29,
'POT' :19,
'S'   :16,
'SD'  :16,
'SG'  :16,
'CLA' :17,
'P'   :15,
'CAL' :20,
'CES' :55,
}

#prefix+'[0-9]'

names2types = {
'H'   :'H'   ,
'HA'  :'HA'  ,
'HB'  :'HB'  ,
'HE'  :'HE'  ,
'HG'  :'HG'  ,
'HH'  :'HH'  ,
'HN'  :'HN'  ,
'HZ'  :'HZ'  ,
'HG11':'HG11',
'HG12':'HG12',
'HG21':'HG21',
'HG22':'HG22',
'QQH' :'QQH' ,
'C'   :'C'   ,
'CZ'  :'CZ'  ,
'CD'  :'CD'  ,
'CE'  :'CE'  ,
'CG'  :'CG'  ,
'CH'  :'CH'  ,
'CY'  :'CY'  ,
'CA'  :'CA'  ,
'CB'  :'CB'  ,
'CG1' :'CG1' ,
'N'   :'N'   ,
'ND'  :'ND'  ,
'NE'  :'NE'  ,
'NG1' :'NG1' ,
'NG2' :'NG2' ,
'NH1' :'NH1' ,
'NH2' :'NH2' ,
'NZ'  :'NZ'  ,
'O'   :'O'   ,
'OPP' :'OPP' ,
'OXT' :'OXT' ,
'OY'  :'OY'  ,
'OH'  :'OH'  ,
'OD'  :'OD'  ,
'OG'  :'OG'  ,
'OE'  :'OE'  ,
'OA'  :'OA'  ,
'OT'  :'OT'  ,
'F'   :'F'   ,
'SOD' :'SOD' ,
'FE'  :'FE'  ,
'CU'  :'CU'  ,
'POT' :'POT' ,
'SD'  :'SD'  ,
'SG'  :'SG'  ,
'CLA' :'CLA' ,
'P'   :'P'   ,
'CAL' :'CAL' ,
'CES' :'CES' ,
}


synonyms = { 'amino_acids':'amino_acids',
             'amino acids':'amino_acids',
             'amino_acid' :'amino_acids',
             'amino acid' :'amino_acids',
             'proteins'   :'amino_acids',
             'protein'    :'amino_acids',
             'enzyme'     :'amino_acids',
             'enzymes'    :'amino_acids',
             'haem'       :'haems',
             'haems'      :'haems',
             'heme'       :'haems',
             'hemes'      :'haems',
             'waters'     :'waters',
             'water'      :'waters',
             'H2O'        :'waters',
             'HOH'        :'waters',
             'solvent'    :'solvents',
             'solvents'   :'solvents',
             'ions'       :'ions',
             'ion'        :'ions',
             'dnas'       :'nucleotides',
             'dna'        :'nucleotides',
             'rnas'       :'nucleotides',
             'rna'        :'nucleotides',
             'nucleobases':'nucleotides',
             'nucleobase' :'nucleotides',
           }


ion_charges = { 'Li' : 1,
                'Na' : 1,
                'K'  : 1,
                'Cs' : 1,
                'Rb' : 1,
                'Mg' : 2,
                'Ca' : 2,
                'Ba' : 2,
                'Zn' : 2,
                'Cd' : 2,
                'Cl' :-1,
                'OH' :-1,
                'SO4':-2,
              }

ion_znums = { 'Na': 11,
              'K' : 19,
              'Cs': 21,
              'Mg': 12,
              'Ca': 20,
              'Zn': 30,
              'Cl': 17,
             }










### FIXME NOT IN USE ###
# TODO: adapt the UFF parameters

# default bond lengths
bonds = zeros((nspecies,nspecies,2), dtype=float16).view(utils.nputils.SymArray)
bonds[1,6,:] = 1.09000,  680.0
bonds[1,8,:] = 0.96000, 1106.0
bonds[6,6,:] = 1.52600,  620.0
bonds[6,8,:] = 1.41000,  640.0
# symmetrise
#print("### bonds in MB =", bonds.nbytes/1024/1024)

# van der Waals
vdw_12_6_xx = zeros((nspecies,nspecies,2), dtype=float32).view(utils.nputils.SymArray)
vdw_12_6_xx[6,6,:] = 1.0430e+06,  6.7560e+02
vdw_12_6_xx[6,8,:] = 7.9150e+05,  6.9310e+02
vdw_12_6_xx[8,8,:] = 5.8180e+05,  6.9970e+02
#print("### xx in MB =", vdw_12_6_xx.nbytes/1024/1024)

# interaction between H and other elements: 1st index means the znum of atom bonded with this hydrogen, for example 6 is 'HC' and 8 is 'HO'
vdw_12_6_hx = zeros((nspecies,nspecies,2), dtype=float32).view(utils.nputils.SymArray)
vdw_12_6_hx[6,6,:] = 9.7170e+04,  1.2690e+02
vdw_12_6_hx[6,8,:] = 6.8280e+04,  1.2530e+02

# interaction between H and H: both indices mean the znums of atoms bonded with the current two hydrogens
vdw_12_6_hh = zeros((nspecies,nspecies,2), dtype=float32).view(utils.nputils.SymArray)
vdw_12_6_hh[6,6,:] = 7.5160e+03,  2.1730e+01

vdw_12_6 = { 'h-h': vdw_12_6_hh,
             'h-x': vdw_12_6_hx,
             'x-x': vdw_12_6_xx }
vdw = { '12-6': vdw_12_6 }

# default bond angles
angles = zeros((nspecies,nspecies,nspecies,2), dtype=float16).view(utils.nputils.SymArray) 
angles[1,6,1,:] = 109.500,  70.0
angles[1,6,6,:] = 109.500, 100.0
angles[6,6,8,:] = 109.500, 100.0
angles[6,8,1,:] = 108.500, 110.0
angles[1,6,8,:] = 109.500,  70.0
angles[6,6,6,:] = 109.500,  80.0
#print("### angles in MB =", angles.nbytes/1024/1024)


# default bond dihedrals
# FIXME: not including heavy atoms for now due to very large arrays in memory
nspecies3 = int(nspecies/3)
dihedrals = zeros((nspecies3,nspecies3,nspecies3,nspecies3,3), dtype=float16).view(utils.nputils.SymArray)
dihedrals[1,6,6,1,:] = 0.0, 0.15556,  3
dihedrals[1,6,6,8,:] = 0.0, 0.15556,  3
dihedrals[1,6,8,1,:] = 0.0, 0.16667,  3
dihedrals[1,8,6,6,:] = 0.0, 0.16667,  3
dihedrals[6,6,6,1,:] = 0.0, 0.15556,  3
dihedrals[6,6,6,6,:] = 0.0, 0.15556,  3
dihedrals[6,6,6,8,:] = 0.0, 0.15556,  3
#print("### dihedrals in MB =", dihedrals.nbytes/1024/1024)

### FIXME END NOT IN USE ###


def getFFObj(ff):
    '''Return the requested forcefield object'''

    # we potentially support other forcefields
    selectcases = { 'charmm': charmm }

    return selectcases.get(ff, charmm)


def getFFDir(ff, dir=''):

    from os.path import abspath, isdir, join
    from ..io    import format
    from ..utils import testutils

    ffObj = getFFObj(ff)

    # preparation
    if not dir:
        if isdir(join(ffObj.default_location)):
            ff_dir = ffObj.default_location
            print("\n Using default {} topology & parameter files in (version: {}):\n {}".format(ff.upper(), ffObj.default_version, ff_dir))
            stdout.flush()
        else:
            format.printwrap(">>> ERROR: No default {} topology/parameter files installed in this ChemShell build.\n".format(ff.upper()))
            # YL 27/02/2023: set unsupported if FF is missing
            testutils.setUnsupported()
            exit(112)
    else:
        ff_dir = abspath(dir)
        if isdir(ff_dir):
            print("\n Using specified topology & parameter files in:\n {}".format(ff_dir))
        else:
            print("\n >>> ERROR: Directory specified by `ff_dir` does NOT exist:\n {}".format(ff_dir))
            exit(999)

    return ff_dir


