# Copyright 2014-2021 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

'''
There are two options to call wannier90 in PySCF.  One is the pyWannier90.py
interface as implemented in this file.

(1)
pyWannier90: Wannier90 for PySCF (https://github.com/hungpham2017/pyWannier90)
Hung Q. Pham
email: pqh3.14@gmail.com

(2)
Another wannier90 python interface is available on the repo:
    https://github.com/zhcui/wannier90
Contact its author "Zhihao Cui" <zcui@caltech.edu> for more details of
installation and implementations.
'''

import os, time
import numpy as np
from scipy.io import FortranFile
import pyscf.data.nist as param
from pyscf import lib
from pyscf.pbc import df
from pyscf.pbc.dft import gen_grid, numint

try:
    import libwannier90
except ImportError:
    print('WARNING: Check the installation of libwannier90 and its path in pyscf/pbc/tools/pywannier90.py')
    print('libwannier90 can be found at: https://github.com/hungpham2017/pyWannier90')
    raise

def save_kmf(kmf, chkfile):
    ''' Save a wavefunction'''
    from pyscf.lib.chkfile import save
    kpts = kmf.kpts
    mo_energy_kpts = kmf.mo_energy_kpts
    mo_coeff_kpts = kmf.mo_coeff_kpts

    scf_dic = {'kpts'          : kpts,
               'mo_energy_kpts': mo_energy_kpts,
               'mo_coeff_kpts' : mo_coeff_kpts}
    save(chkfile, 'scf', scf_dic)

def load_kmf(chkfile):
    ''' Load a wavefunction'''
    from pyscf.lib.chkfile import load
    kmf = load(chkfile, 'scf')
    class fake_kmf:
        def __init__(self, kmf):
            self.kpts = kmf['kpts']
            self.mo_energy_kpts = kmf['mo_energy_kpts']
            self.mo_coeff_kpts = kmf['mo_coeff_kpts']
    kmf = fake_kmf(kmf)
    return kmf

def angle(v1, v2):
    '''
    Return the angle (in radiant between v1 and v2)
    '''

    v1 = np.asarray(v1)
    v2 = np.asarray(v2)
    cosa = v1.dot(v2)/ np.linalg.norm(v1) / np.linalg.norm(v2)
    return np.arccos(cosa)

def transform(x_vec, z_vec):
    '''
    Construct a transformation matrix to transform r_vec to the new coordinate system defined by x_vec and z_vec
    '''

    x_vec = x_vec/np.linalg.norm(np.asarray(x_vec))
    z_vec = z_vec/np.linalg.norm(np.asarray(z_vec))
    assert x_vec.dot(z_vec) == 0    # x and z have to be orthogonal to one another
    y_vec = -np.cross(x_vec,z_vec)
    new = np.asarray([x_vec, y_vec, z_vec])
    original = np.asarray([[1,0,0],[0,1,0],[0,0,1]])

    tran_matrix = np.empty([3,3])
    for row in range(3):
        for col in range(3):
            tran_matrix[row,col] = np.cos(angle(original[row],new[col]))

    return tran_matrix.T

def cartesian_prod(arrays, out=None, order='C'):
    '''
    This function is similar to lib.cartesian_prod of PySCF, except the output can be in Fortran or in C order
    '''
    arrays = [np.asarray(x) for x in arrays]
    dtype = np.result_type(*arrays)
    nd = len(arrays)
    dims = [nd] + [len(x) for x in arrays]

    if out is None:
        out = np.empty(dims, dtype)
    else:
        out = np.ndarray(dims, dtype, buffer=out)
    tout = out.reshape(dims)

    shape = [-1] + [1] * nd
    for i, arr in enumerate(arrays):
        tout[i] = arr.reshape(shape[:nd-i])

    return tout.reshape((nd,-1),order=order).T

