"""
Frank-Wolfe algorithm.
"""
# 2015-04-15
#
# We place this in its own module since using:
#
#   from __future__ import division
#
# seems to cause some issues with `cvxopt.matrix`, causing a TypeError.
#
# >>> from cvxopt import matrix
# >>> x = matrix([[1,2],[3,5]])
# >>> x / 3
# TypeError: unsupported operand type(s) for /: 'cvxopt.base.matrix' and 'int'

from debtcollector import removals

import logging
import numpy as np

from dit.utils import basic_logger

__all__ = (
    'frank_wolfe',
)

@removals.remove(message="Please see methods in dit.algorithms.distribution_optimizers.py.",
                 version='1.0.1')
def frank_wolfe(objective, gradient, A, b, initial_x,
                maxiters=2000, tol=1e-4, clean=True, verbose=None):
    """
    Uses the Frank--Wolfe algorithm to minimize the convex objective.

    Minimization is subject to the linear equality constraint: :math:`A x = b`.

    Assumes x should be nonnegative.

    Parameters
    ----------
    objective : callable
        The objective function. It would receive a ``cvxopt`` matrix for the
        input `x` and return the value of the objective function.
    gradient : callable
        The gradient function. It should receive a ``cvxopt`` matrix for the
        input `x` and return the value of the gradient evaluated at `x`.
    A : matrix
        A ``cvxopt`` matrix specifying the LHS linear equality constraints.
    b : matrix
        A ``cvxopt`` matrix specifying the RHS linear equality constraints.
    initial_x : matrix
        A ``cvxopt`` matrix specifying the initial `x` to use.
    maxiters : int
        The maximum number of iterations to perform. If convergence was not
        reached after the last iteration, a warning is issued and the current
        value of `x` is returned.
    tol : float
        The tolerance used to determine when we have converged to the optimum.
    clean : bool
        Occasionally, the iteration process will take nonnegative values to be
        ever so slightly negative. If ``True``, then we forcibly make such
        values equal to zero and renormalize the vector. This is an application
        specific decision and is probably not more generally useful.
    verbose : int
        An integer representing the logging level ala the ``logging`` module.
        If `None`, then (effectively) the log level is set to `WARNING`. For
        a bit more information, set this to `logging.INFO`. For a bit less,
        set this to `logging.ERROR`, or perhaps 100.

    """
    # Function level import to avoid circular import.
    from dit.algorithms.optutil import op_runner

    # Function level import to keep cvxopt dependency optional.
    # All variables should be cvxopt variables, not NumPy arrays
    from cvxopt.modeling import variable

    # Set up a custom logger.
    logger = basic_logger('dit.frankwolfe', verbose)

    # Set cvx info level based on logging.DEBUG level.
    if logger.isEnabledFor(logging.DEBUG):
        show_progress = True
    else:
        show_progress = False

    assert (A.size[1] == initial_x.size[0])

    n = initial_x.size[0]
    x = initial_x
    xdiff = 0

    TOL = 1e-7
    verbosechunk = maxiters / 10
    for i in range(maxiters):
        obj = objective(x)
        grad = gradient(x)

        xbar = variable(n)

        new_objective = grad.T * xbar
        constraints = []
        constraints.append((xbar >= 0))
        constraints.append((-TOL <= A * xbar - b))
        constraints.append((A * xbar - b <= TOL))

        logger.debug('FW Iteration: {}'.format(i))
        opt = op_runner(new_objective, constraints, show_progress=show_progress)
        if opt.status != 'optimal':
            msg = '\tFrank-Wolfe: Did not find optimal direction on '
            msg += 'iteration {}: {}'
            msg = msg.format(i, opt.status)
            logger.info(msg)

        # Calculate optimality gap
        xbar_opt = opt.variables()[0].value
        opt_bd = grad.T * (xbar_opt - x)

        msg = "i={:6}  obj={:10.7f}  opt_bd={:10.7f}  xdiff={:12.10f}"
        if logger.isEnabledFor(logging.DEBUG):
            logger.debug(msg.format(i, obj, opt_bd[0, 0], xdiff))
            logger.debug("")
        elif i % verbosechunk == 0:
            logger.info(msg.format(i, obj, opt_bd[0, 0], xdiff))

        xnew = (i * x + 2 * xbar_opt) / (i + 2)
        xdiff = np.linalg.norm(xnew - x)
        x = xnew

        if xdiff < tol:
            obj = objective(x)
            break
    else:
        msg = "Only converged to xdiff={:12.10f} after {} iterations. "
        msg += "Desired: {}"
        logger.warn(msg.format(xdiff, maxiters, tol))

    xopt = np.array(x)

    if clean:
        xopt[np.abs(xopt) < tol] = 0
        xopt /= xopt.sum()

    return xopt, obj


"""All main functions for et_module
"""
import os
import sys
import csv
import logging
import numpy as np
import pandas as pd

from et_module import diffusion_functions

def load_curve_assignement(
        curr_yr,
        base_yr,
        yr_until_changed,
        et_service_demand_yh,
        load_profiles,
        regions,
        charging_scenario,
        diffusion='linear'
    ):
    """Assign input electrictiy demand (given as "tranport service"
    for every hour in a year) to an hourly energy demand load profile
    depending (see documentation for more information).

    Arguments
    =========
    curr_yr : int
        Current simulation year
    base_yr : int
        Base year of simulation
    yr_until_changed : int
        Year until changed is fully implemented
    et_service_demand_yh : dict
        Transport energy demand for every region (hourly demand)
    load_profiles : list
        Load profile objects
    regions : list
        All region names
    charging_scenario : str
        Scenario

            'sheduled' :

            'sheduled' :
            TODO

    diffusion : str
        Type of diffusion between base year and end year load profile

            'linear':   Linear change over time towards future load profile

            'sigmoid':  Sigmoid change over time towards future load profile

    Returns
    =========
    et_demand_yh : array
        Houlry demand, np.array(reg_array_nr, 8760 timesteps)
    """
    et_demand_yh = np.zeros((len(regions), 365 * 24), dtype=float)

    # -------------------
    # Calculate diffusion
    # -------------------
    if diffusion == 'linear':
        simulation_year_p = diffusion_functions.linear_diff(
            base_yr=base_yr,
            curr_yr=curr_yr,
            value_start=0,
            value_end=1,
            yr_until_changed=yr_until_changed)

    elif diffusion == 'sigmoid':
        # Default sigmoid parameters
        simulation_year_p = diffusion_functions.sigmoid_diffusion(
            base_yr=base_yr,
            curr_yr=curr_yr,
            end_yr=yr_until_changed,
            sig_midpoint=0,
            sig_steeppness=1)
    else:
        sys.exit("Error: No diffusion option is selected")
    # --------------------------------------------------------------------
    # Calculate current year profile with base year and profile from 2015
    # --------------------------------------------------------------------

    # Get base year load profile
    for load_profile in load_profiles:
        if load_profile.name == 'av_lp_2015.csv':
            profile_yh_by = load_profile.shape_yh

    # Get future year load profile
    for load_profile in load_profiles:

        if charging_scenario == 'unsheduled':

            # Unsheduled load profile (same as base year)
            if load_profile.name == 'av_lp_2015.csv':
                profile_yh_ey = load_profile.shape_yh

        elif charging_scenario == 'sheduled':

            # Sheduled load profile
            if load_profile.name == 'av_lp_2050.csv':
                profile_yh_ey = load_profile.shape_yh

    if base_yr == curr_yr:
        profile_yh_cy = profile_yh_by
    elif curr_yr == yr_until_changed or curr_yr > yr_until_changed:
        profile_yh_cy = profile_yh_ey
    else:

        # Calculate difference between by and ey
        diff_profile = profile_yh_ey - profile_yh_by

        # Calculate difference up to cy
        diff_profile_cy = diff_profile * simulation_year_p

        # Add difference to by
        profile_yh_cy = profile_yh_by + diff_profile_cy

    assert round(np.sum(profile_yh_cy), 3) == 1

    # ----------
    # Plotting
    # ----------
    #from et_module import plotting_functions
    #fig_lp.plot_lp_dh(profile_yh_cy, day=2)

    # ------------------------------------
    # Disaggregate for every region
    # ------------------------------------
    for region_array_nr, region in enumerate(regions):

        # Sum total service demand to annual demand
        et_service_demand_y = np.sum(et_service_demand_yh[region])

        # Multiply the annual total service demand with yh load profile
        reg_profile_yh = et_service_demand_y * profile_yh_cy

        logging.debug(
            "Assinging new shape {}  {} {}".format(
                region, et_service_demand_y, np.sum(profile_yh_cy)))

        # Reshape (365 days, 24hours) into 8760 timesteps
        et_demand_yh[region_array_nr] = reg_profile_yh.reshape(8760)

    return et_demand_yh

def get_load_profiles(path):
    """Read in all load profiles from csv files and store in
    `LoadProfile`.

    Arguments
    =========
    path : str
        Path where load profiles are stored

    Returns
    =======
    load_profiles : list
        All load profiles objects
    """
    load_profiles = []

    # Name of load profiles to load
    names = [
        'av_lp_2015.csv',
        'av_lp_2050.csv']

    for name in names:

        # Create path to csv file
        path_to_csv = os.path.join(path, name)

        # Read in csv load profile
        lp_dh = read_load_shape(path_to_csv)

        lp_dh_p = lp_dh / 100 # convert percentage to fraction

        # Shape for every hour in a year (Assign same profile to every day)
        shape_yd = np.full((365), 1/365)

        # Shape for every hour in a year (365) * (24)
        shape_yh = shape_yd[:, np.newaxis]  * lp_dh_p

        # Create load profile
        load_profile = LoadProfile(
            name=name,
            year=name[-8:-4],
            shape_yd=shape_yd,
            shape_yh=shape_yh)

        load_profiles.append(load_profile)

    return load_profiles

def read_load_shape(path_to_csv):
    """This function reads in a load profile from
    a csv file of a single day.

    Arguments
    =========
    path_to_csv : str
        Path to csv file

    Returns
    =======
    shape_dh : array (24)
        Load profile
    """
    with open(path_to_csv, 'r') as csvfile:
        read_lines = csv.reader(csvfile, delimiter=',')
        _headings = next(read_lines) # Skip first row

        for row in read_lines:
            shape_dh = np.zeros((24), dtype=float)
            for cnt, row_entry in enumerate(row):
                shape_dh[int(_headings[cnt])] = float(row_entry)

    return shape_dh

class LoadProfile(object):
    """Class to store load profiles

    Arguments
    ----------
    name : str
        Name of load profile
    year : int
        Year of load profile
    shape_yd : array
        Yearly load profile
    shape yh : array
        Daily load profile

    Note
    ====
    -   `Yearly load profile (yd)` can be used to derive the energy demand
        of all days in year. This is achieved by multiplying total
        annual demand with this _yd array with the array shape (365)

    -   `Daily load profile (yh)` can be used to derive the energy demand
        of all hours in a year. This is achieved by multiplying total
        annual demand with the this _yh array with the array shape (365, 24).
    """
    def __init__(
            self,
            name,
            year,
            shape_yd,
            shape_yh
        ):
        """Constructor
        """
        self.name = name
        self.year = year
        self.shape_yd = shape_yd
        self.shape_yh = shape_yh


import numpy as np
import sys
import warnings
import argparse
from fit import do_fit
from output import *
from input import get_input,get_imax
warnings.simplefilter("error")

