# AUTOGENERATED! DO NOT EDIT! File to edit: StructureIO.ipynb (unless otherwise specified).

__all__ = ['Arrow3D', 'fancy_quiver3d', 'export_poscar', 'InvokeMaterialsProject', 'get_kpath', 'read_ticks',
           'str2kpath', 'get_kmesh', 'order', 'rotation', 'get_bz', 'splot_bz', 'iplot_bz', 'to_R3', 'kpoints2bz',
           'fix_sites', 'get_pairs', 'iplot_lat', 'splot_lat', 'join_poscars', 'write_poscar', 'scale_poscar',
           'rotate_poscar']

# Cell
import os, re
import json
import math
import numpy as np
from pathlib import Path
import requests as req
from collections import namedtuple
from itertools import product, permutations
from functools import lru_cache

from scipy.spatial import ConvexHull, Voronoi, KDTree
import plotly.graph_objects as go
import matplotlib.pyplot as plt
import matplotlib.colors as mplc #For viewpoint
from mpl_toolkits import mplot3d
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
# Inside packages import to work both with package and jupyter notebook.
try:
    from pivotpy import vr_parser as vp
    from pivotpy import s_plots as sp
except:
    import pivotpy.vr_parser as vp
    import pivotpy.s_plots as sp

# Cell
from matplotlib.patches import FancyArrowPatch
from mpl_toolkits.mplot3d import proj3d

class Arrow3D(FancyArrowPatch):
    """Draw 3D fancy arrow."""
    def __init__(self, x, y, z, u, v, w, *args, **kwargs):
        FancyArrowPatch.__init__(self, (0, 0), (0, 0), *args, **kwargs)
        self._verts3d = [x,x+u], [y,y+v], [z,z+w]

    def draw(self, renderer):
        xs3d, ys3d, zs3d = self._verts3d
        xs, ys, zs = proj3d.proj_transform(xs3d, ys3d, zs3d, renderer.M)
        self.set_positions((xs[0], ys[0]), (xs[1], ys[1]))
        FancyArrowPatch.draw(self, renderer)

    def on(self,ax):
        ax.add_artist(self)

def fancy_quiver3d(X,Y,Z,U,V,W,ax=None,C = 'r',L = 0.7,mutation_scale=10,**kwargs):
    """Plots 3D arrows on a given ax. See [FancyArrowPatch](https://matplotlib.org/3.1.1/api/_as_gen/matplotlib.patches.FancyArrowPatch.html).
    - **Parameters**
        - X, Y, Z : 1D arrays of coordinates of arrows' tail point.
        - U, V, W : 1D arrays of dx,dy,dz of arrows.
        - ax: 3D axes, if not given, auto created.
        - C : 1D colors array mapping for arrows. Could be one color.
        - L : 1D linwidths array mapping for arrows. Could be one linewidth.
        - mutation_scale: Arrow head width/size scale. Default is 10.
        - kwargs: FancyArrowPatch's keyword arguments excluding positions,color, lw and mutation_scale, shrinkA, shrinkB which are already used. An important keyword argument is `arrowstyle` which could be '->','-|>', their inverted forms and many more. See on matplotlib.
    """
    if not ax:
        ax = sp.get_axes(figsize=(3.4,3.4),axes_3d=True) # Same aspect ratio.
    if not isinstance(C,(list,np.ndarray)):
        C = [[*mplc.to_rgb(C)] for x in X]
    C = np.array(C) # Safe for list

    if not isinstance(L,(list,np.ndarray)):
        L = [L for x in X]
    args_dict = dict(mutation_scale=mutation_scale,shrinkA=0, shrinkB=0)
    for x,y,z,u,v,w,c,l in zip(X,Y,Z,U,V,W,C,L):
        Arrow3D(x, y, z, u, v, w, color=c,lw=l,**args_dict,**kwargs).on(ax)

    return ax

# Cell
def export_poscar(path=None,content=None):
    """Export POSCAR file to python objects.
    - **Parameters**
        - path: Path/to/POSCAR file. Auto picks in CWD.
        - contet: POSCAR content as string, This takes precedence to path.
    """
    if content and isinstance(content,str):
        path = content.splitlines() # This acts as slice for islice2array.
    elif not path:
        path = './POSCAR'
        if not os.path.isfile(path):
            raise FileNotFoundError

    SYSTEM = vp.islice2array(path,start=0,nlines=1,raw=True,exclude=None).split()[0]
    scale = vp.islice2array(path,start=1,nlines=1,exclude=None).squeeze()
    basis = scale*vp.islice2array(path,start=2,nlines=3,exclude=None).reshape((-1,3))
    volume = np.linalg.det(basis)
    rec_basis = np.linalg.inv(basis).T # general formula
    out_dict = {'SYSTEM':SYSTEM,'volume':volume,'basis':basis,'rec_basis':rec_basis}

    elems = vp.islice2array(path,raw=True,start=5,nlines=1,exclude=None).split()
    ions = vp.islice2array(path,start=6,nlines=1,exclude=None)
    N = np.sum(ions).astype(int)
    inds = np.cumsum([0,*ions]).astype(int)
    # Two lines are excluded in below command before start. so start = 7-2
    positions = vp.islice2array(path,start=5,exclude="^\s+[a-zA-Z]|^[a-zA-Z]",cols=[0,1,2]).reshape((-1,3))[:N]


    unique_d = {}
    for i,e in enumerate(elems):
        unique_d.update({e:range(inds[i],inds[i+1])})

    elem_labels = []
    for i, name in enumerate(elems):
        for ind in range(inds[i],inds[i+1]):
            elem_labels.append(f"{name} {str(ind - inds[i] + 1)}")
    out_dict.update({'positions':positions,'labels':elem_labels,'unique':unique_d})
    return vp.Dict2Data(out_dict)


# Cell
def _save_mp_API(api_key):
    """
    - Save materials project api key for autoload in functions.
    """
    home = str(Path.home())
    file = os.path.join(home,'.pivotpyrc')
    lines = []
    if os.path.isfile(file):
        with open(file,'r') as fr:
            lines = fr.readlines()
            lines = [line for line in lines if 'MP_API_KEY' not in line]

    with open(file,'w') as fw:
        fw.write("MP_API_KEY = {}".format(api_key))
        for line in lines:
            fw.write(line)

# Cell
def _load_mp_data(formula,api_key=None,mp_id=None,max_sites = None, min_sites = None):
    """
    - Returns fetched data using request api of python form materials project website.
    - **Parameters**
        - formula  : Material formula such as 'NaCl'.
        - api_key  : API key for your account from material project site. Auto picks if you already used `_save_mp_API` function.
        - mp_id     : Optional, you can specify material ID to filter results.
        - max_sites : Maximum number of sites. If None, sets `min_sites + 1`, if `min_sites = None`, gets all data.
        - min_sites : Minimum number of sites. If None, sets `max_sites + 1`, if `max_sites = None`, gets all data.
    """
    if api_key is None:
        try:
            home = str(Path.home())
            file = os.path.join(home,'.pivotpyrc')
            with open(file,'r') as f:
                lines=f.readlines()
                for line in lines:
                    if 'MP_API_KEY' in line:
                        api_key = line.split('=')[1].strip()
        except:
            return print("api_key not given. provide in argument or generate in file using `_save_mp_API(your_mp_api_key)")

    #url must be a raw string
    url = r"https://www.materialsproject.org/rest/v2/materials/{}/vasp?API_KEY={}".format(formula,api_key)
    resp = req.request(method='GET',url=url)
    jl = json.loads(resp.text)

    try: jl['response'] #check if response
    except: return print("Either formula {!r} or API_KEY is incorrect.".format(formula))

    all_res = jl['response']

    if max_sites != None and min_sites != None:
        lower, upper = min_sites, max_sites
    elif max_sites == None and min_sites != None:
        lower, upper = min_sites, min_sites + 1
    elif max_sites != None and min_sites == None:
        lower, upper = max_sites - 1, max_sites
    else:
        lower, upper = '-1', '-1' # Unknown

    if lower != '-1' and upper != '-1':
        sel_res=[]
        for res in all_res:
            if res['nsites'] <= upper and res['nsites'] >= lower:
                sel_res.append(res)
        return sel_res
    # Filter to mp_id at last. more preferred
    if mp_id !=None:
        for res in all_res:
            if mp_id == res['material_id']:
                return [res]
    return all_res

