# coding: utf-8
# Distributed under the terms of the MIT License.

import pdb
import sys
import os.path

from ababe.stru.scaffold import SitesGrid, CStru, GeneralCell
from ababe.stru.element import GhostSpecie, Specie
from itertools import combinations, zip_longest
from progressbar import ProgressBar

import numpy as np
import spglib
from spglib import get_symmetry
import os
import xxhash
from itertools import tee
# Filename sogen is for Site-Occupy-GENerator


def numbers2id(numbers_arr):
    num_hash = xxhash.xxh64(numbers_arr).intdigest()
    return num_hash


class OccupyGenerator(object):
    """
    The class takes a GeneralCell instance as input,
    and can produces a generator GeneralCell instances which are
    nonduplicated to each others.
    """
    def __init__(self, general_cell):
        self.init_cell = general_cell
        self.lattice = general_cell.lattice
        self.positions = general_cell.positions
        self.numbers = general_cell.numbers

        self.symmetry_permutation = general_cell.get_symmetry_permutation()
        self.wyckoffs = general_cell.get_wyckoffs()

    def is_equivalent(self, cell_i, cell_other):
        numbers_i = cell_i.numbers
        numbers_other = cell_other.numbers
        for perm in self.symmetry_permutation:
            new_numbers = numbers_i[perm]
            if np.array_equal(new_numbers, numbers_other):
                return True

        return False

    def gen_dup_unitary(self, n, sp):
        init_numbers = self.init_cell.numbers
        num_count = init_numbers.size   # number of atoms in structure
        i_speckle = sp.Z
        from itertools import combinations
        for comb_index in combinations(range(num_count), n):
            numbers = init_numbers.copy()
            for index in comb_index:
                numbers[index] = i_speckle
            yield GeneralCell(self.lattice, self.positions, numbers)

    def gen_nodup_unitary(self, n, sp):
        dup = self.gen_dup_unitary(n, sp)
        # sym_perm = self.symmetry_permutation

        # isoset = dict()
        # for cell in dup:
        #     cell_id = cell.id
        #     if cell_id not in isoset:
        #         yield cell
        #         # from ababe.stru.io import VaspPOSCAR
        #         # pdb.set_trace()
        #         self._update_isoset(isoset, cell.numbers, sym_perm)
        return self.gen_2nodup_gen(dup)

    def gen_dup(self, wy, n, sp):
        init_numbers = self.init_cell.numbers
        i_speckle = sp.Z
        wyckoffs = self.wyckoffs
        sp_ind = [i for i, w in enumerate(wyckoffs) if w is wy]
        # pdb.set_trace()
        for comb_index in combinations(sp_ind, n):
            numbers = init_numbers.copy()
            for index in comb_index:
                numbers[index] = i_speckle
            yield GeneralCell(self.lattice, self.positions, numbers)

    def gen_nodup(self, wy, n, sp):
        dup = self.gen_dup(wy, n, sp)
        return self.gen_2nodup_gen(dup)

    def gen_dup_of_ele(self, ele, n, sp):
        # ele is a Specie instance

        ele_num = ele.Z
        init_numbers = self.init_cell.numbers
        i_speckle = sp.Z
        sp_ind = [i for i, e in enumerate(init_numbers) if e == ele_num]

        for comb_index in combinations(sp_ind, n):
            numbers = init_numbers.copy()
            for index in comb_index:
                numbers[index] = i_speckle
            yield GeneralCell(self.lattice, self.positions, numbers)

    def gen_nodup_of_ele(self, ele, n, sp):
        dup = self.gen_dup_of_ele(ele, n, sp)
        return self.gen_2nodup_gen(dup)

    def gen_dup_exch(self, sp1, sp2, n):
        init_numbers = self.init_cell.numbers
        n_sp1 = sp1.Z
        n_sp2 = sp2.Z
        sp1_ind = [i for i, e in enumerate(init_numbers) if e == n_sp1]
        sp2_ind = [i for i, e in enumerate(init_numbers) if e == n_sp2]

        for ex_sp1 in combinations(sp1_ind, n):
            for ex_sp2 in combinations(sp2_ind, n):
                numbers = init_numbers.copy()
                for ind_sp1, ind_sp2 in zip(ex_sp1, ex_sp2):
                    numbers[ind_sp1] = n_sp2
                    numbers[ind_sp2] = n_sp1
                yield GeneralCell(self.lattice, self.positions, numbers)


    def gen_nodup_exch(self, sp1, sp2, n):
        dup = self.gen_dup_exch(sp1, sp2, n)
        return self.gen_2nodup_gen(dup)

    def _gen_dup_trinary_alloy(self, sp1, n1, sp2, n2):
        init_numbers = self.init_cell.numbers
        isp1 = sp1.Z
        isp2 = sp2.Z

        sp_ind_origin = [i for i, s in enumerate(init_numbers)]
        for sp1_comb_index in combinations(sp_ind_origin, n1):
            sp_ind_bin = [x for x in sp_ind_origin if x not in sp1_comb_index]
            for sp2_comb_index in combinations(sp_ind_bin, n2):
                numbers = init_numbers.copy()
                for i1, i2 in zip_longest(sp1_comb_index, sp2_comb_index):
                    if i1 is not None:
                        numbers[i1] = isp1
                    if i2 is not None:
                        numbers[i2] = isp2
                yield GeneralCell(self.lattice, self.positions, numbers)
                # pdb.set_trace()

    def gen_nodup_trinary_alloy(self, sp1, n1, sp2, n2):
        dup = self._gen_dup_trinary_alloy(sp1, n1, sp2, n2)
        return self.gen_2nodup_gen(dup)

    def gen_add_one_speckle_unitary(self, gen, sp):
        """
        input a structure generator __ spg_cell(mostly nonduplicate)
        output a generator with one more speckle.
        This a method give duplicate structures which have one more speckle
        than the input structures.
        """
        atom = sp.Z
        id_db = dict()
        for cell in gen:
            for index, val in enumerate(cell.numbers):
                numbers_new = cell.numbers.copy()
                if atom != val:
                    numbers_new[index] = atom
                    num_id = numbers2id(numbers_new)
                    if num_id not in id_db:
                        yield GeneralCell(cell.lattice, cell.positions,
                                          numbers_new)
                        id_db[num_id] = None

    def gen_add_one_speckle(self, gen, wy, sp):
        atom = sp.Z
        id_db = dict()
        for cell in gen:
            wy_ele = zip(self.wyckoffs, cell.numbers)
            for index, tuple_w_e in enumerate(wy_ele):
                numbers_new = cell.numbers.copy()
                if tuple_w_e[0] == wy and atom != tuple_w_e[1]:
                    numbers_new[index] = atom
                    num_id = numbers2id(numbers_new)
                    if num_id not in id_db:
                        yield GeneralCell(cell.lattice, cell.positions,
                                          numbers_new)
                        id_db[num_id] = None

    def gen_add_one_speckle_of_ele(self, gen, ele, sp):
        ele_Z = ele.Z
        atom = sp.Z
        id_db = dict()
        for cell in gen:
            # site_ele = zip(self.wyckoffs, cell.numbers)
            for index, e in enumerate(cell.numbers):
                numbers_new = cell.numbers.copy()
                if e == ele_Z and e != atom:
                    numbers_new[index] = atom
                    num_id = numbers2id(numbers_new)
                    if num_id not in id_db:
                        yield GeneralCell(cell.lattice, cell.positions,