#!/usr/bin/env python
# -*- coding: utf-8 -*-


import ase
from ase import Atom
from ase import data
import numpy as np

# Chemical symbols of new types of atoms
data.chemical_symbols += ['Fe1', 'Fe2']
data.chemical_symbols += ['O1']

# Atomic numbers
Z_Fe = data.atomic_numbers['Fe']
Z_O = data.atomic_numbers['O']
data.atomic_numbers['Fe1'] = Z_Fe
data.atomic_numbers['Fe2'] = Z_Fe
data.atomic_numbers['O1'] = Z_O

# Atomic names
data.atomic_names += ['Iron', 'Iron']
data.atomic_names += ['Oxygen'] 

# Atomic masses
Fe_mass = data.atomic_masses_iupac2016[Z_Fe]
O_mass = data.atomic_masses_iupac2016[Z_O] 
data.atomic_masses_iupac2016 = np.concatenate((data.atomic_masses_iupac2016, np.array([Fe_mass,Fe_mass])))
data.atomic_masses_iupac2016 = np.concatenate((data.atomic_masses_iupac2016, np.array([O_mass])))

data.atomic_masses = data.atomic_masses_iupac2016

Fe_mass = data.atomic_masses_legacy[Z_Fe]
O_mass = data.atomic_masses_legacy[Z_O]
data.atomic_masses_legacy = np.concatenate((data.atomic_masses_legacy, np.array([Fe_mass,Fe_mass])))
data.atomic_masses_legacy = np.concatenate((data.atomic_masses_legacy, np.array([O_mass])))

# Atomic radius
Fe_radius = data.covalent_radii[Z_Fe]
O_radius = data.covalent_radii[Z_O]
data.covalent_radii = np.concatenate((data.covalent_radii, np.array([Fe_radius,Fe_radius])))
data.covalent_radii = np.concatenate((data.covalent_radii, np.array([O_radius])))

# Reference and ground state datas
data.reference_states += [data.reference_states[Z_Fe], data.reference_states[Z_Fe]]
data.reference_states += [data.reference_states[Z_O]]

Fe_gsmm = data.ground_state_magnetic_moments[Z_Fe]
O_gsmm = data.ground_state_magnetic_moments[Z_O]

data.ground_state_magnetic_moments = np.concatenate((data.ground_state_magnetic_moments, np.array([Fe_gsmm,Fe_gsmm])))
data.ground_state_magnetic_moments = np.concatenate((data.ground_state_magnetic_moments, np.array([O_gsmm])))

# VdW raduis
Fe_vdw_r = data.vdw_radii[Z_Fe]
O_vdw_r = data.vdw_radii[Z_O]
data.vdw_radii = np.concatenate((data.vdw_radii, np.array([Fe_vdw_r,Fe_vdw_r])))
data.vdw_radii = np.concatenate((data.vdw_radii, np.array([O_vdw_r])))