# Cell
class InvokeMaterialsProject:
    """Connect to materials project and get data using api_key from their site.
    Usage:
        >>> from pivotpyr.sio import InvokeMaterialsProject # or import pivotpy.InvokeMaterialsProject as InvokeMaterialsProject
        >>> mp = InvokeMaterialsProject(api_key='your_api_key')
        >>> mp.request(formula='NaCl') #returns nothing but saves response
        >>> mp.poscars #returns poscars data
        >>> mp.cifs #returns cifs data
    """
    def __init__(self,api_key=None):
        "Request Materials Project acess. api_key is on their site. Your only need once and it is saved for later."
        self.api_key = api_key
        self.__cifs = None
        self.__poscars = None
        self.__response = None
        self.success = False

    def save_api_key(self,api_key):
        "Save api_key for auto reloading later."
        _save_mp_API(api_key)

    @property
    def poscars(self):
        "Get all loaded poscars as list. Any item in list can be written to file using `item.write`"
        if self.__poscars is None: # empty means still got it
            raise Exception("No poscars loaded. Use `request` function to load data and then access.")
        return self.__poscars

    @property
    def cifs(self):
        "Get all loaded cifs as list. Any item in list can be written to file using `item.write`"
        if self.__cifs is None: # empty means still got it
            raise Exception("No cifs loaded. Use `request` function to load data and then access.")
        return self.__cifs

    @lru_cache(maxsize=2) #cache for 2 calls
    def request(self,formula,mp_id=None,max_sites = None,min_sites=None):
        """Fetch data using request api of python form materials project website. After request, you can access `cifs` and `poscars`.
        - **Parameters**
            - formula  : Material formula such as 'NaCl'.
            - mp_id     : Optional, you can specify material ID to filter results.
            - max_sites : Maximum number of sites. If None, sets `min_sites + 1`, if `min_sites = None`, gets all data.
            - min_sites : Minimum number of sites. If None, sets `max_sites + 1`, if `max_sites = None`, gets all data.
        """
        try:
            self.__response = _load_mp_data(formula = formula,api_key = self.api_key, mp_id = mp_id, max_sites = max_sites,min_sites=min_sites)
        except:
            return req.HTTPError("Error in request. Check your api_key or formula.")
        # get cifs
        cifs = []
        for res in self.__response:
            cif     = res['cif']
            symbol  = res['spacegroup']['symbol']
            crystal = res['spacegroup']['crystal_system']
            unit    = res['unit_cell_formula']
            mp_id   = res['material_id']
            crs = dict(mp_id = mp_id,symbol = symbol, crystal = crystal, unit = unit, content = cif)
            cifs.append(vp.Dict2Data(crs))
        self.__cifs = cifs

        # get poscars
        poscars = []
        for cr in self.__cifs:
            cif = cr.content
            lines = cif.split('\n')
            if '' in lines.copy():
                lines.remove('')
            abc = []
            abc_ang = []
            index = 0
            for ys in lines:
                if '_cell' in ys:
                    if '_length' in ys:
                        abc.append(ys.split()[1])
                    if '_angle' in ys:
                        abc_ang.append(ys.split()[1])
                    if '_volume' in ys:
                        volume = float(ys.split()[1])
                if '_structural' in ys:
                    system = ys.split()[1]
            for i,ys in enumerate(lines):
                if '_atom_site_occupancy' in ys:
                    index = i +1 # start collecting pos.
            poses = lines[index:]
            pos_str = ""
            for pos in poses:
                s_p = pos.split()
                pos_str += "{0:>12}  {1:>12}  {2:>12}  {3}\n".format(*s_p[3:6],s_p[0])

            # ======== Cleaning ===========
            abc_ang = [float(ang) for ang in abc_ang]
            abc     = [float(a) for a in abc]
            a = "{0:>22.16f} {1:>22.16f} {2:>22.16f}".format(1.0,0.0,0.0) # lattic vector a.
            to_rad = 0.017453292519
            gamma = abc_ang[2]*to_rad
            bx,by = abc[1]*np.cos(gamma),abc[1]*np.sin(gamma)
            b = "{0:>22.16f} {1:>22.16f} {2:>22.16f}".format(bx/abc[0],by/abc[0],0.0) # lattic vector b.
            cz = volume/(abc[0]*by)
            cx = abc[2]*np.cos(abc_ang[1]*to_rad)
            cy = (abc[1]*abc[2]*np.cos(abc_ang[0]*to_rad)-bx*cx)/by
            c = "{0:>22.16f} {1:>22.16f} {2:>22.16f}".format(cx/abc[0],cy/abc[0],cz/abc[0]) # lattic vector b.
            unit = cr.unit.to_dict()
            elems = [elem for elem in unit.keys()]
            elems = '\t'.join(elems)
            nums  = [str(int(unit[elem])) for elem in unit.keys()]
            nums  = '\t'.join(nums)
            # Get others important info
            symbol  = cr.symbol
            crystal = cr.crystal
            mp_id   = cr.mp_id
            # ======================
            top = system + "\t# [{0}] Generated by PivotPy using Materials Project Database.".format(symbol)
            content = "{0}\n  {1}\n {2}\n {3}\n {4}\n  {5}\n  {6}\nDirect\n{7}".format(top,abc[0],a,b,c,elems,nums,pos_str)
            # =======================
            net_out = dict(system = system, symbol = symbol, crystal =crystal,mp_id = mp_id,content = content)
            poscars.append(vp.Dict2Data(net_out))
        self.__poscars = poscars

        # Add write method to cifs and poscars
        def write(dd_instance,filename):
            "Write current poscar/cif string to file"
            with open(filename,'w') as f:
                f.write(dd_instance.content)

        for pos in [*self.__poscars, *self.__cifs]:
            pos.write = write.__get__(pos,type(pos)) # Attch method to poscar and cifs

        self.success = True # set success flag


# Cell
def get_kpath(hsk_list=[],labels=[], n = 5,weight= None ,ibzkpt = None,outfile=None):
    """
    Generate list of kpoints along high symmetry path. Options are write to file or return KPOINTS list.
    It generates uniformly spaced point with input `n` as just a scale factor of number of points per unit length.
    You can also specify custom number of kpoints in an interval by putting number of kpoints as 4th entry in left kpoint.
    - **Parameters**
        - hsk_list : N x 3 list of N high symmetry points, if broken path then [[N x 3],[M x 3],...].
                    Optionally you can put a 4 values point where 4th entry will decide number of kpoints in current interval.
                    Make sure that points in a connected path patch are at least two i.e. `[[x1,y1,z1],[x2,y2,z2]]` or `[[x1,y1,z1,N],[x2,y2,z2]]`.
        - n        : int, number per unit length, this makes uniform steps based on distance between points.
        - weight : Float, if None, auto generates weights.
        - ibzkpt : Path to ibzkpt file, required for HSE calculations.
        - labels : Hight symmetry points labels. Good for keeping record of lables and points indices for later use.
                    > Note: If you do not want to label a point, label it as 'skip' at its index and it will be removed.
        - outfile: Path/to/file to write kpoints.

    If `outfile = None`, KPONITS file content is printed.
    """
    if hsk_list:
        try: hsk_list[0][0][0]
        except: hsk_list = [hsk_list] # Make overall 3 dimensions to include breaks in path
        if not labels:
            ['' for ks in hsk_list for k in ks] #make flatten empty labels
    else: return print('Give at least non empty first argument.')

    xs,ys,zs, inds,joinat = [],[],[],[0],[] # 0 in inds list is important
    _labels = []
    _len_prev = 0
    for j,a in enumerate(hsk_list):
        for i in range(len(a)-1):
            _vec = [_a-_b for _a,_b in zip(a[i][:3],a[i+1] )] # restruct point if 4 entries
            _m = np.rint(np.linalg.norm(_vec)*n).astype(int) # Calculate
            try: _m = a[i][3] # number of points given explicitly.
            except: pass

            inds.append(inds[-1]+_m) #Append first then do next
            _labels.append(labels[i+_len_prev])
            if j !=0 and i == 0:
                joinat.append(inds[-2]) # Add previous in joinpath and label
                if 'skip' not in _labels[-2]:
                    _labels[-1] = _labels[-2] + '|' + _labels[-1]
                    _labels = [*_labels[:-2],_labels[-1]] # Drop the label we added before

            xs.append(list(np.linspace(a[i][0],a[i+1][0],_m)))
            ys.append(list(np.linspace(a[i][1],a[i+1][1],_m)))
            zs.append(list(np.linspace(a[i][2],a[i+1][2],_m)))

        _labels.append(labels[len(a) -1 +_len_prev]) # Add last in current interval
        _len_prev += len(a)

    xs = [y for z in xs for y in z] #flatten values.
    ys = [y for z in ys for y in z]
    zs = [y for z in zs for y in z]

    if weight == None and xs:
        weight = 1/len(xs)

    out_str = ["{0:>16.10f}{1:>16.10f}{2:>16.10f}{3:>12.6f}".format(x,y,z,weight) for x,y,z in zip(xs,ys,zs)]
    out_str = '\n'.join(out_str)
    N = np.size(xs)
    if ibzkpt != None:
        if os.path.isfile(ibzkpt):
            with open(ibzkpt,'r') as f:
                lines = f.readlines()

            N = int(lines[1].strip())+N # Update N.
            slines = lines[3:N+4]
            ibz_str = ''.join(slines)
            out_str = "{}\n{}".format(ibz_str,out_str) # Update out_str
    if inds:
        inds[-1] = -1 # last index to -1

    inds = [i for k,i in enumerate(inds) if 'skip' != _labels[k]]
    _labels = [l.replace('|skip','') for l in _labels if l != 'skip']
    top_str = "Automatically generated using PivotPy with HSK-INDS = {}, LABELS = {}, SEG-INDS = {}\n\t{}\nReciprocal Lattice".format(inds,_labels,joinat,N)
    out_str = "{}\n{}".format(top_str,out_str)
    if outfile != None:
        with open(outfile,'w') as f:
            f.write(out_str)
    else:
        print(out_str)

