#!/usr/bin/env python3
#
# This file is part of the GROMACS molecular simulation package.
#
# Copyright 2012- The GROMACS Authors
# and the project initiators Erik Lindahl, Berk Hess and David van der Spoel.
# Consult the AUTHORS/COPYING files and https://www.gromacs.org for details.
#
# GROMACS 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 2.1
# of the License, or (at your option) any later version.
#
# GROMACS 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 GROMACS; if not, see
# https://www.gnu.org/licenses, or write to the Free Software Foundation,
# Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
#
# If you want to redistribute modifications to GROMACS, please
# consider that scientific software is very special. Version
# control is crucial - bugs must be traceable. We will be happy to
# consider code for inclusion in the official distribution, but
# derived work must not be called official GROMACS. Details are found
# in the README & COPYING files - if they are missing, get the
# official version at https://www.gromacs.org.
#
# To help us fund GROMACS development, we humbly ask that you cite
# the research papers on the package. Check out https://www.gromacs.org.

"""Script for generating spherical test configurations."""

import math
import random
import sys

def dot(a, b):
    return sum(x*y for (x, y) in zip(a, b))

def norm(a):
    return math.sqrt(dot(a, a))

def angle(a, b):
    return math.degrees(math.acos(dot(a, b) / (norm(a) * norm(b))))

def minangle(a, list):
    minangle = 180
    minindex = -1
    for index, x in enumerate(list):
        xangle = angle(a, x)
        if xangle < minangle:
            minangle = xangle
            minindex = index
    return (minangle, minindex)

def get_single_vec():
    while True:
        x = random.randint(-1000, 1000) / 1000.0
        y = random.randint(-1000, 1000) / 1000.0
        z = random.randint(-1000, 1000) / 1000.0
        pos = (x, y, z)
        dist = norm(pos)
        if dist <= 1.0 and dist > 0.25:
            return pos

def write_gro_title(fp, title, atomcount):
    fp.write(title + '\n')
    fp.write('%5d\n' % (atomcount))

def write_gro_atom(fp, resnr, resname, atomname, index, x):
    fp.write('%5d%-5s%5s%5d%8.3f%8.3f%8.3f\n' %
            (resnr, resname, atomname, index, x[0], x[1], x[2]))

def write_gro_box(fp, box):
    fp.write('%10.5f%10.5f%10.5f\n' % box)

random.seed(1097)
center = (0, 0, 0)
cutoff = 20
possamples = 500
negsamples = 500
totsamples = 10000

sys.stderr.write("Generating reference points\n")
refpoints = []
refpoints.append((0, 0, -1))
refpoints.append((-0.5, 0.6, 0.1))
refpoints.append((-0.5, -0.5, 0.25))
while len(refpoints) < 30:
    pos = get_single_vec()
    if pos[0] > 0 and pos[1] > pos[0] and pos[2] > 0:
        refpoints.append(pos)

sys.stderr.write("Generating test points\n")
postestpoints = []
negtestpoints = []
hits = 0
samplecount = 0
while samplecount < totsamples or len(postestpoints) < possamples or len(negtestpoints) < negsamples:
    pos = get_single_vec()
    (pangle, index) = minangle(pos, refpoints)
    if pangle < cutoff:
        hits += 1
        if len(postestpoints) < possamples:
            postestpoints.append(pos)
    if pangle > cutoff:
        if len(negtestpoints) < negsamples:
            negtestpoints.append(pos)
    samplecount += 1

cfrac = float(hits) / samplecount
errest = math.sqrt((cfrac - cfrac*cfrac) / samplecount)
sys.stderr.write('Cutoff: %f angles\n' % (cutoff))
sys.stderr.write('Estimated covered fraction: %f +- %f\n' % (cfrac, errest))

debugfp = open('debug.txt', 'w')
fp = sys.stdout
count = 1 + len(refpoints) + len(postestpoints) + len(negtestpoints)
write_gro_title(fp, 'Spherical test case, cutoff %f, cfrac %f +- %f' %
        (cutoff, cfrac, errest) , count)
n = 1
write_gro_atom(fp, 1, 'C', 'C', n, center)
n += 1
for i in range(len(refpoints)):
    write_gro_atom(fp, 2, 'R', 'R', n, refpoints[i])
    n += 1
for i in range(len(postestpoints)):
    write_gro_atom(fp, 3, 'TP', 'TP', n, postestpoints[i])
    x = postestpoints[i]
    xangle, index = minangle(x, refpoints)
    refx = refpoints[index]
    debugfp.write('%3d%8.3f%8.3f%8.3f  %4.1f  %2d%8.3f%8.3f%8.3f\n' %
            (n-1, x[0], x[1], x[2], xangle, index, refx[0], refx[1], refx[2]))
    n += 1
for i in range(len(negtestpoints)):
    write_gro_atom(fp, 4, 'TN', 'TN', n, negtestpoints[i])
    n += 1
write_gro_box(fp, (10, 10, 10))
debugfp.close()
