#! /usr/bin/env python
# -*- coding: utf-8 -*-
#
# Copyright © 2016 Hao Ren <renh.cn@gmail.com>
#
# Distributed under terms of the LGPLv3 license.

"""
Generate a Gaussian fchk file contains Natural Transition Orbital information.
    An old fchk file is required to read system details apart from the 
    eigenenergies and eigenvalues (MO coefficients).

    The Gaussian output file is required to provide the T matrix (see R. L. Martin,
    JCP, 118, 4775, 2003), which should be stored in a higher precision than 
    default [e.g. using IOp(9/47=7) in the route line]. 
"""
from __future__ import print_function
import argparse
import numpy as np

parser = argparse.ArgumentParser()
parser.add_argument('-l', '--log',
                            help = 'Gaussian log file name as input',
                                                default = '')
parser.add_argument('-r', '--root', type=int,
                            help = 'transition of interest, root=N in Gaussian',
                                                default = 1)
parser.add_argument('-k', '--fchk',
                            help = 'formed checkpoint file generated by Gaussian',
                                                default='')
parser.add_argument('-o', '--output',
                            help = 'Output fchk file contains the new NTO coefficients',
                                                default='')
args = parser.parse_args()

if not args.log:
    raise ValueError('Gaussian log file not specified. Use -h for help')
else:
    log_name = args.log

if not args.fchk:
    raise ValueError('Gaussian formmated checkpoint file not specified. Use -h for help')
else:
    fchk_name = args.fchk

root = args.root

if not args.output:
    out_name = "nto-{:d}.fchk".format(root)
else:
    out_name = args.output

print("==================================================")
print("I am gnerating NTO for electronic transition # {}".format(root))
print("==================================================")


def get_NBasis(fname):
    pattern = ' NBasis'
    with open(fname) as fh:
        while True:
            line = fh.readline()
            if line.startswith(pattern):
                break
        NBasis = int(line.split()[1])
        line = fh.readline()
        NBsUse = int(line.split()[1])
    return NBasis, NBsUse

NBasis, NBsUse = get_NBasis(log_name)
print("NBasis = {}, NBsUse = {}".format(NBasis, NBsUse))


def get_restricted(fname):
    pattern = ' Excited State   1'
    with open(fname) as fh:
        while True:
            line = fh.readline()
            if line.startswith(pattern):
                break
        line = fh.readline()
    if 'A' in line or 'B' in line:
        return False
    else:
        return True

restricted = get_restricted(log_name)
if restricted:
    print('\nThis is a restricted-orbital calculation')
    nspin = 1
else:
    print('\nThis is an un-restricted-orbital calculation')
    print(
        'Singlular value decomposition will be performed in each spin-subspace')
    nspin = 2

def get_Nocc(fname):
    pattern = 'alpha electrons'
    with open(fname) as fh:
        while True:
            line = fh.readline()
            if pattern in line:
                break
    info = line.split()[::3]
    Nocc = [int(x) for x in info]
    return np.array(Nocc)

Nocc = get_Nocc(log_name)
Nvir = NBsUse - Nocc
print("  {} occupied alpha orbs, {} virtual alpha orbs".format(Nocc[0], Nvir[0]))
print("  {} occupied  beta orbs, {} virtual  beta orbs".format(Nocc[0], Nvir[0]))

# parse T-matrix from Gaussian log
pattern = ' Excited State{:4d}'.format(root)
#pattern = pattern.encode()
Ta = np.zeros([Nocc[0], Nvir[0]])
print('NBasis =', NBasis)
print('Nocc = ', Nocc)
print('Nvir = ', Nvir)
print(Ta.shape)
print("\nParsing T-matrix for excited state #{}".format(root))
if not restricted:
    Tb = np.zeros([Nocc[1], Nvir[1]])
with open(log_name) as fh:
    while True:
        line = fh.readline()
        if line.startswith(pattern):
            break
    while True:
        line = fh.readline()
        if not '->' in line:
            break
        occ_ends = 8 if restricted else 7
        iocc = int(line[:occ_ends]) - 1
        ivir = int(line[11:16]) - 1
        if 'B' in line:
            ivir -= Nocc[1]
            Tb[iocc, ivir] = float(line[19:])
        else:
            ivir -= Nocc[0]
            Ta[iocc, ivir] = float(line[19:])
print("  T-matrix parsed.")
print("  max element in T: ", np.max(np.abs(Ta)))
print('  index for max element in T: {}\n'.format(np.argmax(np.abs(Ta))))

# parse MO coefficients from fchk
with open(fchk_name) as fh:
    #  alpha orbitals
    while True:
        line = fh.readline()
        if line.startswith('Alpha MO coefficients'):
            break
    Ndata = int(line.split()[-1])
    if (NBasis * NBsUse) != Ndata:
        print("Read MO coefficients error:")
        print("   NBasis * NBsUse = {}".format(NBasis*NBsUse))
        print("     NData in file = {}".format(Ndata))
        raise ValueError
    MOC = []
    if Ndata % 5:
        nblock = Ndata // 5 + 1
    else:
        nblock = Ndata // 5
    for ib in range(nblock):
        line = fh.readline()
        MOC.extend([float(x) for x in line.split()])
    MOC_alpha = np.array(MOC).reshape([NBasis,NBsUse], order='F')
    print("\nalpha MO coeff parsed")

    if not restricted:
        # beta orbitals
        while True:
            line = fh.readline()
            if line.startswith('Beta MO coefficients'):
                break
        # we use the same Ndata, nblock, etc. as the alpha case
        MOC = []
        for ib in range(nblock):
            line = fh.readline()
            MOC.extend([float(x) for x in line.split()])
        MOC_beta = np.array(MOC).reshape([NBasis, NBsUse], order='F')
        print('\nbeta MO coeff parsed')

