import numpy as np
import sys
import re
import copy
from vasppy import cell
from vasppy.units import angstrom_to_bohr
from pymatgen.core import Lattice as pmg_Lattice
from pymatgen.core import Structure as pmg_Structure
from pymatgen.io.cif import CifWriter
from collections import Counter

# Ignore SIG_PIPE and don't throw exceptions on it...
# http://newbebweb.blogspot.co.uk/2012/02/python-head-ioerror-errno-32-broken.html
from signal import signal, SIGPIPE, SIG_DFL

signal(SIGPIPE, SIG_DFL)


def parity(list):
    return sum(list) % 2


def swap_axes(matrix, axes):
    axes_index = {"x": 0, "y": 1, "z": 2}
    matrix[:, [axes_index[axes[0]], axes_index[axes[1]]]] = matrix[
        :, [axes_index[axes[1]], axes_index[axes[0]]]
    ]
    return matrix


class Poscar:
    lines_offset = 9

    def __init__(self):
        self.title = "Title"
        self.scaling = 1.0
        self.cell = cell.Cell(np.identity(3))
        self.atoms = ["A"]
        self.atom_numbers = [1]
        self.coordinate_type = "Direct"
        self.coordinates = np.array([[0.0, 0.0, 0.0]])
        self.selective_dynamics = False

    def coordinates_by_species(self, species):
        return self.coordinates[self.range_by_species(species)]

    def range_by_species(self, species):
        i = 0
        atom_range = {}
        for a, n in zip(self.atoms, self.atom_numbers):
            atom_range[a] = range(i, i + n)
            i += n
        return atom_range[species]

    def atom_number_by_species(self, species):
        atom_numbers = {
            a: n for a, n in zip(self.atoms, self.atom_numbers)
        }
        return atom_numbers[species]

    def sorted(self, species):
        new_poscar = copy.deepcopy(self)
        new_poscar.atoms = species
        new_poscar.atom_numbers = [self.atom_number_by_species(s) for s in species]
        new_poscar.coordinates = np.concatenate(
            [self.coordinates_by_species(s) for s in species], axis=0
        )
        return new_poscar

    def read_from(self, filename):
        with open(filename) as f:
            lines = f.readlines()
        self.title = lines.pop(0).strip()
        self.scaling = float(lines.pop(0).strip())
        self.cell.matrix = np.array(
            [[float(e) for e in lines.pop(0).split()] for i in range(3)]
        )
        self.cell.inv_matrix = np.linalg.inv(self.cell.matrix)
        self.atoms = lines.pop(0).split()
        self.atom_numbers = [int(element) for element in lines.pop(0).split()]
        self.coordinate_type = lines.pop(0).strip()
        if re.match(r"\A[Ss]", self.coordinate_type):  # test for 'Selective dynamics'
            self.selective_dynamics = True
            self.coordinate_type = lines.pop(0).strip()
        self.coordinates = np.array(
            [
                [float(e) for e in lines.pop(0).split()[0:3]]
                for i in range(sum(self.atom_numbers))
            ]
        )
        if self.coords_are_cartesian():  # Convert to direct coordinates
            self.coordinates = self.fractional_coordinates()
            self.coordinate_type = "Direct"

    @classmethod
    def from_file(cls, filename):
        poscar = cls()
        poscar.read_from(filename)
        return poscar

    def in_bohr(self):
        new_poscar = copy.deepcopy(self)
        bohr_to_angstrom = 0.529177211
        new_poscar.scaling *= bohr_to_angstrom
        new_poscar.cell.matrix /= bohr_to_angstrom
        if new_poscar.coords_are_cartesian():
            new_poscar.coordinates /= bohr_to_angstrom
        return new_poscar

    def coords_are_fractional(self):
        return not self.coords_are_cartesian()

    def coords_are_cartesian(self):
        return re.match(r"\A[CcKk]", self.coordinate_type)

    def fractional_coordinates(self):
        return (
            self.coordinates
            if self.coords_are_fractional()
            else self.coordinates.dot(np.linalg.inv(self.cell.matrix))
        )
        # return ( self.coordinates if self.coords_are_fractional() else self.cell.cartesian_to_fractional_coordinates( coordinates )
        # need to check whether this alternative version works.

    def cartesian_coordinates(self):
        return (
            self.coordinates
            if self.coords_are_cartesian()
            else self.coordinates.dot(self.cell.matrix)
        )

    def select_coordinates(self, coordinate_type="Direct"):
        coord_opts = {
            "Direct": self.fractional_coordinates(),
            "Cartesian": self.cartesian_coordinates(),
        }
        return coord_opts[coordinate_type]

    def output_coordinates_only(self, coordinate_type="Direct", opts=None):
        prefix = []
        suffix = []
        for i in range(self.coordinates.shape[0]):
            prefix_string = ""
            suffix_string = ""
            if "selective" in opts:
                if opts["selective"] == "T":
                    suffix_string += " T T T"
                elif opts["selective"] == "F":
                    suffix_string += " F F F"
                elif opts["selective"]:
                    raise ValueError
            if "numbered" in opts:
                if opts["numbered"]:
                    suffix_string += " {}".format(i + 1)
            if "label" in opts:
                if opts["label"]:
                    if opts["label"] == 1:
                        prefix_string += self.labels()[i].ljust(6)
                    elif opts["label"] == 4:
                        suffix_string += " {}".format(self.labels()[i])
                    elif opts["label"]:
                        raise ValueError(opts["label"])
            prefix.append(prefix_string)
            suffix.append(suffix_string)
        for pref, coord, suff in zip(
            prefix, self.select_coordinates(coordinate_type), suffix
        ):
            print(
                pref
                + "".join(["  {: .10f}".format(element) for element in coord])
                + suff
            )

    def output(self, coordinate_type="Direct", opts=None):
        if opts is None:
            opts = {}
        if not opts.get("coordinates_only"):
            self.output_header(coordinate_type=coordinate_type, opts=opts)
        self.output_coordinates_only(coordinate_type=coordinate_type, opts=opts)

    def output_header(self, coordinate_type="Direct", opts=None):
        if opts is None:
            opts = {}
        print(self.title)
        print(self.scaling)
        if opts.get("orthorhombic"):
            matrix = self.cell.matrix * np.eye(3)
        else:
            matrix = self.cell.matrix
        for row in matrix:
            print("".join(["   {: .10f}".format(element) for element in row]))
        print(" ".join(self.atoms))
        print(" ".join([str(n) for n in self.atom_numbers]))
        if opts.get("selective"):
            print("Selective Dynamics")
        print(coordinate_type)

    def write_to(self, filename, coordinate_type="Direct", opts=None):
        if opts is None:
            opts = {}
        with open(filename, "w") as sys.stdout:
            self.output(coordinate_type=coordinate_type, opts=opts)
        sys.stdout = sys.__stdout__  # make sure sys.stdout is reset

    def output_as_xtl(self):
        print(self.title)
        print("CELL")
        cell_lengths = self.cell_lengths()
        cell_angles = self.cell_angles()
        cell_data = cell_lengths + cell_angles
        print("".join(["   {: .8f}".format(element) for element in cell_data]))
        print(" Symmetry label P1\n\nATOMS\nNAME      X       Y     Z")
        output_opts = {"label": True}
        self.output_coordinates_only(coordinate_type="Direct", opts=output_opts)

    def output_as_cif(self, symprec=None):
        print(CifWriter(self.to_pymatgen_structure(), symprec))

    def output_as_pimaim(self, to_bohr=True):
        if to_bohr is True:
            unit_scaling = angstrom_to_bohr
        else:
            unit_scaling = 1.0
        cell_lengths = self.cell.lengths() * self.scaling / unit_scaling
        print("T\nF\nF\nF")
        for row in self.cell_coordinates():
            print(" ".join([str(x / unit_scaling) for x in row]))
        for row in self.cell.unit_vectors().transpose():
            print(" ".join([str(x) for x in row]))
        for length in cell_lengths:
            print(length)

    def cell_coordinates(self):
        return self.coordinates * self.cell.lengths() * self.scaling

    def labels(self):
        return [
            atom_name
            for (atom_name, atom_number) in zip(
                self.atoms, self.atom_numbers
            )
            for __ in range(atom_number)
        ]

    def replicate(self, h, k, l, group=False):
        lattice_scaling = np.array([h, k, l], dtype=float)
        lattice_shift = np.reciprocal(lattice_scaling)
        new_poscar = Poscar()
        new_poscar.title = self.title
        new_poscar.scaling = self.scaling
        # print( lattice_scaling.dot( self.lattice ) )
        new_poscar.cell.matrix = (self.cell.matrix.T * lattice_scaling).T
        new_poscar.coordinate_type = self.coordinate_type
        new_coordinate_list = []
        cell_shift_indices = [
            [a, b, c] for a in range(h) for b in range(k) for c in range(l)
        ]
        # generate grouped / ungrouped atom names and numbers for supercell
        if group:
            new_poscar.atoms = [
                label + group for label in self.atoms for group in ("a", "b")
            ]
            new_poscar.atom_numbers = [
                int(num * h * k * l / 2) for num in self.atom_numbers for __ in (0, 1)
            ]
        else:
            new_poscar.atoms = self.atoms
            new_poscar.atom_numbers = [num * h * k * l for num in self.atom_numbers]
        # generate grouped / ungrouped atoms coordinates for supercell
        if group:
            for row in self.coordinates:
                pos_in_origin_cell = row / lattice_scaling
                for odd_even in (0, 1):
                    for a, b, c in [
                        cell_shift
                        for cell_shift in cell_shift_indices
                        if parity(cell_shift) == odd_even
                    ]:
                        new_coordinate_list.append(
                            [pos_in_origin_cell + np.array([a, b, c] * lattice_shift)][
                                0
                            ].tolist()
                        )
        else:
            for row in self.coordinates:
                pos_in_origin_cell = row / lattice_scaling
                for a, b, c in cell_shift_indices:
                    new_coordinate_list.append(
                        [pos_in_origin_cell + np.array([a, b, c] * lattice_shift)][
                            0
                        ].tolist()
                    )
        new_poscar.coordinates = np.array(new_coordinate_list)
        return new_poscar

    def cell_lengths(self):
        return (self.cell.lengths() * self.scaling).tolist()
        # return [ np.linalg.norm( row ) for row in self.cell.matrix * self.scaling ]

    def cell_angles(self):
        return self.cell.angles()
        # ( a, b, c ) = [ row for row in self.cell.matrix ]
        # return [ angle( b, c ), angle( a, c ), angle( a, b ) ]

    def swap_axes(self, axes):
        new_poscar = copy.deepcopy(self)
        new_poscar.cell = cell.Cell(swap_axes(self.cell.matrix, axes))
        return new_poscar

    def to_pymatgen_structure(self):
        lattice = pmg_Lattice(self.cell.matrix * self.scaling)
        structure = pmg_Structure(lattice, self.labels(), self.coordinates)
        return structure

    @property
    def stoichiometry(self):
        """
        Stoichiometry for this POSCAR, as a Counter.
        e.g. AB_2O_4 -> Counter( { 'A': 1, 'B': 2, O: 4 } )

        Args:
            None

        Returns:
            None
        """
        return Counter(
            {
                label: number
                for label, number in zip(self.atoms, self.atom_numbers)
            }
        )