def get_args():

    args_parser = argparse.ArgumentParser()

    # Data files arguments
    args_parser.add_argument(
        '--vmcfile',
        help='Local path to vmc data',
        required=True,
        type=str
    )

    args_parser.add_argument(
        '--dmcfile',
        help='Local path to dmc data',
        required=True,
        type=str
    )

    args_parser.add_argument(
        '--rfile',
        help='Local path to dmc data',
        required=True,
        type=str
    )

    args_parser.add_argument(
        '--max-pol',
        help='Maximum polynomial order.',
        required=False,
        type=int,
        default=3
    )

    args_parser.add_argument(
        '--min-pol',
        help='Minimum polynomial order.',
        required=False,
        type=int,
        default=3
    )

    args_parser.add_argument(
        '--plot',
        help='BOOL. Produce plot of the PCFs. 0 (default) FALSE. 1 TRUE. ',
        required=False,
        type=int,
        default=0
    )

    args_parser.add_argument(
        '--num-e',
        help='Total Number of electrons',
        required=False,
        type=int,
        default=3
    )

    args_parser.add_argument(
        '--lat-vec',
        help='Maximum polynomial order.',
        required=True,
        type=float
    )

    args_parser.add_argument(
        '--fit-range',
        help='Maximum polynomial order.',
        required=True,
        type=float
    )

    args_parser.add_argument(
        '--opt-method',
	help='lm,trf or dogbox',
	required=False,
        type=str,
        default='lm'
    )

    args_parser.add_argument(
        '--volume',
        help='Volume of the simulation cell.',
        required=True,
        type=float
    )

    args_parser.add_argument(
        '--metal',
        help='System is metallic (1) or not (0). Affects to the scaling of the extrapolated PCF. Default: 0.',
        required=False,
        type=int,
        default=0
    )

    args_parser.add_argument(
        '--verbosity',
        help='Controls the amount of info printed.',
        required=False,
        type=int,
        default=0
    )

    args_parser.add_argument(
	'--omit_pcf',
        help='skip pcfs.',
        required=False,
        nargs='+',
	type=int,
        default=[-1]
    )

    args_parser.add_argument(
        '--corepart',
        help='tau_valence/tau_total.',
        type=float,
        default=1.0
    )

    args_parser.add_argument(
	'--table',
        help='Printout table in latex.',
        required=False,
	type=int,
        default=[-1]
    )

    args_parser.add_argument(
        '--cross-val',
        help='Perform a cross-validation over the fits.',
        required=False,
        type=int,
        default=[-1]
    )

    return args_parser.parse_args()

def fit_statistics(args,fits,g,r):

    # The dimensions of our current experiment
    Nx=len(r)
    Npcf=len(g)
    p_degs=np.arange(args.min_pol,args.max_pol+1,2)
    N_degs = len(p_degs)
    bond_distance=args.fit_range*args.lat_vec

    # Averages of the PCFs and fits
    fit_average=np.zeros((Nx,N_degs))
    g_average=np.zeros((Nx,))
    for i in range(Npcf):
        g_average[:]=g_average[:]+g[i]
        for d in range(N_degs):
            fit_average[:,d]=fit_average[:,d]+fits[:,i,d]

    g_average=g_average/Npcf
    fit_average=fit_average/Npcf

    # Compute the fitting error
    fit_errors=np.zeros((N_degs,))
    fit_sqerrors=np.zeros((N_degs,))
    imax=get_imax(r,2)
    for i in range(N_degs):
        for j in range(imax):
            fit_errors[i]=fit_errors[i]+np.absolute(np.exp(fit_average[j,i])-g_average[j])
            fit_sqerrors[i]=fit_sqerrors[i]+(np.exp(fit_average[j,i])-g_average[j])**2
    fit_errors=fit_errors/imax
    fit_sqerrors=fit_sqerrors/imax

    if args.plot == 1:
        plot_results(p_degs,N_degs,g_average,fit_average,r,imax)

    return fit_errors,fit_sqerrors

def mean_and_error(g):
    '''
    Input: (N_pcf,N_degs)-order matrix of PCF-zeroes.
    Output:
      * Means over separate PCF's.
      * Mean average errors of PCF's.
      * Standard deviations of PCF's.

    '''
    gmean=np.mean(g,axis=0)
    mean_error=np.mean(np.absolute(gmean-g),axis=0)
    standard_deviation=np.std(g,axis=0)/np.sqrt(g.shape[0])#np.sqrt(np.mean((gmean-g)**2,axis=0))/np.sqrt(g.shape[0])

    return gmean,mean_error,standard_deviation

def main():

    args=get_args()

    r,r_ex,gex,glogex=get_input(args)

    r_range=args.fit_range*args.lat_vec

    # Fitting
    fits,opt_pol_coeff=do_fit(r,r_ex,r_range,glogex,args)

    # Fitting statistics
    fe,fsqe=fit_statistics(args,fits,gex,r)

    # Get g(0) values and statistics
    gzeros=np.exp(fits[0,:,:])
    m,e,std=mean_and_error(gzeros)

    # Lifetime statistics
    lifetimes=1000.0*(100.617/2*args.num_e/args.volume*gzeros)**-1
    mt,et,stdt=mean_and_error(lifetimes)

    if(args.table==1):
        make_table(args,m,mt,fe,fsqe,e,std,stdt,gzeros,lifetimes)
    else:
        print_output(args,m,mt,fe,fsqe,e,std,stdt,gzeros,lifetimes)

    sys.exit('All done.')

if __name__ == '__main__':
    main()


#!/usr/bin/python3

"""Calculate the RMSD between structures."""

import argparse

import numpy as np

# Inspired by
# https://github.com/cclib/cclib/blob/master/cclib/parser/utils.py#L159
element = [
    None,
    "H",
    "He",
    "Li",
    "Be",
    "B",
    "C",
    "N",
    "O",
    "F",
    "Ne",
    "Na",
    "Mg",
    "Al",
    "Si",
    "P",
    "S",
    "Cl",
    "Ar",
    "K",
    "Ca",
    "Sc",
    "Ti",
    "V",
    "Cr",
    "Mn",
    "Fe",
    "Co",
    "Ni",
    "Cu",
    "Zn",
    "Ga",
    "Ge",
    "As",
    "Se",
    "Br",
    "Kr",
    "Rb",
    "Sr",
    "Y",
    "Zr",
    "Nb",
    "Mo",
    "Tc",
    "Ru",
    "Rh",
    "Pd",
    "Ag",
    "Cd",
    "In",
    "Sn",
    "Sb",
    "Te",
    "I",
    "Xe",
    "Cs",
    "Ba",
    "La",
    "Ce",
    "Pr",
    "Nd",
    "Pm",
    "Sm",
    "Eu",
    "Gd",
    "Tb",
    "Dy",
    "Ho",
    "Er",
    "Tm",
    "Yb",
    "Lu",
    "Hf",
    "Ta",
    "W",
    "Re",
    "Os",
    "Ir",
    "Pt",
    "Au",
    "Hg",
    "Tl",
    "Pb",
    "Bi",
    "Po",
    "At",
    "Rn",
    "Fr",
    "Ra",
    "Ac",
    "Th",
    "Pa",
    "U",
    "Np",
    "Pu",
    "Am",
    "Cm",
    "Bk",
    "Cf",
    "Es",
    "Fm",
    "Md",
    "No",
    "Lr",
    "Rf",
    "Db",
    "Sg",
    "Bh",
    "Hs",
    "Mt",
    "Ds",
    "Rg",
    "Cn",
    "Nh",
    "Fl",
    "Mc",
    "Lv",
    "Ts",
    "Og",
]

# Inspired by
# https://github.com/cclib/cclib/blob/master/cclib/parser/utils.py#L159
atomic_number = {
    "H": 1,
    "He": 2,
    "Li": 3,
    "Be": 4,
    "B": 5,
    "C": 6,
    "N": 7,
    "O": 8,
    "F": 9,
    "Ne": 10,
    "Na": 11,
    "Mg": 12,
    "Al": 13,
    "Si": 14,
    "P": 15,
    "S": 16,
    "Cl": 17,
    "Ar": 18,
    "K": 19,
    "Ca": 20,
    "Sc": 21,
    "Ti": 22,
    "V": 23,
    "Cr": 24,
    "Mn": 25,
    "Fe": 26,
    "Co": 27,
    "Ni": 28,
    "Cu": 29,
    "Zn": 30,
    "Ga": 31,
    "Ge": 32,
    "As": 33,
    "Se": 34,
    "Br": 35,
    "Kr": 36,
    "Rb": 37,
    "Sr": 38,
    "Y": 39,
    "Zr": 40,
    "Nb": 41,
    "Mo": 42,
    "Tc": 43,
    "Ru": 44,
    "Rh": 45,
    "Pd": 46,
    "Ag": 47,
    "Cd": 48,
    "In": 49,
    "Sn": 50,
    "Sb": 51,
    "Te": 52,
    "I": 53,
    "Xe": 54,
    "Cs": 55,
    "Ba": 56,
    "La": 57,
    "Ce": 58,
    "Pr": 59,
    "Nd": 60,
    "Pm": 61,
    "Sm": 62,
    "Eu": 63,
    "Gd": 64,
    "Tb": 65,
    "Dy": 66,
    "Ho": 67,
    "Er": 68,
    "Tm": 69,
    "Yb": 70,
    "Lu": 71,
    "Hf": 72,
    "Ta": 73,
    "W": 74,
    "Re": 75,
    "Os": 76,
    "Ir": 77,
    "Pt": 78,
    "Au": 79,
    "Hg": 80,
    "Tl": 81,
    "Pb": 82,
    "Bi": 83,
    "Po": 84,
    "At": 85,
    "Rn": 86,
    "Fr": 87,
    "Ra": 88,
    "Ac": 89,
    "Th": 90,
    "Pa": 91,
    "U": 92,
    "Np": 93,
    "Pu": 94,
    "Am": 95,
    "Cm": 96,
    "Bk": 97,
    "Cf": 98,
    "Es": 99,
    "Fm": 100,
    "Md": 101,
    "No": 102,
    "Lr": 103,
    "Rf": 104,
    "Db": 105,
    "Sg": 106,
    "Bh": 107,
    "Hs": 108,
    "Mt": 109,
    "Ds": 110,
    "Rg": 111,
    "Cn": 112,
    "Nh": 113,
    "Fl": 114,
    "Mc": 115,
    "Lv": 116,
    "Ts": 117,
    "Og": 118,
}

def read_xyz(path_file_or_str):
    """Read a xyz file and return structures.

    Parameters
    ----------
    path_file_or_str : str

    Returns
    -------
    atomnos : array-like
    comments : str
    atomcoords : array-like
    """

    def _process(lines):
        natom = int(lines[0])
        m = 2 + natom
        nstruct = len(lines) // m

        atomnos = []
        comments = []
        atomcoords = []
        for i in range(nstruct):
            nos = []
            coords = []
            initial = i * m + 2
            for j in range(initial, initial + natom):
                fields = lines[j].split()
                nos.append(atomic_number[fields[0]])
                coords.append([float(x) for x in fields[1:]])
            atomnos.append(np.array(nos))
            comments.append(lines[i * m + 1].strip("\n"))
            atomcoords.append(np.array(coords))
        return atomnos, comments, atomcoords

    try:
        with open(path_file_or_str, "r") as xyz_file:
            return _process(xyz_file.readlines())
    except TypeError:
        return _process(path_file_or_str.readlines())
    except FileNotFoundError:
        return _process(path_file_or_str.split("\n"))

def write_xyz(atomnos, atomcoords, comment=""):
    """Format a string as xyz coordinates.

    Parameters
    ----------
    atomnos, atomcoords : array-like

    Returns
    -------
    str
    """
    lines = []
    for no, coord in zip(atomnos, atomcoords):
        lines.append(
            f"{element[no]:2s} {coord[0]:12.6f} {coord[1]:12.6f} {coord[2]:12.6f}"
        )

    lines = "\n".join(lines)
    return f"{len(atomnos)}\n{comment}\n{lines}"

