#!/usr/bin/env python3

import os
import os.path
import math
import numpy as np
# to install this lib, run as root: aptitude install python-numpy
from scipy.optimize import minimize_scalar

"""
Novaco physics for pentagonal symmetry, argparse version
"""


#def abs1(x):
#    return np.linalg.norm(x)

abs1=np.linalg.norm

def abs2(x):
    return np.dot(x, x)

def cabs2(c):
    return c.real**2+c.imag**2

def rotation(angle):
    cangle = np.cos(angle)
    sangle = np.sin(angle)
    rotArray = [[cangle, -sangle], [sangle, cangle]]
    return rotArray

def omegaL(q):
    return cL * abs1(q)

def omegaT(q):
    return cT * abs1(q)

def dyn_mat(q):
    sqrt3 = np.sqrt(3)
    qxa = q[0] * acol
    qya = q[1] * acol
    cosxa = np.cos(qxa)
    cosxa2 = np.cos(qxa/2)
    cosyas3 = np.cos(sqrt3*qya/2)
    offdiag = sqrt3*np.sin(qxa/2)*np.sin(sqrt3*qya/2)
    return kappa/M * np.array([[(3 -cosxa2*cosyas3 - 2*cosxa), offdiag],\
                               [offdiag, 3 - 3*cosxa2*cosyas3]])

def sign(number):
    ab=abs(number)
    if ab == 0:
        return 1
    else:
        return number/ab

def omega(q):
    #return eigenfrequencies and eigenvectors
    m = dyn_mat(q)
    w = np.linalg.eigh(m)
    eigenfreq=np.sqrt(abs(w[0])) # potentially dangerous, no failure if negative eigenvals
    eigenvectors=np.transpose(w[1])
    sig = sign(np.dot(eigenvectors[0], epsilonT(q)))
    eigenvectors[0] = eigenvectors[0]*sig

    sig = sign(np.dot(eigenvectors[1], epsilonL(q)))
    eigenvectors[1] = eigenvectors[1]*sig
    # here it may help to do something about the eigenvectors phase continuity...
    return eigenfreq, eigenvectors


def inBZ(q, rotatedtau):
    '''returns True when q is in BZ, and a weight to make it fuzzy around 0'''
    inbz = True
    # the squared length of the b_i vectors /2
    norm2tau_half = abs2(rotatedtau[0])*0.5
    norm2_adjust = [norm2tau_half,norm2tau_half,norm2tau_half,
                   norm2tau_half-1e-15,norm2tau_half-1e-15,norm2tau_half-1e-15]
    for i in range(len(rotatedtau)):
        # the projection of q in the rtau direction, times the rtau length
        aaa = np.dot(q, rotatedtau[i])
        if aaa > norm2_adjust[i]:  # violation of the 1BZ condition
            inbz = False
    cutoff = 1.0
    absq = abs1(q)
    if absq < qmin:
        cutoff = math.pow(absq/qmin, 2)
    return inbz, cutoff

def _inBZ_dangerous(q, rotatedtau):
    # returns True if q is in BZ, and a weight to make it fuzzy around 0
    # POSSIBLY WRONG AT THE EDGE!!!
    inbz = False
    rtaus = rotatedtau/2
    rtaulen = abs1(rtaus[0])

    for rtau in rtaus:
        inbz = inbz or abs(np.dot(q, rtau)/rtaulen) <= rtaulen

    cutoff = 1.0
    absq = abs1(q)
    if absq < qmin:
        cutoff = math.pow(absq/qmin, 2)
    return inbz, cutoff


def epsilonL(q):
    return q/abs1(q)


def epsilonT(q):
    return [-q[1], q[0]]/abs1(q)


def f2omega(q, G, T, R):
    (omegaT, omegaL), (epsT, epsL) = omega(q)
    # epsL = epsilonL(q)
    # epsT = epsilonT(q)
    ff = cabs2(np.dot(G, epsL)*np.exp(1j*np.dot(T, R))/omegaL) + \
        cabs2(np.dot(G, epsT)*np.exp(1j*np.dot(T, R))/omegaT)
    return ff