def read_ticks(kpoints_file_path):
    "Reads ticks values and labels in header of kpoint file. Returns dictionary of `ktick_inds`,`ktick_vals`,`kseg_inds` that can be unpacked to plotting functions. If not exist in header, returns empty values(still valid)."
    out_dict = dict(ktick_inds=[],ktick_vals=[],kseg_inds=[])
    if os.path.isfile(kpoints_file_path):
        head = vp.islice2array(kpoints_file_path,exclude=None,raw=True,nlines=1)
        if 'HSK-INDS' in head:
            hsk = head.split('HSK-INDS')[1].split(']')[0].split('[')[1].split(',')
            out_dict['ktick_inds'] = [int(h) for h in hsk if h]
        if 'LABELS' in head:
            labs = head.split('LABELS')[1].split(']')[0].split('[')[1].split(',')
            out_dict['ktick_vals'] = [l.replace("'","").replace('"','').strip() for l in labs if l]
        if 'SEG-INDS' in head:
            segs = head.split('SEG-INDS')[1].split(']')[0].split('[')[1].split(',')
            out_dict['kseg_inds'] = [int(s) for s in segs if s]
    return out_dict

# Cell
def str2kpath(kpath_str,n = 5, weight = None, ibzkpt = None, outfile = None):
    """Get Kpath from a string of kpoints (Line-Mode like). Useful in Terminal.
    - **Parameters**
        - kpath_str: str, a multiline string similiar to line mode of KPOINTS, initial 4 lines are not required.
            - If you do not want to label a point, label it as 'skip' and it will be removed.
            - You can add an interger at end of a line to customize number of points in a given patch.
            - Each empty line breaks the path, so similar points before and after empty line are useless here.
        - n      : int, number per unit length, this makes uniform steps based on distance between points.
        - weight : Float, if None, auto generates weights.
        - ibzkpt : Path to ibzkpt file, required for HSE calculations.
        - outfile: Path/to/file to write kpoints.

    - **Example**
        > str2kpath('''0 0 0 !$\Gamma$ 3
                    0.25 0.25 0.25 !L''')
        > Automatically generated using PivotPy with HSK-INDS = [0, -1], LABELS = ['$\\Gamma$', 'L'], SEG-INDS = []
	    >   3
        > Reciprocal Lattice
        >   0.0000000000    0.0000000000    0.0000000000    0.333333
        >   0.1250000000    0.1250000000    0.1250000000    0.333333
        >   0.2500000000    0.2500000000    0.2500000000    0.333333
    """
    lines = kpath_str.splitlines()

    skipN = 0
    for _n, line in enumerate(lines[:6]): #Handle strings from AFLOW-like softwares
        if line.strip().isalpha():
            skipN = _n + 1

    where_blanks = [i for i,line in enumerate(lines) if line.strip() == '']

    hsk_list, labels = [],[]
    for j,line in enumerate(lines[skipN:]):
        if line.strip():
            _labs = re.findall('\$\\\\[a-zA-Z]+\$|[a-zA-Z]+|[α-ωΑ-Ω]+|\|', line)
            labels.append(_labs[0] if _labs else '')
            _ks = re.findall('[-+\d]*[.][\d+]+|[-+]*\d+/\d+|[-+]*\d+',line)
            _ks = [[float(k) for k in w.split('/')] if '/' in w else float(w) for w in _ks]

            for i,_k in enumerate(_ks):
                if type(_k) == list and len(_k) == 2:
                    _ks[i] = _k[0]/_k[1]
                elif type(_k) == list and len(_k) != 2:
                    print(f'Check if you provide fraction correctly in line {j+1+skipN}!')

            if len(_ks) == 4:
                try: _ks[3] = int(_ks[3])
                except: print(f'4th number in line {j+1+skipN} should be integer!')

            hsk_list.append(_ks)

    if where_blanks:
        filtered = [w-i for i,w in enumerate(where_blanks)]
        where_blanks = np.unique([0,*filtered,len(hsk_list)]).tolist()

    new_list = [] if where_blanks else hsk_list #Fix up both
    for a,b in zip(where_blanks[:-1], where_blanks[1:]):
        if b - a < 2:
            return print(f"There should be at least two points in a patch of path at line {a+1}!")
        new_list.append(hsk_list[a:b])

    return get_kpath(hsk_list = new_list, labels = labels,n=n,weight=weight,ibzkpt=ibzkpt,outfile=outfile)

# Cell
def _get_basis(path_pos):
    """Returns given(computed) and inverted(without 2π) basis as tuple(given,inverted).
    - **Parameters**
        - path_pos: path/to/POSCAR or 3 given vectors as rows of a matrix."""
    if isinstance(path_pos,(list,np.ndarray)) and np.ndim(path_pos) ==2:
        basis = np.array(path_pos)
    elif isinstance(path_pos,str) or isinstance(path_pos,type(None)):
        basis = export_poscar(path_pos).basis
    else:
        raise FileNotFoundError("{!r} does not exist or not 3 by 3 list.".format(path_pos))
    # Process. 2π is not included in vasp code.
    rec_basis = np.linalg.inv(basis).T # Compact Formula
    Basis = namedtuple('Basis', ['given', 'inverted'])
    return Basis(basis,rec_basis)

# Cell
def get_kmesh(n_xyz=[5,5,5],weight = None, ibzkpt= None,path_pos=None,outfile=None):
    """
    - Generates uniform mesh of kpoints. Options are write to file, or return KPOINTS list.
    - **Parameters**
        - n_xyz  : List of [nx ny nz] or integer. If integer given, it represents numbers of kpoints along smallest side in reciprocal space and kmesh is autoscaled.
        - weight : Float, if None, auto generates weights.
        - ibzkpt : Path to ibzkpt file, required for HSE calculations.
        - path_pos : POSCAR file path or real space lattice vectors, if None, cubic symmetry is used and it is fast.
        - outfile: Path/to/file to write kpoints.

    If `outfile = None`, KPOINTS file content is printed."""
    if type(n_xyz) == int:
        nx,ny,nz = n_xyz,n_xyz,n_xyz
    else:
        nx,ny,nz = [n for n in n_xyz]

    if type(n_xyz) == int and isinstance(path_pos,(str,list,tuple,np.ndarray)):
        Basis = _get_basis(path_pos)
        weights = np.linalg.norm(Basis.inverted,axis=1)
        weights = weights/np.min(weights) # For making smallest side at given n_xyz
        nx, ny, nz = np.rint(weights*n_xyz).astype(int)

    points = []
    for k in np.linspace(0,1,nz):
        for j in np.linspace(0,1,ny):
            for i in np.linspace(0,1,nx):
                points.append([i,j,k])
    points = np.array(points)
    points[np.abs(points) < 1e-10] = 0
    top_info=''

    if len(points) == 0:
        return print('No KPOINTS in BZ from given input. Try larger input!')
    if weight == None and len(points) != 0:
        weight = float(1/len(points))
    out_str = ["{0:>16.10f}{1:>16.10f}{2:>16.10f}{3:>12.6f}".format(x,y,z,weight) for x,y,z in points]
    out_str = '\n'.join(out_str)
    N = len(points)
    if ibzkpt != None:
        if os.path.isfile(ibzkpt):
            with open(ibzkpt,'r') as f:
                lines = f.readlines()

            N = int(lines[1].strip())+N # Update N.
            slines = lines[3:N+4]
            ibz_str = ''.join(slines)
            out_str = "{}\n{}".format(ibz_str,out_str) # Update out_str
    top_str = "Automatically generated uniform mesh using PivotPy with {}x{}x{} grid{}\n\t{}\nReciprocal Lattice".format(nx,ny,nz,top_info,N)
    out_str = "{}\n{}".format(top_str,out_str)
    if outfile != None:
        with open(outfile,'w') as f:
            f.write(out_str)
    else:
        print(out_str)