def periodic_grid(cell, grid=[50,50,50], supercell=[1,1,1], order='C'):
    '''
    Generate a periodic grid for the unit/computational cell in F/C order
    '''
    ngrid = np.asarray(grid)
    qv = cartesian_prod([np.arange(-ngrid[i]*(supercell[i]//2),ngrid[i]*((supercell[i]+1)//2)) for i in range(3)], order=order)
    a_frac = lib.einsum('i,ij->ij', 1./ngrid, cell.lattice_vectors())
    coords = np.dot(qv, a_frac)

    # Compute weight
    ngrids = np.prod(grid)
    ncells = np.prod(supercell)
    weights = np.empty(ngrids*ncells)
    weights[:] = cell.vol / ngrids / ncells
    return coords, weights

def R_r(r_norm, r=1, zona=1):
    r'''
    Radial functions used to compute \Theta_{l,m_r}(\theta,\phi)
    '''

    if r == 1:
        R_r = 2 * zona**(3/2) * np.exp(-zona*r_norm)
    elif r == 2:
        R_r = 1 / 2 / np.sqrt(2) * zona**(3/2) * (2 - zona*r_norm) * np.exp(-zona*r_norm/2)
    else:
        R_r = np.sqrt(4/27) * zona**(3/2) * (1 - 2*zona*r_norm/3 + 2*(zona**2)*(r_norm**2)/27) * np.exp(-zona*r_norm/3)

    return R_r

def theta(func, cost, phi):
    r'''
    Basic angular functions (s,p,d,f) used to compute \Theta_{l,m_r}(\theta,\phi)
    ref: Table 3.1 of the Wannier90 User guide
        Link: https://github.com/wannier-developers/wannier90/raw/v3.1.0/doc/compiled_docs/user_guide.pdf
    '''
    sint = np.sqrt(1 - cost**2)
    if func == 's':
        theta = 1 / np.sqrt(4 * np.pi) * np.ones([cost.shape[0]])
    elif func == 'pz':
        theta = np.sqrt(3 / 4 / np.pi) * cost
    elif func == 'px':
        theta = np.sqrt(3 / 4 / np.pi) * sint * np.cos(phi)
    elif func == 'py':
        theta = np.sqrt(3 / 4 / np.pi) * sint * np.sin(phi)
    elif func == 'dz2':
        theta = np.sqrt(5 / 16 / np.pi) * (3*cost**2 - 1)
    elif func == 'dxz':
        theta = np.sqrt(15 / 4 / np.pi) * sint * cost * np.cos(phi)
    elif func == 'dyz':
        theta = np.sqrt(15 / 4 / np.pi) * sint * cost * np.sin(phi)
    elif func == 'dx2-y2':
        theta = np.sqrt(15 / 16 / np.pi) * (sint**2) * np.cos(2*phi)
    elif func == 'dxy':
        theta = np.sqrt(15 / 16 / np.pi) * (sint**2) * np.sin(2*phi)
    elif func == 'fz3':
        theta = np.sqrt(7) / 4 / np.sqrt(np.pi) * (5*cost**3 - 3*cost)
    elif func == 'fxz2':
        theta = np.sqrt(21) / 4 / np.sqrt(2*np.pi) * (5*cost**2 - 1) * sint * np.cos(phi)
    elif func == 'fyz2':
        theta = np.sqrt(21) / 4 / np.sqrt(2*np.pi) * (5*cost**2 - 1) * sint * np.sin(phi)
    elif func == 'fz(x2-y2)':
        theta = np.sqrt(105) / 4 / np.sqrt(np.pi) * sint**2 * cost * np.cos(2*phi)
    elif func == 'fxyz':
        theta = np.sqrt(105) / 4 / np.sqrt(np.pi) * sint**2 * cost * np.sin(2*phi)
    elif func == 'fx(x2-3y2)':
        theta = np.sqrt(35) / 4 / np.sqrt(2*np.pi) * sint**3 * (np.cos(phi)**2 - 3*np.sin(phi)**2) * np.cos(phi)
    elif func == 'fy(3x2-y2)':
        theta = np.sqrt(35) / 4 / np.sqrt(2*np.pi) * sint**3 * (3*np.cos(phi)**2 - np.sin(phi)**2) * np.sin(phi)

    return theta

def theta_lmr(l, mr, cost, phi):
    r'''
    Compute the value of \Theta_{l,m_r}(\theta,\phi)
    ref: Table 3.1 and 3.2 of the Wannier90 User guide
        Link: https://github.com/wannier-developers/wannier90/raw/v3.1.0/doc/compiled_docs/user_guide.pdf
    '''
    assert l in [0,1,2,3,-1,-2,-3,-4,-5]
    assert mr in [1,2,3,4,5,6,7]

    if l == 0:                           # s
        theta_lmr = theta('s', cost, phi)
    elif (l == 1) and (mr == 1):         # pz
        theta_lmr = theta('pz', cost, phi)
    elif (l == 1) and (mr == 2):         # px
        theta_lmr = theta('px', cost, phi)
    elif (l == 1) and (mr == 3):         # py
        theta_lmr = theta('py', cost, phi)
    elif (l == 2) and (mr == 1):         # dz2
        theta_lmr = theta('dz2', cost, phi)
    elif (l == 2) and (mr == 2):         # dxz
        theta_lmr = theta('dxz', cost, phi)
    elif (l == 2) and (mr == 3):         # dyz
        theta_lmr = theta('dyz', cost, phi)
    elif (l == 2) and (mr == 4):         # dx2-y2
        theta_lmr = theta('dx2-y2', cost, phi)
    elif (l == 2) and (mr == 5):         # dxy
        theta_lmr = theta('dxy', cost, phi)
    elif (l == 3) and (mr == 1):         # fz3
        theta_lmr = theta('fz3', cost, phi)
    elif (l == 3) and (mr == 2):         # fxz2
        theta_lmr = theta('fxz2', cost, phi)
    elif (l == 3) and (mr == 3):         # fyz2
        theta_lmr = theta('fyz2', cost, phi)
    elif (l == 3) and (mr == 4):         # fz(x2-y2)
        theta_lmr = theta('fz(x2-y2)', cost, phi)
    elif (l == 3) and (mr == 5):         # fxyz
        theta_lmr = theta('fxyz', cost, phi)
    elif (l == 3) and (mr == 6):         # fx(x2-3y2)
        theta_lmr = theta('fx(x2-3y2)', cost, phi)
    elif (l == 3) and (mr == 7):         # fy(3x2-y2)
        theta_lmr = theta('fy(3x2-y2)', cost, phi)
    elif (l == -1) and (mr == 1):         # sp-1
        theta_lmr = 1/np.sqrt(2) * (theta('s', cost, phi) + theta('px', cost, phi))
    elif (l == -1) and (mr == 2):         # sp-2
        theta_lmr = 1/np.sqrt(2) * (theta('s', cost, phi) - theta('px', cost, phi))
    elif (l == -2) and (mr == 1):         # sp2-1
        theta_lmr = 1/np.sqrt(3) * theta('s', cost, phi) - 1/np.sqrt(6) *theta('px', cost, phi) + 1/np.sqrt(2) * theta('py', cost, phi)
    elif (l == -2) and (mr == 2):         # sp2-2
        theta_lmr = 1/np.sqrt(3) * theta('s', cost, phi) - 1/np.sqrt(6) *theta('px', cost, phi) - 1/np.sqrt(2) * theta('py', cost, phi)
    elif (l == -2) and (mr == 3):         # sp2-3
        theta_lmr = 1/np.sqrt(3) * theta('s', cost, phi) + 2/np.sqrt(6) *theta('px', cost, phi)
    elif (l == -3) and (mr == 1):         # sp3-1
        theta_lmr = 1/2 * (theta('s', cost, phi) + theta('px', cost, phi) + theta('py', cost, phi) + theta('pz', cost, phi))
    elif (l == -3) and (mr == 2):         # sp3-2
        theta_lmr = 1/2 * (theta('s', cost, phi) + theta('px', cost, phi) - theta('py', cost, phi) - theta('pz', cost, phi))
    elif (l == -3) and (mr == 3):         # sp3-3
        theta_lmr = 1/2 * (theta('s', cost, phi) - theta('px', cost, phi) + theta('py', cost, phi) - theta('pz', cost, phi))
    elif (l == -3) and (mr == 4):         # sp3-4
        theta_lmr = 1/2 * (theta('s', cost, phi) - theta('px', cost, phi) - theta('py', cost, phi) + theta('pz', cost, phi))
    elif (l == -4) and (mr == 1):         # sp3d-1
        theta_lmr = 1/np.sqrt(3) * theta('s', cost, phi) - 1/np.sqrt(6) *theta('px', cost, phi) + 1/np.sqrt(2) * theta('py', cost, phi)
    elif (l == -4) and (mr == 2):         # sp3d-2
        theta_lmr = 1/np.sqrt(3) * theta('s', cost, phi) - 1/np.sqrt(6) *theta('px', cost, phi) - 1/np.sqrt(2) * theta('py', cost, phi)
    elif (l == -4) and (mr == 3):         # sp3d-3
        theta_lmr = 1/np.sqrt(3) * theta('s', cost, phi) + 2/np.sqrt(6) * theta('px', cost, phi)
    elif (l == -4) and (mr == 4):         # sp3d-4
        theta_lmr = 1/np.sqrt(2) * (theta('pz', cost, phi) + theta('dz2', cost, phi))
    elif (l == -4) and (mr == 5):         # sp3d-5
        theta_lmr = 1/np.sqrt(2) * (-theta('pz', cost, phi) + theta('dz2', cost, phi))
    elif (l == -5) and (mr == 1):         # sp3d2-1
        theta_lmr = 1/np.sqrt(6) * theta('s', cost, phi) - 1/np.sqrt(2) *theta('px', cost, phi) - 1/np.sqrt(12) *theta('dz2', cost, phi) \
                    + 1/2 *theta('dx2-y2', cost, phi)
    elif (l == -5) and (mr == 2):         # sp3d2-2
        theta_lmr = 1/np.sqrt(6) * theta('s', cost, phi) + 1/np.sqrt(2) *theta('px', cost, phi) - 1/np.sqrt(12) *theta('dz2', cost, phi) \
                    + 1/2 *theta('dx2-y2', cost, phi)
    elif (l == -5) and (mr == 3):         # sp3d2-3
        theta_lmr = 1/np.sqrt(6) * theta('s', cost, phi) - 1/np.sqrt(2) *theta('py', cost, phi) - 1/np.sqrt(12) *theta('dz2', cost, phi) \
                    - 1/2 *theta('dx2-y2', cost, phi)
    elif (l == -5) and (mr == 4):         # sp3d2-4
        theta_lmr = 1/np.sqrt(6) * theta('s', cost, phi) + 1/np.sqrt(2) *theta('py', cost, phi) - 1/np.sqrt(12) *theta('dz2', cost, phi) \
                    - 1/2 *theta('dx2-y2', cost, phi)
    elif (l == -5) and (mr == 5):         # sp3d2-5
        theta_lmr = 1/np.sqrt(6) * theta('s', cost, phi) - 1/np.sqrt(2) *theta('pz', cost, phi) + 1/np.sqrt(3) *theta('dz2', cost, phi)
    elif (l == -5) and (mr == 6):         # sp3d2-6
        theta_lmr = 1/np.sqrt(6) * theta('s', cost, phi) + 1/np.sqrt(2) *theta('pz', cost, phi) + 1/np.sqrt(3) *theta('dz2', cost, phi)

    return theta_lmr

def g_r(grids_coor, site, l, mr, r, zona, x_axis=[1,0,0], z_axis=[0,0,1], unit='B'):
    r'''
    Evaluate the projection function g(r) or \Theta_{l,m_r}(\theta,\phi) on a grid
    ref: Chapter 3, wannier90 User Guide
    Attributes:
        grids_coor : a grids for the cell of interest
        site       : absolute coordinate (in Borh/Angstrom) of the g(r) in the cell
        l, mr      : l and mr value in the Table 3.1 and 3.2 of the ref
    Return:
        theta_lmr  : an array (ngrid, value) of g(r)

    '''

    unit_conv = 1
    if unit == 'A': unit_conv = param.BOHR

    r_vec = (grids_coor - site)
    r_vec = lib.einsum('iv,uv ->iu', r_vec, transform(x_axis, z_axis))
    r_norm = np.linalg.norm(r_vec,axis=1)
    if (r_norm < 1e-8).any():
        r_vec = (grids_coor - site - 1e-5)
        r_vec = lib.einsum('iv,uv ->iu', r_vec, transform(x_axis, z_axis))
        r_norm = np.linalg.norm(r_vec,axis=1)
    cost = r_vec[:,2]/r_norm

    phi = np.empty_like(r_norm)
    larger_idx = r_vec[:, 0] > 1e-8
    smaller_idx = r_vec[:, 0] < -1e-8
    neither_idx = not larger_idx and not smaller_idx
    phi[larger_idx] = np.arctan(r_vec[larger_idx,1]/r_vec[larger_idx,0])
    phi[smaller_idx] = np.arctan(r_vec[smaller_idx,1]/r_vec[smaller_idx,0])  + np.pi
    phi[neither_idx] = np.sign(r_vec[neither_idx,1]) * 0.5 * np.pi

    return theta_lmr(l, mr, cost, phi) * R_r(r_norm * unit_conv, r = r, zona = zona)

def get_wigner_seitz_supercell(w90, ws_search_size=[2,2,2], ws_distance_tol=1e-6):
    '''
    Return a grid that contains all the lattice within the Wigner-Seitz supercell
    Ref: the hamiltonian_wigner_seitz(count_pts) in wannier90/src/hamittonian.F90
    '''

    real_metric = w90.real_lattice_loc.T.dot(w90.real_lattice_loc)
    dist_dim = np.prod(2 * (np.asarray(ws_search_size) + 1) + 1)
    ndegen = []
    irvec = []
    mp_grid = np.asarray(w90.mp_grid_loc)
    n1_range =  np.arange(-ws_search_size[0] * mp_grid[0], ws_search_size[0]*mp_grid[0] + 1)
    n2_range =  np.arange(-ws_search_size[1] * mp_grid[1], ws_search_size[1]*mp_grid[1] + 1)
    n3_range =  np.arange(-ws_search_size[2] * mp_grid[2], ws_search_size[2]*mp_grid[2] + 1)
    x, y, z = np.meshgrid(n1_range, n2_range, n3_range)
    n_list = np.vstack([z.flatten('F'), x.flatten('F'), y.flatten('F')]).T
    i1 = np.arange(- ws_search_size[0] - 1, ws_search_size[0] + 2)
    i2 = np.arange(- ws_search_size[1] - 1, ws_search_size[1] + 2)
    i3 = np.arange(- ws_search_size[2] - 1, ws_search_size[2] + 2)
    x, y, z = np.meshgrid(i1, i2, i3)
    i_list = np.vstack([z.flatten('F'), x.flatten('F'), y.flatten('F')]).T

    nrpts = 0
    for n in n_list:
        # Calculate |r-R|^2
        ndiff = n - i_list * mp_grid
        dist = (ndiff.dot(real_metric).dot(ndiff.T)).diagonal()

        dist_min = dist.min()
        if abs(dist[(dist_dim + 1)//2 -1] - dist_min) < ws_distance_tol**2:
            temp = 0
            for i in range(0, dist_dim):
                if (abs(dist[i] - dist_min) < ws_distance_tol**2):
                    temp = temp + 1
            ndegen.append(temp)
            irvec.append(n.tolist())
            if (n**2).sum() < 1.e-10: rpt_origin = nrpts
            nrpts = nrpts + 1

    irvec = np.asarray(irvec)
    ndegen = np.asarray(ndegen)

    # Check the "sum rule"
    tot = np.sum(1/np.asarray(ndegen))
    assert tot - np.prod(mp_grid) < 1e-8, "Error in finding Wigner-Seitz points!!!"

    return ndegen, irvec, rpt_origin

def R_wz_sc(w90, R_in, R0, ws_search_size=[2,2,2], ws_distance_tol=1e-6):
    '''
    TODO: document it
    Ref: This is the replication of the R_wz_sc function of ws_distance.F90
    '''
    ndegenx = 8 #max number of unit cells that can touch in a single point (i.e.  vertex of cube)
    R_bz = np.asarray(R_in).reshape(-1, 3)
    nR = R_bz.shape[0]
    R0 = np.asarray(R0)
    ndeg = np.zeros([nR], dtype=np.int32)
    ndeg_ = np.zeros([nR, ndegenx])
    shifts = np.zeros([nR, ndegenx, 3])
    R_out = np.zeros([nR, ndegenx, 3])

    mod2_R_bz = np.sum((R_bz - R0)**2, axis=1)
    R_in_f = R_bz.dot(w90.recip_lattice_loc.T / 2 / np.pi)
    n1_range =  np.arange(-ws_search_size[0] - 1, ws_search_size[0] + 2)
    n2_range =  np.arange(-ws_search_size[1] - 1, ws_search_size[1] + 2)
    n3_range =  np.arange(-ws_search_size[2] - 1, ws_search_size[2] + 2)
    x, y, z = np.meshgrid(n1_range, n2_range, n3_range)
    n_list = np.vstack([z.flatten('F'), x.flatten('F'), y.flatten('F')]).T
    trans_vecs = n_list * w90.mp_grid_loc

    # First loop:
    R_f = np.repeat(R_in_f[:,np.newaxis,:], trans_vecs.shape[0], axis=1) + trans_vecs
    R = R_f.dot(w90.real_lattice_loc)
    mod2_R = np.sum((R - R0)**2, axis=2)
    mod2_R_min = mod2_R.min(axis=1)
    mod2_R_min_idx = np.argmin(mod2_R, axis=1)
    idx = mod2_R_min < mod2_R_bz
    R_bz[idx] = R[idx, mod2_R_min_idx[idx]]
    mod2_R_bz[idx] = mod2_R_min[idx]
    shifts_data = np.repeat(trans_vecs[np.newaxis,:,:], nR, axis=0)[idx, mod2_R_min_idx[idx]]
    shifts[idx] = np.repeat(shifts_data[:,np.newaxis,:], ndegenx, axis=1)

    idx = mod2_R_bz < ws_distance_tol**2
    ndeg[idx] = 1
    R_out[idx, 0] = R0

    # Second loop:
    R_in_f = R_bz.dot(w90.recip_lattice_loc.T / 2 / np.pi)
    R_f = np.repeat(R_in_f[:,np.newaxis,:], trans_vecs.shape[0], axis=1) + trans_vecs
    R = R_f.dot(w90.real_lattice_loc)
    mod2_R = np.sum((R - R0)**2, axis=2)
    mod2_R_bz = np.repeat(mod2_R_bz[:,np.newaxis], trans_vecs.shape[0], axis=1)
    abs_diff = abs(np.sqrt(mod2_R) - np.sqrt(mod2_R_bz))
    idx = abs_diff < ws_distance_tol
    ndeg = idx.sum(axis=1)
    assert (ndeg <= 8).all(), "The degeneracy cannot be larger than 8"
    for i in range(nR):
        R_out[i, :ndeg[i]] = R[i, idx[i]]
        shifts[i, :ndeg[i]] = shifts[i, :ndeg[i]] + trans_vecs[idx[i]]
        ndeg_[i, :ndeg[i]] = 1.0

    return ndeg_, ndeg, R_out, shifts

def ws_translate_dist(w90, irvec, ws_search_size=[2,2,2], ws_distance_tol=1e-6):
    '''
    TODO: document it
    Ref: This is the replication of the ws_translate_dist function of ws_distance.F90
    '''
    nrpts = irvec.shape[0]
    ndegenx = 8 #max number of unit cells that can touch in a single point (i.e.  vertex of cube)
    num_wann = w90.num_wann
    assert ndegenx*num_wann*nrpts > 0, "Unexpected dimensions in ws_translate_dist"

    irvec_ = []
    wann_centres_i = []
    wann_centres_j = []
    for i in range(3):
        x, y, z = np.meshgrid(irvec[:,i], np.zeros(num_wann), np.zeros(num_wann), indexing='ij')
        irvec_.append(x.flatten())
        x, y, z = np.meshgrid(np.zeros(nrpts), np.zeros(num_wann), w90.wann_centres[:,i], indexing='ij')
        wann_centres_i.append(z.flatten())
        x, y, z = np.meshgrid(np.zeros(nrpts), w90.wann_centres[:,i], np.zeros(num_wann), indexing='ij')
        wann_centres_j.append(y.flatten())

    irvec_list = np.vstack(irvec_).T
    irvec_cart_list = irvec_list.dot(w90.real_lattice_loc)
    wann_centres_i_list = np.vstack(wann_centres_i).T
    wann_centres_j_list = np.vstack(wann_centres_j).T
    R_in = irvec_cart_list - wann_centres_i_list + wann_centres_j_list
    wdist_ndeg_, wdist_ndeg, R_out, shifts = w90.R_wz_sc(R_in, [0,0,0], ws_search_size, ws_distance_tol)
    ndegenx = wdist_ndeg_.shape[1]
    irdist_ws = np.repeat(irvec_list[:,np.newaxis,:], ndegenx, axis=1) + shifts
    crdist_ws = irdist_ws.dot(w90.real_lattice_loc)

    # Reformat the matrices for the computational convenience in lib.einsum
    wdist_ndeg = wdist_ndeg.reshape(nrpts, num_wann, num_wann)
    wdist_ndeg_ = wdist_ndeg_.reshape(nrpts, num_wann, num_wann, ndegenx).transpose(3,0,1,2)
    irdist_ws = irdist_ws.reshape(nrpts, num_wann, num_wann, ndegenx, 3).transpose(3,0,1,2,4)
    crdist_ws = crdist_ws.reshape(nrpts, num_wann, num_wann, ndegenx, 3).transpose(3,0,1,2,4)

    return wdist_ndeg, wdist_ndeg_, irdist_ws, crdist_ws

'''Main class of pyWannier90'''
class W90:
    def __init__(self, kmf, cell, mp_grid, num_wann, gamma=False, spinors=False, spin_up=None, other_keywords=None):

        if isinstance(kmf, str):
            self.kmf = load_kmf(kmf)
        else:
            self.kmf = kmf
        self.cell = cell
        self.num_wann = num_wann
        self.keywords = other_keywords

        # Collect the pyscf calculation info
        nao_kpts = []
        for mo_energy in kmf.mo_energy_kpts:
            nao_kpts.append(mo_energy.shape[0])

        self.num_bands_tot = np.min(nao_kpts)
        if self.num_bands_tot < cell.nao_nr():
            print(('The number of bands at different k-point are not the same. '
                   'The first %d bands are used.') % (self.num_bands_tot) )

        self.num_kpts_loc = self.kmf.kpts.shape[0]
        self.mp_grid_loc = mp_grid
        assert self.num_kpts_loc == np.asarray(self.mp_grid_loc).prod()
        self.real_lattice_loc = self.cell.lattice_vectors() * param.BOHR
        self.recip_lattice_loc = self.cell.reciprocal_vectors() / param.BOHR
        self.kpt_latt_loc = self.cell.get_scaled_kpts(self.kmf.kpts)
        self.num_atoms_loc = self.cell.natm
        self.atom_symbols_loc = [atom[0] for atom in self.cell._atom]
        self.atom_atomic_loc = [int(self.cell._atm[atom][0] + self.cell.atom_nelec_core(atom))
                                for atom in range(self.num_atoms_loc)]
        self.atoms_cart_loc = np.asarray([(np.asarray(atom[1])* param.BOHR).tolist()
                                          for atom in self.cell._atom])
        self.gamma_only, self.spinors = (0 , 0)
        if gamma: self.gamma_only = 1
        if spinors: self.spinors = 1

        # Wannier90_setup outputs
        self.num_bands_loc = None
        self.num_wann_loc = None
        self.nntot_loc = None
        self.nn_list = None
        self.proj_site = None
        self.proj_l = None
        self.proj_m = None
        self.proj_radial = None
        self.proj_z = None
        self.proj_x = None
        self.proj_zona = None
        self.exclude_bands = None
        self.proj_s = None
        self.proj_s_qaxis = None

        # Input for Wannier90_run
        self.band_included_list = None
        self.A_matrix_loc = None
        self.M_matrix_loc = None
        self.eigenvalues_loc = None

        # Wannier90_run outputs
        self.U_matrix = None
        self.U_matrix_opt = None
        self.lwindow = None
        self.wann_centres = None
        self.wann_spreads = None
        self.spread = None

        # Others
        self.use_bloch_phases = False
        self.spin_up = spin_up
        self.mo_energy_kpts = []
        self.mo_coeff_kpts = []
        if np.mod(self.cell.nelectron,2) !=0:
            if spin_up:
                for kpt in range(self.num_kpts_loc):
                    self.mo_energy_kpts.append(self.kmf.mo_energy_kpts[0][kpt][:self.num_bands_tot])
                    self.mo_coeff_kpts.append(self.kmf.mo_coeff_kpts[0][kpt][:,:self.num_bands_tot])
            else:
                for kpt in range(self.num_kpts_loc):
                    self.mo_energy_kpts.append(self.kmf.mo_energy_kpts[1][kpt][:self.num_bands_tot])
                    self.mo_coeff_kpts.append(self.kmf.mo_coeff_kpts[1][kpt][:,:self.num_bands_tot])
        else:

            for kpt in range(self.num_kpts_loc):
                self.mo_energy_kpts.append(self.kmf.mo_energy_kpts[kpt][:self.num_bands_tot])
                self.mo_coeff_kpts.append(self.kmf.mo_coeff_kpts[kpt][:,:self.num_bands_tot])

    def kernel(self, external_AME=None):
        '''
        Main kernel for pyWannier90
        '''
        self.make_win()
        self.setup()
        if external_AME is not None:
            self.M_matrix_loc = self.read_M_mat(external_AME + '.mmn')
            self.A_matrix_loc = self.read_A_mat(external_AME + '.amn')
            self.eigenvalues_loc = self.read_epsilon_mat(external_AME + '.eig')
        else:
            self.M_matrix_loc = self.get_M_mat()
            self.A_matrix_loc = self.get_A_mat()
            self.eigenvalues_loc = self.get_epsilon_mat()
        self.run()

    def make_win(self):
        '''
        Make a basic *.win file for wannier90
        '''

        win_file = open('wannier90.win', "w")
        win_file.write('! Basic input generated by the pyWannier90. Date: %s\n' % (time.ctime()))
        win_file.write('\n')
        win_file.write('num_bands       = %d\n' % (self.num_bands_tot))
        win_file.write('num_wann       = %d\n' % (self.num_wann))
        win_file.write('\n')
        win_file.write('Begin Unit_Cell_Cart\n')
        for row in range(3):
            win_file.write('%10.7f  %10.7f  %10.7f\n' %
                           (self.real_lattice_loc[0, row], self.real_lattice_loc[1, row],
                            self.real_lattice_loc[2, row]))
        win_file.write('End Unit_Cell_Cart\n')
        win_file.write('\n')
        win_file.write('Begin atoms_cart\n')
        for atom in range(len(self.atom_symbols_loc)):
            win_file.write('%s  %7.7f  %7.7f  %7.7f\n' %
                           (self.atom_symbols_loc[atom], self.atoms_cart_loc[atom,0],
                            self.atoms_cart_loc[atom,1], self.atoms_cart_loc[atom,2]))
        win_file.write('End atoms_cart\n')
        win_file.write('\n')
        if self.use_bloch_phases: win_file.write('use_bloch_phases = T\n\n')
        if self.keywords is not None:
            win_file.write('!Additional keywords\n')
            win_file.write(self.keywords)
        win_file.write('\n\n\n')
        win_file.write('mp_grid        = %d %d %d\n' %
                       (self.mp_grid_loc[0], self.mp_grid_loc[1], self.mp_grid_loc[2]))
        if self.gamma_only == 1: win_file.write('gamma_only : true\n')
        win_file.write('begin kpoints\n')
        for kpt in range(self.num_kpts_loc):
            win_file.write('%7.7f  %7.7f  %7.7f\n' %
                           (self.kpt_latt_loc[kpt][0], self.kpt_latt_loc[kpt][1], self.kpt_latt_loc[kpt][2]))
        win_file.write('End Kpoints\n')
        win_file.close()

    def get_M_mat(self):
        r'''
        Construct the ovelap matrix: M_{m,n}^{(\mathbf{k,b})}
        Equation (25) in MV, Phys. Rev. B 56, 12847
        '''

        M_matrix_loc = np.empty([self.num_kpts_loc, self.nntot_loc,
                                 self.num_bands_loc, self.num_bands_loc],
                                dtype = np.complex128)

        for k_id in range(self.num_kpts_loc):
            for nn in range(self.nntot_loc):
                k1 = self.cell.get_abs_kpts(self.kpt_latt_loc[k_id])
                k_id2 = self.nn_list[nn, k_id, 0] - 1
                k2_ = self.kpt_latt_loc[k_id2]
                k2_scaled = k2_ + self.nn_list[nn, k_id, 1:4]
                k2 = self.cell.get_abs_kpts(k2_scaled)
                s_AO = df.ft_ao.ft_aopair(self.cell, -k2+k1, kpti_kptj=[k2,k1], q = np.zeros(3))[0]
                Cm = self.mo_coeff_kpts[k_id][:,self.band_included_list]
                Cn = self.mo_coeff_kpts[k_id2][:,self.band_included_list]
                M_matrix_loc[k_id, nn,:,:] = lib.einsum('nu,vm,uv->nm', Cn.T.conj(), Cm, s_AO).conj()

        return M_matrix_loc

    def read_M_mat(self, filename=None):
        r'''
        Read the ovelap matrix: M_{m,n}^{(\mathbf{k,b})} from seedname.mnn
        '''
        if filename is None: filename = 'wannier90.mmn'
        assert os.path.exists(filename), "Cannot find " + filename

        with open(filename, 'r') as f:
            data = f.readlines()
            num_bands_loc, num_kpts_loc, nntot_loc = np.int64(data[1].split())
            data = data[2:]
            nn_list = []
            nline = num_bands_loc**2 + 1
            M_matrix_loc = np.empty([num_kpts_loc, nntot_loc, num_bands_loc, num_bands_loc], dtype = np.complex128)
            jump = 0
            for kpt in range(num_kpts_loc):
                for nn in range(nntot_loc):
                    temp = data[jump : jump + nline]
                    nn_list.append(np.int64(temp[0].split()))
                    val_in_float = np.float64(" ".join(temp[1:]).split()).reshape(-1,2)
                    val_in_complex = val_in_float[:,0] + 1j * val_in_float[:,1]
                    M_matrix_loc[kpt, nn] = val_in_complex.reshape(num_bands_loc, num_bands_loc)
                    jump += nline

        return M_matrix_loc

    def get_A_mat(self):
        r'''
        Construct the projection matrix: A_{m,n}^{\mathbf{k}}
        Equation (62) in MV, Phys. Rev. B 56, 12847 or equation (22) in SMV, Phys. Rev. B 65, 035109
        '''

        A_matrix_loc = np.empty([self.num_kpts_loc, self.num_wann_loc, self.num_bands_loc], dtype = np.complex128)

        if self.use_bloch_phases:
            Amn = np.zeros([self.num_wann_loc, self.num_bands_loc])
            np.fill_diagonal(Amn, 1)
            A_matrix_loc[:,:,:] = Amn
        else:
            from pyscf.dft import numint as mol_numint
            from pyscf.dft import gen_grid as mol_gen_grid
            grids = mol_gen_grid.Grids(self.cell).build()
            coords = grids.coords
            weights = grids.weights
            for ith_wann in range(self.num_wann_loc):
                frac_site = self.proj_site[ith_wann]
                abs_site = frac_site.dot(self.real_lattice_loc) / param.BOHR
                l = self.proj_l[ith_wann]
                mr = self.proj_m[ith_wann]
                r = self.proj_radial[ith_wann]
                zona = self.proj_zona[ith_wann]
                x_axis = self.proj_x[ith_wann]
                z_axis = self.proj_z[ith_wann]
                gr = g_r(coords, abs_site, l, mr, r, zona, x_axis, z_axis, unit = 'B')
                ao_L0 = mol_numint.eval_ao(self.cell, coords)
                s_aoL0_g = lib.einsum('i,i,iv->v', weights, gr, ao_L0)
                for k_id in range(self.num_kpts_loc):
                    kpt = self.cell.get_abs_kpts(self.kpt_latt_loc[k_id])
                    mo_included = self.mo_coeff_kpts[k_id][:,self.band_included_list]
                    s_kpt = self.cell.pbc_intor('int1e_ovlp', hermi=1, kpts=kpt, pbcopt=lib.c_null_ptr())
                    A_matrix_loc[k_id,ith_wann,:] = lib.einsum('v,vu,um->m', s_aoL0_g, s_kpt, mo_included,
                                                               optimize=True).conj()

        return A_matrix_loc

    def read_A_mat(self, filename=None):
        r'''
        Read the ovelap matrix: M_{m,n}^{(\mathbf{k,b})} from seedname.mnn
        '''
        if filename is None: filename = 'wannier90.amn'
        assert os.path.exists(filename), "Cannot find " + filename

        with open(filename, 'r') as f:
            data = f.readlines()
            num_bands_loc, num_kpts_loc, num_wann_loc = np.int64(data[1].split())
            data = data[2:]
            A_matrix_loc = np.empty([num_kpts_loc, num_wann_loc, num_bands_loc], dtype = np.complex128)
            val_in_float = np.float64(" ".join(data).split()).reshape(-1,5)
            A_matrix_loc = (val_in_float[:,3] + 1j * val_in_float[:,4]).reshape(num_kpts_loc, num_wann_loc, num_bands_loc)

        return A_matrix_loc

    def get_epsilon_mat(self):
        r'''
        Construct the eigenvalues matrix: \epsilon_{n}^(\mathbf{k})
        '''

        return np.asarray(self.mo_energy_kpts, dtype=np.float64)[:,self.band_included_list] * param.HARTREE2EV

    def read_epsilon_mat(self, filename=None):
        r'''
        Read the eigenvalues matrix: \epsilon_{n}^(\mathbf{k})
        '''
        if filename is None: filename = 'wannier90.eig'
        assert os.path.exists(filename), "Cannot find " + filename
        with open(filename, 'r') as f:
            data = f.read()
            temp = np.float64(data.split()).reshape(-1, 3)
            nbands = int(temp[:,0].max())
            nkpts = int(temp[:,1].max())
            eigenvals = temp[:,2].reshape(nkpts, nbands)

        return eigenvals

    def setup(self):
        '''
        Execute the Wannier90_setup
        '''

        real_lattice_loc = self.real_lattice_loc.T.flatten()
        recip_lattice_loc = self.recip_lattice_loc.T.flatten()
        kpt_latt_loc = self.kpt_latt_loc.flatten()
        atoms_cart_loc = self.atoms_cart_loc.flatten()

        (bands_wann_nntot, nn_list, proj_site, proj_l, proj_m, proj_radial,
         proj_z, proj_x, proj_zona, exclude_bands, proj_s, proj_s_qaxis) = \
                libwannier90.setup(self.mp_grid_loc, self.num_kpts_loc, real_lattice_loc,
                                   recip_lattice_loc, kpt_latt_loc,
                                   self.num_bands_tot, self.num_atoms_loc,
                                   self.atom_atomic_loc, atoms_cart_loc, self.gamma_only, self.spinors)

        # Convert outputs to the correct data type
        self.num_bands_loc, self.num_wann_loc, self.nntot_loc = np.int32(bands_wann_nntot)
        self.nn_list = np.int32(nn_list)
        self.proj_site = proj_site
        self.proj_l = np.int32(proj_l)
        self.proj_m = np.int32(proj_m)
        self.proj_radial = np.int32(proj_radial)
        self.proj_z = proj_z
        self.proj_x = proj_x
        self.proj_zona = proj_zona
        self.exclude_bands = np.int32(exclude_bands)
        self.band_included_list = [i for i in range(self.num_bands_tot) if (i + 1) not in self.exclude_bands]
        self.proj_s = np.int32(proj_s)
        self.proj_s_qaxis = proj_s_qaxis

    def run(self):
        '''
        Execute the Wannier90_run
        '''

        assert self.num_wann_loc is not None
        assert isinstance(self.M_matrix_loc, np.ndarray)
        assert isinstance(self.A_matrix_loc, np.ndarray)
        assert isinstance(self.eigenvalues_loc, np.ndarray)

        real_lattice_loc = self.real_lattice_loc.T.flatten()
        recip_lattice_loc = self.recip_lattice_loc.T.flatten()
        kpt_latt_loc = self.kpt_latt_loc.flatten()
        atoms_cart_loc = self.atoms_cart_loc.flatten()
        M_matrix_loc = self.M_matrix_loc.flatten()
        A_matrix_loc = self.A_matrix_loc.flatten()
        eigenvalues_loc = self.eigenvalues_loc.flatten()

        U_matrix, U_matrix_opt, lwindow, wann_centres, wann_spreads, spread = \
                libwannier90.run(self.mp_grid_loc, self.num_kpts_loc, real_lattice_loc,
                                 recip_lattice_loc, kpt_latt_loc, self.num_bands_loc,
                                 self.num_wann_loc, self.nntot_loc, self.num_atoms_loc,
                                 self.atom_atomic_loc, atoms_cart_loc, self.gamma_only,
                                 M_matrix_loc, A_matrix_loc, eigenvalues_loc)

        # Convert outputs to the correct data typ
        self.U_matrix = U_matrix
        self.U_matrix_opt = U_matrix_opt
        lwindow = np.int32(np.abs(lwindow.real))
        self.lwindow = (lwindow == 1)
        self.wann_centres = wann_centres.real
        self.wann_spreads = wann_spreads.real
        self.spread = spread.real

    get_wigner_seitz_supercell = get_wigner_seitz_supercell
    R_wz_sc = R_wz_sc
    ws_translate_dist = ws_translate_dist

    def get_hamiltonian_kpts(self):
        '''Get the Hamiltonian in k-space, this should be identical to Fock matrix from PySCF'''

        assert self.U_matrix is not None, "You must wannierize first, then you can run this function"
        eigenvals_in_window = []
        for k_id in range(self.num_kpts_loc):
            mo_included = self.mo_energy_kpts[k_id][self.band_included_list]
            orbs_in_win = self.lwindow[k_id]
            mo_in_window = mo_included[orbs_in_win]
            U_matrix_opt = self.U_matrix_opt[k_id][ :, orbs_in_win].T
            eigenvals = lib.einsum('m,mo,mo->o', mo_in_window, U_matrix_opt.conj(), U_matrix_opt)
            eigenvals_in_window.append(eigenvals)

        hamiltonian_kpts = lib.einsum('kso,ko,kto->kst', self.U_matrix.conj(), eigenvals_in_window, self.U_matrix)
        return hamiltonian_kpts

    def get_hamiltonian_Rs(self, Rs, ham_kpts=None):
        '''Get the R-space Hamiltonian H(R0, R) centered at R0 or the first R in Rs list
        '''

        assert self.U_matrix is not None, "You must wannierize first, then you can run this function"
        nkpts = self.kpt_latt_loc.shape[0]
        if ham_kpts is not None:
            hamiltonian_kpts = ham_kpts
        else:
            hamiltonian_kpts = self.get_hamiltonian_kpts()

        # Find the center either R(0,0,0) or the first R in the Rs list
        ngrid = len(Rs)
        center = np.arange(ngrid)[(np.asarray(Rs)**2).sum(axis=1) < 1e-10]
        if center.shape[0] == 1:
            center = center[0]
        else:
            center = 0

        # The phase factor is computed using the exp(1j*R.dot(k)) rather than exp(-1j*R.dot(k)) in wannier90
        phase = 1/np.sqrt(nkpts) * np.exp(1j* 2*np.pi * np.dot(Rs, self.kpt_latt_loc.T))
        hamiltonian_R0 = lib.einsum('k,kst,Rk->Rst', phase[center], hamiltonian_kpts, phase.conj())

        return hamiltonian_R0

    def interpolate_ham_kpts(self, frac_kpts, ham_kpts=None,
                             use_ws_distance=True, ws_search_size=[2,2,2],
                             ws_distance_tol=1e-6):
        ''' Interpolate the band structure using the Slater-Koster scheme
            Return:
                eigenvalues and eigenvectors at the desired kpts
        '''

        assert self.U_matrix is not None, "You must wannierize first, then you can run this function"
        ndegen, Rs, center = self.get_wigner_seitz_supercell(ws_search_size, ws_distance_tol)
        hamiltonian_R0 = self.get_hamiltonian_Rs(Rs, ham_kpts)

        # Interpolate H(kpts) at the desired k-pts
        if use_ws_distance:
            wdist_ndeg, wdist_ndeg_, irdist_ws, crdist_ws = self.ws_translate_dist(Rs)
            temp = lib.einsum('iRstx,kx->iRstk', irdist_ws, frac_kpts)
            phase = lib.einsum('iRstk,iRst->Rstk', np.exp(1j* 2*np.pi * temp), wdist_ndeg_)
            inter_hamiltonian_kpts = \
                    lib.einsum('R,Rst,Rts,Rstk->kst', 1/ndegen, 1/wdist_ndeg, hamiltonian_R0, phase)
        else:
            phase = np.exp(1j* 2*np.pi * np.dot(Rs, frac_kpts.T))
            inter_hamiltonian_kpts = \
                    lib.einsum('R,Rst,Rk->kst', 1/ndegen, hamiltonian_R0, phase)

        return inter_hamiltonian_kpts

    def interpolate_band(self, frac_kpts, ham_kpts=None, use_ws_distance=True,
                         ws_search_size=[2,2,2], ws_distance_tol=1e-6):
        ''' Interpolate the band structure using the Slater-Koster scheme
            Return:
                eigenvalues and eigenvectors at the desired kpts
        '''

        assert self.U_matrix is not None, (
            "You must wannierize first, then you can run this function")
        inter_hamiltonian_kpts = self.interpolate_ham_kpts(
            frac_kpts, ham_kpts, use_ws_distance, ws_search_size, ws_distance_tol)
        # Diagonalize H(kpts) to get eigenvalues and eigenvector
        nkpts = frac_kpts.shape[0]
        eigvals, eigvecs = np.linalg.eigh(inter_hamiltonian_kpts)
        idx_kpts = eigvals.argsort()
        eigvals = np.asarray([eigvals[kpt][idx_kpts[kpt]] for kpt in range(nkpts)])
        eigvecs = np.asarray([eigvecs[kpt][:,idx_kpts[kpt]] for kpt in range(nkpts)])

        return eigvals, eigvecs

    def is_real(self, threshold=1.e-6):
        '''
        Fourier transform the mo coefficients to real space and check if it is real
        '''

        assert self.U_matrix is not None, "You must wannierize first, then you can run this function"
        eigenvecs_in_window = []
        for k_id in range(self.num_kpts_loc):
            mo_included = self.mo_coeff_kpts[k_id][:,self.band_included_list]
            orbs_in_win = self.lwindow[k_id]
            mo_in_window = mo_included[:, orbs_in_win].dot(self.U_matrix_opt[k_id][ :, orbs_in_win].T)
            eigenvecs_in_window.append(mo_in_window)

        # Rotate the mo(kpts) into localized basis
        rotated_mo_coeff_kpts = lib.einsum('kum,ksm->kus', eigenvecs_in_window, self.U_matrix)

        # Fourier transform the localized mo
        nkx, nky, nkz = self.mp_grid_loc
        Ts = lib.cartesian_prod((np.arange(nkx), np.arange(nky), np.arange(nkz)))
        nkpts = self.kpt_latt_loc.shape[0]
        phase = 1/np.sqrt(nkpts) * np.exp(1j* 2*np.pi * np.dot(Ts, self.kpt_latt_loc.T))
        mo_coeff_Rs = lib.einsum('k,kus,Rk->Rus', phase[0], rotated_mo_coeff_kpts, phase.conj())

        return mo_coeff_Rs.imag.max() < threshold

    def export_unk(self, grid=[50,50,50]):
        '''
        Export the periodic part of BF in a real space grid for plotting with wannier90
        '''

        grids_coor, weights = periodic_grid(self.cell, grid, order = 'F')

        for k_id in range(self.num_kpts_loc):
            if self.spin_up:
                spin = '.1'
            else:
                spin = '.2'
            kpt = self.cell.get_abs_kpts(self.kpt_latt_loc[k_id])
            ao = numint.eval_ao(self.cell, grids_coor, kpt = kpt)
            u_ao = lib.einsum('x,xi->xi', np.exp(-1j*np.dot(grids_coor, kpt)), ao)
            unk_file = FortranFile('UNK' + "%05d" % (k_id + 1) + spin, 'w')
            unk_file.write_record(np.asarray([grid[0], grid[1], grid[2], k_id + 1, self.num_bands_loc], dtype = np.int32))
            mo_included = self.mo_coeff_kpts[k_id][:,self.band_included_list]
            u_mo = lib.einsum('xi,in->xn', u_ao, mo_included)
            for band in range(len(self.band_included_list)):
                unk_file.write_record(np.asarray(u_mo[:,band], dtype = np.complex128))
            unk_file.close()

    def export_AME(self, grid=[50,50,50]):
        r'''
        Export A_{m,n}^{\mathbf{k}} and M_{m,n}^{(\mathbf{k,b})} and \epsilon_{n}^(\mathbf{k})
        '''

        if self.A_matrix_loc is None:
            self.make_win()
            self.setup()
            self.M_matrix_loc = self.get_M_mat()
            self.A_matrix_loc = self.get_A_mat()
            self.eigenvalues_loc = self.get_epsilon_mat()
            self.export_unk(self, grid = grid)

        with open('wannier90.mmn', 'w') as f:
            f.write('Generated by the pyWannier90. Date: %s\n' % (time.ctime()))
            f.write('    %d    %d    %d\n' % (self.num_bands_loc, self.num_kpts_loc, self.nntot_loc))

            for k_id in range(self.num_kpts_loc):
                for nn in range(self.nntot_loc):
                    k_id1 = k_id + 1
                    k_id2 = self.nn_list[nn, k_id, 0]
                    nnn, nnm, nnl = self.nn_list[nn, k_id, 1:4]
                    f.write('    %d  %d    %d  %d  %d\n' % (k_id1, k_id2, nnn, nnm, nnl))
                    for m in range(self.num_bands_loc):
                        for n in range(self.num_bands_loc):
                            f.write('    %22.18e  %22.18e\n' % (self.M_matrix_loc[k_id, nn,m,n].real,
                                                                self.M_matrix_loc[k_id, nn,m,n].imag))

        with open('wannier90.amn', 'w') as f:
            f.write('Generated by the pyWannier90. Date: %s\n' % (time.ctime()))
            f.write('    %d    %d    %d\n' % (self.num_bands_loc, self.num_kpts_loc, self.num_wann_loc))

            for k_id in range(self.num_kpts_loc):
                for ith_wann in range(self.num_wann_loc):
                    for band in range(self.num_bands_loc):
                        f.write('    %d    %d    %d    %22.18e    %22.18e\n' %
                                (band+1, ith_wann+1, k_id+1,
                                 self.A_matrix_loc[k_id,ith_wann,band].real,
                                 self.A_matrix_loc[k_id,ith_wann,band].imag))

        with open('wannier90.eig', 'w') as f:
            for k_id in range(self.num_kpts_loc):
                for band in range(self.num_bands_loc):
                    f.write('    %d    %d    %22.18e\n' % (band+1, k_id+1, self.eigenvalues_loc[k_id,band]))

    def get_wannier(self, supercell=[1,1,1], grid=[50,50,50]):
        '''
        Evaluate the MLWF using a periodic grid
        '''

        grids_coor, weights = periodic_grid(self.cell, grid, supercell = [1,1,1], order = 'C')
        kpts = self.cell.get_abs_kpts(self.kpt_latt_loc)

        u_mo  = []
        for k_id in range(self.num_kpts_loc):
            mo_included = self.mo_coeff_kpts[k_id][:,self.band_included_list]
            mo_in_window = self.lwindow[k_id]
            C_opt = mo_included[:,mo_in_window].dot(self.U_matrix_opt[k_id][ :, mo_in_window].T)
            C_tildle = C_opt.dot(self.U_matrix[k_id].T)
            kpt = kpts[k_id]
            ao = numint.eval_ao(self.cell, grids_coor, kpt = kpt)
            u_ao = lib.einsum('x,xi->xi', np.exp(-1j*np.dot(grids_coor, kpt)), ao)
            u_mo.append(lib.einsum('xi,in->xn', u_ao, C_tildle))

        u_mo = np.asarray(u_mo)
        WF0 = libwannier90.get_WF0s(self.kpt_latt_loc.shape[0],self.kpt_latt_loc, supercell, grid, u_mo)

        # Fix the global phase following the pw2wannier90 procedure
        max_index = (WF0*WF0.conj()).real.argmax(axis=0)
        norm_wfs = np.diag(WF0[max_index,:])
        norm_wfs = norm_wfs/np.absolute(norm_wfs)
        WF0 = WF0/norm_wfs/self.num_kpts_loc

        # Check the 'reality' following the pw2wannier90 procedure
        for WF_id in range(self.num_wann_loc):
            ratio_max = np.abs(WF0[np.abs(WF0[:,WF_id].real) >= 0.01,WF_id].imag /
                               WF0[np.abs(WF0[:,WF_id].real) >= 0.01,WF_id].real).max(axis=0)
            print('The maximum imag/real for wannier function ', WF_id,' : ', ratio_max)
        return WF0

    def get_guess_orb(self, frac_site=[0,0,0], l=0, mr=1, r=1,
                      zona=1.0, x_axis=[1,0,0], z_axis=[0,0,1],
                      supercell=[1,1,1], grid=[50,50,50]):
        '''
        Evaluate a guess orbital using a periodic uniform grid
        '''
        grids_coor, weights = periodic_grid(self.cell, grid, supercell = supercell, order = 'C')
        frac_site = np.asarray(frac_site)
        abs_site = frac_site.dot(self.real_lattice_loc) / param.BOHR
        gr = g_r(grids_coor, abs_site, l, mr, r, zona, x_axis, z_axis, unit = 'B')
        return gr

    def plot_wf(self, outfile='MLWF', wf_list=None, supercell=[1,1,1], grid=[50,50,50]):
        '''
        Export Wannier function at cell R
        xsf format: http://web.mit.edu/xcrysden_v1.5.60/www/XCRYSDEN/doc/XSF.html
        Attributes:
            wf_list        : a list of MLWFs to plot
            supercell    : a supercell used for plotting
        '''

        if wf_list is None:
            wf_list = list(range(self.num_wann_loc))

        grid = np.asarray(grid)
        origin = np.asarray([-(grid[i]*(supercell[i]//2) + 1)/grid[i] for i in range(3)]).dot(
            self.cell.lattice_vectors().T)* param.BOHR
        real_lattice_loc = (grid*supercell-1)/grid * self.cell.lattice_vectors() * param.BOHR
        nx, ny, nz = grid*supercell
        WF0 = self.get_wannier(supercell = supercell, grid = grid)

        for wf_id in wf_list:
            assert wf_id in list(range(self.num_wann_loc))
            WF = WF0[:,wf_id].reshape(nx,ny,nz).real

            with open(outfile + '-' + str(wf_id) + '.xsf', 'w') as f:
                f.write('Generated by the pyWannier90. Date: %s\n\n' % (time.ctime()))
                f.write('CRYSTAL\n')
                f.write('PRIMVEC\n')
                for row in range(3):
                    f.write('%10.7f  %10.7f  %10.7f\n' %
                            (self.real_lattice_loc[row,0], self.real_lattice_loc[row,1],
                             self.real_lattice_loc[row,2]))
                f.write('CONVVEC\n')
                for row in range(3):
                    f.write('%10.7f  %10.7f  %10.7f\n' %
                            (self.real_lattice_loc[row,0], self.real_lattice_loc[row,1],
                             self.real_lattice_loc[row,2]))
                f.write('PRIMCOORD\n')
                f.write('%3d %3d\n' % (self.num_atoms_loc, 1))
                for atom in range(len(self.atom_symbols_loc)):
                    f.write('%s  %7.7f  %7.7f  %7.7f\n' %
                            (self.atom_symbols_loc[atom], self.atoms_cart_loc[atom][0],
                             self.atoms_cart_loc[atom][1], self.atoms_cart_loc[atom][2]))
                f.write('\n\n')
                f.write('BEGIN_BLOCK_DATAGRID_3D\n3D_field\nBEGIN_DATAGRID_3D_UNKNOWN\n')
                f.write('   %5d     %5d  %5d\n' % (nx, ny, nz))
                f.write('   %10.7f  %10.7f  %10.7f\n' % (origin[0],origin[1],origin[2]))
                for row in range(3):
                    f.write('   %10.7f  %10.7f  %10.7f\n' %
                            (real_lattice_loc[row,0], real_lattice_loc[row,1], real_lattice_loc[row,2]))

                fmt = ' %13.5e' * nx + '\n'
                for iz in range(nz):
                    for iy in range(ny):
                        f.write(fmt % tuple(WF[:,iy,iz].tolist()))
                f.write('END_DATAGRID_3D\nEND_BLOCK_DATAGRID_3D')

    def plot_guess_orbs(self, outfile='guess_orb', frac_site=[0,0,0], l=0, mr=1, r=1,
                        zona=1.0, x_axis=[1,0,0], z_axis=[0,0,1],
                        supercell=[1,1,1], grid=[50,50,50]):
        '''
        Export Wannier function at cell R
        xsf format: http://web.mit.edu/xcrysden_v1.5.60/www/XCRYSDEN/doc/XSF.html
        Attributes:
            wf_list        : a list of MLWFs to plot
            supercell    : a supercell used for plotting
        '''

        grid = np.asarray(grid)
        origin = np.asarray([-grid[i]*(supercell[i]//2)/grid[i] for i in range(3)]).dot(
            self.cell.lattice_vectors().T) * param.BOHR
        real_lattice_loc = (grid*supercell-1)/grid * self.cell.lattice_vectors() * param.BOHR
        nx, ny, nz = grid*supercell
        guess_orb = self.get_guess_orb(frac_site=frac_site, l=l, mr=mr, r=r,
                                       zona=zona, x_axis=x_axis, z_axis=z_axis,
                                       supercell=supercell, grid=grid)
        guess_orb = guess_orb.reshape(nx,ny,nz).real

        with open(outfile + '.xsf', 'w') as f:
            f.write('Generated by the pyWannier90\n\n')
            f.write('CRYSTAL\n')
            f.write('PRIMVEC\n')
            for row in range(3):
                f.write('%10.7f  %10.7f  %10.7f\n' % (self.real_lattice_loc[row,0], self.real_lattice_loc[row,1],
                                                      self.real_lattice_loc[row,2]))
            f.write('CONVVEC\n')
            for row in range(3):
                f.write('%10.7f  %10.7f  %10.7f\n' % (self.real_lattice_loc[row,0], self.real_lattice_loc[row,1],
                                                      self.real_lattice_loc[row,2]))
            f.write('PRIMCOORD\n')
            f.write('%3d %3d\n' % (self.num_atoms_loc, 1))
            for atom in range(len(self.atom_symbols_loc)):
                f.write('%s  %7.7f  %7.7f  %7.7f\n' % (self.atom_symbols_loc[atom], self.atoms_cart_loc[atom][0],
                                                       self.atoms_cart_loc[atom][1], self.atoms_cart_loc[atom][2]))
            f.write('\n\n')
            f.write('BEGIN_BLOCK_DATAGRID_3D\n3D_field\nBEGIN_DATAGRID_3D_UNKNOWN\n')
            f.write('   %5d     %5d  %5d\n' % (nx, ny, nz))
            f.write('   %10.7f  %10.7f  %10.7f\n' % (origin[0],origin[1],origin[2]))
            for row in range(3):
                f.write('   %10.7f  %10.7f  %10.7f\n' % (real_lattice_loc[row,0], real_lattice_loc[row,1],
                                                         real_lattice_loc[row,2]))

            fmt = ' %13.5e' * nx + '\n'
            for iz in range(nz):
                for iy in range(ny):
                    f.write(fmt % tuple(guess_orb[:,iy,iz].tolist()))
            f.write('END_DATAGRID_3D\nEND_BLOCK_DATAGRID_3D')

if __name__ == '__main__':
    from pyscf.pbc import gto as pgto
    from pyscf.pbc import scf as pscf
    import pywannier90

    # build cell object
    cell = pgto.Cell()
    cell.a = [[0.0, 2.7155, 2.7155], [2.7155, 0.0, 2.7155], [2.7155, 2.7155, 0.0]]
    cell.atom = [['Si',[0.0,0.0,0.0]], ['Si',[1.35775, 1.35775, 1.35775]]]
    cell.basis = 'gth-dzv'
    cell.pseudo = 'gth-pade'
    cell.exp_to_discard = 0.1
    cell.build()

    # build and run scf object
    kmesh = [3, 1, 1]
    kpts = cell.make_kpts(kmesh)
    kmf = pscf.KKS(cell, kpts)
    kmf.xc = 'pbe'
    kmf.run()

    # build and run w90 object
    num_wann = 8
    keywords = '''
    begin projections
    Si:sp3
    end projections
    '''
    w90 = pywannier90.W90(kmf, cell, kmesh, num_wann, other_keywords=keywords)
    w90.kernel()


# -*- coding: utf-8 -*-
##########################################################################
# NSAp - Copyright (C) CEA, 2016
# Distributed under the terms of the CeCILL-B license, as published by
# the CEA-CNRS-INRIA. Refer to the LICENSE file or to
# http://www.cecill.info/licences/Licence_CeCILL-B_V1-en.html
# for details.
##########################################################################

"""
Compute the connectome of a given parcellation, like the FreeSurfer aparc+aseg
segmentation, using MRtrix, FSL Probtrackx2 or MITK Gibbs Tracking.
"""

# Standard
import os
import subprocess

# Package
from pyconnectome import DEFAULT_FSL_PATH
from pyconnectome.tractography.probabilist import probtrackx2
from pyconnectome.utils.segtools import fix_freesurfer_subcortical_parcellation
from pyconnectome.utils.filetools import convert_mitk_vtk_fibers_to_tck
from pyconnectome.utils.filetools import convert_trk_fibers_to_tck
from pyconnectome.utils.filetools import convert_probtrackx2_saved_paths_to_tck
from pyconnectome.utils.regtools import freesurfer_bbregister_t1todif

# Third-party
import numpy
import nibabel
from pyfreesurfer import DEFAULT_FREESURFER_PATH
from pyfreesurfer.wrapper import FSWrapper
from pyfreesurfer.utils.filetools import (get_or_check_freesurfer_subjects_dir,
                                          load_look_up_table)

def connectome_snapshot(connectome, snapshot, labels=None, transform=None,
                        colorbar_title="", dpi=200, labels_size=4,
                        vmin=None, vmax=None):
    """
    Create a PNG snapshot of the connectome (i.e. connectivity matrix).

    Parameters
    ----------
    connectome: str
        Path to txt file storing the connectivity matrix.
    snapshot: str
        Path to the output snapshot.
    labels: str, default None
        Path to txt file listing the label names. By default no labels.
        Should be ordered like the rows/cols of the connectivity matrix.
    transform: callable, default None
        A Callable function to apply on the matrix (e.g. numpy.log1p).
        By default no transformation is applied.
    colorbar_title: str, default ""
        How to interpret the values of the connectivity,
        e.g. "Log(# of tracks)" or "% of tracks"
    dpi: int, default 200
        "Dot Per Inch", set higher for better resolution.
    labels_size: int, default 4
        The label font size.
    vmin, vmax: float, default None
        The display range.

    Returns
    -------
    snapshot: str
        Path to the output connectome snapshot.
    """
    # Import in function, so that the rest of the module can be used even
    # if matplotlib is not available
    import matplotlib.pyplot as plt
    # Load the connectivity matrix
    matrix = numpy.loadtxt(connectome)

    # Check connectivity matrix dimensions
    if matrix.ndim != 2 or matrix.shape[0] != matrix.shape[1]:
        raise ValueError("Connectivity matrix should be a square matrix."
                         "Shape of matrix: {}".format(matrix.shape))

    # Apply transformation if requested
    if transform is not None:
        matrix = transform(matrix)

    # -------------------------------------------------------------------------
    # Create the figure with matplotlib
    fig, ax = plt.subplots()
    ax.invert_yaxis()
    ax.xaxis.tick_top()

    ax.set_xticks(numpy.arange(0.5, matrix.shape[0]))
    ax.set_yticks(numpy.arange(0.5, matrix.shape[1]))

    ax.tick_params(which="both", axis="both", width=0, length=0)

    # Add the labels if passed
    if labels is not None:
        labels_array = numpy.loadtxt(labels, dtype=str)

        if len(labels_array) != matrix.shape[0]:
            raise ValueError(
                "Wrong number of labels: {}. Should be {}.".format(
                    len(labels_array), matrix.shape[0]))

        ax.set_xticklabels(labels_array, size=labels_size, rotation=90)
        ax.set_yticklabels(labels_array, size=labels_size)

    ax.set_aspect("equal")
    kwargs = {}
    if vmin is not None:
        kwargs["vmin"] = vmin
    if vmax is not None:
        kwargs["vmax"] = vmax
    heatmap = ax.pcolor(matrix, cmap=plt.cm.Reds, **kwargs)
    colorbar = fig.colorbar(heatmap)
    colorbar.set_label(colorbar_title, rotation=270, labelpad=20)

    fig.tight_layout()
    # -------------------------------------------------------------------------

    # Save to PNG file
    if not snapshot.endswith(".png"):
        snapshot += ".png"
    fig.savefig(snapshot, dpi=dpi)

    # Release memory
    fig.clear()
    plt.close()

    return snapshot

def voxel_to_node_connectivity(probtrackx2_dir, nodes, connectome_lut, outdir,
                               basename="connectome"):
    """
    When using the --omatrix3 option in Probtrackx2, the result is a
    VOXELxVOXEL connectivity matrix. This function creates the NODExNODE
    (i.e. ROIxROI) connectivity matrix for the given parcellation.

    Parameters
    ----------
    probtrackx2_dir: str
        Path to dir where to find the files created by probtrackx2 when using
        --omatrix3 option, i.e "fdt_matrix3.dot" and "coords_for_fdt_matrix3"
    nodes: str
        Path to parcellation defining the nodes of the connectome,
        e.g. FreeSurfer aparc+aseg parcellation with only the regions
        (i.e. labels) to keep in the connectome.
    connectome_lut: str
        Path to the Look Up Table of the given parcellation in the
        FreeSurfer LUT format.
    outdir: str
        Path to directory where output.
    basename: str, default "connectome"
        Basename of output files (<outdir>/<basename>.[mat|labels]).

    Returns
    -------
    out_connectome: str
        The generated connectome.
    out_labels: str
        The coonectome associated labels.
    """
    # Check input and output dirs
    for directory in (probtrackx2_dir, outdir):
        if not os.path.isdir(directory):
            raise ValueError("Directory does not exist: %s" % directory)

    # Coords: map x,y,z coordinates to voxel index
    # <x> <y> <z> <voxel index>
    path_coords = os.path.join(probtrackx2_dir, "coords_for_fdt_matrix3")
    coords = numpy.loadtxt(path_coords, dtype=int, usecols=[0, 1, 2, 4])

    # Load parcellation volume with node labels
    nodes_vol = nibabel.load(nodes).get_data().astype(dtype=int)

    # Load LUT to get the node names and set of int labels
    node_labels, node_names, _ = load_look_up_table(connectome_lut)
    set_labels = set(node_labels)

    # Connectivity matrix
    nb_nodes = len(node_names)
    connectome = numpy.zeros((nb_nodes, nb_nodes), dtype=int)

    # fdt_matrix3.dot: voxel to voxel connectivity
    # <index voxel 1> <index voxel 2> <nb connections>
    path_fdt_matrix3 = os.path.join(probtrackx2_dir, "fdt_matrix3.dot")

    # Since the fdt_matrix3.dot file can be very large, we parse it line by
    # line without loading it completely in memory
    with open(path_fdt_matrix3) as f:
        for line in f:

            # Get indexes of connected voxels and nb of connections
            v1_idx, v2_idx, nb_connections = map(int, line.strip().split())
            if nb_connections == 0:
                continue

            # Volume coordinates of connected voxels
            x1, y1, z1, _v1_idx = coords[v1_idx - 1, :]
            x2, y2, z2, _v2_idx = coords[v2_idx - 1, :]

            # Checking assumption that a voxel's info is stored at row=index-1
            assert v1_idx == _v1_idx, "%i %i" % (v1_idx, _v1_idx)
            assert v2_idx == _v2_idx, "%i %i" % (v2_idx, _v2_idx)

            # Labels of the 2 connected voxels
            label1, label2 = nodes_vol[x1, y1, z1], nodes_vol[x2, y2, z2]

            # Ignore pairs of voxels which labels are not in the connectome
            if not {label1, label2}.issubset(set_labels):
                continue

            # Update counts
            i, j = label1 - 1, label2 - 1  # 0-indexed in python
            connectome[i, j] += nb_connections
            connectome[j, i] += nb_connections

    # Write output connectome
    out_connectome = os.path.join(outdir, basename + ".mat")
    numpy.savetxt(out_connectome, connectome, fmt="%i")

    # Write nodes names
    out_labels = os.path.join(outdir, basename + ".labels")
    numpy.savetxt(out_labels, node_names, fmt="%s")

    return out_connectome, out_labels

def probtrackx2_connectome(
        outdir,
        tempdir,
        subject_id,
        t1_parc,
        t1_parc_lut,
        connectome_lut,
        nodif_brain,
        nodif_brain_mask,
        bedpostx_dir,
        nsamples,
        nsteps,
        steplength,
        fix_freesurfer_subcortical=False,
        subjects_dir=None,
        loopcheck=True,
        cthr=0.2,
        fibthresh=0.01,
        distthresh=0.0,
        sampvox=0.0,
        snapshots=True,
        fs_sh=DEFAULT_FREESURFER_PATH,
        fsl_sh=DEFAULT_FSL_PATH):
    """
    Compute the connectome of a given parcellation, like the FreeSurfer
    aparc+aseg segmentation, using ProbTrackx2.

    Requirements:
        - brain masks for the preprocessed DWI: nodif_brain and
          nodif_brain_mask.
        - FreeSurfer: result of recon-all on the T1.
        - FSL Bedpostx: computed for the preprocessed DWI.
        - a T1 parcellation that defines the nodes of the connectome, it has
          to be in the FreeSurfer space (i.e. aligned with
          <subjects dir>/<subject>/mri/brain.mgz), e.g. aparc+aseg from
          FreeSurfer.

    Connectome construction strategy:
        - Pathways are constructed from 'constitutive points' and not from
          endpoints. A pathway is the result of 2 samples propagating in
          opposite directions from a seed point. It is done using the
          --omatrix3 option of Probtrackx2.
        - The seed mask is the mask of WM voxels that are neighbors
          (12-connexity) of nodes.
        - The stop mask is the inverse of white matter, i.e. a sample stops
          propagating as soon as it leaves the white matter.

    Parameters
    ----------
    outdir: str
        Directory where to output.
    tempdir: str
        Path to the directory where temporary directories should be written.
        It should be a partition with 5+ GB available.
    subject_id: str
        Subject id used with FreeSurfer 'recon-all' command.
    t1_parc: str
        Path to the parcellation that defines the nodes of the connectome, e.g.
        aparc+aseg.mgz from FreeSurfer.
    t1_parc_lut: str
        Path to the Look Up Table for the passed parcellation in the
        FreeSurfer LUT format. If you T1 parcellation is from FreeSurfer, this
        will most likely be <$FREESURFER_HOME>/FreeSurferColorLUT.txt.
    connectome_lut: str
        Path to a Look Up Table in the FreeSurfer LUT format, listing the
        regions from the parcellation to use as nodes in the connectome. The
        region names should match the ones used in the <t1_parc_lut> LUT and
        the integer labels should be the row/col positions in the connectome.
        Alternatively it can be set to 'Lausanne2008' to use the predefined
        LUT for the Lausanne 2008 atlas, which is based on the FreeSurfer
        aparc+aseg parcellation.
    nodif_brain: str
        Path to the preprocessed brain-only DWI volume.
    nodif_brain_mask: str
        Path to the brain binary mask.
    bedpostx_dir: str
        Bedpostx output directory.
    nsamples: int
        Number of samples per voxel to initiate in the seed mask.
    nsteps: int
        Maximum number of steps for a given sample.
    steplength: int
        Step size in mm.
    fix_freesurfer_subcortical: bool, default False
        If the <t1_parc> is aparc+aseg or aparc.a2009s+aseg from FreeSurfer,
        set this option to True, to recompute the subcortical segmentations
        of the 5 structures that are uncorrectly segmented by FreeSurfer,
        using FSL FIRST.
    subjects_dir: str or None, default None
        Path to the FreeSurfer subjects directory. Required if the FreeSurfer
        environment variable (i.e. $SUBJECTS_DIR) is not set.
    cthr: float, optional
        Probtrackx2 option.
    fibthresh, distthresh, sampvox: float, optional
        Probtrackx2 options.
    loopcheck: bool, optional
        Probtrackx2 option.
    snapshots: bool, default True
        If True, create PNG snapshots for QC.
    fs_sh: str, default NeuroSpin path
        Path to the Bash script setting the FreeSurfer environment
    fsl_sh: str, default NeuroSpin path
        Path to the Bash script setting the FSL environment.

    Returns
    -------
    connectome_file: str
        The generated connectome.
    labels_file: str
        The coonectome associated labels.
    connectome_snap_file: str
        A grphical representation of the connectome.
    """
    # -------------------------------------------------------------------------
    # STEP 0 - Check arguments

    # FreeSurfer subjects_dir
    subjects_dir = get_or_check_freesurfer_subjects_dir(subjects_dir)

    if connectome_lut.lower() == "lausanne2008":
        module_dir = os.path.dirname(os.path.abspath(__file__))
        connectome_lut = os.path.join(module_dir, "Lausanne2008LUT.txt")

    # Check input paths
    paths_to_check = [t1_parc, t1_parc_lut, connectome_lut, nodif_brain,
                      nodif_brain_mask, bedpostx_dir, fs_sh, fsl_sh]
    for p in paths_to_check:
        if not os.path.exists(p):
            raise ValueError("File or directory does not exist: %s" % p)

    # Create <outdir> and/or <tempdir> if not existing
    for directory in [outdir, tempdir]:
        if not os.path.isdir(directory):
            os.makedirs(directory)

    # -------------------------------------------------------------------------
    # STEP 1 - Compute T1 <-> DWI rigid transformation

    # FreeSurfer T1 to Nifti
    fs_t1_brain = os.path.join(subjects_dir, subject_id, "mri", "brain.mgz")
    t1_brain = os.path.join(outdir, "t1_brain.nii.gz")
    cmd_1a = ["mri_convert", fs_t1_brain, t1_brain]
    FSWrapper(cmd_1a, shfile=fs_sh)()

    # Register diffusion to T1
    _, dif2anat_dat, dif2anat_mat = freesurfer_bbregister_t1todif(
            outdir=outdir,
            subject_id=subject_id,
            nodif_brain=nodif_brain,
            subjects_dir=subjects_dir,
            fs_sh=fs_sh,
            fsl_sh=fsl_sh)

    # Invert dif2anat transform
    m = numpy.loadtxt(dif2anat_mat)
    m_inv = numpy.linalg.inv(m)
    anat2dif_mat = os.path.join(outdir, "anat2dif.mat")
    numpy.savetxt(anat2dif_mat, m_inv)

    # -------------------------------------------------------------------------
    # STEP 2 - Convert LUT
    # Change integer labels in the LUT so that the each label corresponds
    # to the row/col position in the connectome
    nodes = os.path.join(outdir, "nodes.nii.gz")
    cmd_2 = ["labelconvert", t1_parc, t1_parc_lut, connectome_lut, nodes,
             "-nthreads", "0", "-failonwarn"]
    subprocess.check_call(cmd_2)

    # -------------------------------------------------------------------------
    # STEP 3 - If the T1 parcellation is aparc+aseg or aparc.a2009s+aseg
    # from FreeSurfer, this option allows the recompute the subcortical
    # segmentations of 5 structures that are uncorrectly segmented by
    # FreeSurfer, using FSL FIRST
    if fix_freesurfer_subcortical:
        fixed_nodes = os.path.join(outdir, "nodes_fixSGM.nii.gz")
        nodes = fix_freesurfer_subcortical_parcellation(parc=nodes,
                                                        t1_brain=t1_brain,
                                                        lut=connectome_lut,
                                                        output=fixed_nodes,
                                                        tempdir=tempdir,
                                                        nb_threads=0,
                                                        fsl_sh=fsl_sh)

    # -------------------------------------------------------------------------
    # STEP 4 - Create the masks for Probtrackx2

    # White matter mask
    aparc_aseg = os.path.join(subjects_dir, subject_id, "mri",
                              "aparc+aseg.mgz")
    wm_mask = os.path.join(outdir, "wm_mask.nii.gz")
    cmd_4a = ["mri_binarize",
              "--i", aparc_aseg,
              "--o", wm_mask,
              "--wm"]
    FSWrapper(cmd_4a, shfile=fs_sh)()

    # Stop mask is inverse of white matter mask
    stop_mask = os.path.join(outdir, "inv_wm_mask.nii.gz")
    cmd_4b = ["mri_binarize",
              "--i", aparc_aseg,
              "--o", stop_mask,
              "--wm", "--inv"]
    FSWrapper(cmd_4b, shfile=fs_sh)()

    # Create target mask: a mask of all nodes
    target_mask = os.path.join(outdir, "target_mask.nii.gz")
    cmd_4c = ["mri_binarize",
              "--i",   nodes,
              "--o",   target_mask,
              "--min", "1"]
    FSWrapper(cmd_4c, shfile=fs_sh)()

    # Dilate target mask by one voxel (12-connexity)
    target_mask_dil = os.path.join(outdir, "target_mask_dilated.nii.gz")
    cmd_4d = ["mri_morphology", target_mask, "dilate", "1", target_mask_dil]
    FSWrapper(cmd_4d, shfile=fs_sh)()

    # Create seed mask: white matter voxels near nodes (target regions)
    # Intersect dilated target mask and white matter mask
    # -> white matter voxels neighbor (12-connectivity) to node voxels
    seed_mask = os.path.join(outdir, "wm_nodes_interface_mask.nii.gz")
    cmd_4e = ["mri_and", wm_mask, target_mask_dil, seed_mask]
    FSWrapper(cmd_4e, shfile=fs_sh)()

    # -------------------------------------------------------------------------
    # STEP 7 - Run Probtrackx2
    probtrackx2_dir = os.path.join(outdir, "probtrackx2")
    probtrackx2(dir=probtrackx2_dir,
                forcedir=True,
                seedref=t1_brain,
                xfm=anat2dif_mat,
                invxfm=dif2anat_mat,
                samples=os.path.join(bedpostx_dir, "merged"),
                mask=nodif_brain_mask,
                seed=seed_mask,
                omatrix3=True,
                target3=nodes,
                stop=stop_mask,
                nsamples=nsamples,
                nsteps=nsteps,
                steplength=steplength,
                loopcheck=loopcheck,
                cthr=cthr,
                fibthresh=fibthresh,
                distthresh=distthresh,
                sampvox=sampvox,
                shfile=fsl_sh)

    # ------------------------------------------------------------------------
    # STEP 8 - Create NODExNODE connectivity matrix for nodes from <t1_parc>
    connectome_file, labels_file = voxel_to_node_connectivity(
        probtrackx2_dir=probtrackx2_dir,
        nodes=nodes,
        connectome_lut=connectome_lut,
        outdir=outdir)

    # ------------------------------------------------------------------------
    # STEP 9 - Create a connectome snapshot if requested
    if snapshots:
        connectome_snap_file = os.path.join(outdir, "connectome.png")
        connectome_snapshot(connectome_file, connectome_snap_file,
                            labels=labels_file, transform=numpy.log1p, dpi=300,
                            labels_size=4, colorbar_title="log(# of tracks)")

    return connectome_file, labels_file, connectome_snap_file

def mrtrix_connectomes(
        outdir,
        tempdir,
        tractogram,
        t1_brain,
        nodif_brain,
        t1_parc,
        t1_parc_lut,
        connectome_lut,
        tractogram_weights=None,
        tractogram_type="mrtrix",
        dif2anat_dat=None,
        dif2anat_mat=None,
        fix_freesurfer_subcortical=False,
        radial_search_dist=2.,
        forward_search_dist=5.,
        snapshots=True,
        fs_sh=DEFAULT_FREESURFER_PATH,
        fsl_sh=DEFAULT_FSL_PATH):
    """ Compute the reduced connectome from a parcellation using MRtrix.

    Parameters
    ----------
    outdir: str
        Directory where to output.
    tempdir: str
        Path to the directory where temporary directories should be written.
        It should be a partition with 5+ GB available.
    tractogram: str or list of str
        The tractogram to be used in VTK, TRK, TXT or TRK format. It is
        possible to provide a list of tractograms only for Connectomist and
        Tracula.
    t1_brain: str
        The anatomical image.
    nodif_brain: str, default None
        Diffusion brain-only Nifti volume with bvalue ~ 0. If not passed, it is
        generated automatically by averaging all the b0 volumes of the DWI.
    t1_parc: str
        Path to the parcellation that defines the nodes of the connectome, e.g.
        aparc+aseg.mgz from FreeSurfer in the 't1_brain' space.
    t1_parc_lut: str
        Path to the Look Up Table for the passed parcellation in the
        FreeSurfer LUT format. If you T1 parcellation is from FreeSurfer, this
        will most likely be <$FREESURFER_HOME>/FreeSurferColorLUT.txt.
    connectome_lut: str
        Path to a Look Up Table in the FreeSurfer LUT format, listing the
        regions from the parcellation to use as nodes in the connectome. The
        region names should match the ones used in the <t1_parc_lut> LUT and
        the integer labels should be the row/col positions in the connectome.
        Alternatively it can be set to 'Lausanne2008' to use the predefined
        LUT for the Lausanne 2008 atlas, which is based on the FreeSurfer
        aparc+aseg parcellation.
    tractogram_weights: str, default None
        The weight associated to each fiber: one weight per line.
    tractogram_type: str, default 'mrtrix'
        The software used to generate the tractogram. This parameter is used
        for format conversion purposes.
    dif2anat_dat: str, default None
        The diffusion to T1 FreeSurfer registration '.dat' file.
    dif2anat_mat: str
        The diffusion to T1 FSL registration '.mat' file.
    fix_freesurfer_subcortical: bool, default False
        If the <t1_parc> is aparc+aseg or aparc.a2009s+aseg from FreeSurfer,
        set this option to True, to recompute the subcortical segmentations
        of the 5 structures that are uncorrectly segmented by FreeSurfer,
        using FSL FIRST.
    radial_search_dist: float, default 2.0
        Multiple connectomes are generated depending on the streamline-node
        association strategy. The radial search assigns the nearest
        node from the streamline endpoint within this radius (in mm).
    forward_search_dist: float, default 5.0
        Multiple connectomes are generated depending on the streamline-node
        association strategy. The forward assignment projects the
        streamline forward from the endpoint to find a node, within this
        distance (in mm).
    snapshots: bool, default True
        If True, create PNG snapshots for QC.
    fs_sh: str, default NeuroSpin path
        Path to the Bash script setting the FreeSurfer environment
    fsl_sh: str, default NeuroSpin path
        Path to the Bash script setting the FSL environment.

    Returns
    -------
    connectome_endvox, connectome_radial, connectome_forward: str
        The generated reduced connectomes.
    """
    # -------------------------------------------------------------------------
    # STEP 0 - Check arguments

    # Get the default module LUT if parameter not provided
    if connectome_lut.lower() == "lausanne2008":
        module_dir = os.path.dirname(os.path.abspath(__file__))
        connectome_lut = os.path.join(module_dir, "Lausanne2008LUT.txt")

    # Check input paths
    paths_to_check = [t1_parc, t1_parc_lut, connectome_lut, t1_brain,
                      nodif_brain]
    paths_to_check.extend(tractogram)
    for p in paths_to_check:
        if not os.path.exists(p):
            raise ValueError("File or directory does not exist: %s" % p)

    # Check supported tractogram
    if tractogram_type not in ["mrtrix", "mitk", "connectomist", "fsl",
                               "tracula"]:
        raise ValueError("Unsupported tractogram: {0}".format(tractogram_type))

    # Create <outdir> and/or <tempdir> if not existing
    for directory in [outdir, tempdir]:
        if not os.path.isdir(directory):
            os.makedirs(directory)

    # -------------------------------------------------------------------------
    # STEP 1 - Align T1 parcellation to diffusion without downsampling
    parc_name = os.path.basename(t1_parc).split(".nii")[0].split(".mgz")[0]
    t1_parc_to_dif = os.path.join(outdir, parc_name + "_to_dif.nii.gz")
    if dif2anat_dat is not None:
        cmd_1c = ["mri_vol2vol",
                  "--mov",  nodif_brain,
                  "--targ", t1_parc,
                  "--inv",
                  "--no-resample",
                  "--interp", "nearest",
                  "--o",   t1_parc_to_dif,
                  "--reg", dif2anat_dat,
                  "--no-save-reg"]
        FSWrapper(cmd_1c, shfile=fs_sh)()
    elif dif2anat_mat is not None:
        raise NotImplementedError(
            "This code only supports 'dif2anat_dat' parameter.")
    else:
        raise ValueError("One transformation matrix is mandatory.")

    # -------------------------------------------------------------------------
    # STEP 2 - Convert LUT
    # Change integer labels in the LUT so that the each label corresponds
    # to the row/col position in the connectome
    nodes = os.path.join(outdir, "nodes.nii.gz")
    cmd_2 = ["labelconvert", t1_parc_to_dif, t1_parc_lut, connectome_lut,
             nodes, "-nthreads", "0", "-failonwarn"]
    subprocess.check_call(cmd_2)

    # -------------------------------------------------------------------------
    # STEP 3 - If the T1 parcellation is aparc+aseg or aparc.a2009s+aseg
    # from FreeSurfer, this option allows the recompute the subcortical
    # segmentations of 5 structures that are uncorrectly segmented by
    # FreeSurfer, using FSL FIRST
    if fix_freesurfer_subcortical:
        fixed_nodes = os.path.join(outdir, "nodes_fixSGM.nii.gz")
        nodes = fix_freesurfer_subcortical_parcellation(
            parc=nodes,
            t1_brain=t1_brain,
            lut=connectome_lut,
            output=fixed_nodes,
            tempdir=tempdir,
            nb_threads=0,
            fsl_sh=fsl_sh)

    # -------------------------------------------------------------------------
    # STEP 4 - Create connectomes with labels by combining fibers and nodes

    # Convert streamlines so that they can be used in MRtrix
    tck_tractogram = os.path.join(outdir, "fibers.tck")
    if tractogram_type == "mitk":
        if len(tractogram) != 1:
            raise ValueError("A one-file tractogram is expected.")
        convert_mitk_vtk_fibers_to_tck(tractogram[0], tck_tractogram)
    elif tractogram_type in ("connectomist", "tracula"):
        convert_trk_fibers_to_tck(
            nodif_brain, tractogram, tck_tractogram, tempdir)
    elif tractogram_type == "fsl":
        convert_probtrackx2_saved_paths_to_tck(
            nodif_brain, tractogram, tck_tractogram, tempdir)
    else:
        if len(tractogram) != 1:
            raise ValueError("A one-file tractogram is expected.")
        tck_tractogram = tractogram[0]

    # Read labels from LUT and create a list of labels: labels.txt
    labels_array = numpy.loadtxt(connectome_lut, dtype=str, usecols=[1])
    path_labels = os.path.join(outdir, "labels.txt")
    numpy.savetxt(path_labels, labels_array, fmt="%s")

    # Create connectome with end-voxel assignment
    connectome_endvox = os.path.join(outdir, "connectome_endvox.txt")
    cmd_8a = ["tck2connectome", tck_tractogram, nodes, connectome_endvox,
              "-keep_unassigned",  # Keep 'Unknown' label
              "-assignment_end_voxels"]
    if tractogram_weights is not None:
        cmd_8a += ["-tck_weights_in", tractogram_weights]
    subprocess.check_call(cmd_8a)

    # Create connectome with radial search assignment
    connectome_radial = os.path.join(
        outdir, "connectome_radial_{:.2f}mm.txt".format(radial_search_dist))
    cmd_8b = ["tck2connectome", tck_tractogram, nodes, connectome_radial,
              "-keep_unassigned",  # Keep 'Unknown' label
              "-assignment_radial_search", "{0}".format(radial_search_dist)]
    if tractogram_weights is not None:
        cmd_8b += ["-tck_weights_in", tractogram_weights]
    subprocess.check_call(cmd_8b)

    # Create connectome assigning the streamline to all nodes it intersects
    connectome_forward = os.path.join(
        outdir, "connectome_forward_{:.2f}mm.txt".format(forward_search_dist))
    cmd_8c = ["tck2connectome", tck_tractogram, nodes, connectome_forward,
              "-keep_unassigned",  # Keep 'Unknown' label
              "-assignment_forward_search", "{0}".format(forward_search_dist)]
    if tractogram_weights is not None:
        cmd_8c += ["-tck_weights_in", tractogram_weights]
    subprocess.check_call(cmd_8c)

    # If snapshots are requested
    if snapshots:
        snapshot_endvox = os.path.join(outdir, "connectome_endvox.png")
        connectome_snapshot(connectome_endvox, snapshot_endvox,
                            labels=path_labels, transform=numpy.log1p, dpi=300,
                            labels_size=4, colorbar_title="log(# of tracks)")

        snapshot_radial = os.path.join(
            outdir, "connectome_radial_{:.2f}mm.png".format(
                radial_search_dist))
        connectome_snapshot(connectome_radial, snapshot_radial,
                            labels=path_labels, transform=numpy.log1p, dpi=300,
                            labels_size=4, colorbar_title="log(# of tracks)")

        snapshot_forward = os.path.join(
            outdir, "connectome_forward_{:.2f}mm.png".format(
                forward_search_dist))
        connectome_snapshot(connectome_forward, snapshot_forward,
                            labels=path_labels, transform=numpy.log1p, dpi=300,
                            labels_size=4, colorbar_title="log(# of tracks)")

    return connectome_endvox, connectome_radial, connectome_forward


from typing import Any, Tuple
import pkgutil
import importlib

import numpy as np
import pandas as pd
import scipy
import scipy.stats as stats
from scipy.special import gamma, gammaln

from pvrpm.core.enums import ConfigKeys as ck

# override to getattr to get modules case insensitve
def getattr_override(obj: Any, attr: str) -> Any:
    for a in dir(obj):
        if a.lower() == attr.lower():
            return getattr(obj, a)

# TODO: there has to be a better way to do this...
def load_pysam_modules():
    """
    Loads ALL of PySAM's modules manually and globalizes them

    This is needed because PySAM is a wrapper for the ssc and sdk of SAM, which includes dynamic modules that are not properly defined for pybind, so using pkgutil's walk_packages function does not work (import error). Since the modules need to be loaded in order for getattr to find it, this must be done once when the program starts
    """
    global pysam
    import PySAM as pysam

    for loader, module_name, is_pkg in pkgutil.walk_packages(pysam.__path__):
        try:
            importlib.import_module(f"{pysam.__name__}.{module_name}")
        except:
            pass

def filename_to_module(filename: str) -> object:
    """
    Takes the filename of an exported json file from SAM, extracts the module name, and returns a callback to that module that can be used to create an object

    Args:
        filename (str): Filename of the exported case

    Returns:
        :obj:`PySAM`: PySAM object the file represents
    """
    # for certain modules the name from SAM doesnt match up with the module name (extra spaces in module name)
    broken_modules = ["host_developer"]
    for mod in broken_modules:
        if mod in filename:
            module_str = filename.strip().split("_")[-2:]
            module_str = "".join(module_str).split(".")[0].strip()
            return getattr_override(pysam, module_str)

    # if not a broken module:
    # SAM case file exporting should be underscores, with the last word being the module type
    module_str = filename.strip().split("_")[-1].split(".")[0].strip()
    return getattr_override(pysam, module_str)

def summarize_dc_energy(dc_power_output: tuple, split: int) -> np.array:
    """
    Calculates the DC energy (kWh) based on an input array of timeseries DC power (kW) for the system lifetime (likely the 'dc_net' output from SAM)

    Can be used to summarize similar hourly, daily data to yearly

    Args:
        dc_power_output (:obj:`tuple`): Tuple output from SAM simulation
        split (int): The frequency to split the data too, typically this is the number of years the system was simulated for (system_lifetime_yrs)

    Returns:
        :obj:`np.array`: Numpy array of length system_lifetime_yrs containing the yearly energy in kWh
    """
    data = np.array(dc_power_output)
    data = np.reshape(data, (int(split), int(len(dc_power_output) / split)))
    return np.sum(data, axis=1)

def component_degradation(percent_per_day: float, t: int) -> float:
    """
    Calculate the degradation of a component given the time since last replacement

    Args:
        percent_per_day (float): The percent degradation per day of the module
        t (int): Time since the module was last replaced, or if its a new module, installed

    Returns:
        float: The performance of the module, between 0 and 1

    Note:
        This gives the overall module performance based on degradation, so if the module has degraded 2 percent so far, this function returns 0.98
    """

    return 1 / np.power((1 + percent_per_day / 100), t)

def sample(distribution: str, parameters: dict, num_samples: int) -> np.array:
    """
    Sample data from a distribution. If distribution is a supported distribution, parameters should be a dictionary with keys "mean" and "std". Otherwise, distribution should be a scipy stats function and parameters be the kwargs for the distribution.

    Supported Distributions (only requires mean and std):
        - lognormal
        - normal
        - uniform (one std around mean)
        - weibull
        - exponential

    Args:
        distribution (str): Name of the distribution function
        parameters (:obj:`dict`): Kwargs for the distribution (for a supported distribution should only be the mean and std)
        num_samples (int): Number of samples to return from distribution

    Returns:
        :obj:(list): List of floats containing samples from the distribution
    """
    distribution = distribution.lower().strip()

    if distribution == "lognormal":
        # lognormal uses the mean and std of the underlying normal distribution of log(X)
        # so they must be normalized first
        mu, sigma = parameters[ck.MEAN], parameters[ck.STD]
        normalized_std = np.sqrt(np.log(1 + (sigma / mu) ** 2))
        normalized_mean = np.log(mu) - normalized_std ** 2 / 2
        dist = stats.lognorm(s=normalized_std, scale=np.exp(normalized_mean))
    elif distribution == "normal":
        dist = stats.norm(loc=parameters[ck.MEAN], scale=parameters[ck.STD])
    elif distribution == "uniform":
        a = parameters[ck.MEAN] - parameters[ck.STD]
        b = parameters[ck.STD] * 2
        dist = stats.uniform(loc=a, scale=b)
    elif distribution == "weibull":
        # for weibull, we have to solve for c and the scale parameter
        # this fails for certain parameter ranges, raising a runtime error
        # see https://github.com/scipy/scipy/issues/12134 for reference
        def _h(c):
            r = np.exp(gammaln(2 / c) - 2 * gammaln(1 / c))
            return np.sqrt(1 / (2 * c * r - 1))

        if ck.STD in parameters:
            mean, std = parameters[ck.MEAN], parameters[ck.STD]
            c0 = 1.27 * np.sqrt(mean / std)
            c, info, ier, msg = scipy.optimize.fsolve(lambda t: _h(t) - (mean / std), c0, xtol=1e-10, full_output=True,)

            # Test residual rather than error code.
            if np.abs(info["fvec"][0]) > 1e-8:
                raise RuntimeError(f"with mean={mean} and std={std}, solve failed: {msg}")

            c = c[0]
        else:
            mean, c = parameters[ck.MEAN], parameters[ck.SHAPE]

        scale = mean / gamma(1 + 1 / c)
        dist = stats.weibull_min(c=c, scale=scale)
    elif distribution == "exponential":
        dist = stats.expon(scale=parameters[ck.MEAN])
    else:
        # else, we don't know this distribution, pass the distribution directly to scipy
        dist = getattr(stats, distribution)
        if not dist:
            raise AttributeError(f"Scipy stats doesn't have a distribution '{distribution}'")
        dist = dist(**parameters)

    # scipy rvs uses rou sampling method
    return dist.rvs(size=num_samples)

def get_higher_components(
    top_level: str, start_level: str, case, start_level_df: pd.DataFrame = None,
) -> Tuple[np.array, np.array, int]:
    """
    Calculates the indicies of the top level that correspond to the given level df indicies and returns the given level indicies count per top level component and the total number of start_level components per top_level component

    Args:
        top_level (str): The string name of the component level to calculate indicies for
        start_level (str): The string name of the component level to start at
        case (SamCase): The case object for this simulation
        start_level_df (:obj:`pd.DataFrame`, Optional): The dataframe of the component level for which to find the corresponding top level indicies for

    Returns:
        tuple(:obj:`np.array`, :obj:`np.array`, int): If start_level_df is given, returns the top level indicies, the number of start_level components in start_level_df per top level index, and the total number of start_level components per top_level component. If start_level_df is None this only returns the total number of start_level components per top_level component.
    """
    # the number of disconnects equals the number of inverters, if this every changes this would need to be changed
    # otherwise, the inverter per trans is the same for disconnects
    # dictionaries to make is easier transitioning between levels
    component_per = {
        ck.TRANSFORMER: case.config[ck.INVERTER_PER_TRANS],
        ck.DISCONNECT: 1,
        ck.INVERTER: case.config[ck.COMBINER_PER_INVERTER],
        ck.COMBINER: case.config[ck.STR_PER_COMBINER],
        ck.STRING: case.config[ck.MODULES_PER_STR],
        ck.MODULE: 1,
    }

    # hierarchy of component levels:
    component_hier = {
        ck.MODULE: 0,
        ck.STRING: 1,
        ck.COMBINER: 2,
        ck.INVERTER: 3,
        ck.DISCONNECT: 4,
        ck.TRANSFORMER: 5,
    }

    above_levels = [
        c
        for c in component_hier.keys()
        if component_hier[c] > component_hier[start_level] and component_hier[c] <= component_hier[top_level]
    ]

    total_comp = 1
    # above levels is ordered ascending
    for level in above_levels:
        total_comp *= component_per[level]

    if start_level_df is not None:
        indicies = start_level_df.index.copy()
        indicies = np.floor(indicies / total_comp)
        # sum up the number of occurences for each index and return with total number of components at start level per top level
        indicies, counts = np.unique(indicies, return_counts=True)
        return indicies.astype(np.int64), counts.astype(np.int64), int(total_comp)
    else:
        return int(total_comp)