#!/usr/bin/env python

import numpy as np

#######################################################################
# Using the local Atoms-like class (BSD license) where a small set of #
# ASE Atoms features is compatible but enough for this example.       #
#######################################################################
from atoms import Atoms

import spglib

#######################################
# Uncomment to use Atoms class in ASE #
#######################################
# from ase import Atoms


def show_symmetry(symmetry):
    for i in range(symmetry["rotations"].shape[0]):
        print("  --------------- %4d ---------------" % (i + 1))
        rot = symmetry["rotations"][i]
        trans = symmetry["translations"][i]
        print("  rotation:")
        for x in rot:
            print("     [%2d %2d %2d]" % (x[0], x[1], x[2]))
        print("  translation:")
        print("     (%8.5f %8.5f %8.5f)" % (trans[0], trans[1], trans[2]))


def show_lattice(lattice):
    print("Basis vectors:")
    for vec, axis in zip(lattice, ("a", "b", "c")):
        print(
            "%s %10.5f %10.5f %10.5f"
            % (
                tuple(
                    axis,
                )
                + tuple(vec)
            )
        )


def show_cell(lattice, positions, numbers):
    show_lattice(lattice)
    print("Atomic points:")
    for p, s in zip(positions, numbers):
        print("%2d %10.5f %10.5f %10.5f" % ((s,) + tuple(p)))


silicon_ase = Atoms(
    symbols=["Si"] * 8,
    cell=[(4, 0, 0), (0, 4, 0), (0, 0, 4)],
    scaled_positions=[
        (0, 0, 0),
        (0, 0.5, 0.5),
        (0.5, 0, 0.5),
        (0.5, 0.5, 0),
        (0.25, 0.25, 0.25),
        (0.25, 0.75, 0.75),
        (0.75, 0.25, 0.75),
        (0.75, 0.75, 0.25),
    ],
    pbc=True,
)

silicon = (
    [(4, 0, 0), (0, 4, 0), (0, 0, 4)],
    [
        (0, 0, 0),
        (0, 0.5, 0.5),
        (0.5, 0, 0.5),
        (0.5, 0.5, 0),
        (0.25, 0.25, 0.25),
        (0.25, 0.75, 0.75),
        (0.75, 0.25, 0.75),
        (0.75, 0.75, 0.25),
    ],
    [
        14,
    ]
    * 8,
)

silicon_dist = (
    [(4.01, 0, 0), (0, 4, 0), (0, 0, 3.99)],
    [
        (0.001, 0, 0),
        (0, 0.5, 0.5),
        (0.5, 0, 0.5),
        (0.5, 0.5, 0),
        (0.25, 0.25, 0.251),
        (0.25, 0.75, 0.75),
        (0.75, 0.25, 0.75),
        (0.75, 0.75, 0.25),
    ],
    [
        14,
    ]
    * 8,
)

silicon_prim = (
    [(0, 2, 2), (2, 0, 2), (2, 2, 0)],
    [(0, 0, 0), (0.25, 0.25, 0.25)],
    [14, 14],
)

rutile = (
    [(4, 0, 0), (0, 4, 0), (0, 0, 3)],
    [
        (0, 0, 0),
        (0.5, 0.5, 0.5),
        (0.3, 0.3, 0.0),
        (0.7, 0.7, 0.0),
        (0.2, 0.8, 0.5),
        (0.8, 0.2, 0.5),
    ],
    [14, 14, 8, 8, 8, 8],
)

rutile_dist = (
    [(3.97, 0, 0), (0, 4.03, 0), (0, 0, 3.0)],
    [
        (0, 0, 0),
        (0.5001, 0.5, 0.5),
        (0.3, 0.3, 0.0),
        (0.7, 0.7, 0.002),
        (0.2, 0.8, 0.5),
        (0.8, 0.2, 0.5),
    ],
    [14, 14, 8, 8, 8, 8],
)

a = 3.07
c = 3.52
MgB2 = (
    [(a, 0, 0), (-a / 2, a / 2 * np.sqrt(3), 0), (0, 0, c)],
    [(0, 0, 0), (1.0 / 3, 2.0 / 3, 0.5), (2.0 / 3, 1.0 / 3, 0.5)],
    [12, 5, 5],
)

a = [3.0, 0.0, 0.0]
b = [-3.66666667, 3.68178701, 0.0]
c = [-0.66666667, -1.3429469, 1.32364995]
niggli_lattice = np.array([a, b, c])

# For VASP case
# import vasp
# bulk = vasp.read_vasp(sys.argv[1])

print("[get_spacegroup]")
print("  Spacegroup of Silicon is %s." % spglib.get_spacegroup(silicon))
print("")
print("[get_spacegroup]")
print(
    "  Spacegroup of Silicon (ASE Atoms-like format) is %s."
    % spglib.get_spacegroup(silicon_ase)
)
print("")
print("[get_spacegroup]")
print("  Spacegroup of Rutile is %s." % spglib.get_spacegroup(rutile))
print("")
print("[get_spacegroup]")
print("  Spacegroup of MgB2 is %s." % spglib.get_spacegroup(MgB2))
print("")
print("[get_symmetry]")
print("  Symmetry operations of Rutile unitcell are:")
print("")
symmetry = spglib.get_symmetry(rutile)
show_symmetry(symmetry)
print("")
print("[get_symmetry]")
print("  Symmetry operations of MgB2 are:")
print("")
symmetry = spglib.get_symmetry(MgB2)
show_symmetry(symmetry)
print("")
print("[get_pointgroup]")
print("  Pointgroup of Rutile is %s." % spglib.get_pointgroup(symmetry["rotations"])[0])
print("")