# Cell
def _tan_inv(vy,vx):
    """
    - Returns full angle from x-axis counter clockwise.
    - **Parameters**
        - vy : Perpendicular componet of vector including sign.
        - vx : Base compoent of vector including sign.
    """
    angle = 0  # Place hodler to handle exceptions
    if vx == 0 and vy == 0:
        angle = 0
    elif vx == 0 and np.sign(vy) == -1:
        angle = 3*np.pi/2
    elif vx == 0 and np.sign(vy) == 1:
        angle = np.pi/2
    else:
        theta = abs(np.arctan(vy/vx))
        if np.sign(vx) == 1 and np.sign(vy) == 1:
            angle = theta
        if np.sign(vx) == -1 and np.sign(vy) == 1:
            angle = np.pi - theta
        if np.sign(vx) == -1 and np.sign(vy) == -1:
            angle = np.pi + theta
        if np.sign(vx) == 1 and np.sign(vy) == -1:
            angle = 2*np.pi - theta
        if np.sign(vx) == -1 and vy == 0:
            angle = np.pi
        if np.sign(vx) == 1 and vy == 0:
            angle = 2*np.pi
    return angle

def order(points,loop=True):
    """
    - Returns indices of counterclockwise ordered vertices of a plane in 3D.
    - **Parameters**
        - points: numpy array of shape (N,3) or List[List(len=3)].
        - loop  : Default is True and appends start point at end to make a loop.
    - **Example**
        > pts = np.array([[1,0,3],[0,0,0],[0,1,2]])
        > inds = order(pts)
        > pts[inds]
        ```
        array([[1, 2, 3],
               [0, 0, 0],
               [1, 0, 3]
               [0, 1, 2]])
        ```
    """
    points = np.array(points) # Make array.
    # Fix points if start point is zero.
    if np.sum(points[0]) == 0:
        points = points + 0.5

    center = np.mean(points,axis=0) # 3D cent point.
    vectors = points - center # Relative to center

    ex = vectors[0]/np.linalg.norm(vectors[0])  # i
    ey = np.cross(center,ex)
    ey = ey/np.linalg.norm(ey)  # j

    angles= []
    for i, v in enumerate(vectors):
        vx = np.dot(v,ex)
        vy = np.dot(v,ey)
        angle = _tan_inv(vy,vx)
        angles.append([i,angle])

    s_angs = np.array(angles)
    ss = s_angs[s_angs[:,1].argsort()] #Sort it.

    if loop: # Add first at end for completing loop.
        ss = np.concatenate((ss,[ss[0]]))

    return ss[:,0].astype(int) # Order indices.


def _out_bz_plane(test_point,plane):
    """
    - Returns True if test_point is between plane and origin. Could be used to sample BZ mesh in place of ConvexHull.
    - **Parameters**
        - test_points: 3D point.
        - plane      : List of at least three coplanar 3D points.
    """
    outside = True
    p_test = np.array(test_point)
    plane = np.unique(plane,axis=0) #Avoid looped shape.
    c = np.mean(plane,axis=0) #center
    _dot_ = np.dot(p_test-c,c)
    if _dot_ < -1e-5:
        outside = False
    return outside


def _rad_angle(v1,v2):
    """
    - Returns interier angle between two vectors.
    - **Parameters**
        - v1,v2 : Two vectors/points in 3D.
    """
    v1 = np.array(v1)
    v2 = np.array(v2)
    norm  = np.linalg.norm(v1)*np.linalg.norm(v2)
    dot_p = np.round(np.dot(v1,v2)/norm,12)
    angle = np.arccos(dot_p)
    return angle

from scipy.spatial.transform import Rotation
def rotation(angle_deg,axis_vec):
    """Get a scipy Rotation object at given `angle_deg` around `axis_vec`.
    Usage:
        rot = rotation(60,[0,0,1])
        rot.apply([1,1,1])
        [-0.3660254  1.3660254  1.] #give this
    """
    axis_vec = np.array(axis_vec)/np.linalg.norm(axis_vec) # Normalization
    angle_rad = np.deg2rad(angle_deg)
    return Rotation.from_rotvec(angle_rad * axis_vec)


# Cell
def get_bz(path_pos = None,loop = True,digits=8,primitive=False):
    """
    - Return required information to construct first Brillouin zone in form of tuple (basis, normals, vertices, faces).
    - **Parameters**
        - path_pos : POSCAR file path or list of 3 vectors in 3D aslist[list,list,list].
        - loop   : If True, joins the last vertex of a BZ plane to starting vertex in order to complete loop.
        - digits : int, rounding off decimal places, no effect on intermediate calculations, just for pretty final results.
        - primitive: Defualt is False and returns Wigner-Seitz cell, If True returns parallelipiped in rec_basis.
    - **Attributes**
        - basis   : get_bz().basis, recprocal lattice basis.
        - normals : get_bz().normals, all vertors that are perpendicular BZ faces/planes.
        - vertices: get_bz().vertices, all vertices of BZ, could be input into ConvexHull for constructing 3D BZ.
        - faces   : get_bz().faces, vertices arranged into faces, could be input to Poly3DCollection of matplotlib for creating BZ from faces' patches.
        - specials : get_bz().specials, Data with attributes `coords`,`kpoints` and `near` in on-one correspondence for high symmetry KPOINTS in recirprocal coordinates space. `near` gives indices of nearest special points around a vertex. All vertices with z > 0 are included.
    """
    basis = _get_basis(path_pos).inverted # Reads
    b1, b2, b3 = basis # basis are reciprocal basis
    # Get all vectors for BZ
    if primitive:
        b0 = np.array([0,0,0])
        bd = b1+b2+b3 #Diagonal point
        faces = np.array([
                 [b0, b1, b1+b2, b2],
                 [b0,b2,b2+b3,b3],
                 [b0,b3,b3+b1,b1],
                 [b1,b1+b2,b1+b3,bd],
                 [b2,b2+b1,b2+b3,bd],
                 [b3,b3+b1,b3+b2,bd]
                ])
    else:
        vectors = []
        for i,j,k in product([0,1,-1],[0,1,-1],[0,1,-1]):
            vectors.append(i*b1+j*b2+k*b3)
        vectors = np.array(vectors)
        # Generate voronoi diagram
        vor = Voronoi(vectors)
        faces = []
        vrd = vor.ridge_dict
        for r in vrd:
            if r[0] == 0 or r[1] == 0:
                verts_in_face = np.array([vor.vertices[i] for i in vrd[r]])
                faces.append(verts_in_face)
        #faces = np.array(faces) # should be a list instead as not regular shape.

    verts = [v for vs in faces for v in vs]
    verts = np.unique(verts,axis=0)

    face_vectors = []
    for f in faces:
        face_vectors.append(np.mean(f,axis=0)) # In primitive point at face center
    if primitive == False:
        face_vectors = [2*f for f in face_vectors] # In regular, cross plane as well.

    # Order Faces.
    faces = [face[order(face,loop=loop)] for face in faces] # order based on given value of loop

    # High symmerty KPOINTS in primitive BZ (positive only)
    mid_faces = np.array([np.mean(np.unique(face,axis=0),axis=0) for face in faces])
    mid_edges = []
    for f in faces:
        for i in range(len(f)-1):
            # Do not insert point between unique vertices
            if np.isclose(np.linalg.norm(f[i]),np.linalg.norm(f[i+1])):
                mid_edges.append(np.mean([f[i],f[i+1]],axis=0))
    if mid_edges!=[]:
        mid_edges = np.unique(mid_edges,axis=0) # because faces share edges
        mid_faces = np.concatenate([mid_faces,mid_edges])
    # Bring all high symmetry points together.
    mid_all = np.concatenate([[[0,0,0]],mid_faces,verts]) # Coords
    mid_basis_all = np.array([np.linalg.solve(basis.T,v) for v in mid_all]) # Kpoints

    # Round off results
    mid_all_p    = np.round(mid_all,digits) # Coordinates
    mid_basis_p  = np.round(mid_basis_all,digits) # Relative points
    bais         = np.round(basis,digits)
    face_vectors = np.round(face_vectors,digits)
    verts        = np.round(verts,digits)
    faces        = tuple([np.round(face,digits) for face in faces])

    #Order special points near each vertex for z > 0.
    _arrs = []
    for v in verts[verts[:,2]>=0]: # Only upper hemisphere.
        _arr = []
        for i,c in enumerate(mid_all_p): # coordinates.
            _arr.append([i, np.linalg.norm(v-c)])
        _arr = np.array(_arr)
        _arr = _arr[_arr[:,1].argsort()][:,0].astype(int)
        upto = np.where(_arr == 0)[0][0]
        _arrs.append([0,*_arr[:upto]])

    one2one  = {'coords': mid_all_p ,'kpoints': mid_basis_p,'near': _arrs}
    out_dict = {'basis':basis, 'normals':face_vectors, 'vertices':verts,
                'faces':faces,'specials':one2one}
    return vp.dict2tuple('BZ',out_dict)