def calc_rmsd(P, Q, translate=True):
    """Calculate the RMSD between P and Q using Kabsch algorithm.

    Parameters
    ----------
    P, Q : array-like
    translate : bool

    Returns
    -------
    float

    Notes
    -----
    If structures are not comparable, np.inf is returned.

    Examples
    --------
    >>> P = [[0, 0, 0],
    ...      [1, 1, 1]]
    >>> calc_rmsd(P, P)
    0.0
    >>> Q = [[1, 1, 1],
    ...      [2, 2, 2]]
    >>> calc_rmsd(P, Q)
    0.0
    >>> P = [[-3.652796902, 0.000000000, -4.445975658],
    ...      [-3.527558151, 0.000000000, -3.430150234],
    ...      [-4.501637325, 0.000000000, -3.833697320]]
    >>> Q = [[ 1.885538972, 0.000000000, -0.577489796],
    ...      [ 0.911459798, 0.000000000,  0.201773543],
    ...      [ 1.165525626, 1.078706033, -0.490031274]]
    >>> calc_rmsd(P, Q)
    0.140663071340813
    """
    P, Q = np.asanyarray(P), np.asanyarray(Q)
    if translate:
        P = P - P.mean(axis=0)
        Q = Q - Q.mean(axis=0)

    try:
        V, s, W = np.linalg.svd(P.T @ Q)
    except ValueError:
        return np.inf
    if np.linalg.det(V) * np.linalg.det(W) < 0:
        s[-1] = -s[-1]
        V[:, -1] = -V[:, -1]

    U = V @ W
    P = P @ U
    return np.sqrt(np.sum((P - Q) ** 2) / len(P))

def main():
    """Run main procedure."""
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "xyz_files", type=argparse.FileType("r"), default="-", nargs="+"
    )
    args = parser.parse_args()

    coords = []
    names = []
    for xyz_file in args.xyz_files:
        coords.append(read_xyz(xyz_file)[2])
        names.append(xyz_file.name)

    print("RMSD:")
    for i in range(len(coords)):
        for j in range(i):
            rmsd = calc_rmsd(coords[j][-1], coords[i][-1])
            print(f"{names[j]:14s} ~ {names[i]:14s} = {rmsd:6.4f} Å")

if __name__ == "__main__":
    main()


#!/usr/bin/env python
# coding: utf-8

# In[1]:

import os
import sys
import glob
import numpy as np
import itk
from itk import TubeTK as ttk
#from itkwidgets import view

# In[2]:

if len(sys.argv) != 2:
  print("ctp-head-CombinedScript.py <ctp-directory>")
  print("   <ctp-directory> format = C:/Data/unc/HighRes-005-ctp")
  sys.exit()

# NRRD Study Name
print("Directory = ", sys.argv[1])
studyname = sys.argv[1] #'C:/Users/steph/Desktop/Data/unc/HighRes-005-ctp'

# NRRD Files
directory = (studyname + '/')

# Saved NRRD Files
directory2 = (studyname + '-Reg/')
try:
    os.mkdir(directory2)
except OSError as error:
    print(error)

# Mask Creation and Location
directory3 = (studyname + '-MinMax/')
try:
    os.mkdir(directory3)
except OSError as error:
    print(error)

pic_folder = os.listdir(directory)
pic_folder = [pic_folder for pic_folder in pic_folder if ".nii" in pic_folder]
pic_folder.sort()
print(pic_folder)
num_images = len(pic_folder)

im0Tmp = itk.imread(directory + pic_folder[int(num_images/2)], itk.F)

resample = ttk.ResampleImage.New(Input=im0Tmp,MakeIsotropic=True)
resample.Update()
im0 = resample.GetOutput()
immath = ttk.ImageMath.New(Input=im0)
immath.Blur(1)
im0Blur = immath.GetOutput()

immath.Threshold(150, 800, 1, 0)
immath.Dilate(10, 1, 0)
mask0 = immath.GetOutputUChar()
mask0Tmp = itk.GetArrayViewFromImage(mask0)
mask0Tmp[0:4,:,:] = 0
sizeZ = mask0Tmp.shape[0]
mask0Tmp[sizeZ-4:sizeZ,:,:] = 0   #No need to update mask0 since mask0Tmp is a view of mask0 (shared memory)

itk.imwrite(mask0, directory3 + 'mask.mha', compression=True)
maskObj = itk.ImageMaskSpatialObject[3].New()
maskObj.SetImage(mask0)
maskObj.Update()

# In[3]:

#view(mask0)

# In[ ]:

Dimension = 3
PixelType = itk.ctype('float')
ImageType = itk.Image[PixelType, Dimension]

imdatamax = itk.GetArrayFromImage(im0)
imdatamin = imdatamax
imdatamax2 = imdatamax
imdatamin2 = imdatamax
imdatamax3 = imdatamax
imdatamin3 = imdatamax

imFixedBlur = im0Blur

for imNum in range(num_images):
    imMoving = itk.imread( directory + pic_folder[imNum], itk.F )

    immath.SetInput(imMoving)
    immath.Blur(1)
    imMovingBlur = immath.GetOutput()

    imreg = ttk.RegisterImages[ImageType].New()
    imreg.SetFixedImage(imFixedBlur)
    imreg.SetMovingImage(imMovingBlur)

    imreg.SetRigidMaxIterations(3000)
    imreg.SetRegistration("RIGID")
    imreg.SetExpectedOffsetMagnitude(20)
    imreg.SetExpectedRotationMagnitude(0.3)
    imreg.SetMetric("MEAN_SQUARED_ERROR_METRIC")

    imreg.SetFixedImageMaskObject(maskObj)
    #imreg.SetSampleFromOverlap(True)

    imreg.SetReportProgress(True)
    imreg.Update()

    tfm = imreg.GetCurrentMatrixTransform()
    #imFixedBlur = imreg.GetFinalMovingImage("LINEAR_INTERPOLATION", -1024)
    imMovingReg = imreg.ResampleImage("LINEAR_INTERPOLATION", imMoving, tfm, -1024)

    itk.imwrite( imMovingReg, directory2 + pic_folder[imNum], compression=True )

    print(tfm)

    imdataTmp = itk.GetArrayFromImage(imMovingReg)

    imdatamax = np.maximum(imdatamax,imdataTmp)
    imdatamin = np.minimum(imdatamin,imdataTmp)
    imdataTmp[np.where(imdataTmp==imdatamax)] = 0
    imdataTmp[np.where(imdataTmp==imdatamin)] = 0
    imdatamax2 = np.maximum(imdatamax2,imdataTmp)
    imdatamin2 = np.minimum(imdatamin2,imdataTmp)
    imdataTmp[np.where(imdataTmp==imdatamax)] = 0
    imdataTmp[np.where(imdataTmp==imdatamin)] = 0
    imdatamax3 = np.maximum(imdatamax3,imdataTmp)
    imdatamin3 = np.minimum(imdatamin3,imdataTmp)

    #out = itk.GetImageFromArray(imdatamax)
    #out.CopyInformation(im0)
    #itk.imwrite(out, (directory3 + 'max_' + str(imNum) + '.nrrd'))

    #out = itk.GetImageFromArray(imdatamax3)
    #out.CopyInformation(im0)
    #itk.imwrite(out, (directory3 + 'max3_' + str(imNum) + '.nrrd'))

    percent = (imNum + 1) / num_images * 100
    print(str(round(percent)) + '% : ' + pic_folder[imNum])

print('Done')

# In[ ]:

out = itk.GetImageFromArray(imdatamax3)
out.CopyInformation(im0)
itk.imwrite(out, (directory3 + 'max3.nrrd'), compression=True)

out = itk.GetImageFromArray(imdatamin3)
out.CopyInformation(im0)
itk.imwrite(out, (directory3 + 'min3.nrrd'), compression=True)

out = itk.GetImageFromArray(imdatamax3 - imdatamin3)
out.CopyInformation(im0)
itk.imwrite(out, (directory3 + 'diff3.nrrd'), compression=True)

print('Done3')

# In[ ]:

out = itk.GetImageFromArray(imdatamax)
out.CopyInformation(im0)
itk.imwrite(out, (directory3 + 'max.nrrd'), compression=True)

out = itk.GetImageFromArray(imdatamin)
out.CopyInformation(im0)
itk.imwrite(out, (directory3 + 'min.nrrd'), compression=True)

out = itk.GetImageFromArray(imdatamax - imdatamin)
out.CopyInformation(im0)
itk.imwrite(out, (directory3 + 'diff.nrrd'), compression=True)

print('Done')

# In[ ]:

#out = itk.GetImageFromArray(imdatamax)
#view(out)

# In[ ]:

#!/usr/bin/env python
# coding: utf-8

# This notebook is intended to demonstrate how select registration, segmentation, and image mathematical methods of ITKTubeTK can be combined to perform multi-channel brain extraction (aka. skull stripping for patient data containing multiple MRI sequences).
#
# There are many other (probably more effective) brain extraction methods available as open-source software such as BET and BET2 in the FSL package (albeit such methods are only for single channel data).   If you need to perform brain extraction for a large collection of scans that do not contain major pathologies, please use one of those packages.   This notebook is meant to show off the capabilities of specific ITKTubeTK methods, not to demonstration how to "solve" brain extraction.

# In[1]:

#import itk
#from itk import TubeTK as ttk

#from itkwidgets import view

#import numpy as np

# In[2]:

ImageType = itk.Image[itk.F, 3]

InputBaseName = studyname + "-MinMax/max3"

filename = InputBaseName + ".nrrd"
im1iso = itk.imread(filename, itk.F)

# In[3]:

N = 8
readerList = ["003", "010", "026", "034", "045", "056", "063", "071"]

imBase = []
imBaseB = []
for i in range(0,N):
    name = "../Data/Normal"+readerList[i]+"-FLASH.mha"
    nameB = "../Data/Normal"+readerList[i]+"-FLASH-Brain.mha"
    imBaseTmp = itk.imread(name, itk.F)
    imBaseBTmp = itk.imread(nameB, itk.F)
    imBase.append(imBaseTmp)
    imBaseB.append(imBaseBTmp)

# In[4]:

#view(im1iso)

# In[5]:

#view(imBase[0])

# In[6]:

imMath = ttk.ImageMath.New(Input=im1iso)
#imMath.Threshold(-4000,-500,1,0)
#headMask = imMath.GetOutput()
imMath.SetInput(im1iso)
#imMath.IntensityWindow(0,1000,1000,0)
#imMath.ReplaceValuesOutsideMaskRange(headMask,-0.5,0.5,-500)
imMath.Blur(1)
imMath.NormalizeMeanStdDev()
imMath.IntensityWindow(-5,5,-500,500)
im1isoBlur = imMath.GetOutput()
#view(im1isoBlur)

# In[7]:

RegisterImagesType = ttk.RegisterImages[ImageType]
regB = []
regBB = []
for i in range(0,N):
    imMath.SetInput(imBase[i])
    imMath.Blur(1)
    imMath.NormalizeMeanStdDev()
    imMath.IntensityWindow(-5,5,-500,500)
    imBaseBlur = imMath.GetOutput()

    #regBTo1 = RegisterImagesType.New(FixedImage=im1isoBlur, MovingImage=imBaseBlur)
    regBTo1 = RegisterImagesType.New(FixedImage=imBaseBlur, MovingImage=im1isoBlur)
    regBTo1.SetReportProgress(True)

    regBTo1.SetRigidMaxIterations(3000)
    regBTo1.SetAffineMaxIterations(3000)

    regBTo1.SetExpectedRotationMagnitude(0.2)
    regBTo1.SetExpectedScaleMagnitude(0.25)
    regBTo1.SetExpectedSkewMagnitude(0.01)
    regBTo1.SetExpectedOffsetMagnitude(40)

    regBTo1.SetRigidSamplingRatio(0.1)
    regBTo1.SetAffineSamplingRatio(0.1)

    regBTo1.SetSampleFromOverlap(True)

    regBTo1.SetInitialMethodEnum("INIT_WITH_IMAGE_CENTERS")
    regBTo1.SetRegistration("PIPELINE_AFFINE")
    regBTo1.SetMetric("MATTES_MI_METRIC")

    regBTo1.Update()

    tfm = regBTo1.GetCurrentMatrixTransform()
    tfmInv = tfm.GetInverseTransform()
    print(tfm)

    resm = ttk.ResampleImage.New(Input=imBase[i])
    resm.SetMatchImage(im1iso)
    resm.SetTransform(tfmInv)
    resm.SetLoadTransform(True)
    resm.Update()
    img = resm.GetOutput()
    regB.append( img )

    resm = ttk.ResampleImage.New(Input=imBaseB[i])
    resm.SetMatchImage(im1iso)
    resm.SetTransform(tfmInv)
    resm.SetLoadTransform(True)
    resm.Update()
    img = resm.GetOutput()
    regBB.append( img )

