#!/usr/bin/env python
from ase.io import read, write
from ase import Atoms
from math import pi, sin, cos
from numpy.random import random
from ase.visualize import view
from numpy.linalg import norm
from ase.utils.geometry import sort
import numpy as np
import sys

def str2bool(v):
    if v.lower() in ('yes', 'true', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected.')

def cartesianPBC(r, cell):
    icell = np.linalg.inv(cell)
    vdir  = np.dot(r, icell)
    vdir  = (vdir % 1.0 + 1.5) % 1.0 - 0.5
    newr  = np.dot(vdir, cell)
    return newr

def random_water(center=None):
    if center == None:
        center = (0.,0.,0.)
    d = 0.96
    angle = 105.0*pi/180.
    pos = [ (0, 0, 0),
            (d, 0, 0),
            (cos(angle), sin(angle), 0) ]
    water = Atoms('OHH', positions=pos)
    phi = 2*pi*random()
    theta = pi*random()
    psi = pi*random()
    water.rotate_euler('COM', phi, theta, psi)
    water.translate(center)
    return water

def clips(atoms, bounding_box):
    for atom in atoms:
        clips_box = False
        for i in range(3): 
            if atom.position[i] < bounding_box[i][0]: break
            if atom.position[i] > bounding_box[i][1]: break
        else:
            clips_box = True
        if clips_box:
            break
    return clips_box

def overlap(particle, water):
    for atom_i in particle:
        if atom_i.symbol not in ('H', 'O'): continue
        cell=particle.get_cell()
        for atom_j in water:
            dr=cartesianPBC(atom_j.position - atom_i.position, cell)
            d = norm(dr)
            if d < 1.5: return True
    return False

def add_water(particle, bounding_box):
    while True:
        pos = random(3)*particle.get_cell().diagonal()
        water = random_water(pos)
        if not clips(water, bounding_box):
            if not overlap(particle, water):
                break
    particle += water

def solvate(particle, thickness, padding, molecules, lsort):
    vacuum = thickness/2.0;
    #padding = 1.0
    particle.center(vacuum, axis=2)
    pos = particle.get_positions()
    cell = particle.get_cell()
    bounding_box = ((min(pos[:,0]), cell[0][0]),
                    (min(pos[:,1]), cell[1][1]),
                    (min(pos[:,2])-padding, max(pos[:,2])+padding))
   
    cell_novac=cell.copy()
    cell_novac[2,:] = cell_novac[2,:]-(thickness-2.*padding)*cell_novac[2,:]/np.linalg.norm(cell_novac[2,:])
    vol = np.linalg.det(cell) - np.linalg.det(cell_novac)
    print 'Your cell is %.3f Angstrom^3' % vol 
    print 'It can fit %i(ice) to %i(water) water molecules' % ( int(vol*0.030642831608863625) ,int(vol*0.03342732803410454))
    print 'You chose to add %i water molecules' % molecules

    print 'Adding %i H2O molecules' % molecules
    sys.stdout.write('Progress: ')
    for i in range(molecules):
        add_water(particle, bounding_box)
        sys.stdout.write('%i '%(i+1))
        sys.stdout.flush()
    print
    if lsort:
        particle=sort(particle)

    return particle

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('structure', help='file with structure to solvate')
    parser.add_argument('thickness', type=float, help='thickness of solvent layer')
    parser.add_argument('padding', type=float, help='padding on structure bounding box')
    parser.add_argument('molecules', type=int, help='number of solvent molecules')
    parser.add_argument('lsort', type=str2bool, help='sort or not')
    args = parser.parse_args()
    particle = read(args.structure)
    print args.lsort
    solvated_particle = solvate(particle, args.thickness, args.padding, args.molecules, args.lsort)
    #solvated_particle.wrap()
    write('POSCAR_%iH2O'%args.molecules, solvated_particle, 'vasp')