# Cell
def splot_bz(path_pos_bz = None, ax = None, plane=None,color='blue',fill=True,vectors=True,v3=False,vname='b',colormap='plasma',light_from=(1,1,1),alpha=0.4):
    """
    - Plots matplotlib's static figure.
    - **Parameters**
        - path_pos_bz: Path/to/POSCAR or List of 3 basis vectors or output of `get_zb`. Auto picks POSCAR from current directory.
        - fill       : True by defult, determines whether to fill surface of BZ or not.
        - color      : color to fill surface and stroke color.
        - vectors    : Plots basis vectors, default is True.
        - v3         : Plots 3rd vector as well. Only works in 2D and when `vectors=True`.
        - plane      : Default is None and plots 3D surface. Can take 'xy','yz','zx' to plot in 2D.
        - ax         : Auto generated by default, 2D/3D axes, auto converts in 3D on demand as well.
        - vname      : Default is `b` for reciprocal space, can set `a` for plotting cell as after `get_bz(get_bz().basis)` you get real space lattice back if `primitive=True` both times.
        - colormap  : If None, single color is applied, only works in 3D and `fill=True`. Colormap is applied along z.
        - light_from: Point from where light is thrown on BZ planes, default is (1,1,1). Only works on plane in 3D.
        - alpha    : Opacity of filling in range [0,1]. Increase for clear viewpoint.
    - **Returns**
        - ax   : Matplotlib's 2D axes if `plane=None`.
        - ax3d : Matplotlib's 2D axes if `plane` is given.

    > Tip: `splot_bz(rec_basis,primitive=True)` will plot cell in real space.
    """
    try:
        path_pos_bz.faces #See if it is BZ
        bz = path_pos_bz
    except AttributeError:
        bz = get_bz(_get_basis(path_pos_bz).given) # If it is a Matrix or picks POSCAR from CWD.
    label = r"$k_{}$" if vname=='b' else "{}"
    if not ax: #For both 3D and 2D, initialize 2D axis.
        ax = sp.get_axes(figsize=(3.4,3.4)) #For better display

    _label = r'\vec{}'.format(vname) # For both

    valid_planes = 'xyzxzyx' # cylic
    if plane and plane not in valid_planes:
        return print(f"`plane` expects value in 'xyzxzyx' or None, got {plane!r}")
    elif plane and plane in valid_planes: #Project 2D
        faces = bz.faces
        ind = valid_planes.index(plane)
        arr = [0,1,2,0,2,1,0]
        i, j = arr[ind], arr[ind+1]
        _ = [ax.plot(f[:,i],f[:,j],color=(color),lw=0.7) for f in faces]

        if vectors:
            if v3:
                s_basis = bz.basis
                ijk = [0,1,2]
            else:
                s_basis = bz.basis[[i,j]]# Only two.
                ijk = [i,j]

            for k,y in zip(ijk,s_basis):
                l = "\n" + r" ${}_{}$".format(_label,k+1)
                ax.text(0.8*y[i],0.8*y[j], l, va='center',ha='left')
                ax.scatter([y[i]],[y[j]],color='w',s=0.0005) # Must be to scale below arrow.

            s_zero = [0 for s_b in s_basis] # either 3 or 2.
            ax.quiver(s_zero,s_zero,*s_basis.T[[i,j]],lw=0.9,color='k',angles='xy', scale_units='xy', scale=1)

        ax.set_xlabel(label.format(valid_planes[i]))
        ax.set_ylabel(label.format(valid_planes[j]))
        ax.set_aspect(1) # Must for 2D axes to show actual lengths of BZ
        return ax
    else: # Plot 3D
        if ax and ax.name == "3d":
            ax3d = ax
        else:
            pos = ax.get_position()
            fig = ax.get_figure()
            ax.remove()
            ax3d = fig.add_axes(pos,projection='3d',azim=45,elev=30)

        if fill:
            if colormap:
                colormap = colormap if colormap in plt.colormaps() else 'viridis'
                cz = [np.mean(np.unique(f,axis=0),axis=0)[2] for f in bz.faces]
                levels = (cz - np.min(cz))/np.ptp(cz) # along Z.
                colors = plt.cm.get_cmap(colormap)(levels)
            else:
                colors = np.array([[*mplc.to_rgb(color)] for f in bz.faces]) # Single color.
            if light_from:
                intensity = bz.normals.dot(light_from) #Plane facing light
                intensity = (intensity - np.min(intensity) + 0.2)/np.ptp(intensity)
                intensity = intensity.clip(0,1)
                colors = np.array([i*c[:3] for i, c in zip(intensity,colors)])

            poly = Poly3DCollection(bz.faces,edgecolors=[color,],facecolors=colors, alpha=alpha)
            ax3d.add_collection3d(poly)
            ax3d.autoscale_view()
        else:
            _ = [ax3d.plot3D(f[:,0],f[:,1],f[:,2],color=(color),lw=0.7) for f in bz.faces]

        if vectors:
            for k,v in enumerate(0.35*bz.basis):
                ax3d.text(*v,r"${}_{}$".format(_label,k+1),va='center',ha='center')

            XYZ,UVW = [[0,0,0],[0,0,0],[0,0,0]], 0.3*bz.basis.T
            fancy_quiver3d(*XYZ,*UVW,C='k',L=0.7,ax=ax3d,arrowstyle="-|>",mutation_scale=7)

        l_ = np.min(bz.vertices,axis=0)
        h_ = np.max(bz.vertices,axis=0)
        ax3d.set_xlim([l_[0],h_[0]])
        ax3d.set_ylim([l_[1],h_[1]])
        ax3d.set_zlim([l_[2],h_[2]])

        # Set aspect to same as data.
        ax3d.set_box_aspect(np.ptp(bz.vertices,axis=0))

        ax3d.set_xlabel(label.format('x'))
        ax3d.set_ylabel(label.format('y'))
        ax3d.set_zlabel(label.format('z'))
        return ax3d