# In[8]:

imMath.SetInput(regB[1])
imMath.AddImages(im1iso,20,1)
img = imMath.GetOutput()
#view( img )

# In[9]:

regBBT = []
for i in range(0,N):
    imMath.SetInput(regBB[i])
    imMath.Threshold(0,1,0,1)
    img = imMath.GetOutput()
    if i==0:
        imMath.SetInput( img )
        imMath.AddImages( img, 1.0/N, 0 )
        sumBBT = imMath.GetOutput()
    else:
        imMath.SetInput( sumBBT )
        imMath.AddImages( img, 1, 1.0/N )
        sumBBT = imMath.GetOutput()

#view(sumBBT)

# In[10]:

imMath.SetInput(sumBBT)
imMath.Threshold(0.85,1.1,1,0)
imMath.Dilate(5,1,0)
imMath.Erode(25,1,0)
brainInside = imMath.GetOutput()

imMath.SetInput( sumBBT )
imMath.Threshold(0,0,1,0)
imMath.Erode(1,1,0)
brainOutsideAll = imMath.GetOutput()
imMath.Erode(20,1,0)
imMath.AddImages(brainOutsideAll, -1, 1)
brainOutside = imMath.GetOutput()

imMath.AddImages(brainInside,1,2)
brainCombinedMask = imMath.GetOutputUChar()
brainCombinedMaskF = imMath.GetOutput()

# In[11]:

imMath.SetInput(brainCombinedMaskF)
imMath.AddImages(im1iso, 100, 1)
brainCombinedMaskView = imMath.GetOutput()
#view(brainCombinedMaskView)

# In[12]:

LabelMapType = itk.Image[itk.UC,3]

segmenter = ttk.SegmentConnectedComponentsUsingParzenPDFs[ImageType,LabelMapType].New()
segmenter.SetFeatureImage( im1iso )
segmenter.SetInputLabelMap( brainCombinedMask )
segmenter.SetObjectId( 2 )
segmenter.AddObjectId( 1 )
segmenter.SetVoidId( 0 )
segmenter.SetErodeDilateRadius( 10 )
segmenter.SetHoleFillIterations( 40 )
segmenter.Update()
segmenter.ClassifyImages()
brainCombinedMaskClassified = segmenter.GetOutputLabelMap()

# In[13]:

#view(brainCombinedMaskClassified)

# In[14]:

cast = itk.CastImageFilter[LabelMapType, ImageType].New()
cast.SetInput(brainCombinedMaskClassified)
cast.Update()
brainMaskF = cast.GetOutput()

brainMath = ttk.ImageMath[ImageType,ImageType].New(Input = brainMaskF)
brainMath.Threshold(2,2,1,0)
brainMath.Erode(1,1,0)
brainMaskD = brainMath.GetOutput()
brainMath.SetInput( im1iso )
brainMath.ReplaceValuesOutsideMaskRange( brainMaskD, 1, 1, 0)
brain = brainMath.GetOutput()

# In[15]:

#view(brain)

# In[16]:

writer = itk.ImageFileWriter[ImageType].New(Input = brain)
filename = InputBaseName + "-Brain.nrrd"
writer.SetFileName(filename)
writer.SetUseCompression(True)
writer.Update()

# In[ ]:

#!/usr/bin/env python
# coding: utf-8

# This notebook is intended to demonstrate how select registration, segmentation, and image mathematical methods of ITKTubeTK can be combined to perform multi-channel brain extraction (aka. skull stripping for patient data containing multiple MRI sequences).
#
# There are many other (probably more effective) brain extraction methods available as open-source software such as BET and BET2 in the FSL package (albeit such methods are only for single channel data).   If you need to perform brain extraction for a large collection of scans that do not contain major pathologies, please use one of those packages.   This notebook is meant to show off the capabilities of specific ITKTubeTK methods, not to demonstration how to "solve" brain extraction.

# In[1]:

#import itk
#from itk import TubeTK as ttk

#from itkwidgets import view

#import numpy as np

# In[2]:

InputBaseDir = studyname

CTPMaxFilename = InputBaseDir + "-MinMax/max.nrrd"
CTPMinFilename = InputBaseDir + "-MinMax/min.nrrd"
CTPBrainFilename = InputBaseDir + "-MinMax/max3-Brain.nrrd"

imMax = itk.imread(CTPMaxFilename, itk.F)
imMin = itk.imread(CTPMinFilename, itk.F)
imBrain = itk.imread(CTPBrainFilename, itk.F)

# In[3]:

#view(imBrain)

# In[4]:

ImageType = itk.Image[itk.F, 3]

imMath = ttk.ImageMath.New(Input=imBrain)
imMath.Threshold( 0.00001, 2000, 1, 0)
imMath.Erode(10,1,0)
imBrainMaskErode = imMath.GetOutput()

imMath.SetInput(imMax)
imMath.AddImages(imMin,1,-1)
imDiff = imMath.GetOutput()
imMath.ReplaceValuesOutsideMaskRange(imBrain, 0.0001, 2000, 0)
imDiffBrain = imMath.GetOutput()
imMath.ReplaceValuesOutsideMaskRange(imBrainMaskErode, 0.5, 1.5, 0)
imDiffBrainErode = imMath.GetOutput()

# In[5]:

tmpA = itk.GetArrayViewFromImage(imDiffBrain)
tmpAE = itk.GetArrayViewFromImage(imDiffBrainErode)
zMax = tmpA.shape[0]
clip = 0
while((np.amax(tmpA[clip:clip+1,:,:])>1000) | (np.amax(tmpA[clip:clip+1,:,:])==0)):
    clip += 1
if(clip>0):
    tmpA[0:clip,:,:]=0
    tmpAE[0:clip,:,:]=0
clip = 1
while((np.amax(tmpA[zMax-clip:zMax-clip+1,:,:])>1000) | (np.amax(tmpA[zMax-clip:zMax-clip+1,:,:])==0)):
    clip += 1
print(clip, np.amax(tmpA[zMax-clip:zMax-clip+1,:,:]))
clip = clip - 1
if(clip>0):
    tmpA[zMax-clip:zMax,:,:]=0  #Happens to imDiffBrain since this array is a view of an itk image
    tmpAE[zMax-clip:zMax,:,:]=0  #Happens to imDiffBrain since this array is a view of an itk image

# In[6]:

#view(imDiffBrain)

# In[7]:

imMath = ttk.ImageMath[ImageType,ImageType].New()
imMath.SetInput(imDiffBrainErode)
imMath.Blur(1.5)
imBlur = imMath.GetOutput()
imBlurArray = itk.GetArrayViewFromImage(imBlur)

numSeeds = 15
seedCoverage = 20
seedCoord = np.zeros([numSeeds,3])
for i in range(numSeeds):
    seedCoord[i] = np.unravel_index(np.argmax(imBlurArray, axis=None), imBlurArray.shape)
    indx = [int(seedCoord[i][0]),int(seedCoord[i][1]),int(seedCoord[i][2])]
    minX = max(indx[0]-seedCoverage,0)
    maxX = max(indx[0]+seedCoverage,imBlurArray.shape[0])
    minY = max(indx[1]-seedCoverage,0)
    maxY = max(indx[1]+seedCoverage,imBlurArray.shape[1])
    minZ = max(indx[2]-seedCoverage,0)
    maxZ = max(indx[2]+seedCoverage,imBlurArray.shape[2])
    imBlurArray[minX:maxX,minY:maxY,minZ:maxZ]=0
    indx.reverse()
    seedCoord[:][i] = imDiffBrain.TransformIndexToPhysicalPoint(indx)
print(seedCoord)

# In[8]:

# Manually extract a few vessels to form an image-specific training set
vSeg = ttk.SegmentTubes.New(Input=imDiffBrain)
vSeg.SetVerbose(True)
vSeg.SetMinRoundness(0.4)
vSeg.SetMinCurvature(0.002)
vSeg.SetRadiusInObjectSpace( 1 )
for i in range(numSeeds):
    print("**** Processing seed " + str(i) + " : " + str(seedCoord[i]))
    vSeg.ExtractTubeInObjectSpace( seedCoord[i], i )

tubeMaskImage = vSeg.GetTubeMaskImage()

# In[9]:

imMath.SetInput(tubeMaskImage)
imMath.AddImages(imDiffBrain, 200, 1)
blendIm = imMath.GetOutput()
#view(blendIm)

# In[10]:

LabelMapType = itk.Image[itk.UC,3]

trMask = ttk.ComputeTrainingMask[ImageType,LabelMapType].New()
trMask.SetInput( tubeMaskImage )
trMask.SetGap( 4 )
trMask.SetObjectWidth( 1 )
trMask.SetNotObjectWidth( 1 )
trMask.Update()
fgMask = trMask.GetOutput()

# In[11]:

#view(fgMask)

# In[12]:

enhancer = ttk.EnhanceTubesUsingDiscriminantAnalysis[ImageType,LabelMapType].New()
enhancer.AddInput( imDiff )
enhancer.SetLabelMap( fgMask )
enhancer.SetRidgeId( 255 )
enhancer.SetBackgroundId( 128 )
enhancer.SetUnknownId( 0 )
enhancer.SetTrainClassifier(True)
enhancer.SetUseIntensityOnly(True)
enhancer.SetScales([0.43,1.29,3.01])
enhancer.Update()
enhancer.ClassifyImages()

# In[13]:

im1vess = itk.SubtractImageFilter( Input1=enhancer.GetClassProbabilityImage(0), Input2=enhancer.GetClassProbabilityImage(1))

imMath.SetInput(imDiffBrain)
imMath.Threshold(0.0001,2000,1,0)
imMath.Erode(2,1,0)
imBrainE = imMath.GetOutput()

imMath.SetInput(im1vess)
imMath.ReplaceValuesOutsideMaskRange(imBrainE, 1, 1, -0.001)
im1vessBrain = imMath.GetOutput()
#view(enhancer.GetClassProbabilityImage(0))
#view(im1vessBrain)

# In[14]:

itk.imwrite( im1vess, InputBaseDir + "-MinMax/diff-VesselEnhanced.nrrd", compression=True)

itk.imwrite( im1vessBrain, InputBaseDir + "-MinMax/Brain-VesselEnhanced.nrrd", compression=True)

# In[ ]:

#!/usr/bin/env python
# coding: utf-8

# This notebook is intended to demonstrate how vessel segmentation methods of ITKTubeTK can be applied to multi-channel MRI (MRA + T1, T2, etc).

# In[1]:

#import itk
#from itk import TubeTK as ttk

#from itkwidgets import view

#import numpy as np

# In[2]:

ImageType = itk.Image[itk.F, 3]

imDir = studyname + "-MinMax/"

im1iso = itk.imread(imDir + "diff3.nrrd")
im1BrainVess = itk.imread(imDir + "Brain-VesselEnhanced.nrrd")

# In[3]:

imMath = ttk.ImageMath.New(im1BrainVess)
imMath.MedianFilter(1)
imMath.Threshold(0.000001, 1, 1, 0)
im1VessMask = imMath.GetOutputShort()

ccSeg = ttk.SegmentConnectedComponents.New(im1VessMask)
ccSeg.SetMinimumVolume(10)
ccSeg.Update()
im1VessMaskCC = ccSeg.GetOutput()

# In[4]:

#view(im1VessMaskCC)

# In[5]:

imMathSS = ttk.ImageMath.New(im1VessMaskCC)
imMathSS.Threshold(0,0,1,0)
im1VessMaskInv = imMathSS.GetOutputFloat()

distFilter = itk.DanielssonDistanceMapImageFilter.New(im1VessMaskInv)
distFilter.Update()
dist = distFilter.GetOutput()

imMath.SetInput(dist)
imMath.Blur(0.4)
tmp = imMath.GetOutput()
imMath.ReplaceValuesOutsideMaskRange(tmp, 0.1, 10, 0)
im1SeedRadius = imMath.GetOutput()