# parse orbital energies from fchk
with open(fchk_name) as fh:
    # alpha
    while True:
        line = fh.readline()
        if line.startswith('Alpha Orbital Energies'):
            break
    Ndata = int(line.split()[-1])
    if Ndata != NBsUse:
        print('No. energies not equal to NBsUse')
        print(' check diffuse basis function related issues')
    if Ndata % 5:
        nblock = Ndata // 5 + 1
    else:
        nblock = Ndata // 5
    MOE = []
    for ib in range(nblock):
        line = fh.readline()
        MOE.extend([float(x) for x in line.split()])
    MOE_alpha = np.array(MOE)

    if not restricted:
        # for beta orb energies
        # use the same block info
        while True:
            line = fh.readline()
            if line.startswith('Beta Orbital Energies'):
                break
        MOE = []
        for ib in range(nblock):
            line = fh.readline()
            MOE.extend([float(x) for x in line.split()])
        MOE_beta = np.array(MOE)

print('\nOrbital energies read')


    
# Then perform SVD, generate NTO
#   alpha
Ua, sa, VHa = np.linalg.svd(Ta)
Va = VHa.T

#   rearrange eigenvalues for alpha spin
MOE_alpha[:Nocc[0]] = sa[::-1]
MOE_alpha[Nocc[0] : Nocc[0]*2] = sa

#   generate NTO
MOC_alpha_occ = MOC_alpha[:, :Nocc[0]]
MOC_alpha_vir = MOC_alpha[:, Nocc[0]:]
NTO_alpha_occ = np.dot(MOC_alpha_occ, Ua)
NTO_alpha_vir = np.dot(MOC_alpha_vir, Va)
#   reverse the order for occupied (hole) NTO states, with corresponding
#   eigenvalues in ascending order (reverse the column vectors)
NTO_alpha_occ = NTO_alpha_occ[:, ::-1]
#   concatenate occ(hole) and vir(electron) states by columns
NTO_alpha = np.concatenate(
        (NTO_alpha_occ, NTO_alpha_vir), axis = 1
        )



if not restricted:
    Ub, sb, VHb = np.linalg.svd(Tb)
    Vb = VHb.T
    MOE_beta[:Nocc[1]] = sb[::-1]
    MOE_beta[Nocc[1] : Nocc[1]*2] = sb

    MOC_beta_occ = MOC_beta[:, :Nocc[1]]
    MOC_beta_vir = MOC_beta[:, Nocc[1]:]
    NTO_beta_occ = np.dot(MOC_beta_occ, Ub)
    NTO_beta_vir = np.dot(MOC_beta_vir, Vb)

    NTO_beta_occ = NTO_beta_occ[:, ::-1]
    NTO_beta = np.concatenate(
        (NTO_beta_occ, NTO_beta_vir), axis = 1
        )


# Then generate the fchk file by only modifying the 'orbital energies' and 
#   the 'MO coefficients' sections.

#   first read the section before orbital energies 'Alpha Orbital Energies'
out_fh = open(out_name, 'w')
with open(fchk_name, 'r') as fh:
    while True:
        line = fh.readline()
        if line.startswith('Alpha Orbital Energies'):
            out_fh.write(line)
            break
        out_fh.write(line)

    #   write modified alpha orbital energies
    NEnergies = len(MOE_alpha)
    if NEnergies % 5:
        nblock_E = NEnergies // 5  + 1
    else:
        nblock_E = NEnergies // 5
    for ib in range(nblock_E):
        fh.readline()
        start = ib * 5
        end = min(start + 5, NEnergies)
        fmt = "{:16.8E}" * (end - start) + "\n"
        out_fh.write(fmt.format(*(MOE_alpha[start : end + 1])))

    if not restricted:
        line = fh.readline()
        out_fh.write(line)
        for ib in range(nblock_E):
            fh.readline()
            start = ib * 5
            end = min(start + 5, NEnergies)
            fmt = "{:16.8E}" * (end - start) + '\n'
            out_fh.write(fmt.format(
                *(MOE_beta[start : end + 1])
                ))

    # write NTO coefficients for alpha spin
    line = fh.readline()
    assert(line.startswith('Alpha MO'))
    out_fh.write(line)
    ntoca = NTO_alpha.flatten('F') # Fortran style, column wise
    Nmoc = len(ntoca)

    if Nmoc % 5:
        nblock_C = Nmoc // 5 + 1
    else:
        nblock_C = Nmoc // 5
    for ib in range(nblock_C):
        fh.readline()
        start = ib * 5
        end = min(start + 5, Nmoc)
        fmt = "{:16.8E}" * (end - start) + '\n'
        out_fh.write(fmt.format(
            *(ntoca[start : end + 1])
            ))

    if not restricted:
        # write beta NTO coefficients
        line = fh.readline()
        assert(line.startswith('Beta MO'))
        out_fh.write(line)

        ntocb = NTO_beta.flatten('F')

        for ib in range(nblock_C):
            fh.readline()
            start = ib * 5
            end = min(start + 5, Nmoc)
            fmt = "{:16.6E}" * (end - start) + '\n'
            out_fh.write(fmt.format(
                *(ntocb[start : end + 1])
                ))
    while True:
        line = fh.readline()
        if not line: break
        out_fh.write(line)

out_fh.close()








