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

import sys
import numpy as np
import ase
from ase import io
from ase.io import read
from ase.io import write
from ase import Atom
from ase import build
from ase.build import add_adsorbate
from ase import constraints
from ase.constraints import FixAtoms
from ase import neighborlist
from ase.neighborlist import NeighborList
from ase import utils
import ase_Fe1_Fe2

if len(sys.argv) < 7 :
    print "This script has to be used with at least 6 arguments."
    print "1. Name of the file containing the surface"
    print "2. Atom number of the atom of the surface where the adsorbate will be adsorbed"
    print "3. Name of the file containing the molecule to be adsorbed"
    print "4. Atom number of the atom of the molecule that will be adsorbed of the surface"
    print "5. Distance between the molecule and the surface"
    print "6. Name of the output file"
    print "Other arguments can be added before the name of the output file :"
    print "  Second atom of the molecule"
    print "  phi angle (rotation around z)"
    print "  theta angle (rotation around new x)"
    print "  psi angle (rotation around new z) "
    sys.exit(1)


#############################################################
# Définition de la fonction permettant de calculer les angles
#############################################################

def get_proper_angle(v1, v2, degrees=True):
    norm_dot = np.dot(v1/np.linalg.norm(v1), v2/np.linalg.norm(v2))
    angle = np.arccos(np.sign(norm_dot) if np.abs(norm_dot) >= 1 else norm_dot)
    return(angle*180/np.pi if degrees else angle)


##########################################
# Lecture des différents arguments de base
##########################################

surface_file = sys.argv[1]
atom_surface = int(sys.argv[2])
molecule_file = sys.argv[3]
atom_molecule = int(sys.argv[4])
distance = int(sys.argv[5])
if len(sys.argv) < 8 :
    output = sys.argv[6]

surface = read(surface_file)
molecule = read(molecule_file)

#################################
# Lecture des positions de atomes
#################################

atom_surf_1 = surface[atom_surface].position
x_atom_surf_1 = atom_surf_1[0]
y_atom_surf_1 = atom_surf_1[1]
z_atom_surf_1 = atom_surf_1[2]

atom_mol_1 = molecule[atom_molecule].position
x_atom_mol_1 = atom_mol_1[0]
y_atom_mol_1 = atom_mol_1[1]
z_atom_mol_1 = atom_mol_1[2]

nb_atom_mol = int(len(molecule))


#######################################################
# Cas avec un plus grand nombre d'arguments (rotations)
#######################################################


if len(sys.argv) > 7 :

    #######################
    # Lecture des arguments
    #######################

    atom_molecule_2 = int(sys.argv[6])
    phi = int(sys.argv[7])
    theta = int(sys.argv[8])
    psi = int(sys.argv[9])
    output = sys.argv[10]

    theta_min = theta - 0.01
    theta_max = theta + 0.01
    phi_min = phi - 0.01
    phi_max = phi + 0.01
    psi_min = psi - 0.01
    psi_max = psi + 0.01

    atom_mol_2 = molecule[atom_molecule_2].position
    x_atom_mol_2 = atom_mol_2[0]
    y_atom_mol_2 = atom_mol_2[1]
    z_atom_mol_2 = atom_mol_2[2]

    ##################################
    # Rotation avec les angles d'Euler
    ##################################

    #Initialisation : placement du second atome de la molecule sur Oy
    vect_y = [0, 1, 0]
    vect_atom_mol_2 = molecule[atom_molecule_2].position - molecule[atom_molecule].position
    angle_1 = get_proper_angle(vect_y, vect_atom_mol_2)
    if (angle_1 != 0) :
        vect_normal = np.cross(vect_y, vect_atom_mol_2)
        molecule.rotate(-angle_1, vect_normal, center=(x_atom_mol_1,y_atom_mol_1,z_atom_mol_1))
        vect_atom_mol_2_verif = molecule[atom_molecule_2].position - molecule[atom_molecule].position
        angle_1_verif = get_proper_angle(vect_atom_mol_2_verif, vect_y)
        angle_max = 0.01
        angle_min = -0.01
        if (angle_1_verif < angle_min or angle_1_verif > angle_max) :
            print 'Error in initialisation'

    #Rotation Euler
    molecule.euler_rotate(phi, theta, psi, center=(x_atom_mol_1,y_atom_mol_1,z_atom_mol_1))

    #Correction des collisions entre la molécule et la surface
    z_atom_surf_max = z_atom_surf_1 + 1

    collision = 0.5
    rotation_tot = 0
    nb_atom_collision_min = 0
    rotation_opt = 0
    rotation_modif = 5
    
    for i in range(0,nb_atom_mol) :
        atom_test = molecule[i].position
        z_atom_test = atom_test[2]
        z_atom_test_final = z_atom_test - z_atom_mol_1 + z_atom_surf_1 + distance
        if (z_atom_test_final < z_atom_surf_max) :
            collision = 1
            nb_atom_collision_min += 1
    if collision == 1 :
        print 'Collision between the molecule and the surface - modification of theta angle'

    vect_z = [0, 0, 1]

    while (collision == 1 and rotation_tot < 354) :
        vect_atom_mol_2_post_euler = molecule[atom_molecule_2].position - molecule[atom_molecule].position
        
        if (vect_atom_mol_2_post_euler[0] != 0 and vect_atom_mol_2_post_euler[1] != 0) :
            vect_rotation_theta = np.cross(vect_atom_mol_2_post_euler, vect_z)
        else :
            vect_rotation_theta = [0.0001, 0.0001 , 0]
        molecule.rotate(rotation_modif, vect_rotation_theta, center=(x_atom_mol_1,y_atom_mol_1,z_atom_mol_1))
        rotation_tot += 5
        
        collision = 0
        nb_atom_collision = 0

        for i in range(0,nb_atom_mol) :
            atom_test = molecule[i].position
            z_atom_test = atom_test[2]
            z_atom_test_final = z_atom_test - z_atom_mol_1 + z_atom_surf_1 + distance
            if (z_atom_test_final < z_atom_surf_max) :
                nb_atom_collision += 1
                collision = 1
        if (nb_atom_collision < nb_atom_collision_min) :
            nb_atom_collision_min = nb_atom_collision
            rotation_opt = rotation_tot

        if nb_atom_collision_min != 0 :
            rotation_tot_2 = 0
            rotation_opt_2 = 0

            while (collision == 1 and rotation_tot_2 < 354) :
                vect_atom_mol_2_post_euler = molecule[atom_molecule_2].position - molecule[atom_molecule].position
                molecule.rotate(rotation_modif, vect_atom_mol_2_post_euler, center=(x_atom_mol_1,y_atom_mol_1,z_atom_mol_1))
                rotation_tot_2 += 5
                
                collision = 0
                nb_atom_collision = 0

                for i in range(0,nb_atom_mol) :
                    atom_test = molecule[i].position
                    z_atom_test = atom_test[2]
                    z_atom_test_final = z_atom_test - z_atom_mol_1 + z_atom_surf_1 + distance
                    if (z_atom_test_final < z_atom_surf_max) :
                        nb_atom_collision += 1
                        collision = 1
                if (nb_atom_collision < nb_atom_collision_min) :
                    nb_atom_collision_min = nb_atom_collision
                    rotation_opt_2 = rotation_tot_2

    if collision == 0 :
        print 'Collision corrected'
    elif collision == 1 :
        print 'Error: the collision could not be corrected'
        pos1 = output.rfind('/')
        pos = pos1 + 1
        num_coll=output[pos:]
        j = output[:pos1]
        pos2 = j.rfind('/')
        molecule_directory = j[:pos2]
        fichier = open("%s/errors" % molecule_directory , "a")
        fichier.write("Adsorption %s : ERROR the collision could not be corrected\n" % num_coll )
        fichier.close()