itk.imwrite(im1SeedRadius, imDir+"VesselsSeedRadius.mha")

# In[6]:

#view(im1SeedRadius)

# In[7]:

imMath.SetInput(im1iso)
imMath.ReplaceValuesOutsideMaskRange(im1BrainVess, 0, 1000, 0)
imMath.Blur(0.4)
imMath.IntensityWindow(0.5,1000,0,1000)
im1Input = imMath.GetOutput()

itk.imwrite(im1iso, imDir+"VesselsInput.mha")

#view(im1Input)

# In[8]:

numSeeds = 40

vSeg = ttk.SegmentTubes.New(Input=im1Input)
#vSeg.SetVerbose(True)
vSeg.SetMinCurvature(0)#.0001)
vSeg.SetMinRoundness(0.02)
vSeg.SetMinRidgeness(0.5)
vSeg.SetMinLevelness(0.0)
vSeg.SetRadiusInObjectSpace( 0.8 )
vSeg.SetBorderInIndexSpace(3)
vSeg.SetSeedMask( im1SeedRadius )
vSeg.SetSeedRadiusMask( im1SeedRadius )
vSeg.SetOptimizeRadius(False)
vSeg.SetUseSeedMaskAsProbabilities(True)
vSeg.SetSeedExtractionMinimumProbability(0.4)
#vSeg.SetSeedMaskMaximumNumberOfPoints( numSeeds )
vSeg.ProcessSeeds()

# In[9]:

tubeMaskImage = vSeg.GetTubeMaskImage()
#view(tubeMaskImage)

# In[10]:

SOWriter = itk.SpatialObjectWriter[3].New()
SOWriter.SetInput(vSeg.GetTubeGroup())
SOWriter.SetBinaryPoints(True)
SOWriter.SetFileName( imDir+"Vessels.tre" )
SOWriter.Update()

# In[ ]:


#!/usr/bin/env python
# -*- coding: utf-8 -*

"""
tools module
"""

__author__ = 'Dr. Janus Juul Eriksen, University of Bristol, UK'
__maintainer__ = 'Dr. Janus Juul Eriksen'
__email__ = 'janus.eriksen@bristol.ac.uk'
__status__ = 'Development'

import sys
import os
import copy
import numpy as np
try:
    import opt_einsum as oe
    OE_AVAILABLE = True
except ImportError:
    OE_AVAILABLE = False
from subprocess import Popen, PIPE
from pyscf import gto, scf, dft, symm, lib
from pyscf import tools as pyscf_tools
from typing import Tuple, List, Dict, Union

MAX_CYCLE = 100
NATORB_THRES = 1.e-12

class Logger(object):
        """
        this class pipes all write statements to both stdout and output_file
        """
        def __init__(self, output_file, both=True) -> None:
            """
            init Logger
            """
            self.terminal = sys.stdout
            self.log = open(output_file, 'a')
            self.both = both

        def write(self, message) -> None:
            """
            define write
            """
            self.log.write(message)
            if self.both:
                self.terminal.write(message)

        def flush(self) -> None:
            """
            define flush
            """
            pass

def git_version() -> str:
        """
        this function returns the git revision as a string
        """
        def _minimal_ext_cmd(cmd):
            env = {}
            for k in ['SYSTEMROOT', 'PATH', 'HOME']:
                v = os.environ.get(k)
                if v is not None:
                    env[k] = v
            # LANGUAGE is used on win32
            env['LANGUAGE'] = 'C'
            env['LANG'] = 'C'
            env['LC_ALL'] = 'C'
            out = Popen(cmd, stdout=PIPE, env=env, \
                        cwd=os.path.dirname(__file__)).communicate()[0]
            return out

        try:
            out = _minimal_ext_cmd(['git', 'rev-parse', 'HEAD'])
            GIT_REVISION = out.strip().decode('ascii')
        except OSError:
            GIT_REVISION = "Unknown"

        return GIT_REVISION