dataset = spglib.get_symmetry_dataset(rutile)
print("[get_symmetry_dataset] ['international']")
print(
    "  Spacegroup of Rutile is %s (%d)." % (dataset["international"], dataset["number"])
)
print("")
print("[get_symmetry_dataset] ['pointgroup']")
print("  Pointgroup of Rutile is %s." % (dataset["pointgroup"]))
print("")
print("[get_symmetry_dataset] ['hall']")
print("  Hall symbol of Rutile is %s (%d)." % (dataset["hall"], dataset["hall_number"]))
print("")
print("[get_symmetry_dataset] ['wyckoffs']")
alphabet = "abcdefghijklmnopqrstuvwxyz"
print("  Wyckoff letters of Rutile are: ", dataset["wyckoffs"])
print("")
print("[get_symmetry_dataset] ['equivalent_atoms']")
print("  Mapping to equivalent atoms of Rutile are: ")
for i, x in enumerate(dataset["equivalent_atoms"]):
    print("  %d -> %d" % (i + 1, x + 1))
print("")
print("[get_symmetry_dataset] ['rotations'], ['translations']")
print("  Symmetry operations of Rutile unitcell are:")
for i, (rot, trans) in enumerate(zip(dataset["rotations"], dataset["translations"])):
    print("  --------------- %4d ---------------" % (i + 1))
    print("  rotation:")
    for x in rot:
        print("     [%2d %2d %2d]" % (x[0], x[1], x[2]))
    print("  translation:")
    print("     (%8.5f %8.5f %8.5f)" % (trans[0], trans[1], trans[2]))
print("")

print("[refine_cell]")
print(" Refine distorted rutile structure")
lattice, positions, numbers = spglib.refine_cell(rutile_dist, symprec=1e-1)
show_cell(lattice, positions, numbers)
print("")

print("[find_primitive]")
print(" Fine primitive distorted silicon structure")
lattice, positions, numbers = spglib.find_primitive(silicon_dist, symprec=1e-1)
show_cell(lattice, positions, numbers)
print("")

print("[standardize_cell]")
print(" Standardize distorted rutile structure:")
print(" (to_primitive=0 and no_idealize=0)")
lattice, positions, numbers = spglib.standardize_cell(
    rutile_dist, to_primitive=0, no_idealize=0, symprec=1e-1
)
show_cell(lattice, positions, numbers)
print("")

print("[standardize_cell]")
print(" Standardize distorted rutile structure:")
print(" (to_primitive=0 and no_idealize=1)")
lattice, positions, numbers = spglib.standardize_cell(
    rutile_dist, to_primitive=0, no_idealize=1, symprec=1e-1
)
show_cell(lattice, positions, numbers)
print("")

print("[standardize_cell]")
print(" Standardize distorted silicon structure:")
print(" (to_primitive=1 and no_idealize=0)")
lattice, positions, numbers = spglib.standardize_cell(
    silicon_dist, to_primitive=1, no_idealize=0, symprec=1e-1
)
show_cell(lattice, positions, numbers)
print("")

print("[standardize_cell]")
print(" Standardize distorted silicon structure:")
print(" (to_primitive=1 and no_idealize=1)")
lattice, positions, numbers = spglib.standardize_cell(
    silicon_dist, to_primitive=1, no_idealize=1, symprec=1e-1
)
show_cell(lattice, positions, numbers)
print("")

symmetry = spglib.get_symmetry(silicon)
print("[get_symmetry]")
print("  Number of symmetry operations of silicon conventional")
print("  unit cell is %d (192)." % len(symmetry["rotations"]))
show_symmetry(symmetry)
print("")

symmetry = spglib.get_symmetry_from_database(525)
print("[get_symmetry_from_database]")
print("  Number of symmetry operations of silicon conventional")
print("  unit cell is %d (192)." % len(symmetry["rotations"]))
show_symmetry(symmetry)
print("")

reduced_lattice = spglib.niggli_reduce(niggli_lattice)
print("[niggli_reduce]")
print("  Original lattice")
show_lattice(niggli_lattice)
print("  Reduced lattice")
show_lattice(reduced_lattice)
print("")


mapping, grid = spglib.get_ir_reciprocal_mesh(
    [11, 11, 11], silicon_prim, is_shift=[0, 0, 0]
)
num_ir_kpt = len(np.unique(mapping))
print("[get_ir_reciprocal_mesh]")
print("  Number of irreducible k-points of primitive silicon with")
print("  11x11x11 Monkhorst-Pack mesh is %d (56)." % num_ir_kpt)
print("")

mapping, grid = spglib.get_ir_reciprocal_mesh([8, 8, 8], rutile, is_shift=[1, 1, 1])
num_ir_kpt = len(np.unique(mapping))
print("[get_ir_reciprocal_mesh]")
print("  Number of irreducible k-points of Rutile with")
print("  8x8x8 Monkhorst-Pack mesh is %d (40)." % num_ir_kpt)
print("")

mapping, grid = spglib.get_ir_reciprocal_mesh([9, 9, 8], MgB2, is_shift=[0, 0, 1])
num_ir_kpt = len(np.unique(mapping))
print("[get_ir_reciprocal_mesh]")
print("  Number of irreducible k-points of MgB2 with")
print("  9x9x8 Monkhorst-Pack mesh is %s (48)." % num_ir_kpt)
print("")