# Cell
def iplot_bz(path_pos_bz = None,fill = True,color = 'rgba(168,204,216,0.4)',background = 'rgb(255,255,255)',vname = 'b', alpha=0.4,ortho3d=True,fig=None):
    """
    - Plots interactive figure showing axes,BZ surface, special points and basis, each of which could be hidden or shown.
    - **Parameters**
        - path_pos_bz: Path/to/POSCAR or List of 3 basis vectors or output of `get_zb`. Auto picks POSCAR from current directory.
        - fill       : True by defult, determines whether to fill surface of BZ or not.
        - color      : color to fill surface 'rgba(168,204,216,0.4)` by default.
        - background : Plot background color, default is 'rgb(255,255,255)'.
        - vname      : Default is `b` for reciprocal space, can set `a` for plotting cell as after `get_bz(get_bz().basis)` you get real space lattice back if `primitive=True` both times.
        - alpha      : Opacity of BZ planes.
        - ortho3d    : Default is True, decides whether x,y,z are orthogonal or perspective.
        - fig        : (Optional) Plotly's `go.Figure`. If you want to plot on another plotly's figure, provide that.
    - **Returns**
        - fig   : plotly.graph_object's Figure instance.

    > Tip: `iplot_bz(rec_basis,primitive=True)` will plot cell in real space.
    """
    try:
        path_pos_bz.faces #See if it is BZ object
        bz = path_pos_bz
    except AttributeError:
        bz = get_bz(_get_basis(path_pos_bz).given) # If it is a Matrix or picks POSCAR from CWD.

    if not fig:
        fig = go.Figure()
    # Name fixing
    axes_text = ["<b>k</b><sub>x</sub>","","<b>k</b><sub>y</sub>","","<b>k</b><sub>z</sub>"]
    s_name = 'BZ'
    a_name = 'Axes'
    if vname == 'a':
        axes_text = ["<b>x</b>","","<b>y</b>","","<b>z</b>"] # Real space
        s_name = 'Lattice'
        a_name = 'RealAxes'

    # Axes
    _len = 0.5*np.mean(bz.basis)
    fig.add_trace(go.Scatter3d(x=[_len,0,0,0,0],y=[0,0,_len,0,0],z=[0,0,0,0,_len],
        mode='lines+text',
        text= axes_text,
        line_color='green', legendgroup=a_name,name=a_name))
    fig.add_trace(go.Cone(x=[0.95*_len,0,0],y=[0,0.95*_len,0],z=[0,0,0.95*_len],
        u=[0.2*_len,0,0],v=[0,0.2*_len,0],w=[0,0,0.2*_len],showscale=False,
        colorscale='Greens',legendgroup=a_name,name=a_name))
    # Basis
    for i,b in enumerate(bz.basis):
        fig.add_trace(go.Scatter3d(x=[0,b[0]], y=[0,b[1]],z=[0,b[2]],
            mode='lines+text',legendgroup="{}<sub>{}</sub>".format(vname,i+1), line_color='red',
            name="<b>{}</b><sub>{}</sub>".format(vname,i+1),text=["","<b>{}</b><sub>{}</sub>".format(vname,i+1)]))
        fig.add_trace(go.Cone(x=[0.95*b[0]],y=[0.95*b[1]],z=[0.95*b[2]],
            u=[0.2*b[0]],v=[0.2*b[1]],w=[0.2*b  [2]],showscale=False,colorscale='Reds',
            legendgroup="{}<sub>{}</sub>".format(vname,i+1),name="<b>{}</b><sub>{}</sub>".format(vname,i+1)))

    # Faces
    legend = True
    for pts in bz.faces:
        fig.add_trace(go.Scatter3d(x=pts[:,0], y=pts[:,1],z=pts[:,2],
            mode='lines',line_color=color, legendgroup=s_name,name=s_name,
            showlegend=legend))
        legend = False # Only first legend to show for all

    if fill:
        xc = bz.vertices[ConvexHull(bz.vertices).vertices]
        fig.add_trace(go.Mesh3d(x=xc[:, 0], y=xc[:, 1], z=xc[:, 2],
                        color=color,
                        opacity=alpha,
                        alphahull=0,
                        lighting=dict(diffuse=0.5),
                        legendgroup=s_name,name=s_name))

    # Special Points only if in reciprocal space.
    if vname == 'b':
        texts,values =[],[]
        norms = np.round(np.linalg.norm(bz.specials.coords,axis=1),5)
        sps = bz.specials
        for key,value, (i,norm) in zip(sps.kpoints, sps.coords, enumerate(norms)):
            texts.append("P{}</br>d = {}</br> Index = {}".format(key,norm,i))
            values.append([[*value,norm]])

        values = np.array(values).reshape((-1,4))
        norm_max = np.max(values[:,3])
        c_vals = np.array([int(v*255/norm_max) for v in values[:,3]])
        colors = [0 for i in c_vals]
        _unique = np.unique(np.sort(c_vals))[::-1]
        _lnp = np.linspace(0,255,len(_unique)-1)
        _u_colors = ["rgb({},0,{})".format(r,b) for b,r in zip(_lnp,_lnp[::-1])]
        for _un,_uc in zip(_unique[:-1],_u_colors):
            _index = np.where(c_vals == _un)[0]
            for _ind in _index:
                colors[_ind]=_uc

        colors[0]= "rgb(255,215,0)" # Gold color at Gamma!.
        fig.add_trace(go.Scatter3d(x=values[:,0], y=values[:,1],z=values[:,2],
                hovertext=texts,name="HSK",marker=dict(color=colors,size=4),mode='markers'))
    proj = dict(projection=dict(type = "orthographic")) if ortho3d else {}
    camera = dict(center=dict(x=0.1, y=0.1, z=0.1),**proj)
    fig.update_layout(scene_camera=camera,paper_bgcolor=background, plot_bgcolor=background,
        font_family="Times New Roman",font_size= 14,
        scene = dict(aspectmode='data',xaxis = dict(showbackground=False,visible=False),
                        yaxis = dict(showbackground=False,visible=False),
                        zaxis = dict(showbackground=False,visible=False)),
                        margin=dict(r=10, l=10,b=10, t=30))
    return fig

# Cell
def to_R3(basis,points):
    """Transforms coordinates of points (relative to non-othogonal basis) into orthogonal space.
    - **Parameters**
        - basis : Non-orthogonal basis of real or reciprocal space.
        - points: 3D points relative to basis, such as KPOINTS and Lattice Points.
    """
    rec_basis = np.array(basis)
    points = np.array(points)
    # Formula to make coordinates from relative points.
    # kx, ky, kz = n1*b1 + n2*b2 +n3*b3
    #            = [n1, n2, n3].dot(rec_basis)
    coords = points.dot(rec_basis)
    return coords

# Cell
def kpoints2bz(bz,kpoints,primitive=False):
    """Brings KPOINTS inside BZ. Applies `to_R3` only if `primitive=True`.
    - **Parameters**
        - bz       : Output of get_bz(), make sure use same value of `primitive` there and here.
        - kpoints  : List or array of KPOINTS to transorm into BZ or R3.
        - primitive: Default is False and brings kpoints into regular BZ. If True, returns `to_R3()`.
    """
    # Shift to 1st octant for both cases of primitive argument.
    kpoints = kpoints - np.min(kpoints,axis=0)
    if primitive:
        return to_R3(bz.basis,kpoints)

    # Work ahead if not primitive
    kpoints = kpoints - 0.5 # For regular BZ only
    cent_planes = [np.mean(np.unique(face,axis=0),axis=0) for face in bz.faces]

    out_coords = np.empty(np.shape(kpoints)) # To store back

    def inside(coord,cent_planes):
        _dots_ = np.max([np.dot(coord-c, c) for c in cent_planes]) #max in all planes
        #print(_dots_)
        if np.max(_dots_) > 1e-8: # Outside
            return [] # empty for comparison
        else: # Inside
            return list(coord) # Must be in list form


    from itertools import product
    for i,p in enumerate(kpoints):
        for q in product([0,1,-1],[0,1,-1],[0,1,-1]):
            # First translate, then make coords, then feed it back
            #print(q)
            pos = to_R3(bz.basis,p + np.array(q))
            r = inside(pos,cent_planes)
            if r:
                #print(p,'-->',r)
                out_coords[i] = r
                StopIteration

    return out_coords # These may have duplicates, apply np.unique(out_coords,axis=0). do this in surface plots

# Cell
def fix_sites(poscar,tol=1e-2,eqv_sites=True,translate=None):
    """Add equivalent sites to make a full data shape of lattice. Returns same data after fixing.
    - **Parameters**
        - poscar: Output of `export_poscar` or `export_vasprun().poscar`.
        - tol   : Tolerance value. Default is 0.01.
        - eqv_sites: If True, add sites on edges and faces. If False, just fix coordinates, i.e. `pos > 1 - tol -> pos - 1`, useful for merging poscars to make slabs.
        - translate: A number(+/-) or list of three numbers to translate in x,y,z directions.
    """
    sys,vol,basis,rec_basis,pos,labels,unique = poscar.to_tuple()
    pos = pos.copy() # Must to avoid chnaging outside.
    out_dict = poscar.to_dict() # For output

    if translate and isinstance(translate,(int,float)):
        pos = pos + (translate - int(translate)) # Only translate in 0 - 1
    elif translate and len(translate) == 3:
        txyz = np.array([translate])
        pos = pos + (txyz - txyz.astype(int))

    if eqv_sites:
        inds, new_pos = [],[]
        for p in product([-1,0,1],[-1,0,1],[-1,0,1]):
            t_pos = pos + [[*p]] # test position.
            t_inds = [i for i,t_p in enumerate(
                (t_pos > -tol) & (t_pos < tol+1)
                ) if not False in t_p]
            new_pos = [*new_pos, *t_pos[t_inds]]
            inds = [*inds,*t_inds]

        new_labs = [labels[i] for i in inds]
        new_pos = np.array(new_pos)

        # order things.
        _ord = np.array([[i,l] for i,l in enumerate(new_labs)])
        _ord = _ord[_ord[:,1].argsort()]

        oinds, olabels = _ord[:,0].T.astype(int), _ord[:,1].T
        new_pos = new_pos[oinds] # Order points too
        new_labs = [new_labs[i] for i in oinds] # Order Labels

        ul, ui = np.unique([ol.split()[0] for ol in olabels],
                    return_counts=True)
        eleminds = np.cumsum([0,*ui]) #Element ranges.
        uelems = {l:range(eleminds[i],eleminds[i+1]) for i,l in enumerate(ul)}
        # Update Things
        out_dict['positions'] = new_pos
        out_dict['labels'] = new_labs
        out_dict['unique'] = uelems
    else:
        xinds,yinds = np.where(pos > 1-tol)
        for i,j in zip(xinds,yinds):
            pos[i,j] = pos[i,j] -1

        # Update Things
        out_dict['positions'] = pos

    return vp.Dict2Data(out_dict)

def get_pairs(basis, positions, r, eps=1e-2):
    """Returns a tuple of Lattice (coords,pairs), so coords[pairs] given nearest site bonds.
    - **Parameters**
        - basis: Real space lattice basis.
        - positions: Array(N,3) of fractional positions of lattice sites. If coordinates positions, provide unity basis.
        - r        : Cartesian distance between the pairs in units of Angstrom e.g. 1.2 -> 1.2E-10.
        - eps      : Tolerance value. Default is 10^-2.
    """
    coords = to_R3(basis,positions)
    tree = KDTree(coords)
    inds = np.array([[*p] for p in tree.query_pairs(r,eps=eps)])
    return vp.dict2tuple('Lattice',{'coords':coords,'pairs':inds})