def dim(mo_occ: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """
        determine molecular dimensions
        """
        return np.where(np.abs(mo_occ[0]) > 0.)[0], np.where(np.abs(mo_occ[1]) > 0.)[0]

def mf_info(mf: Union[scf.hf.SCF, dft.rks.KohnShamDFT]) -> Tuple[Tuple[np.ndarray, np.ndarray], \
                                                                 Tuple[np.ndarray, np.ndarray]]:
        """
        retrieve mf information (mo coefficients & occupations)
        """
        # mo occupations
        if np.asarray(mf.mo_occ).ndim == 1:
            mo_occ = (np.ones(np.count_nonzero(0. < mf.mo_occ)), np.ones(np.count_nonzero(1. < mf.mo_occ)))
        else:
            mo_occ = (mf.mo_occ[0][np.nonzero(mf.mo_occ[0])], mf.mo_occ[1][np.nonzero(mf.mo_occ[1])])
        # dimensions
        alpha, beta = dim(mo_occ)
        # mo coefficients
        if np.asarray(mf.mo_coeff).ndim == 2:
            mo_coeff = (mf.mo_coeff[:, alpha], mf.mo_coeff[:, beta])
        else:
            mo_coeff = (mf.mo_coeff[0][:, alpha], mf.mo_coeff[1][:, beta])

        return mo_coeff, mo_occ

def orbsym(mol, mo_coeff):
        """
        this functions returns orbital symmetries
        """
        if isinstance(mo_coeff, np.ndarray):
            if mo_coeff.ndim == 2:
                try:
                    orbsymm = symm.label_orb_symm(mol, mol.irrep_name, mol.symm_orb, mo_coeff)
                except:
                    orbsymm = np.array(['A'] * mo_coeff.shape[1])
            else:
                try:
                    orbsymm = np.array([symm.label_orb_symm(mol, mol.irrep_name, mol.symm_orb, c) for c in mo_coeff])
                except:
                    orbsymm = np.array([['A'] * c.shape[1] for c in mo_coeff])
        else:
            try:
                orbsymm = np.array([symm.label_orb_symm(mol, mol.irrep_name, mol.symm_orb, c) for c in mo_coeff])
            except:
                orbsymm = np.array([['A'] * c.shape[1] for c in mo_coeff])

        return orbsymm

def make_rdm1(mo: np.ndarray, occup: np.ndarray) -> np.ndarray:
        """
        this function returns an 1-RDM (in ao basis) corresponding to given mo(s)
        """
        return contract('ip,jp->ij', occup * mo, mo)

def make_natorb(mol: gto.Mole, mo_coeff: np.ndarray, \
                rdm1: np.ndarray, thres: float = NATORB_THRES) -> Tuple[Tuple[np.ndarray, np.ndarray], \
                                                                        Tuple[np.ndarray, np.ndarray]]:
        """
        this function returns no coefficients and occupations corresponding
        to given mo coefficients and rdm1
        """
        # reshape mo_coeff and rdm1
        if mo_coeff.ndim == 2:
            c = np.asarray((mo_coeff,) * 2)
        else:
            c = mo_coeff
        if rdm1.ndim == 2:
            d = np.array([rdm1, rdm1]) * .5
        else:
            d = rdm1
        # overlap matrix
        s = mol.intor_symmetric('int1e_ovlp')
        # ao to mo transformation of dm
        rdm1_mo = contract('xpi,pq,xqr,rs,xsj->xij', c, s, d, s, c)
        # diagonalize rdm1_mo
        occ_no, u = np.linalg.eigh(rdm1_mo)
        # transform to no basis
        mo_no = contract('xip,xpj->xij', c, u)
        # retain only significant nos
        return (mo_no[0][:, np.where(np.abs(occ_no[0]) >= thres)[0]], mo_no[1][:, np.where(np.abs(occ_no[1]) >= thres)[0]]), \
               (occ_no[0][np.where(np.abs(occ_no[0]) >= thres)], occ_no[1][np.where(np.abs(occ_no[1]) >= thres)])

def write_rdm1(mol: gto.Mole, part: str, \
               mo_coeff: np.ndarray, mo_occ: np.ndarray, fmt: str, \
               weights: List[np.ndarray], \
               suffix: str = '') -> None:
        """
        this function writes a 1-RDM as a numpy or cube (default) file
        """
        # assertion
        assert part == 'atoms', '`write_rdm1` function only implemented for `atoms` partitioning'
        assert fmt in ['cube', 'numpy'], 'fmt arg to `write_rdm1` must be `cube` or `numpy`'
        # molecular dimensions
        alpha, beta = dim(mo_occ)
        # compute total 1-RDM (AO basis)
        rdm1_tot = np.array([make_rdm1(mo_coeff[0], mo_occ[0]), make_rdm1(mo_coeff[1], mo_occ[1])])
        # loop over atoms
        for a in range(mol.natm):
            # atom-specific rdm1
            rdm1_atom = np.zeros_like(rdm1_tot)
            # loop over spins
            for i, spin_mo in enumerate((alpha, beta)):
                # loop over spin-orbitals
                for m, j in enumerate(spin_mo):
                    # get orbital(s)
                    orb = mo_coeff[i][:, j].reshape(mo_coeff[i].shape[0], -1)
                    # orbital-specific rdm1
                    rdm1_orb = make_rdm1(orb, mo_occ[i][j])
                    # weighted contribution to rdm1_atom
                    rdm1_atom[i] += rdm1_orb * weights[i][m][a]
            if fmt == 'cube':
                # write rdm1_atom as cube file
                pyscf_tools.cubegen.density(mol, f'atom_{mol.atom_symbol(a).upper():s}{a:d}_rdm1{suffix:}.cube', \
                                            np.sum(rdm1_atom, axis=0))
            else:
                # write rdm1_atom as numpy file
                np.save(f'atom_{mol.atom_symbol(a).upper():s}{a:d}_rdm1{suffix:}.npy', np.sum(rdm1_atom, axis=0))

def res_add(res_a, res_b):
        """
        this function adds two result dictionaries
        """
        return {key: res_a[key] + res_b[key] for key in res_a.keys()}

def res_sub(res_a, res_b):
        """
        this function subtracts two result dictionaries
        """
        return {key: res_a[key] - res_b[key] for key in res_a.keys()}

def contract(eqn, *tensors):
        """
        interface to optimized einsum operation
        """
        if OE_AVAILABLE:
            return oe.contract(eqn, *tensors)
        else:
            return np.einsum(eqn, *tensors, optimize=True)


import numpy as np
from cppdefs import *
from pom.modules_old import RLEN, bfm_lwp, LOGUNIT, seconds_per_day

class TimeInfo:

    def __init__(self, date0, time0, timeEnd, step0, timestep, stepnow, stepEnd):
        self.date0 = date0
        self.time0 = time0
        self.timeEnd = timeEnd
        self.step0 = step0
        self.timestep = timestep
        self.stepnow = stepnow
        self.stepEnd = stepEnd

# PUBLIC DATA MEMBERS:
timestr = ' ' * 19
start = '2000-01-01 00:00:00'
stop = ' ' * 19

# PRIVATE DATA MEMBERS
jul0 = -1
secs0 = -1

# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
#
# # ROUTINE: INITIALIZE THE TIME SYSTEM
#
# DESCRIPTION:  The subroutine {\tt init\_time()} initialises the time module by reading
#               a namelist and take actions according to the specifications.
#               On exit from this subroutine the two variables MinN and MaxN have well
#               defined values and can be used in the time loop.
#
# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=

def initialize_time_system(MinN,MaxN,timestep,timefmt,simdays,start,stop):

    """
    Description: Initializes the time module by reading a namelist.
                 Takes actions accrding to the specifications.

    :param MinN: minimum number of time steps
    :param MaxN: maximum number of time steps
    :param timestep: size of time step
    :param timefmt: time format
    :param simdays: number of days for simulation
    :param start: start date
    :param stop: end date
    :return:
    """

    jul1 = -1
    secs1 = -1

    # -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
    #   READ TIME SPECIFIC THINGS FROM THE NAMELIST
    # -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=

    LEVEL1(); print('initialize_time_system')

    # CALCULATE MaxN -> MinN IS 1 IF NOT CHANGED BY HOTSTART
    MinN = 1
    LEVEL2(); print('Time step:      ', timestep, ' seconds')
    LEVEL2(); print('Time format:    ', timefmt)

    if timefmt == 1:
        HasRealTime = False
        LEVEL2(); print('# of timesteps: ', MaxN)
        # start = '2000-01-01 00:00:00'
        LEVEL2(); print('Fake start:     ', start)

    elif timefmt == 2:
        HasRealTime = True
        LEVEL2(); print('Start:          ', start)
        LEVEL2(); print('Stop:           ', stop)
        jul1, secs1 = read_time_string(start)
        jul2, secs2 = read_time_string(stop)

        nsecs = time_diff(jul2, secs2, jul1, secs1)
        MaxN = round(nsecs / timestep)

        ndays = jul2 - jul1
        if nsecs < 86400 and jul1 != jul2:
            ndays = ndays - 1

        nsecs = nsecs - 86400 * ndays
        string = '  ==> ' + str(ndays) + ' day(s) and ' + str(nsecs) + ' seconds ==> ' + str(MaxN) + ' time steps'
        STDERR(string)

    elif timefmt == 3:
        HasRealTime = True
        LEVEL2(); print('Start:          ', start)
        LEVEL2(); print('# of timesteps: ', MaxN)

        jul1, secs1 = read_time_string(start)

        nsecs = np.rint(MaxN * timestep) + secs1
        ndays = nsecs / 86400
        jul2 = jul1 + ndays
        secs2 = nsecs % 86400

        write_time_string(jul2, secs2)
        LEVEL2(); print('Stop:           ', stop)

    elif timefmt == 4:
        HasRealTime = False
        nsecs = simdays * 86400
        MaxN = np.rint(nsecs / timestep)
        LEVEL2(); print('# of timesteps: ', MaxN)
        # start = '2000-01-01 00:00:00'
        LEVEL2(); print('Fake start:     ', start)

    else:
        STDERR('Fatal error: A non valid input format has been chosen')
        return

    jul0 = jul1
    secs0 = secs1

    julianday = jul0
    secondsofday = secs0

    simtime = timestep * (MaxN - MinN + 1)

    # SET BFM TIME
    jday = float(jul0)
    yy, mm, dd, hh, nn = calendar_date(jday)

    bfmtime = TimeInfo(str(yy) + '-' + str(mm) + '-' + str(dd) + ' ' + str(hh) + ':' + str(nn),
                       jday,
                       jday + (float(MaxN) * timestep) / seconds_per_day,
                       MinN - 1,
                       timestep,
                       MinN - 1,
                       MaxN)

    LEVEL2(); print('bfmtime : ', bfmtime)

    return bfmtime

# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
#
# # ROUTINE: CONVERT TRUE JULIAN DAY TO CALENDAR DATE
#
# DESCRIPTION:  Converts a Julian day to a calendar date --- year, month and day.
#               Based on a similar routine in \emph{Numerical Recipes}.
#
# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=

def calendar_date(julian):

    """
    Description: Converts a Julian day to a calender date --- year, month, and day

    :param julian: Julian day
    :return: calender date
    """

    # LOCAL VARIABLES
    IGREG = 2299161

    # -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=

    jday = np.floor(julian)
    if jday >= IGREG:
        x = ((jday - 1867216) - 0.25) / 36524.25
        ja = jday + 1 + int(x) - int(0.25 * x)
    else:
        ja = jday

    jb = ja + 1524
    jc = int(6680 + ((jb - 2439870) - 122.1) / 365.25)
    jd = int(365 * jc + (0.25 * jc))
    je = int((jb - jd) / 30.6001)

    dd = jb - jd - int(30.6001 * je)
    mm = je - 1

    if mm > 12:
        mm = mm - 12

    yyyy = jc - 4715

    if mm > 2:
        yyyy = yyyy - 1

    if yyyy <= 0:
        yyyy = yyyy - 1

    res = julian - float(jday)
    hh = np.floor(res * 24)
    nn = np.floor(((res * 24) - float(hh)) * 60)

    return yyyy, mm, dd, hh, nn

# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
#
# # ROUTINE: CONVERT CALENDAR DATE TO JULIAN DAY
#
# DESCRIPTION:  Converts a calendar date to a Julian day.
#               Based on a similar routine in \emph{Numerical Recipes}.
#
# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=

def julian_day(yyyy,mm,dd,hh,nn):

    """
    Description: Converts a calender date to a Julian day.

    :param yyyy: year
    :param mm: month
    :param dd: day
    :param hh: hour
    :param nn: minute
    :return: Julian day
    """

    # LOCAL VARIABLES
    IGREG = 15 + 31 * (10 + 12 * 1582)

    # -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=

    jy = yyyy
    if jy < 0:
        jy = jy + 1

    if mm > 2:
        jm = mm + 1
    else:
        jy = jy - 1
        jm = mm + 13

    jday = int(np.floor(365.25 * jy) + np.floor(30.6001 * jm) + dd + 1720995)
    if dd + 31 * (mm + 12 * yyyy) >= IGREG:
        ja = int(0.01 * jy)
        jday = jday + 2 - ja + int(0.25 * ja)

    jh = hh
    jn = nn

    if jn >= 60:
        jh = jh + 1
        jn = jn - 60

    if jh >= 24:
        jday = jday + 1
        jh = jh - 24

    julian = float(jday) + float(jh) / 24 + float(jn) / (24 * 60)

    return julian

# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
#
# # ROUTINE: KEEP TRACK OF TIME (JULIAN DAYS AND SECONDS)
#
# DESCRIPTION:  Based on a starting time this routine calculates the actual time
#               in a model integration using the number of time steps, {\tt n},
#               and the size of the time step, {\tt timestep}. More public variables
#               can be updated here if necessary.
#
# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=

def update_time(n,timestep,secs0):

    """
    Description: Based on a starting time this routine calculates the actual time in a model integration
                 using the number of time steps and the size of the time step.

    :param n: number of time steps
    :param timestep: size of time step
    :param secs0:
    :return: actual time in a model integration
    """

    nsecs = np.rint(n * timestep) + secs0
    fsecs = n * timestep + secs0
    julianday = jul0 + nsecs / 86400
    secondsofday = nsecs % 86400

    return fsecs, julianday, secondsofday

# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
#
# # ROUTINE: CONVERT A TIME STRING TO JULIAN DAY AND SECONDS
#
# DESCRIPTION:  Converts a time string to the true Julian day and seconds of that day.
#               The format of the time string must be: {\tt yyyy-mm-dd hh:hh:ss }.
#
# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=

def read_time_string(timestr):

    """
    Description: Converts a time string to the true Julian day and the seconds of that day.

    :param timestr: time string
    :return: Julian day and seconds of that day
    """

    yy = timestr[0:4]
    mm = timestr[5:7]
    dd = timestr[8:10]
    hh = timestr[11:13]
    mins = timestr[14:16]
    ss = timestr[17:19]

    julian, jday = julian_day(yy,mm,dd,0,0)
    jul = int(jday)
    secs = 3600*hh + 60*mins + ss

    return jul, secs

# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
#
# # ROUTINE: CONVERT JULIAN DAY AND SECONDS TO A TIME STRING
#
# DESCRIPTION:  Formats Julian day and seconds of that day to a nice looking
#               character string.
#
# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=

def write_time_string(jul,secs):

    """
    Description: Formats Julian day and seconds of that day into a character string.

    :param jul: Julian day
    :param secs: seconds of Julian day
    :return: time string
    """

    jday = float(jul)
    yy, mm, dd, hh, nn = calendar_date(jday)

    hh = secs / 3600
    mins = (secs - hh * 3600) / 60
    ss = secs - 3600 * hh - 60 * mins

    timestr = str(yy) + '-' + str(mm) + '-' + str(dd) + ' ' + str(hh) + ':' + str(mins) + ':' + str(ss)

    return timestr

# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
#
# # ROUTINE: RETURN THE TIME DIFFERENCE IN SECONDS
#
# DESCRIPTION:  This functions returns the time difference between two
#               dates in seconds. The dates are given as Julian day and seconds
#               of that day.
#
# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=

def time_diff(jul2,secs2,jul1,secs1):

    """
    Description: Returns the time difference between two dates in seconds.

    :param jul2: Julian day 2
    :param secs2: seconds of julian day 2
    :param jul1: Julian day 1
    :param secs1: seconds of julian day 1
    :return: difference between Julian day 2 and Julian day 1
    """

    time_diff = 86400*(jul2-jul1) + (secs2-secs1)

    return time_diff

# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
#
# # ROUTINE: CONVERT A CALENDAR DATE TO TRUE JULIAN DAY
#
# DESCRIPTION:  Converts a Julian day to the day number in the current year
#
# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=

def dayofyear(julian,yy):

    """
    Description: Converts a Julian day to the day number in the curent year.

    :param julian: Julian day
    :param yy: current year
    :return: day number of the current year
    """

    jday = float(julian)
    calendar_date(jday)
    julian0, jday = julian_day(yy,1,1,0,0)

    ddyear = julian - int(julian0) + 1

    return ddyear

# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
#
# # ROUTINE:
#
# DESCRIPTION:
#
# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=

def eomdays(Year,Month):

    """
    Description: Provides system with the number of days in a given month.
                 Accounts for leap year using Year input.

    :param Year: current year
    :param Month: current month
    :return: numbr of days in current month
    """

    eomdays = 0
    if Month < 1 or Month > 12:
        print('eomdays: Invalid Month!!')
    elif Month in [1,3,4,7,8,10,12]:  # Jan, Mar, May, July, Aug, Oct, Dec
        eomdays = 31
    elif Month == 2:  # Feb
        eomdays = 28
        if Year % 4 == 0:  # Leap Year
            eomdays = 29
    else:  # Apr, June, Sep, Nov
        eomdays = 30

    return eomdays

# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
#
# # ROUTINE:
#
# DESCRIPTION:
#
# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=

def yeardays(Year):

    """
    Description: Provides system with the number of days in the current year.

    :param Year: current year
    :return: number of days in current year
    """

    yeardays = 0
    for i in range(0,12):
        yeardays = yeardays + float(eomdays(Year,i))

    if yeardays == 0 or yeardays > 366:
        print('yeardays out of bounds!!')

    return yeardays

# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
#   Copyright by the GOTM-team under the GNU Public License - www.gnu.org
# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=


"""Inception_ResNet 1D models in Tensorflow-Keras.
Reference - [Rethinking the Inception Architecture for Computer Vision](http://arxiv.org/abs/1512.00567)
Inception_ResNet Review: https://towardsdatascience.com/review-inception-v4-evolved-from-googlenet-merged-with-resnet-idea-image-classification-5e8c339d18bc
"""

import numpy as np
import tensorflow as tf

def Conv_1D_Block(x, model_width, kernel, strides=1, padding="same"):
    # 1D Convolutional Block with BatchNormalization
    x = tf.keras.layers.Conv1D(model_width, kernel, strides=strides, padding=padding, kernel_initializer="he_normal")(x)
    x = tf.keras.layers.BatchNormalization()(x)
    x = tf.keras.layers.Activation('relu')(x)

    return x

def classifier(inputs, class_number):
    # Construct the Classifier Group
    # inputs       : input vector
    # class_number : number of output classes
    out = tf.keras.layers.Dense(class_number, activation='softmax')(inputs)
    return out

def regressor(inputs, feature_number):
    # Construct the Regressor Group
    # inputs         : input vector
    # feature_number : number of output features
    out = tf.keras.layers.Dense(feature_number, activation='linear')(inputs)
    return out

def SE_Block(inputs, num_filters, ratio):
    squeeze = tf.keras.layers.GlobalAveragePooling1D()(inputs)

    excitation = tf.keras.layers.Dense(units=num_filters/ratio)(squeeze)
    excitation = tf.keras.layers.Activation('relu')(excitation)
    excitation = tf.keras.layers.Dense(units=num_filters)(excitation)
    excitation = tf.keras.layers.Activation('sigmoid')(excitation)
    excitation = tf.keras.layers.Reshape([1, num_filters])(excitation)

    scale = inputs * excitation

    return scale

def Inception_ResNet_Module_A(inputs, filterB1_1, filterB2_1, filterB2_2, filterB3_1, filterB3_2, filterB3_3, filterB4_1, i):
    # Inception ResNet Module A - Block i
    branch1x1 = Conv_1D_Block(inputs, filterB1_1, 1)

    branch3x3 = Conv_1D_Block(inputs, filterB2_1, 1)
    branch3x3 = Conv_1D_Block(branch3x3, filterB2_2, 3)

    branch3x3dbl = Conv_1D_Block(inputs, filterB3_1, 1)
    branch3x3dbl = Conv_1D_Block(branch3x3dbl, filterB3_2, 3)
    branch3x3dbl = Conv_1D_Block(branch3x3dbl, filterB3_3, 3)

    branch_concat = tf.keras.layers.concatenate([branch1x1, branch3x3, branch3x3dbl], axis=-1)
    branch1x1_ln = tf.keras.layers.Conv1D(filterB4_1, 1, activation='linear', strides=1, padding='same', kernel_initializer="he_normal")(branch_concat)

    x = tf.keras.layers.Add(name='Inception_ResNet_Block_A'+str(i))([inputs, branch1x1_ln])
    x = tf.keras.layers.BatchNormalization()(x)
    out = tf.keras.layers.Activation('relu')(x)

    return out

def Inception_ResNet_Module_B(inputs, filterB1_1, filterB2_1, filterB2_2, filterB2_3, filterB3_1, i):
    # Inception ResNet Module B - Block i
    branch1x1 = Conv_1D_Block(inputs, filterB1_1, 1)

    branch7x7 = Conv_1D_Block(inputs, filterB2_1, 1)
    branch7x7 = Conv_1D_Block(branch7x7, filterB2_2, 7)
    branch7x7 = Conv_1D_Block(branch7x7, filterB2_3, 7)

    branch_concat = tf.keras.layers.concatenate([branch1x1, branch7x7], axis=-1)
    branch1x1_ln = tf.keras.layers.Conv1D(filterB3_1, 1, activation='linear', strides=1, padding='same', kernel_initializer="he_normal")(branch_concat)

    x = tf.keras.layers.Add(name='Inception_ResNet_Block_B'+str(i))([inputs, branch1x1_ln])
    x = tf.keras.layers.BatchNormalization()(x)
    out = tf.keras.layers.Activation('relu')(x)

    return out

def Inception_ResNet_Module_C(inputs, filterB1_1, filterB2_1, filterB2_2, filterB2_3, filterB3_1, i):
    # Inception ResNet Module C - Block i
    branch1x1 = Conv_1D_Block(inputs, filterB1_1, 1)

    branch3x3 = Conv_1D_Block(inputs, filterB2_1, 1)
    branch3x3 = Conv_1D_Block(branch3x3, filterB2_2, 3)
    branch3x3 = Conv_1D_Block(branch3x3, filterB2_3, 3)

    branch_concat = tf.keras.layers.concatenate([branch1x1, branch3x3], axis=-1)
    branch1x1_ln = tf.keras.layers.Conv1D(filterB3_1, 1, activation='linear', strides=1, padding='same', kernel_initializer="he_normal")(branch_concat)

    x = tf.keras.layers.Add(name='Inception_ResNet_Block_C'+str(i))([inputs, branch1x1_ln])
    x = tf.keras.layers.BatchNormalization()(x)
    out = tf.keras.layers.Activation('relu')(x)

    return out

def Reduction_Block_A(inputs, filterB1_1, filterB2_1, filterB2_2, filterB2_3):
    # Reduction Block A
    branch_pool = tf.keras.layers.MaxPooling1D(pool_size=3, strides=2)(inputs)

    branch3x3 = Conv_1D_Block(inputs, filterB1_1, 3, strides=2, padding='valid')

    branch3x3dbl = Conv_1D_Block(inputs, filterB2_1, 1)
    branch3x3dbl = Conv_1D_Block(branch3x3dbl, filterB2_2, 3)
    branch3x3dbl = Conv_1D_Block(branch3x3dbl, filterB2_3, 3, strides=2, padding='valid')

    x = tf.keras.layers.concatenate([branch_pool, branch3x3, branch3x3dbl], axis=-1, name='Reduction_Block_A')
    x = tf.keras.layers.BatchNormalization()(x)
    out = tf.keras.layers.Activation('relu')(x)

    return out

def Reduction_Block_B(inputs, filterB1_1, filterB1_2, filterB2_1, filterB2_2, filterB3_1, filterB3_2, filterB3_3):
    # Reduction Block B
    branch_pool = tf.keras.layers.MaxPooling1D(pool_size=3, strides=2)(inputs)

    branch3x3 = Conv_1D_Block(inputs, filterB1_1, 1)
    branch3x3 = Conv_1D_Block(branch3x3, filterB1_2, 3, strides=2, padding='valid')

    branch3x3_2 = Conv_1D_Block(inputs, filterB2_1, 1)
    branch3x3_2 = Conv_1D_Block(branch3x3_2, filterB2_2, 3, strides=2, padding='valid')

    branch3x3dbl = Conv_1D_Block(inputs, filterB3_1, 1)
    branch3x3dbl = Conv_1D_Block(branch3x3dbl, filterB3_2, 3)
    branch3x3dbl = Conv_1D_Block(branch3x3dbl, filterB3_3, 3, strides=2, padding='valid')

    x = tf.keras.layers.concatenate([branch_pool, branch3x3, branch3x3_2, branch3x3dbl], axis=-1)
    x = tf.keras.layers.BatchNormalization()(x)
    out = tf.keras.layers.Activation('relu')(x)

    return out

class SEInception_ResNet:
    def __init__(self, length, num_channel, num_filters, ratio=4, problem_type='Regression',
                 output_nums=1, pooling='avg', dropout_rate=False, auxilliary_outputs=False):
        # length: Input Signal Segments Length
        # model_depth: Depth of the Model
        # model_width: Width of the Model
        # kernel_size: Kernel or Filter Size of the Input Convolutional Layer
        # num_channel: Number of Channels of the Input Predictor Signals
        # problem_type: Regression or Classification
        # output_nums: Number of Output Classes in Classification mode and output features in Regression mode
        # pooling: Choose either 'max' for MaxPooling or 'avg' for Averagepooling
        # dropout_rate: If turned on, some layers will be dropped out randomly based on the selected proportion
        # auxilliary_outputs: Two extra Auxullary outputs for the Inception models, acting like Deep Supervision
        self.length = length
        self.num_channel = num_channel
        self.num_filters = num_filters
        self.ratio = ratio
        self.problem_type = problem_type
        self.output_nums = output_nums
        self.pooling = pooling
        self.dropout_rate = dropout_rate
        self.auxilliary_outputs = auxilliary_outputs

    def MLP(self, x):
        if self.pooling == 'avg':
            x = tf.keras.layers.GlobalAveragePooling1D()(x)
        elif self.pooling == 'max':
            x = tf.keras.layers.GlobalMaxPooling1D()(x)
        x = tf.keras.layers.Flatten()(x)
        if self.dropout_rate:
            x = tf.keras.layers.Dropout(self.dropout_rate)(x)
        outputs = tf.keras.layers.Dense(self.output_nums, activation='linear')(x)
        if self.problem_type == 'Classification':
            outputs = tf.keras.layers.Dense(self.output_nums, activation='softmax')(x)

        return outputs

    def SEInception_ResNet_v1(self):
        inputs = tf.keras.Input((self.length, self.num_channel))  # The input tensor
        # Stem
        x = Conv_1D_Block(inputs, 32, 3, strides=2, padding='valid')
        x = Conv_1D_Block(x, 32, 3, padding='valid')
        x = Conv_1D_Block(x, 64, 3)
        x = tf.keras.layers.MaxPooling1D(3, strides=2)(x)
        x = Conv_1D_Block(x, 80, 1)
        x = Conv_1D_Block(x, 192, 3, padding='valid')
        x = Conv_1D_Block(x, 256, 3, strides=2, padding='valid')

        # 5x Inception ResNet A Blocks - 35 x 35 x 256
        for i in range(5):
            x = Inception_ResNet_Module_A(x, 32, 32, 32, 32, 32, 32, 256, i)
            x = SE_Block(x, int(np.shape(x)[-1]), self.ratio)

        aux_output_0 = []
        if self.auxilliary_outputs:
            # Auxilliary Output 0
            aux_pool = tf.keras.layers.AveragePooling1D(pool_size=5, strides=3)(x)
            aux_conv = Conv_1D_Block(aux_pool, 128, 1)
            aux_conv = Conv_1D_Block(aux_conv, 768, 5, padding='valid')
            aux_output_0 = self.MLP(aux_conv)

        x = Reduction_Block_A(x, 384, 192, 224, 256)  # Reduction Block A: 17 x 17 x 768

        # 10x Inception ResNet B Blocks - 17 x 17 x 768
        for i in range(10):
            x = Inception_ResNet_Module_B(x, 128, 128, 128, 128, 896, i)
            x = SE_Block(x, int(np.shape(x)[-1]), self.ratio)

        aux_output_1 = []
        if self.auxilliary_outputs:
            # Auxilliary Output 1
            aux_pool = tf.keras.layers.AveragePooling1D(pool_size=5, strides=3, padding='valid')(x)
            aux_conv = Conv_1D_Block(aux_pool, 128, 1)
            aux_conv = Conv_1D_Block(aux_conv, 768, 5, padding='valid')
            aux_output_1 = self.MLP(aux_conv)

        x = Reduction_Block_B(x, 256, 384, 256, 256, 256, 256, 256)  # Reduction Block B: 8 x 8 x 1280

        # 5x Inception ResNet C Blocks - 8 x 8 x 1280
        for i in range(5):
            x = Inception_ResNet_Module_C(x, 128, 192, 192, 192, 1792, i)
            x = SE_Block(x, int(np.shape(x)[-1]), self.ratio)

        # Final Dense MLP Layer for the outputs
        final_output = self.MLP(x)
        # Create model.
        model = tf.keras.Model(inputs, final_output, name='Inception_v4')
        if self.auxilliary_outputs:
            model = tf.keras.Model(inputs, outputs=[final_output, aux_output_0, aux_output_1], name='Inception_ResNet_v1')

        return model

    def SEInception_ResNet_v2(self):
        inputs = tf.keras.Input((self.length, self.num_channel))  # The input tensor
        # Stem
        x = Conv_1D_Block(inputs, 32, 3, strides=2, padding='valid')
        x = Conv_1D_Block(x, 32, 3, padding='valid')
        x = Conv_1D_Block(x, 64, 3)
        #
        branch1 = Conv_1D_Block(x, 96, 3, strides=2, padding='valid')
        branch2 = tf.keras.layers.MaxPooling1D(3, strides=2)(x)
        x = tf.keras.layers.concatenate([branch1, branch2], axis=-1)
        #
        branch1 = Conv_1D_Block(x, 64, 1)
        branch1 = Conv_1D_Block(branch1, 96, 3, padding='valid')
        branch2 = Conv_1D_Block(x, 64, 1)
        branch2 = Conv_1D_Block(branch2, 64, 7)
        branch2 = Conv_1D_Block(branch2, 96, 3, padding='valid')
        x = tf.keras.layers.concatenate([branch1, branch2], axis=-1)
        #
        branch1 = Conv_1D_Block(x, 192, 3, padding='valid')
        branch2 = tf.keras.layers.MaxPooling1D(3, strides=1)(x)
        x = tf.keras.layers.concatenate([branch1, branch2], axis=-1)

        # 5x Inception ResNet A Blocks - 35 x 35 x 256
        for i in range(10):
            x = Inception_ResNet_Module_A(x, 32, 32, 32, 32, 48, 64, 384, i)
            x = SE_Block(x, int(np.shape(x)[-1]), self.ratio)

        aux_output_0 = []
        if self.auxilliary_outputs:
            # Auxilliary Output 0
            aux_pool = tf.keras.layers.AveragePooling1D(pool_size=5, strides=3, padding='valid')(x)
            aux_conv = Conv_1D_Block(aux_pool, 96, 1)
            aux_output_0 = self.MLP(aux_conv)

        x = Reduction_Block_A(x, 384, 192, 224, 256)  # Reduction Block A: 17 x 17 x 768

        # 10x Inception ResNet B Blocks - 17 x 17 x 768
        for i in range(20):
            x = Inception_ResNet_Module_B(x, 192, 128, 160, 192, 1024, i)
            x = SE_Block(x, int(np.shape(x)[-1]), self.ratio)

        aux_output_1 = []
        if self.auxilliary_outputs:
            # Auxilliary Output 1
            aux_pool = tf.keras.layers.AveragePooling1D(pool_size=5, strides=3, padding='valid')(x)
            aux_conv = Conv_1D_Block(aux_pool, 128, 1)
            aux_conv = Conv_1D_Block(aux_conv, 768, 5)
            aux_output_1 = self.MLP(aux_conv)

        x = Reduction_Block_B(x, 256, 384, 256, 288, 256, 288, 320)  # Reduction Block B: 8 x 8 x 1280

        # 5x Inception ResNet C Blocks - 8 x 8 x 1280
        for i in range(10):
            x = Inception_ResNet_Module_C(x, 192, 192, 224, 256, 2016, i)
            x = SE_Block(x, int(np.shape(x)[-1]), self.ratio)

        # Final Dense MLP Layer for the outputs
        final_output = self.MLP(x)
        # Create model.
        model = tf.keras.Model(inputs, final_output)
        if self.auxilliary_outputs:
            model = tf.keras.Model(inputs, outputs=[final_output, aux_output_0, aux_output_1])

        return model

if __name__ == '__main__':
    # Configurations
    length = 1024  # Length of each segment
    model_name = 'SEInceptionResNetV1'  # DenseNet Models
    model_width = 64 # Width of the Initial Layer, subsequent layers start from here
    num_channel = 1  # Number of Input Channels in the Model
    problem_type = 'Regression' # Classification or Regression
    output_nums = 1  # Number of Class for Classification Problems, always '1' for Regression Problems
    reduction_ratio = 8  # Reduction Ratio or Cardinality for the Squeeze and Excite Block
    #
    Model = SEInception_ResNet(length, num_channel, model_width, ratio=reduction_ratio, problem_type=problem_type, output_nums=output_nums,
                      pooling='avg', dropout_rate=False, auxilliary_outputs=False).SEInception_ResNet_v1()
    Model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.0003), loss=tf.keras.losses.MeanAbsoluteError(), metrics=tf.keras.metrics.MeanSquaredError())
    Model.summary()