##########################################
# Adsorption de la molecule sur la surface
##########################################

add_adsorbate(surface, molecule, distance, (x_atom_surf_1,y_atom_surf_1), mol_index=atom_molecule)

out=output+".xyz"
write(out, surface)


############################
# Dissociation si necessaire
############################

if collision == 0 :
    surface = read(surface_file)
    molecule_reminder = molecule
    
    list_cutoffs = utils.natural_cutoffs(molecule)
    nl=NeighborList(list_cutoffs, self_interaction=False, bothways=True)

    nl.update(molecule)
    neighbor_indices, offsets = nl.get_neighbors(atom_molecule)

    symbols = molecule.get_chemical_symbols()

    neighbors_symbols=[]
    for i in neighbor_indices:
            neighbors_symbols.append(symbols[i])

    nb_neighbors = len(neighbor_indices)

    diss=0
    for i in range(0,nb_neighbors):
        if neighbors_symbols[i] == "H":
            diss=1
            H_atom=neighbor_indices[i] #nb of the H atom
            print "Dissociation possible"

    if diss == 1:
        list_cutoffs_surface = utils.natural_cutoffs(surface)
        nl2=NeighborList(list_cutoffs_surface, self_interaction=False, bothways=True)
        nl2.update(surface)
        neighbor_indices_surf, offsets_surf = nl2.get_neighbors(atom_surface)
        symbols_surf = surface.get_chemical_symbols()
        for i in neighbor_indices_surf:
            surface = read(surface_file)
            molecule = molecule_reminder
            neighbor_surf_coord = surface[i].position
            coord_H_atom=molecule[H_atom].position
            if (neighbor_surf_coord[2] > z_atom_surf_1-1 and symbols_surf[i] == "O"):
                vector_H_diss=[neighbor_surf_coord[0]-coord_H_atom[0], neighbor_surf_coord[1]-coord_H_atom[1], neighbor_surf_coord[2]+1-coord_H_atom[2]]
                translation_matrix_dissociation=[]
                for k in range (0, nb_atom_mol):
                    if k != H_atom:
                        translation_matrix_dissociation.append([0,0,0])
                    else:
                        translation_matrix_dissociation.append(vector_H_diss)
                molecule.translate(translation_matrix_dissociation)
                add_adsorbate(surface, molecule, distance, (x_atom_surf_1,y_atom_surf_1), mol_index=atom_molecule)
                output_diss = output + "_diss_" + str(i) + ".xyz"
                write(output_diss, surface)