def test_dispersion(b1, firstshelltau):
    with open('omega.dat', 'w') as fomega:
        print('# step diag_omegaT diag_omegaL exact_omegaT exact_omegaL', file=fomega)
        for i in range(1,200):
            q = np.array([4*i*1e-3, i*1e-2])
            sqrt3 = np.sqrt(3)
            qxa = q[0] * acol
            qya = q[1] * acol
            cosxa = np.cos(qxa)
            cosxa2 = np.cos(qxa/2)
            cosya = np.cos(qya)
            cosyas3 = np.cos(sqrt3*qya/2)
            a = 3 - cosxa - 2 * cosxa2*cosyas3
            b = cosxa**2 + (cosxa2*cosyas3)**2 - 2 * cosxa*cosxa2*cosyas3 + 3 \
                * (np.sin(qxa/2)*np.sin(sqrt3*qya/2))**2
            if inBZ(q, firstshelltau):
                (oT, oL), (epsT, epsL) = omega(q)
                print(q[0], q[1], oT, oL, np.sqrt(kappa/M * (a - np.sqrt(b))), np.sqrt(kappa/M * (a + np.sqrt(b))), file=fomega)


def novaco_energy(alpha, gvectors, tauvectorsu, firstshelltau):

    rot = rotation(-alpha*twopi/360)
    rotatedg = np.dot(gvectors, np.transpose(rot))

    en_1ph = 0.
    for gg in rotatedg:
        for tau in tauvectorsu:
            q = gg-tau
            inbz, cutoff = inBZ(q, firstshelltau)
            if inbz:
                en_1ph -= f2omega(q, gg, tau, R0)*cutoff
                (omegaT, omegaL), (epsT, epsL) = omega(q)
                if file_q:
                    file_q.write('{} {} {} {} {} {} {} {} {} {} {}\n'.format(
                        alpha, q[0], q[1], gg[0], gg[1] , tau[0], tau[1], epsT[0], epsT[1], epsL[0], epsL[1]))
    if file_q:
        file_q.write('\n\n')
    return en_1ph*Vnorm**2/2./M

def minimize_en(args, gvectors, tauvectorsu, firstshelltau):
    bound=(args.minalpha,args.maxalpha)
    res=minimize_scalar(novaco_energy, bounds=bound, args=(gvectors, tauvectorsu, firstshelltau), tol=1e-8)
    if not res.success:
        print("ERROR: minimization failure, with params", args)
    return res.x, res.fun
    
def loop_overangles(args, gvectors, tauvectorsu, firstshelltau):
    for ialpha in range(int(round(args.minalpha/args.deltaalpha)),int(round(1.+args.maxalpha/args.deltaalpha))):
        alpha = ialpha*args.deltaalpha
        
        en_1ph =novaco_energy(alpha, gvectors, tauvectorsu, firstshelltau)
        print(alpha, en_1ph)


def novaco_preliminaries(args):
    global acol, apot
    global pi, twopi
    global qmin
    global cT, cL
    global kappa	# the spring constant of the lattice
    global M		# the mass of the particles in the lattice 
    global R0
    pi = math.pi
    twopi = 2*pi

    rcatino = 1170.
#   values from commandline:
    nsymm=args.nsymm

    apot = args.apot
    acol = args.acol

# for these coefficients see Table 1 of SI of [PNAS 109, 16429-16433 (2012)]:
    coefs=[0,0,0.5,2./3.,1/np.sqrt(2.),1,2/np.sqrt(3.),1,1,1,1,1,1,1,1,1]
    gscalingfactor=coefs[nsymm]

    # NOTE: LAMMPS harmonic potential is K_{lammps}(d_i - d_{i-1})^2
    #      in the notes the harmonic potential is K_{notes}/2 (d_i - d_{i-1})^2
    #
    # K_{lammps} = K_{notes}/2 => K_{notes} = 2*K_{lammps}
    kappa = 2e-7		# in fkg / µs² (K_{notes})
    kappa = kappa * 1e6		# in fkg / ms²
    M = 31.06                   # in fkg
    cT = np.sqrt(3*kappa/M) * acol / math.pow(2, 3/2)   # µm/ms
    cL = np.sqrt(kappa/M) * 3 * acol / math.pow(2, 3/2) # µm/ms
    V = args.V0      # in zJ
    print("# Santoro/Novaco analysis for a_pot=", apot, "a_col=", acol," (micrometers); nsymm=",nsymm)

    print("# V0:", V, "zJ", "K acol^2:", kappa*acol*acol,"zJ","; g=V0/(kappa*acol^2):", V/(kappa*acol*acol))

    Vnorm = V/float(nsymm**2)   # note the factor nsymm**2 in the proper definition of V!

    qmin = twopi/rcatino
    klength = gscalingfactor*twopi/apot

    print("# cutoffing parameters qmin:", qmin," klength:", klength)

    R0 = [args.X0, args.Y0]
    print("# center position R0:", R0)

    firstkappa = [klength, 0.]
    kvectors = []
    for i in range(nsymm):
        gg = np.dot(rotation(twopi*i/nsymm), firstkappa)
        kvectors.append(gg)

    gvectors = []
    for ki in kvectors:
        for kj in kvectors:
            gg = ki-kj
            if abs1(gg) > 0:
                gvectors.append(gg)

    if args.debug:
        print("debug1", gvectors)

    b1 = np.array([twopi/acol, twopi/acol/math.sqrt(3)])
    b2 = np.array([0., 2*twopi/acol/(math.sqrt(3))])

    maxlength = 3.1*klength
    maxradius = maxlength*math.sqrt(2)

    nmax = 7
    tauvectorsu = []
    for n1 in range(-nmax, nmax):
        for n2 in range(-nmax, nmax):
            tau = float(n1)*b1+float(n2)*b2
            if args.debug:
                print("debug0", tau)
            if abs1(tau) < maxradius:
                tauvectorsu.append(tau)

    # first shell of reciprocal tau vectors, to be used in inBZ
    firstshelltau = []
    for ishell in range(6):
        tau = np.dot(rotation(twopi*ishell/6.), b1)
        firstshelltau.append(tau)
    firstshelltau = np.array(firstshelltau)

    if args.writefile:    # file to store q at each angle
        file_q = open('file_q.dat', 'w')
        file_q.write('# angle qx qy ggx ggy taux tauy epsTx, epsTy, epsLx, epsLy\n')
    else:
        file_q=False

    return file_q, Vnorm, gvectors, tauvectorsu, firstshelltau