def _get_bond_length(poscar,given=None,eps=1e-2):
    "eps is add to calculated bond length in order to fix small differences, paramater `given` in range [0,1] which is scaled to V^(1/3)."
    if given != None:
        return given*poscar.volume**(1/3) + eps
    else:
        _coords = to_R3(poscar.basis,poscar.positions)
        _arr = sorted(np.linalg.norm(_coords[1:] - _coords[0],axis=1)) # Sort in ascending. returns list
        return np.mean(_arr[:2]) + eps if _arr else 1 #Between nearest and second nearest.

# Cell
def iplot_lat(poscar,sizes=10,colors='blue',
              bond_length=None,tol=1e-1,eps=1e-2,eqv_sites=True,
              translate = None,
              line_width=4,edge_color = 'black',
              fill=False,alpha=0.4, ortho3d=True,fig=None):
    """Interactive plot of lattice.
    - **Main Parameters**
        - poscar     : Output of export_poscar or export_vasprun().poscar.
        - sizes      : Size of sites. Either one int/float or list equal to type of ions.
        - colors     : Colors of sites. Either one colors or list equal to type of ions.
        - bond_length: Length of bond in fractional unit [0,1]. It is scaled to V^1/3 and auto calculated if not provides.
    Other parameters just mean what they seem to be.
    """
    poscar = fix_sites(poscar=poscar,tol=tol,eqv_sites=eqv_sites,translate=translate)
    bond_length = _get_bond_length(poscar,given=bond_length,eps=eps)
    coords, pairs = get_pairs(basis=poscar.basis,
                        positions =poscar.positions,
                        r=bond_length,eps=eps)
    if not fig:
        fig = go.Figure()

    uelems = poscar.unique.to_dict()
    if not isinstance(sizes,(list,np.ndarray)):
        sizes = [sizes for elem in uelems.keys()]

    if not isinstance(colors, str):
        colors = np.array(colors)
    else:
        colors = [colors for elem in uelems.keys()]

    h_text = np.array(poscar.labels)

    _colors = []
    for i,vs in enumerate(uelems.values()):
        for v in vs:
            _colors.append(colors[i])
    _colors = np.array(_colors)  # Full List

    if np.any(pairs):
        coords_p = coords[pairs] #paired points
        _colors = _colors[pairs] # Colors at pairs
        coords_n = []
        colors_n = []
        for c_p, _c in zip(coords_p,_colors):
            mid = np.mean(c_p,axis=0)
            arr = np.concatenate([c_p[0],mid,mid,c_p[1]]).reshape((-1,2,3))
            coords_n = [*coords_n,*arr] # Same shape
            colors_n = [*colors_n,*_c] # same shape.

        coords_n = np.array(coords_n)
        colors_n = np.array(colors_n)

        for (i, cp),c in zip(enumerate(coords_n),colors_n):
            showlegend = True if i==0 else False
            fig.add_trace(go.Scatter3d(
                x = cp[:,0].T,
                y = cp[:,1].T,
                z = cp[:,2].T,
                mode='lines',line_color=c,
                legendgroup='Bonds',showlegend=showlegend,
                name='Bonds',line_width=line_width))

    for (k,v),c,s in zip(uelems.items(),colors,sizes):
        fig.add_trace(go.Scatter3d(
            x = coords[v][:,0].T,
            y = coords[v][:,1].T,
            z = coords[v][:,2].T,
            mode='markers',marker_color=c,
            hovertext = h_text[v],
            line_color='rgba(1,1,1,0)',line_width=0.001,
            marker_size = s,opacity=1,name=k))
    bz = get_bz(path_pos=poscar.rec_basis, primitive=True)
    _ = iplot_bz(bz,fig=fig,vname='a',color=edge_color,
                fill=fill,alpha=alpha,ortho3d=ortho3d)
    return fig

# Cell
def splot_lat(poscar,sizes=50,colors=[],colormap=None,
              bond_length=None,tol=1e-1,eps=1e-2,eqv_sites=True,
              translate = None,
              line_width=1,edge_color=((1,0.5,0,0.4)),
              vectors=True,v3=False,plane=None,
              light_from=(1,1,1),
              fill=False,alpha=0.4,ax=None):
    """Static plot of lattice.
    - **Main Parameters**
        - poscar     : Output of export_poscar or export_vasprun().poscar.
        - sizes      : Size of sites. Either one int/float or list equal to type of ions.
        - bond_length: Length of bond in fractional unit [0,1]. It is scaled to V^1/3 and auto calculated if not provides.
        - colors: List of colos. If given, preffered over colormap, should have same length as type of ions.
    Other parameters just mean what they seem to be.

    > Tip: Use `plt.style.use('ggplot')` for better 3D perception.
    """
    #Plane fix
    if plane and plane not in 'xyzxzyx':
        return print("plane expects in 'xyzxzyx' or None.")
    if plane:
        ind = 'xyzxzyx'.index(plane)
        arr = [0,1,2,0,2,1,0]
        ix,iy = arr[ind], arr[ind+1]
    poscar = fix_sites(poscar=poscar,tol=tol,eqv_sites=eqv_sites,translate=translate)
    bond_length = _get_bond_length(poscar,given=bond_length,eps=eps)
    coords, pairs = get_pairs(basis=poscar.basis,
                        positions =poscar.positions,
                        r=bond_length,eps=eps)
    bz = get_bz(poscar.rec_basis, primitive=True)
    ax = splot_bz(bz,ax=ax,vname='a',
                color=edge_color,colormap=colormap,
                fill=fill,alpha=alpha,plane=plane,v3=v3,
                vectors=vectors,light_from=light_from)

    uelems = poscar.unique.to_dict()
    if not isinstance(sizes,(list,np.ndarray)):
        sizes = [sizes for elem in uelems.keys()]

    if not colormap in plt.colormaps():
        colormap = 'brg'

    if not colors or len(colors) != len(uelems):
        colors = plt.cm.get_cmap(colormap)(np.linspace(0.1,0.9,len(uelems)))

    _colors = []
    for i,vs in enumerate(uelems.values()):
        for v in vs:
            _colors.append(colors[i])
    _colors = np.array(_colors)  # Full List

    if np.any(pairs):
        coords_p = coords[pairs] #paired points
        _colors = _colors[pairs] # Colors at pairs
        coords_n = []
        colors_n = []
        for c_p, _c in zip(coords_p,_colors):
            mid = np.mean(c_p,axis=0)
            arr = np.concatenate([c_p[0],mid,mid,c_p[1]]).reshape((-1,2,3))
            coords_n = [*coords_n,*arr] # Same shape
            colors_n = [*colors_n,*_c] # same shape.

        coords_n = np.array(coords_n)
        colors_n = np.array(colors_n)

        if not plane:
            _ = [ax.plot(*c.T,c=_c,lw=line_width) for c,_c in zip(coords_n,colors_n)]
        elif plane in 'xyzxzyx':
            _ = [ax.plot(c[:,ix],c[:,iy],c=_c,lw=line_width) for c,_c in zip(coords_n,colors_n)]

    for (k,v),c,s in zip(uelems.items(),colors,sizes):
        if not plane:
            ax.scatter(coords[v][:,0],coords[v][:,1],coords[v][:,2],color = c ,s =s,label=k,depthshade=False)
        elif plane in 'xyzxzyx':
            ax.scatter(coords[v][:,ix],coords[v][:,iy],color = c ,s =s,label=k,zorder=3)
    ax.set_axis_off()
    sp.add_legend(ax)
    return ax