# MIT License
#
# Copyright (c) 2016-2017 Anders Steen Christensen, Lars Andersen Bratholm
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import numpy as np
import copy

from arad import ARAD
from aras import ARAS

from data import NUCLEAR_CHARGE

from representations import fgenerate_coulomb_matrix
from representations import fgenerate_unsorted_coulomb_matrix
from representations import fgenerate_local_coulomb_matrix
from representations import fgenerate_atomic_coulomb_matrix

HOF_DFTB3 = dict()
HOF_DFTB3["H"] = -172.3145
HOF_DFTB3["C"] = -906.4342
HOF_DFTB3["N"] = -1327.2991
HOF_DFTB3["O"] = -1936.6161
HOF_DFTB3["S"] = -1453.3907

class Molecule(object):
    """ Implemented molecular descriptors:
        1. Sorted Coulomb Matrix
        2. Unsorted Coulomb Matrix
        3. Sorted Coulomb Matrix Eigenvalues
        4. Reduced Coulomb Matrix
        5. Bag of Bonds

        Implemented local descriptors:
        1. Local Coulomb Matrix
        2. Atomic Coulomb Matrix
        3. ARAD
        4. ARAS

        Implemented periodic descriptors
        1. ARAD
        2. ARAS
    """

    def __init__(self):
        self.natoms = -1
        self.energy = float("nan")
        self.molid = -1
        self.name = None
        self.dftb3_energy = float("nan")
        self.dftb3_hof = float("nan")

        self.atomtypes = []
        self.nuclear_charges = []
        self.coordinates = []
        self.active_atoms = []
        self.unit_cell = None

        # Container for misc properties
        self.properties = []
        self.properties2 = []

    def generate_coulomb_matrix(self, size=23):
        self.descriptor = fgenerate_coulomb_matrix(self.nuclear_charges, \
                self.coordinates, self.natoms, size)

    def generate_bob(self, size = 23, asize = {"O":3, "C":7, "N":3, "H":16, "S":1}):
        coulomb_matrix = fgenerate_unsorted_coulomb_matrix(self.nuclear_charges,
                self.coordinates, self.natoms, size)
        coulomb_matrix = vector_to_matrix(coulomb_matrix)
        self.descriptor = []
        self.atomtypes = np.asarray(self.atomtypes)
        for atom1, size1 in asize.items():
            pos1 = np.where(self.atomtypes == atom1)[0]
            feature_vector = np.zeros(size1)
            feature_vector[:pos1.size] = np.diag(coulomb_matrix)[pos1]
            feature_vector.sort()
            self.descriptor.append(feature_vector[:])
            for atom2, size2 in asize.items():
                if atom1 > atom2:
                    continue
                if atom1 == atom2:
                    size = size1*(size1-1)/2
                    feature_vector = np.zeros(size)
                    sub_matrix = coulomb_matrix[np.ix_(pos1,pos1)]
                    feature_vector[:pos1.size*(pos1.size-1)/2] = sub_matrix[np.triu_indices(pos1.size, 1)]
                    feature_vector.sort()
                    self.descriptor.append(feature_vector[:])
                else:
                    pos2 = np.where(self.atomtypes == atom2)[0]
                    feature_vector = np.zeros(size1*size2)
                    feature_vector[:pos1.size*pos2.size] = coulomb_matrix[np.ix_(pos1,pos2)].ravel()
                    feature_vector.sort()
                    self.descriptor.append(feature_vector[:])

        self.descriptor = np.concatenate(self.descriptor)

    def generate_eigenvalue_coulomb_matrix(self, size=23):
        coulomb_matrix = fgenerate_coulomb_matrix(self.nuclear_charges, \
                self.coordinates, self.natoms, size)
        self.descriptor = np.linalg.eigh(vector_to_matrix(coulomb_matrix))[0]

    def generate_reduced_coulomb_matrix(self, size=23):
        coulomb_matrix = fgenerate_coulomb_matrix(self.nuclear_charges, \
                self.coordinates, self.natoms, size)
        coulomb_matrix = vector_to_matrix(coulomb_matrix)
        self.descriptor = np.concatenate([np.diag(coulomb_matrix), coulomb_matrix[1:,0]])

    def generate_unsorted_coulomb_matrix(self, size=23):
        self.descriptor = fgenerate_unsorted_coulomb_matrix(self.nuclear_charges, \
                self.coordinates, self.natoms, size)

    def generate_local_coulomb_matrix(self, calc="all",size=23):
        self.local_descriptor = fgenerate_local_coulomb_matrix( \
                self.nuclear_charges, self.coordinates, self.natoms, size)

    def generate_atomic_coulomb_matrix(self, size = 23, cutoff = 8.0):
        self.local_descriptor = fgenerate_atomic_coulomb_matrix( \
                self.nuclear_charges, self.coordinates, self.natoms, size, cutoff)

    def generate_atomic_unsorted_coulomb_matrix(self, size = 23, cutoff = 8.0):
        self.local_descriptor = fgenerate_atomic_unsorted_coulomb_matrix( \
                self.nuclear_charges, self.coordinates, self.natoms, size, cutoff)

    def generate_arad_descriptor(self, size=23):
        arad_object = ARAD(maxMolSize=size,maxAts=size)
        self.arad_descriptor = arad_object.describe(np.array(self.coordinates), \
                np.array(self.nuclear_charges))

        assert (self.arad_descriptor).shape[0] == size, "ERROR: Check ARAD descriptor size!"
        assert (self.arad_descriptor).shape[2] == size, "ERROR: Check ARAD descriptor size!"

    def generate_arad_descriptor_periodic(self, size=23, unit_cell=None):

        if unit_cell is None:
            unit_cell = self.unit_cell

        arad_object = ARAD(maxMolSize=size,maxAts=size)
        self.arad_descriptor = arad_object.describe(np.array(self.coordinates), \
                np.array(self.nuclear_charges), cell=unit_cell)

        assert (self.arad_descriptor).shape[0] == size, "ERROR: Check ARAD descriptor size!"
        assert (self.arad_descriptor).shape[2] == size, "ERROR: Check ARAD descriptor size!"

    def generate_aras_descriptor(self, size=23):
        aras_object = ARAS(maxMolSize=size,maxAts=size)
        self.aras_descriptor = aras_object.describe(np.array(self.coordinates), \
                np.array(self.nuclear_charges))

        assert (self.aras_descriptor).shape[0] == size, "ERROR: Check ARAS descriptor size!"
        assert (self.aras_descriptor).shape[2] == size, "ERROR: Check ARAS descriptor size!"

    def read_xyz(self, filename):

        f = open(filename, "r")
        lines = f.readlines()
        f.close()

        self.natoms = int(lines[0])

        for line in lines[2:]:
            tokens = line.split()

            if len(tokens) < 4:
                break

            self.atomtypes.append(tokens[0])
            self.nuclear_charges.append(NUCLEAR_CHARGE[tokens[0]])

            x = float(tokens[1])
            y = float(tokens[2])
            z = float(tokens[3])

            self.coordinates.append(np.array([x, y, z]))

        self.coordinates = np.array(self.coordinates)

