#!/usr/bin/env python
import numpy

AA2AU = 1.0 / 0.529
eps=78.39

max_iterations = 50
diis_start_from_iter = 1
diis_vector_threshold = 1.0e-6
inverse = False
surface_area_minimum = 1.0e-6
threshold = 1.0e-9
#k = 1.07
# from Tomasi J, Mennucci B, Cammi R, Chem. Rev. 2005, 105, 2999-3093
# on pp. 3013
k = 1.0694
scal_charges_iter = 0.05 # empirical factor to slowly turn on the charges

def CMAT(areas, coords):
    n = len(areas)
    C = numpy.zeros((n,n))
    for its in range(n):
        C[its,its] = k * numpy.sqrt( 4 * numpy.pi / areas[its] )
        for jts in range(its):
            dr = coords[its] - coords[jts]
            R = numpy.sqrt(dr.dot(dr))
            Ri = 1.0 / R
            C[its,jts] = Ri
            C[jts,its] = Ri

    return C

f = open("water.surface", "r")
areas = []
coordinates = []
for line in f:
    data = line.split()

    area = float(data[2])
    if area < surface_area_minimum:
        area = surface_area_minimum
    areas.append( area ) #bohr**2

    coordinates.append( map(float, data[3:6] ) ) # bohr
f.close()

areas = numpy.array(areas)
coordinates = numpy.array(coordinates)
nts = len(areas)
vtes = numpy.zeros(nts)

# potential on the surface of TIP3P water
qpos = numpy.array([
[1.0849871266,  -0.0223279991,   0.0246861362],
[2.0528740514,   0.0086017465,  -0.0099440794],
[0.8061379337,   0.5958520619,  -0.6674536909]
]) * AA2AU
nat = len(qpos)
#qval = numpy.array([-0.834, 0.417, 0.417])
#efp water
qval = numpy.array([
    -8.5999937120+8.00000,
    -0.7000031440+1.00000,
    -0.7000031440+1.00000
    ])
for iat in range(nat):
    for its in range(nts):
        dr = qpos[iat] - coordinates[its]
        R = numpy.sqrt(dr.dot(dr))
        vtes[its] = vtes[its] + qval[iat]/R

vtes *= -(eps-1)/eps

# get the C-matrix. Let's just invert it for now.
C = CMAT(areas, coordinates)
if inverse:
    Ci = numpy.linalg.inv(C)
    qout = Ci.dot( vtes )
else:
    d0 = 1.0 / numpy.diag(C)
    qin = scal_charges_iter * vtes * d0

    q_storage = []
    q_diff_storage = []

    for k in range(1, max_iterations):
        vpot = numpy.zeros(nts)
        for its in range(nts):
            for jts in range(its):
                dr = coordinates[its] - coordinates[jts]
                R = numpy.sqrt(dr.dot(dr))
                vpot[its] = vpot[its] + qin[jts] / R
                vpot[jts] = vpot[jts] + qin[its] / R

        qout = (vtes - vpot) * d0
        dq = qout - qin

        if k >= diis_start_from_iter:
            q_storage.append( qout )
            q_diff_storage.append( dq )

        if len(q_diff_storage) > 1:
            n = len(q_diff_storage)
            rhs = numpy.zeros(n+1)
            rhs[n] = -1

            # setup system of linear equations
            B = numpy.zeros((n + 1, n + 1))
            B[:][n] = -1
            B = B.transpose()
            B[:][n] = -1
            B[n][n] = 0

            for i, qi in enumerate(q_diff_storage):
                for j, qj in enumerate(q_diff_storage):
                    B[i, j] = qi.dot(qj)

            # solve the system of linear equations
            c = numpy.linalg.solve(B, rhs)[:n]

            # better guess
            qout = numpy.zeros(numpy.shape(dq))
            for i in range(n):
                qout += c[i] * q_storage[i]

            # remove, if needed, items from the DIIS space
            sel = list(numpy.where(numpy.abs(c) < diis_vector_threshold)[0])
            sel.reverse()
            for i in sel:
                q_storage.pop(i)
                q_diff_storage.pop(i)

        abs_err = numpy.sqrt(dq.dot(dq)/nts)
        e_q = -0.5 * vtes.dot(qout)
        print("iter = {0:3d}, energy = {1:12.8f}, Q = {2:9.5f}, error = {3:16.9f}".format(k, e_q, sum(qout), abs_err))
        if abs_err < threshold:
            break
        qin = qout[:]

Q = sum(qout)
G = -0.5 * vtes.dot(qout)
print("ASC    = {0:16.9f}".format(Q))
print("Energy = {0:16.9f}".format(G))