# Cell
def join_poscars(poscar1,poscar2,direction='z',tol=1e-2):
    """Joins two POSCARs in a given direction. In-plane lattice parameters are kept from `poscar1` and basis of `poscar2` parallel to `direction` is modified while volume is kept same.
    - **Parameters**
        - poscar1, poscar2:  Base and secondary POSCARs respectivly. Output of `export_poscar` or similar object from other functions.
        - direction: The joining direction. It is general and can join in any direction along basis. Expect one of ['a','b','c','x','y','z'].
        - tol: Default is 0.01. It is used to bring sites near 1 to near zero in order to complete sites in plane. Vasp relaxation could move a point, say at 0.00100 to 0.99800 which is not useful while merging sites.
    """
    _poscar1 = fix_sites(poscar1,tol=tol,eqv_sites=False)
    _poscar2 = fix_sites(poscar2,tol=tol,eqv_sites=False)
    pos1 = _poscar1.positions.copy()
    pos2 = _poscar2.positions.copy()

    s1,s2 = 0.5, 0.5 # Half length for each.
    a1,b1,c1 = np.linalg.norm(_poscar1.basis,axis=1)
    a2,b2,c2 = np.linalg.norm(_poscar2.basis,axis=1)
    basis = _poscar1.basis.copy() # Must be copied, otherwise change outside.

    # Processing in orthogonal space
    if direction in ['z','c']:
        c2 = (a2*b2)/(a1*b1)*c2
        netc = c1+c2
        s1, s2 = c1/netc, c2/netc
        pos1[:,2] = s1*pos1[:,2]
        pos2[:,2] = s2*pos2[:,2] + s1
        basis[2] = netc*basis[2]/np.linalg.norm(basis[2]) #Update 3rd vector

    elif direction in ['y','b']:
        b2 = (a2*c2)/(a1*c1)*b2
        netb = b1+b2
        s1, s2 = b1/netb, b2/netb
        pos1[:,1] = s1*pos1[:,1]
        pos2[:,1] = s2*pos2[:,1] + s1
        basis[1] = netb*basis[1]/np.linalg.norm(basis[1]) #Update 2nd vector

    elif direction in ['x','a']:
        a2 = (b2*c2)/(b1*c1)*a2
        neta = a1+a2
        s1, s2 = a1/neta, a2/neta
        pos1[:,0] = s1*pos1[:,0]
        pos2[:,0] = s2*pos2[:,0] + s1
        basis[0] = neta*basis[0]/np.linalg.norm(basis[0]) #Update 1st vector

    else:
        return print("direction expects one of ['a','b','c','x','y','z']")

    rec_basis = np.linalg.inv(basis).T # Transpose is must
    volume = np.dot(basis[0],np.cross(basis[1],basis[2]))
    u1 = _poscar1.unique.to_dict()
    u2 = _poscar2.unique.to_dict()
    u_all = np.unique([*u1.keys(),*u2.keys()])


    pos_all = []
    i_all = []
    for u in u_all:
        _i_ = 0
        if u in u1.keys():
            _i_ = len(u1[u])
            pos_all = [*pos_all,*pos1[u1[u]]]
        if u in u2.keys():
            _i_ = _i_ + len(u2[u])
            pos_all = [*pos_all,*pos2[u2[u]]]
        i_all.append(_i_)

    labels = [f"{u} {i+1}" for u,ii in zip(u_all,i_all) for i in range(ii)]
    i_all = np.cumsum([0,*i_all]) # Do it after labels
    uelems = {_u:range(i_all[i],i_all[i+1]) for i,_u in enumerate(u_all)}
    sys = ''.join(uelems.keys()) + "   # Created by Pivotpy"
    out_dict = {'SYSTEM':sys,'volume':volume,'basis':basis,'rec_basis':rec_basis,'positions':np.array(pos_all),'labels':labels,'unique':uelems}
    return vp.Dict2Data(out_dict)

# Cell
def write_poscar(poscar,sd_list=None,outfile=None,overwrite=False):
    """Writes poscar data object to a file or returns string
    - **Parameters**
        - poscar   : Output of `export_poscar`,`join_poscars` etc.
        - sd_list  : A list ['T T T','F F F',...] strings to turn on selective dynamics at required sites. len(sd_list)==len(sites) should hold.
        - outfile  : str,file path to write on.
        - overwrite: bool, if file already exists, overwrite=True changes it.
    """
    out_str = poscar.SYSTEM + "  # Created by Pivotpy"
    scale = np.linalg.norm(poscar.basis[0])
    out_str += "\n  {:<20.14f}\n".format(scale)
    out_str += '\n'.join(["{:>22.16f}{:>22.16f}{:>22.16f}".format(*a) for a in poscar.basis/scale])
    uelems = poscar.unique.to_dict()
    out_str += "\n  " + '\t'.join(uelems.keys())
    out_str += "\n  " + '\t'.join([str(len(v)) for v in uelems.values()])
    if sd_list:
        out_str += "\nSelective Dynamics"
    out_str += "\nDirect\n"
    pos_list = ["{:>21.16f}{:>21.16f}{:>21.16f}".format(*a) for a in poscar.positions]
    if sd_list:
        if len(pos_list) != len(sd_list):
            return print("len(sd_list) != len(sites).")
        pos_list = [f"{p}   {s}" for p,s in zip(pos_list,sd_list)]
    out_str += '\n'.join(pos_list)
    if outfile:
        if not os.path.isfile(outfile):
            with open(outfile,'w') as f:
                f.write(out_str)
            return
        elif overwrite and os.path.isfile(outfile):
            with open(outfile,'w') as f:
                f.write(out_str)
            return
        else:
            return print(f"{outfile!r} exists, can not overwrite, \nuse overwrite=True if you want to chnage.")
    # Return string anyway
    return out_str

# Cell
def scale_poscar(path_poscar,scale=(1,1,1),tol=1e-2):
    """Create larger/smaller cell from a given POSCAR.
    - **Parameters**
        - path_poscar: Path/to/POSCAR or `poscar` data object.
        - scale: Tuple of three values along (a,b,c) vectors. int or float values. If number of sites are not as expected in output, tweak `tol` instead of `scale`. You can put a minus sign with `tol` to get more sites and plus sign to reduce sites.
        - tol: It is used such that site positions are blow `1 - tol`, as 1 belongs to next cell, not previous one.
    > Tip: scale = (2,2,2) enlarges a cell and next operation of (1/2,1/2,1/2) should bring original cell back.
    """
    if isinstance(path_poscar,str):
        poscar = export_poscar(path_poscar)
    else:
        poscar = path_poscar

    ii, jj, kk = np.ceil(scale).astype(int) # Need int for joining.
    _orig_poscar = poscar # Assign base for x
    if ii > 1:
        for i in range(1,ii): # 1 poscar is already there.
            poscar = join_poscars(_orig_poscar,poscar,direction='x',tol=tol)
        _orig_poscar = poscar # Reassign base for y

    if jj > 1:
        for i in range(1,jj):
            poscar = join_poscars(_orig_poscar,poscar,direction='y',tol=tol)
        _orig_poscar = poscar # Reassign base for z

    if kk > 1:
        for i in range(1,kk):
            poscar = join_poscars(_orig_poscar,poscar,direction='z',tol=tol)

    new_poscar = poscar.to_dict() # Update in it

    # Get clip fraction
    fi, fj, fk = scale[0]/ii, scale[1]/jj, scale[2]/kk

    # Clip at end according to scale, change length of basis as fractions.
    pos   = poscar.positions.copy()/np.array([fi,fj,fk]) # rescale for clip
    basis = poscar.basis.copy()
    for i,f in zip(range(3),[fi,fj,fk]):
        basis[i] = f*basis[i] # Basis rescale for clip

    new_poscar['basis']      = basis
    new_poscar['volume']     = np.linalg.det(basis)
    new_poscar['rec_basis']  = np.linalg.inv(basis).T

    uelems = poscar.unique.to_dict()
    # Minus in below for block is because if we have 0-2 then 1 belongs to next cell not original.
    positions,drop = [],0
    for key,value in uelems.items():
        i0 = np.where(pos[value][:,0] <= 1 - tol)
        i1 = np.where(pos[value][:,1] <= 1 - tol)
        i2 = np.where(pos[value][:,2] <= 1 - tol)
        all_i = np.intersect1d(np.intersect1d(i0,i1),i2)
        uelems[key] = range(drop,drop + len(all_i))
        positions = [*positions,*pos[value][all_i]] # Pick positions
        drop = len(all_i) #Update for next element

    labels = ["{} {}".format(k,_v - vs[0] + 1) for k,vs in uelems.items() for _v in vs]
    new_poscar['labels']    = labels
    new_poscar['unique']    = uelems
    new_poscar['positions'] = np.array(positions)
    return vp.Dict2Data(new_poscar)

def rotate_poscar(path_poscar,angle_deg,axis_vec):
    """Rotate a given POSCAR.
    - **Parameters**
        - path_poscar: Path/to/POSCAR or `poscar` data object.
        - angle_deg: Rotation angle in degrees.
        - axis_vec : (x,y,z) of axis about which rotation takes place. Axis passes through origin.
    """
    if isinstance(path_poscar,str):
        poscar = export_poscar(path_poscar)
    else:
        poscar = path_poscar
    rot = rotation(angle_deg=angle_deg,axis_vec=axis_vec)
    p_dict = poscar.to_dict()
    p_dict['basis'] = rot.apply(p_dict['basis']) # Rotate basis so that they are transpose
    p_dict['rec_basis'] = np.linalg.inv(p_dict['basis']).T # invert rotated basis
    return vp.Dict2Data(p_dict)