def get_lines(filename):

    f = open(filename, "r")
    lines = f.readlines()
    f.close()

    return lines

def parse_molecules(filename):

    lines = get_lines(filename)

    mols = []

    mol = Molecule()

    for line in lines:

        tokens = line.split()

        if len(tokens) == 1:

            if mol.natoms > 0:
                mols.append(mol)

            mol = Molecule()
            mol.natoms = int(tokens[0])

        if len(tokens) == 2:
            mol.molid = int(tokens[0])
            mol.energy = float(tokens[1])
            mol.dftb3_energy = parse_dft3_energy(mol.molid)

        if len(tokens) == 7:

            atom_type = tokens[0]
            mol.atomtypes.append(atom_type)
            mol.nuclear_charges.append(NUCLEAR_CHARGE[atom_type])
            x = float(tokens[4])
            y = float(tokens[5])
            z = float(tokens[6])

            mol.coordinates.append(np.array([x, y, z]))

            mol.dftb3_hof = 0.0
            mol.dftb3_hof += mol.dftb3_energy

            for atom in ["H", "C", "N", "O", "S"]:

                n = mol.atomtypes.count(atom)
                mol.dftb3_hof -= n * HOF_DFTB3[atom]

    mol.atomtypes = np.asarray(mol.atomtypes)
    mol.coordinates = np.asarray(mol.coordinates)
    mol.nuclear_charges = np.asarray(mol.nuclear_charges)

    return mols

def parse_dft3_energy(molid):

    filename = "../logfiles/" + str(molid) + ".log"
    f = open(filename, "r")
    lines = f.readlines()
    f.close()

    energy = float("nan")
    for line in lines:
        if "Total Energy" in line:
            tokens = line.split()
            energy = float(tokens[2]) * 627.51

    return energy

def vector_to_matrix(v):
    if not (np.sqrt(8*v.shape[0]+1) == int(np.sqrt(8*v.shape[0]+1))):
        print "ERROR: Can not make a square matrix."
        exit(1)

    n = v.shape[0]
    l = (-1 + int(np.sqrt(8*n+1)))/2
    M = np.empty((l,l))

    index = 0
    for i in range(l):
        for j in range(l):
            if j > i:
                continue

            M[i,j] = v[index]
            M[j,i] = M[i,j]

            index += 1
    return M