# -*- coding:utf-8 -*-
"""
========================================================================
Provide iteration-related file class which do operations on these files.
========================================================================
Written by PytLab <shaozhengjiang@gmail.com>, August 2015
Updated by PytLab <shaozhengjiang@gmail.com>, March 2017
========================================================================

"""
import re
from collections import namedtuple
from string import whitespace

import numpy as np
try:
    import matplotlib.pyplot as plt
    plt_installed = True
except ImportError:
    print('Warning: Module matplotlib.pyplot is not installed')
    plt_installed = False

from vaspy import VasPy, PY2
from vaspy import LazyProperty
from vaspy.atomco import PosCar, XyzFile
from vaspy.functions import line2list
# Copy a XdatCar from atomco.
from vaspy.atomco import XdatCar


class OsziCar(VasPy):
    def __init__(self, filename='OSZICAR'):
        """
        Create a OSZICAR file class.

        Example:

        >>> a = OsziCar(filename='OSZICAR')

        Class attributes descriptions
        =======================================================
          Attribute      Description
          ============  =======================================
          filename       string, name of the SPLITED DOS file
          vars           list of strings, 每次迭代得到的数据
          esort()        method, 对数据进行排序
          plot()         method, 对数据绘图
          ============  =======================================
        """
        super(self.__class__, self).__init__(filename)

        #set regex patterns
        float_regex = r'[\+|-]?\d*\.\d*(?:[e|E][\+|-]?\d+)?'
        eq_regex = r'\s*([\w|\d|\s]+)\=\s*(' + float_regex + r')\s*'
        split_regex = r'^\s*(\d+)\s*((' + eq_regex + r')+)$'  # 将step和其余部分分开

        self.eq_regex = re.compile(eq_regex)
        self.split_regex = re.compile(split_regex)

        self.load()

    def match(self, line):
        "匹配每一步迭代的数据"
        m = self.split_regex.search(line)
        if m:
            # Get step
            step = int(m.group(1))

            # Get other data
            resid = m.group(2)
            eq_tuples = self.eq_regex.findall(resid)  # list of tuples
            if PY2:
                names, numbers = zip(*eq_tuples)
            else:
                names, numbers = list(zip(*eq_tuples))

            # Remove space in names.
            names = [name.replace(' ', '') for name in names]

            #convert string to float
            numbers = [float(number) for number in numbers]
            if PY2:
                eq_tuples = [('step', step)] + zip(names, numbers)
            else:
                eq_tuples = [('step', step)] + list(zip(names, numbers))
            return eq_tuples
        else:
            return None

    def load(self):
        "加载文件数据信息"
        with open(self.filename, 'r') as f:
            content = ''
            for line in f:
                eq_tuples = self.match(line)
                if eq_tuples:  # if matched
                    if not hasattr(self, 'vars'):
                        if PY2:
                            self.vars, numbers = zip(*eq_tuples)
                        else:
                            self.vars, numbers = list(zip(*eq_tuples))

                    for name, number in eq_tuples:
                        if not hasattr(self, name):
                            setattr(self, name, [number])
                        else:
                            getattr(self, name).append(number)
                    content += line
            self.content = content
            #convert list to numpy array
            for var in self.vars:
                data = getattr(self, var)
                setattr(self, var, np.array(data))

        return

    def esort(self, var, n, reverse=False):
        '''
        进行数据var排序, 获取排序后的前n个值.

        Example:
        >>> esort('E0', 10, reverse=True)

        Parameters
        ----------
        var: string, data to be sorted.

        n: int, top numbers of sorted data.

        '''
        if PY2:
            zipped = zip(getattr(self, var), self.step)  # (E0, step)
        else:
            zipped = list(zip(getattr(self, var), self.step))
        dtype = [('var', float), ('step', int)]
        zipped = np.array(zipped, dtype=dtype)
        srted = np.sort(zipped, order='var')

        if reverse:
            return srted[-n:]
        else:
            return srted[:n]

    def plot(self, var, mode='show'):
        "绘制不同变量随step的变化曲线"
        fig = plt.figure()
        ax = fig.add_subplot(111)
        ax.set_ylabel(var)
        ax.set_xlabel('step')
        ax.plot(self.step, getattr(self, var), linewidth=2.5)
        if mode == 'show':
            plt.show()
        elif mode == 'save':
            fname = "%s_vs_%s.png" % (var, 'step')
            fig.savefig(fname)
        else:
            raise ValueError('Unrecognized show mode parameter : ' + mode)

        return fig