if __name__ == "__main__":
    import sys
    import argparse
    commandname=sys.argv[0]

    desc="""Novaco energetics physics for pentagonal symmetry in k space
    OUTPUT: the angular dependence of the energy per particle
    by Emanuele Panizon & Nicola Manini & Mario Forzanini
 
    """

    epil="""            v. 0.9  by Nicola Manini, 07/11/2024"""

    parser = argparse.ArgumentParser( formatter_class=argparse.ArgumentDefaultsHelpFormatter
                                    , description=desc, epilog=epil)

    parser.add_argument( 'filenames', nargs='*', default=['-'],
                         help='UNUSED Files to be processed. If none given, stdin is used')

    parser.add_argument( '-a', 
                         dest='acol', type=float, default=5.8,
                         help='lattice spacing of particles (microm)')

    parser.add_argument( '-b', 
                         dest='apot', type=float, default=5.2,
                         help='lattice spacing of corrugation (microm)')

    parser.add_argument( '-d', 
                         dest='debug', type=int, default=0,
                         help='raise the debug level')

    parser.add_argument( '-f', 
                         dest='writefile', action='store_true',
                         help='write a file with all q vectors')

    parser.add_argument( '-i', 
                         dest='deltaalpha', type=float, default=0.1,
                         help='the angular increment (degrees)')

    parser.add_argument( '--min', 
                         dest='minimize', action='store_true',
                         help='search for energy minimum')

    parser.add_argument( '-z', 
                         dest='minalpha', type=float, default=0.,
                         help='the initial angle (degrees)')

    parser.add_argument( '-m', 
                         dest='maxalpha', type=float, default=6.,
                         help='the maximum angle (degrees)')

    parser.add_argument( '--nsymm', 
                         dest='nsymm', type=int, default=5,
                         help='the symmetry of the potential corrugation (e.g. 3->hexagonal, 5->decagonal)')

    parser.add_argument( '-V', 
                         dest='V0', type=float, default=1.,
                         help='the potential corrugation amplitude (zJ)')

    parser.add_argument( '-x', 
                         dest='X0', type=float, default=0.,
                         help='relative x displacement (microm)')

    parser.add_argument( '-y', 
                         dest='Y0', type=float, default=0.,
                         help='relative y displacement (microm)')

## End arg parser definition
    args=parser.parse_args(sys.argv[1:])
    d = vars(args)	 # adding prog to args (for unknown reasons it's not there already...)
    d['prog']=parser.prog

#   here the actual functions doing the job are called:
    file_q, Vnorm, gvectors, tauvectorsu, firstshelltau = novaco_preliminaries(args)
    if args.minimize:
        alpha, en = minimize_en(args, gvectors, tauvectorsu, firstshelltau)
        print("#minimum_twist_angle	min_energy	apot/acoll")
        print(alpha,en,args.apot/args.acol)
    else:
        loop_overangles(args, gvectors, tauvectorsu, firstshelltau)
        
#        novaco_energy(args)