'''
Random structure generation using PyXtal (https://github.com/qzhu2017/PyXtal)
'''

import collections
from contextlib import redirect_stdout, redirect_stderr
from multiprocessing import Process, Queue
import os
import random
import sys

import numpy as np
from pymatgen.core import Structure, Molecule
from pymatgen.core.periodic_table import DummySpecie
from pyxtal import pyxtal
from pyxtal.database.collection import Collection
from pyxtal.tolerance import Tol_matrix

from ...util.struc_util import check_distance, sort_by_atype, out_poscar
from ...util.struc_util import get_atype_dummy, scale_cell_mol, rot_mat
from ...IO import read_input as rin


class Rnd_struc_gen_pyxtal:
    '''
    Random structure generation using pyxtal
    '''

    def __init__(self, mindist):
        self.mindist = mindist
        self.spg_error = []

    def set_mol(self):
        '''
        set molecule files and number of molecules
        '''
        # ----------
        mol_data = []
        pyxtal_mol_data = Collection('molecules')
        pyxtal_mol_names = list(Collection('molecules'))
        for i, mf in enumerate(rin.mol_file):
            if os.path.isfile(mf):
                mol = Molecule.from_file(mf)
            elif mf in pyxtal_mol_names:
                mol = pyxtal_mol_data[mf]
            else:
                raise ValueError('no molecular files')
            mol_data.append(mol)
        # ---------- self.xxx
        self.mol_data = mol_data

    def gen_struc(self, nstruc, id_offset=0, init_pos_path=None):
        '''
        Generate random structures for given space groups

        # ---------- args
        nstruc (int): number of generated structures

        id_offset (int): default: 0
                         structure ID starts from id_offset
                         e.g. nstruc = 3, id_offset = 10
                              you obtain ID 10, ID 11, ID 12

        init_pos_path (str): default: None
                             specify a path of file
                             if you write POSCAR data of init_struc_data
                             ATTENSION: data are appended to the specified file

        # ---------- comment
        generated structure data are saved in self.init_struc_data
        '''
        # ---------- initialize
        self.init_struc_data = {}
        # ---------- Tol_matrix
        tolmat = self._set_tol_mat(rin.atype, self.mindist)
        # ---------- loop for structure generattion
        while len(self.init_struc_data) < nstruc:
            # ------ spgnum --> spg
            if rin.spgnum == 'all':
                spg = random.randint(1, 230)
            else:
                spg = random.choice(rin.spgnum)
            if spg in self.spg_error:
                continue
            # ------ vol_factor
            rand_vol = random.uniform(rin.vol_factor[0], rin.vol_factor[1])
            # ------ generate structure
            tmp_crystal = pyxtal()
            try:
                with redirect_stdout(sys.stderr):
                    tmp_crystal.from_random(dim=3, group=spg, species=rin.atype,
                                            numIons=rin.nat, factor=rand_vol,
                                            conventional=False, tm=tolmat)
            except Exception as e:
                print(e, ':spg = {} retry.'.format(spg), file=sys.stderr, flush=True)
                self.spg_error.append(spg)
                continue
            if tmp_crystal.valid:
                tmp_struc = tmp_crystal.to_pymatgen(resort=False)    # pymatgen Structure format
                # -- check the number of atoms
                if not self._check_nat(tmp_struc):
                    # (pyxtal 0.1.4) cryspy adopts "conventional=False",
                    #     which is better for DFT calculation
                    # pyxtal returns conventional cell, that is, too many atoms
                    tmp_struc = tmp_struc.get_primitive_structure()
                    # recheck nat
                    if not self._check_nat(tmp_struc):    # failure
                        continue
                # -- sort, just in case
                tmp_struc = sort_by_atype(tmp_struc, rin.atype)
                # -- scale volume
                if rin.vol_mu is not None:
                    vol = random.gauss(mu=rin.vol_mu, sigma=rin.vol_sigma)
                    tmp_struc.scale_lattice(volume=vol)
                # -- check minimum distance
                #    from CrySPY 0.10.4
                #    Tol_matrix is used for mindist
                #
                #if self.mindist is not None:
                #    success, mindist_ij, dist = check_distance(tmp_struc,
                #                                               self.atype,
                #                                               self.mindist)
                #    if not success:
                #        print('mindist in gen_struc: {} - {}, {}. retry.'.format(
                #            self.atype[mindist_ij[0]],
                #            self.atype[mindist_ij[1]],
                #            dist), file=sys.stderr)
                #        continue    # failure
                # -- check actual space group
                try:
                    spg_sym, spg_num = tmp_struc.get_space_group_info(
                        symprec=rin.symprec)
                except TypeError:
                    spg_num = 0
                    spg_sym = None
                # -- register the structure in pymatgen format
                cid = len(self.init_struc_data) + id_offset
                self.init_struc_data[cid] = tmp_struc
                print('Structure ID {0:>6} was generated.'
                      ' Space group: {1:>3} --> {2:>3} {3}'.format(
                       cid, spg, spg_num, spg_sym), flush=True)
                # -- save init_POSCARS
                if init_pos_path is not None:
                    out_poscar(tmp_struc, cid, init_pos_path)
            else:
                self.spg_error.append(spg)

    def gen_struc_mol(self, nstruc, id_offset=0, init_pos_path=None):
        '''
        Generate random molecular crystal structures for given space groups
        one have to run self.set_mol() in advance
        # ---------- args
        nstruc (int): number of generated structures

        id_offset (int): default: 0
                         structure ID starts from id_offset
                         e.g. nstruc = 3, id_offset = 10
                              you obtain ID 10, ID 11, ID 12

        init_pos_path (str): default: None
                             specify a path of file
                             if you write POSCAR data of init_struc_data
                             ATTENSION: data are appended to the specified file

        # ---------- comment
        generated structure data are saved in self.init_struc_data
        '''
        # ---------- check args
        if not (type(nstruc) is int and nstruc > 0):
            raise ValueError('nstruc must be positive int')
        if type(id_offset) is not int:
            raise TypeError('id_offset must be int')
        if init_pos_path is None or type(init_pos_path) is str:
            pass
        else:
            raise ValueError('init_pos_path is wrong.'
                             ' init_pos_path = {}'.format(init_pos_path))
        # ---------- initialize
        self.init_struc_data = {}
        if rin.algo == 'EA' and rin.struc_mode in ['mol', 'mol_bs']:
            self.struc_mol_id = {}
        # ---------- Tol_matrix
        tolmat = self._set_tol_mat(rin.atype, self.mindist)
        # ---------- loop for structure generattion
        while len(self.init_struc_data) < nstruc:
            # ------ spgnum --> spg
            if rin.spgnum == 'all':
                spg = random.randint(1, 230)
            else:
                spg = random.choice(rin.spgnum)
            if spg in self.spg_error:
                continue
            rand_vol = random.uniform(rin.vol_factor[0], rin.vol_factor[1])
            # ------ generate structure
            # -- multiprocess for measures against hangup
            q = Queue()
            p = Process(target=self._mp_mc, args=(tolmat, spg, rin.nmol, rand_vol, q, rin.algo))
            p.start()
            p.join(timeout=rin.timeout_mol)
            if p.is_alive():
                p.terminate()
                p.join()
            if sys.version_info.minor >= 7:
                # Process.close() available from python 3.7
                p.close()
            if q.empty():
                print('timeout for molecular structure generation. retry.',
                      file=sys.stderr, flush=True)
                continue
            else:
                # -- get struc data from _mp_mc
                if rin.algo == 'EA':
                    tmp_q = q.get()
                    tmp_struc = tmp_q[0]    # structure data
                    dums = tmp_q[1]         # dummy element
                    in_dists = tmp_q[2]     # interatomic distance in a molecule
                    mol_dists = tmp_q[3]    # distance between molecules
                else:
                    tmp_struc = q.get()
                tmp_valid = q.get()
                if tmp_struc == 'error':
                    self.spg_error.append(spg)
                    continue
            if tmp_valid:
                # -- scale volume
                if rin.vol_mu is not None:
                    vol = random.gauss(mu=rin.vol_mu, sigma=rin.vol_sigma)
                    vol = vol * tmp_struc.num_sites / rin.natot    # for conv. cell
                    tmp_struc = scale_cell_mol(tmp_struc, self.mol_data, vol)
                    if not tmp_struc:    # case: scale_cell_mol returns False
                        print('failed scale cell. retry.', file=sys.stderr, flush=True)
                        continue
                # -- check nat
                if not self._check_nat(tmp_struc):
                    # cryspy adopts conventional=True
                    # pyxtal returns conventional cell,
                    # too many atoms if centering
                    tmp_struc = tmp_struc.get_primitive_structure()
                    # recheck nat
                    if not self._check_nat(tmp_struc):    # failure
                        print('different num. of atoms. retry.', file=sys.stderr, flush=True)
                        continue
                # -- grouping atoms for molecule using interatomic distance
                if rin.algo == 'EA':
                    loopcnt = 0
                    while loopcnt < 10:
                        loopcnt += 1
                        mol_group = []
                        group_id = 0
                        # check atom distance from dummy atom
                        for i, ts in enumerate(tmp_struc.species):
                            for dn, dummyatom in enumerate(dums):
                                if ts == dummyatom:
                                    mol_indx = int(ts.symbol[1:2])
                                    for j, tc_fc in enumerate(tmp_struc.frac_coords):
                                        if i == j:
                                            continue
                                        dist = tmp_struc.get_distance(i, j)
                                        for dist_indx, true_dist in enumerate(in_dists[dn]):
                                            if round(true_dist, 2) == round(dist, 2):
                                                mol_group.append([tc_fc, group_id, mol_indx, dist_indx])
                                                break
                                    group_id += 1
                        success_cnt = 0
                        # check number of atom per group
                        # if failed, shuffle struc data, and continue(max 10times)
                        mol_cnt = collections.Counter([x[2] for x in mol_group])
                        for mc_key, mc_value in mol_cnt.items():
                            if mc_value % len(self.mol_data[mc_key]) == 0:
                                success_cnt += 1
                        if success_cnt == len(self.mol_data):
                            break
                        random.shuffle(tmp_struc)
                    if loopcnt == 10:
                        continue
                    tmp_struc.remove_species(dums)
                # -- sort, necessary in molecular crystal
                tmp_struc = sort_by_atype(tmp_struc, rin.atype)
                # -- record group_id, mol_id, and dist_index
                if rin.algo == 'EA':
                    tmp_id = []
                    tmp_mol_indx = []
                    tmp_dist_indx = []
                    for atom_fc in tmp_struc.frac_coords:
                        for gatom in mol_group:
                            if np.array_equal(atom_fc, gatom[0]):
                                tmp_id.append(gatom[1])
                                tmp_mol_indx.append(gatom[2])
                                tmp_dist_indx.append(gatom[3])
                # -- check minimum distance
                #    from CrySPY 0.10.4
                #    Tol_matrix is used for mindist
                #
                #if self.mindist is not None:
                #    success, mindist_ij, dist = check_distance(tmp_struc,
                #                                               self.atype,
                #                                               self.mindist)
                #    if not success:
                #        print('mindist in gen_struc_mol: {} - {}, {}. retry.'.format(
                #            self.atype[mindist_ij[0]],
                #            self.atype[mindist_ij[1]],
                #            dist), file=sys.stderr)
                #        continue    # failure
                # -- check actual space group
                try:
                    spg_sym, spg_num = tmp_struc.get_space_group_info(
                        symprec=rin.symprec)
                except TypeError:
                    spg_num = 0
                    spg_sym = None
                # -- register the structure in pymatgen format
                cid = len(self.init_struc_data) + id_offset
                self.init_struc_data[cid] = tmp_struc
                if rin.algo == 'EA' and rin.struc_mode in ['mol', 'mol_bs']:
                    self.struc_mol_id.update({cid: [tmp_mol_indx, tmp_id, mol_dists]})
                print('Structure ID {0:>6} was generated.'
                      ' Space group: {1:>3} --> {2:>3} {3}'.format(
                       cid, spg, spg_num, spg_sym), flush=True)
                # -- save init_POSCARS
                if init_pos_path is not None:
                    out_poscar(tmp_struc, cid, init_pos_path)
            else:
                self.spg_error.append(spg)

    def gen_struc_mol_break_sym(self, nstruc, mindist_dummy,
                                id_offset=0, init_pos_path=None):
        '''
        Generate random molecular crystal structures
        one have to run self.set_mol() in advance
        molecules are put a Wyckoff position without consideration of symmetry

        # ---------- args
        nstruc (int): number of generated structures

        mindist_dummy: mindist for dummy atoms

        id_offset (int): default: 0
                         structure ID starts from id_offset
                         e.g. nstruc = 3, id_offset = 10
                              you obtain ID 10, ID 11, ID 12

        init_pos_path (str): default: None
                             specify a path of file
                             if you write POSCAR data of init_struc_data
                             ATTENSION: data are appended to the specified file

        # ---------- comment
        generated structure data are saved in self.init_struc_data
        '''
        # ---------- check args
        if not (type(nstruc) is int and nstruc > 0):
            raise ValueError('nstruc must be positive int')
        if type(id_offset) is not int:
            raise TypeError('id_offset must be int')
        if init_pos_path is None or type(init_pos_path) is str:
            pass
        else:
            raise ValueError('init_pos_path is wrong.'
                             ' init_pos_path = {}'.format(init_pos_path))
        # ---------- initialize
        self.init_struc_data = {}
        if rin.algo == 'EA' and rin.struc_mode in ['mol', 'mol_bs']:
            self.struc_mol_id = {}
        # ------ dummy atom type
        atype_dummy = get_atype_dummy()
        # ---------- Tol_matrix for dummy atoms
        tolmat = self._set_tol_mat(atype_dummy, mindist_dummy)
        # ---------- loop for structure generattion
        while len(self.init_struc_data) < nstruc:
            # ------ spgnum --> spg
            if rin.spgnum == 'all':
                spg = random.randint(1, 230)
            else:
                spg = random.choice(rin.spgnum)
            if spg in self.spg_error:
                continue
            # ------ vol_factor
            rand_vol = random.uniform(rin.vol_factor[0], rin.vol_factor[1])
            # ------ generate structure
            tmp_crystal = pyxtal()
            try:
                with redirect_stdout(sys.stderr):
                    tmp_crystal.from_random(dim=3, group=spg, species=atype_dummy,
                                            numIons=rin.nmol, factor=rand_vol,
                                            conventional=False, tm=tolmat)
            except Exception as e:
                print(e, ':spg = {} retry.'.format(spg), file=sys.stderr, flush=True)
                self.spg_error.append(spg)
                continue
            if tmp_crystal.valid:
                # -- each wyckoff position --> dummy atom
                dums = []        # dummy atoms
                dum_pos = []     # internal position of dummy
                dum_type = {}    # type of dummy
                                 #  e.g. {DummySpecie X00+: 'Rn',
                                 #        DummySpecie X10+: 'Rn',
                                 #        DummySpecie X20+: 'Xe',
                                 #        DummySpecie X30+: 'Xe'}
                for i, site in enumerate(tmp_crystal.atom_sites):
                    dum = DummySpecie("X{}".format(i))
                    dums.append(dum)
                    dum_pos.append(site.position)
                    dum_type[dum] = site.specie
                # -- tmp_struc with dummy atoms
                lattice = tmp_crystal.lattice.get_matrix()
                tmp_struc = Structure.from_spacegroup(spg, lattice, dums, dum_pos)
                tmp_struc = tmp_struc.get_primitive_structure()
                # -- scale volume
                if rin.vol_mu is not None:
                    vol = random.gauss(mu=rin.vol_mu, sigma=rin.vol_sigma)
                    tmp_struc.scale_lattice(volume=vol)
                # -- save dummy coords
                dum_species = tmp_struc.species
                dum_coords = tmp_struc.cart_coords
                if rin.algo == 'EA':
                    tmp_mol_indx = []
                    tmp_id = []
                    tmp_id_cnt = 0
                # -- remove dummy
                tmp_struc.remove_sites(range(0, len(dum_species)))
                tmp_struc_ori = tmp_struc.copy()
                rot_success = False
                # -- rotate molecules
                for nrel in range(rin.nrot):
                    tmp_struc = tmp_struc_ori.copy()
                    for (dum_specie, dum_coord) in zip(dum_species, dum_coords):
                        mol_index = atype_dummy.index(dum_type[dum_specie])
                        mol = self.mol_data[mol_index]
                        # rotation option
                        if rin.rot_mol is None:
                            rot_mol_coord = mol.cart_coords
                        if rin.rot_mol == 'random':
                            angle = 2 * np.pi * np.random.rand(3)
                            R = rot_mat(angle)
                            rot_mol_coord = np.matmul(mol.cart_coords, R)
                        if rin.rot_mol == 'random_mol':
                            # each mol_data
                            mol_angles = []    # [angles of mol 1, angles of mol 2, ...]
                            for i in range(len(self.mol_data)):
                                mol_angles.append(2 * np.pi * np.random.rand(3))
                            angle = mol_angles[mol_index]
                            R = rot_mat(angle)
                            rot_mol_coord = np.matmul(mol.cart_coords, R)
                        if rin.rot_mol == 'random_wyckoff':
                            # each Wyckoff
                            dum_angles = {}    # e.g.
                                               # {DummySpecie X00+: array([ , , ]),
                                               #  DummySpecie X10+: array([ , , ]),
                                               #  DummySpecie X20+: array([ , , ]),
                                               #  DummySpecie X30+: array([ , , ])}
                            angles = 2 * np.pi * np.random.rand(len(dums), 3)
                            for (dum, angle) in zip(dums, angles):
                                dum_angles[dum] = angle
                            angle = dum_angles[dum_specie]
                            R = rot_mat(angle)
                            rot_mol_coord = np.matmul(mol.cart_coords, R)
                        # rotate coord
                        coord = rot_mol_coord + dum_coord
                        # append mol
                        for i, ms in enumerate(mol.species):
                            tmp_struc.append(ms, coord[i], coords_are_cartesian=True)
                            if rin.algo == 'EA':
                                tmp_id.append(tmp_id_cnt)
                                tmp_mol_indx.append(mol_index)
                        if rin.algo == 'EA':
                            tmp_id_cnt += 1
                    # -- check nat
                    if not self._check_nat(tmp_struc):
                        # pyxtal returns conventional cell,
                        # too many atoms if centering
                        tmp_struc = tmp_struc.get_primitive_structure()
                        # recheck nat
                        if not self._check_nat(tmp_struc):    # failure
                            if rin.rot_mol is None:
                                break    # go back to the while loop
                            continue
                    # -- sort, necessary in molecular crystal
                    tmp_struc = sort_by_atype(tmp_struc, rin.atype)
                    # -- check minimum distance
                    if self.mindist is not None:
                        success, mindist_ij, dist = check_distance(tmp_struc,
                                                                   rin.atype,
                                                                   self.mindist)
                        if not success:
                            print('mindist: {} - {}, {}. retry.'.format(
                                rin.atype[mindist_ij[0]],
                                rin.atype[mindist_ij[1]],
                                dist), file=sys.stderr, flush=True)
                            if rin.rot_mol is None:
                                break    # go back to the while loop
                            continue    # failure
                    # -- check actual space group (success)
                    try:
                        spg_sym, spg_num = tmp_struc.get_space_group_info(
                            symprec=rin.symprec)
                    except TypeError:
                        spg_num = 0
                        spg_sym = None
                    rot_success = True
                    break
                # -- reach maximum times to rotate (failure)
                if not rot_success:
                    continue    # go back to the while loop
                # -- register the structure in pymatgen format
                cid = len(self.init_struc_data) + id_offset
                self.init_struc_data[cid] = tmp_struc
                if rin.algo == 'EA' and rin.struc_mode in ['mol', 'mol_bs']:
                    # calculate interatomic distance of mol_data
                    mol_dists = []
                    for j, mol in enumerate(self.mol_data):
                        tmp_dists = []
                        for n, m in enumerate(mol):
                            tmp_dists.append(mol.get_distance(0, n))
                        mol_dists.append(tmp_dists)
                    self.struc_mol_id.update({cid: [tmp_mol_indx, tmp_id, mol_dists]})
                print('Structure ID {0:>6} was generated.'
                      ' Space group: {1:>3} --> {2:>3} {3}'.format(
                       cid, spg, spg_num, spg_sym), flush=True)
                # -- save init_POSCARS
                if init_pos_path is not None:
                    out_poscar(tmp_struc, cid, init_pos_path)
            else:
                self.spg_error.append(spg)

    def _set_tol_mat(self, atype, mindist):
        tolmat = Tol_matrix()
        for i, itype in enumerate(atype):
            for j, jtype in enumerate(atype):
                if i <= j:
                    tolmat.set_tol(itype, jtype, mindist[i][j])
        return tolmat

    def _check_nat(self, struc):
        # ---------- count number of atoms in each element for check
        species_list = [a.species_string for a in struc]
        for i in range(len(rin.atype)):
            if species_list.count(rin.atype[i]) != rin.nat[i]:
                return False    # failure
        return True

    def _mp_mc(self, tolmat, spg, nmol, rand_vol, q, algo):
        '''
        multiprocess part
        here cannot use rin.xxx
        '''
        try:
            np.random.seed(random.randint(0, 100000000))
            tmp_crystal = pyxtal(molecular=True)
            with open('err_cryspy', 'a') as f:
                with redirect_stdout(f):
                    with redirect_stderr(f):
                        tmp_crystal.from_random(dim=3, group=spg,
                                                species=self.mol_data, numIons=nmol,
                                                factor=rand_vol, conventional=False, tm=tolmat)
            if algo == 'EA':
                tmp_struc = tmp_crystal.to_pymatgen(resort=False)
                tmp_lattice = tmp_struc.lattice
                dums = []        # dummy atoms
                dum_pos = []     # internal position of dummy
                in_dists = []
                mol_dists = []
                for i, site in enumerate(tmp_crystal.mol_sites):
                    for j, mol in enumerate(self.mol_data):
                        if mol.species == site.mol.species:
                            dum = DummySpecie("X{}{}".format(j, i))
                            dums.append(dum)
                            dum_pos.append(site.position)
                            # -- calculate atom distance by pyxtal
                            in_dists.append(np.linalg.norm(site.position.dot(tmp_crystal.lattice.get_matrix()) -
                                                           site.get_coords_and_species(absolute=True)[0]
                                                           [0:len(site.mol.species)], axis=1))
                # -- calculate atom distance from self.mol_data
                for j, mol in enumerate(self.mol_data):
                    tmp_dists = []
                    for n, m in enumerate(mol):
                        tmp_dists.append(mol.get_distance(0, n))
                    mol_dists.append(tmp_dists)
                dum_struc = Structure.from_spacegroup(spg, tmp_lattice, dums, dum_pos)
                for ds, dc in zip(dum_struc.species, dum_struc.cart_coords):
                    tmp_struc.append(ds, dc, coords_are_cartesian=True)
            # ---------- queue
            if tmp_crystal.valid:
                if algo == 'EA':
                    q.put([tmp_struc, dums, in_dists, mol_dists])
                else:
                    q.put(tmp_crystal.to_pymatgen(resort=False))
                q.put(tmp_crystal.valid)
            else:
                q.put(None)
                q.put(tmp_crystal.valid)
        except Exception as e:
            with open('err_cryspy', 'a') as f:
                f.write(e.args[0])
                f.write('    spg = {} retry.\n'.format(spg))
            q.put('error')
            q.put(None)