class OutCar(VasPy):

    # 抽取原子受力信息的正则表达式
    # Regular expression for forces information.
    force_regex = re.compile(r"^\s*POSITION\s+TOTAL-FORCE\s*\(eV\/Angst\)\s*$")
    force_info = ("ion_step", "coordinates", "forces")

    # 抽取频率信息的正则表达式
    # Regular expression for frequency information.
    float_regex = r"(\d+\.\d+)"
    freq_regex = (r"^\s*(\d+)\s*(f|f\/i)\s*=\s*" +
                  float_regex + r"\s*THz\s*" +
                  float_regex + r"\s*2PiTHz\s*" +
                  float_regex + r"\s*cm-1\s*" +
                  float_regex + r"\s*meV\s*$")
    freq_regex = re.compile(freq_regex)
    freq_info = ("index", "freq_type", "THz", "2PiTHz",
                 "cm-1", "meV", "coordinates", "deltas")
    title_regex = re.compile(r"\s*X\s*Y\s*Z\s*dx\s*dy\s*dz\s*")

    def __init__(self, filename="OUTCAR", poscar="POSCAR"):
        """
        Create a OUTCAR file class.

        Parameters:
        -----------
        filename: File name of OUTCAR, default name is "OUTCAR"(OUTCAR in current path).

        poscar: File name of POSCAR, default value is "POSCAR"(POSCAR in current path).

        Example:

        >>> a = OutCar(filename='OUTCAR', poscar="POSCAR")

        Class attributes descriptions
        =================================================================
          Attribute           Description
          ===============    ============================================
          filename            string, name of OUTCAR file
          max_forces          list of float, 每个离子步迭代的最大原子受力
          last_max_force      float, 最后一步的最大原子受力
          last_max_atom       int, 最后一步受力最大原子序号
          zpe                 float, 零点能
          freq_types          list of str, 频率类型列表
          ===============    ============================================
        """
        VasPy.__init__(self, filename)

        # Get PosCar object.
        self.poscar = PosCar(poscar)

        # Check parameter validity.
        self.__check()

    def __check(self):
        """
        Private helper function to check consistency of POSCAR and OUTCAR.
        """
        # Coordinates in OUTCAR.
        _, coords, _ = next(self.iforces)
        coords_outcar = np.array(coords)
        shape_outcar = coords_outcar.shape

        # Coordinates in POSCAR
        coords_poscar = self.poscar.data
        shape_poscar = coords_poscar.shape

        if shape_poscar != shape_outcar:
            msg = "Shape of data in POSCAR({}) and OUTCAR({}) are different."
            msg = msg.format(shape_poscar, shape_outcar)
            raise ValueError(msg)

    @property
    def iforces(self):
        """
        返回每个离子步迭代的步数，坐标和每个原子受力信息。
        Return a generator yield ionic_step, coordinates, forces on atoms.

        NOTE: ionic step starts from 1 **NOT 0**.
        """
        # Define namedtuple for item in force iteration.
        ForceItem = namedtuple('ForceItem', ['step', 'coordinates', 'forces'])

        with open(self.filename, "r") as f:
            ion_step = 0

            # Force data collection flags.
            collect_begin = False
            collecting = False

            # Collect force data for each ionic step and yield.
            for line in f:
                if not collect_begin:
                    if self.force_regex.match(line):
                        collect_begin = True
                        ion_step += 1
                elif not collecting:
                    if "-"*6 in line:
                        collecting = True
                        coordinates = []
                        forces = []
                else:
                    if "-"*6 in line:
                        collecting = False
                        collect_begin = False
                        yield ForceItem._make([ion_step, coordinates, forces])
                    else:
                        x, y, z, fx, fy, fz = line2list(line)
                        coordinates.append([x, y, z])
                        forces.append([fx, fy, fz])

    def __mask_forces(self, atom_forces, tfs):
        """
        Private helper function to use F/T info to mask forces.

        Returns:
        --------
        Masked forces 2D array.
        """
        # Check atom forces.
        if len(tfs) != len(atom_forces):
            msg = "Length of atom forces({}) must be equal to length of atoms({})."
            msg = msg.format(len(atom_forces), len(tfs))
            raise ValueError(msg)

        masked_forces = []
        for tf_vector, forces in zip(tfs, atom_forces):
            masked_force = []
            for tf, force in zip(tf_vector, forces):
                if tf == "F":
                    masked_force.append(0.0)
                else:
                    masked_force.append(force)
            masked_forces.append(masked_force)

        return masked_forces

    def fmax(self, atom_forces):
        """
        Function to get the max forces vector and atom index.

        Parameters:
        -----------
        atom_forces: 2D array of floats, forces on each atom.

        Return:
        -------
        The atom number with max force and force vector.
        NOTE: atom number start from **1 NOT 0**
        """
        # Mask forces.
        masked_forces = self.__mask_forces(atom_forces, self.poscar.tf)

        # Get max forces.
        max_force = max(masked_forces, key=lambda x: sum([i**2 for i in x]))
        index = masked_forces.index(max_force)

        return index + 1, max_force

    def forces(self, step=-1):
        """
        获取特定离子步的原子受力信息
        Function to get forces info for a specific step.

        Parameters:
        -----------
        step: The step number, int.

        Return:
        -------
        Coordinates and forces for that step.
        """
        for i, coord, forces in self.iforces:
            if step != -1 and i == step:
                return coord, forces

        if step == -1:
            return coord, forces
        elif step > i:
            raise ValueError("Illegal step {} (> {})".format(step, i))

    @LazyProperty
    def total_forces(self):
        """
        获取离子步迭代的原子最大受力（合力）列表。
        Function to get max force for every ionic step.
        """
        max_forces = []
        for _, _, forces in self.iforces:
            _, fvector = self.fmax(forces)
            max_force = np.linalg.norm(fvector)
            max_forces.append(max_force)

        return max_forces

    @LazyProperty
    def last_forces(self):
        """
        最后离子步所有原子受力矩阵。
        Function to get forces info of last ionic step.
        """
        _, forces = self.forces(-1)
        return forces

    @LazyProperty
    def last_max_force(self):
        """
        最后一步离子步中原子最大受力（合力）。
        Function to get the max force in last ionic step.

        Returns:
        --------
        The max force value of last ionic step.
        """
        _, fvector = self.fmax(self.last_forces)
        f = np.linalg.norm(fvector)

        return f

    @LazyProperty
    def last_max_atom(self):
        """
        最后一个离子步受力最大原子的原子序号。
        Function to get atom number with max force
        in the last ionic step.
        """
        atom_number, _ = self.fmax(self.last_forces)
        return atom_number

    @property
    def ifreq(self):
        """
        返回频率信息字典的迭代器。
        Return frequency iterator to generating frequency related data.
        """
        with open(self.filename, "r") as f:
            collecting = False

            for line in f:
                freq = self.freq_regex.match(line)
                title = self.title_regex.match(line)
                empty_line = (line.strip(whitespace) == "")

                if freq:
                    freq_data = list(freq.groups())

                # Collect start.
                if title and not collecting:
                    collecting = True
                    coords, deltas = [], []
                # Collect stop.
                elif empty_line and collecting:
                    collecting = False
                    freq_data.append(coords)
                    freq_data.append(deltas)
                    freq_dict = dict(zip(self.freq_info, freq_data))
                    yield freq_dict
                # Collect data.
                elif collecting:
                    x, y, z, dx, dy, dz = line2list(line)
                    coord = (x, y, z)
                    delta = (dx, dy, dz)
                    coords.append(coord)
                    deltas.append(delta)

    def check_freq_exists(func):
        """
        Decorator to check if frequency information exists.
        """
        def wrapper(self):
            try:
                next(self.ifreq)
            except StopIteration:
                msg = "'{}' has no attribtue '{}'".format(self.__class__.__name__, "zpe")
                raise AttributeError(msg)
            return func(self)

        return wrapper


    @LazyProperty
    @check_freq_exists
    def zpe(self):
        """
        从OUTCAR的频率信息返回计算的零点能。
        Function to get Zero Point Energy(ZPE) in eV.
        """
        E = [float(freq_dict["meV"])
             for freq_dict in self.ifreq if freq_dict["freq_type"] == "f"]

        return sum(E)/2000.0

    @LazyProperty
    @check_freq_exists
    def freq_types(self):
        """
        获取频率类型列表。
        Function to get frequency types.
        """
        freq_types = [freq_dict["freq_type"] for freq_dict in self.ifreq]

        # Reshape.
        try:
            freq_types = np.array(freq_types).reshape((-1, 3)).tolist()
        except ValueError:
            raise ValueError("Number of frequency can not be divided by 3.")

        return freq_types


class AniFile(VasPy):
    """ Class for file containing multiple xyz content.
        e.g. *.ANI file for molden animation visualization.

    Example:

    >>> x = MultiXyz(filename="OUT.ANI")
    """
    def __init__(self, filename="OUT.ANI"):
        VasPy.__init__(self, filename)

        # Get the total atom number for data reading.
        with open(self.filename, "r") as f:
            natom = f.readline()

        self.natom = int(natom)

    def __iter__(self):
        """ Make MultiXyz object iterable.
        """
        with open(self.filename, "r") as f:
            content_list = []
            for i, line in enumerate(f):
                if (i + 1) % (self.natom + 2) == 0:
                    content_list.append(line)
                    yield XyzFile(content_list=content_list)
                    content_list = []
                else:
                    content_list.append(line)

