#!/usr/bin/env python
# -*- coding: utf-8 -*-

# This file is part of the udkm1Dsimpy module.
#
# udkm1Dsimpy is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, see <http://www.gnu.org/licenses/>.
#
# Copyright (C) 2019 Daniel Schick

__all__ = ["UnitCell"]

__docformat__ = "restructuredtext"

import numpy as np
from inspect import isfunction
from sympy import integrate, Symbol
from sympy.utilities.lambdify import lambdify
from tabulate import tabulate
from . import u, Q_


class UnitCell:
    """UnitCell

    The unitCell class hold different structural properties of real
    physical unit cells and also an array of atoms at different postions
    in the unit cell.

    Args:
        id (str): id of the UnitCell
        name (str): name of the UnitCell
        c_axis (float): c-axis of the UnitCell

    Keyword Args:
        a_axis (float): a-axis of the UnitCell
        b_axis (float): b-axis of the UnitCell
        deb_wal_fac (float): Debye Waller factor
        sound_vel (float): sound velocity
        phonon_damping (float): phonon damping
        opt_pen_depth (float): optical penetration depth
        opt_ref_index (float): refractive index
        opt_ref_index_per_strain (float): change of refractive index per
           strain
        heat_capacity (float): heat capacity
        therm_cond (float): thermal conductivity
        lin_therm_exp (float): linear thermal expansion
        sub_system_coupling (float): sub-system coupling

    Attributes:
        id (str): id of the unit cell
        name (str): name of the unit cell
        atoms (list[atom, @lambda]): list of atoms and funtion handle
           for strain dependent displacement
        num_atoms (int): number of atoms in unit cell
        deb_wal_fac (float): Debye Waller factor <u>² [m²]
        spring_const (ndarray[float]): spring constant of the unit cell
           [kg/s²] and higher orders
        opt_ref_index (ndarray[float]): optical refractive index - real
           and imagenary part :math:`n + i\kappa`
        opt_ref_index_per_strain (ndarray[float]): optical refractive
           index change per strain - real and imagenary part
           :math:`\\frac{d n}{d \eta} + i\\frac{d \kappa}{d \eta}`
        therm_cond (list[@lambda]): list of HANDLES T-dependent thermal
           conductivity [W/(m K)]
        lin_therm_exp (list[@lambda]): list of HANDLES T-dependent
           linear thermal expansion coefficient (relative)
        int_lin_therm_exp (list[@lambda]): list of HANDLES T-dependent
           integrated linear thermal expansion coefficient
        heat_capacity (list[@lambda]): list of HANDLES T-dependent heat
           capacity function [J/(kg K)]
        int_heat_capacity (list[@lambda]): list of HANDLES T-dependent
           integrated heat capacity function
        sub_system_coupling (list[@lambda]): list of HANDLES of coupling
           functions of different subsystems [W/m³]
        num_sub_systems (int): number of subsystems for heat and phonons
           (electrons, lattice, spins, ...)

    """

    def __init__(self, id, name, c_axis, **kwargs):
        self.id = id
        self.name = name
        self.c_axis = c_axis
        self.a_axis = kwargs.get('a_axis', self.c_axis)
        self.b_axis = kwargs.get('b_axis', self.a_axis)
        self.atoms = []
        self.num_atoms = 0
        self.mass = 0*u.kg
        self.density = 0*u.kg/u.m**2
        self.spring_const = np.array([0])
        self.deb_wal_fac = kwargs.get('deb_wal_fac', 0*u.m**2)
        self.sound_vel = kwargs.get('sound_vel', 0)
        self.phonon_damping = kwargs.get('phonon_damping', 0*u.kg/u.s)
        self.opt_pen_depth = kwargs.get('opt_pen_depth', 0*u.nm)
        self.opt_ref_index = kwargs.get('opt_ref_index', 0)
        self.opt_ref_index_per_strain = kwargs.get('opt_ref_index_per_strain', 0)
        self.heat_capacity, self.heat_capacity_str = self.check_cell_array_input(
                kwargs.get('heat_capacity', 0))
        self.therm_cond, self.therm_cond_str = self.check_cell_array_input(
                kwargs.get('therm_cond', 0))
        self.lin_therm_exp, self.lin_therm_exp_str = self.check_cell_array_input(
                kwargs.get('lin_therm_exp', 0))
        self.sub_system_coupling, self.sub_system_coupling_str = self.check_cell_array_input(
                kwargs.get('sub_system_coupling', 0))

        if (len(self.heat_capacity) == len(self.therm_cond)
            and len(self.heat_capacity) == len(self.lin_therm_exp)
                and len(self.heat_capacity) == len(self.sub_system_coupling)):
            self.num_sub_systems = len(self.heat_capacity)
        else:
            raise ValueError('Heat capacity, thermal conductivity, linear'
                             'thermal expansion and subsystem coupling have not'
                             'the same number of elements!')

        self.area = self.a_axis * self.b_axis
        self.volume = self.area * self.c_axis

    def __str__(self):
        """String representation of this class"""
        output = [['id', self.id],
                  ['name', self.name],
                  ['a-axis', '{:.4~P}'.format(self.a_axis)],
                  ['b-axis', '{:.4~P}'.format(self.b_axis)],
                  ['c-axis', '{:.4~P}'.format(self.c_axis)],
                  ['area', '{:.4~P}'.format(self.area.to('angstrom**2'))],
                  ['volume', '{:.4~P}'.format(self.volume.to('angstrom**3'))],
                  ['mass', '{:.4~P}'.format(self.mass)],
                  ['density', '{:.4~P}'.format(self.density.to('kg/meter**3'))],
                  ['Debye Waller Factor', self.deb_wal_fac.to('meter**2')],
                  ['sound velocity', '{:.4~P}'.format(self.sound_vel.to('meter/s'))],
                  ['spring constant', self.spring_const * u.kg/u.s**2],
                  ['phonon damping', self.phonon_damping.to('kg/s')],
                  ['opt. pen. depth', self.opt_pen_depth.to('nm')],
                  ['opt. refractive index', self.opt_ref_index],
                  ['opt. ref. index/strain', self.opt_ref_index_per_strain],
                  ['thermal conduct.', ' W/(m K)\n'.join(self.therm_cond_str) + ' W/(m K)'],
                  ['linear thermal expansion', '\n'.join(self.lin_therm_exp_str)],
                  ['heat capacity', ' J/(kg K)\n'.join(self.heat_capacity_str) + ' J/(kg K)'],
                  ['subsystem coupling', ' W/m³\n'.join(self.sub_system_coupling_str) + ' W/m³']]

        class_str = 'Unit Cell with the following properties\n\n'
        class_str += tabulate(output, headers=['parameter', 'value'], tablefmt="rst",
                              colalign=('right',), floatfmt=('.2f', '.2f'))
        class_str += '\n\n' + str(self.num_atoms) + ' Constituents:\n'

        atoms_str = []
        for i in range(self.num_atoms):
            atoms_str.append([self.atoms[i][0].name,
                              '{:0.2f}'.format(self.atoms[i][1](0)),
                              self.atoms[i][2]])
        class_str += tabulate(atoms_str, headers=['atom', 'position', 'position function'],
                              tablefmt="rst")
        return class_str

    def clone_2_multiple(self, N, *args):
        """clone_2_multiple"""
#        %% clone2multiple
#        % Returns a cloned unit cell N time repeated. Accordingly, all
#        % physical properties are adapted, *despite of any specific function
#        % handle of the atomic position in the unit cell and higher orders
#        % of spring constants!*
#        % If no _ID_ or _name_ is given for the cloned unit cell, the name
#        % of the current unit cell is taken plus $N$.
#        function clone = clone2multiple(obj,N,varargin)
#            % initialize input parser and define defaults and validators
#            p = inputParser;
#            p.addRequired('N'                                       , @isnumeric);
#            p.addParamValue('ID'            , [obj.ID num2str(N)]   , @ischar);
#            p.addParamValue('name'          , [obj.name num2str(N)] , @ischar);
#            % parse the input
#            p.parse(N,varargin{:});
#            % assign parser results
#            N       = p.Results.N;
#            newID   = p.Results.ID;
#            newName = p.Results.name;
#            % initialize the cloned object with required inputs
#            clone = unitCell(obj.ID,obj.name,obj.cAxis);
#            % copy all properties from the current to the cloned object
#            props = properties(obj);
#            for i = 1:length(props)
#                clone.(props{i}) = obj.(props{i});
#            end%for
#            % reset some _N_-dependent properties
#            clone.ID = newID;
#            clone.name = newName;
#            clone.cAxis = N*obj.cAxis;
#            clone.volume = N*obj.volume;
#            % reset all atom-dependent properties
#            clone.atoms = {};
#            clone.mass = 0;
#            clone.density = 0;
#            clone.springConst = 0;
#            clone.numAtoms = 0;
#            % add all atoms to cloned unit cell, but the position handle is
#            % transfered!
#            for i = 0:N-1
#               for j = 1:obj.numAtoms
#                  clone.addAtom(obj.atoms{j,1},obj.atoms{j,2}(0)/N+i/N);
#               end%for
#            end%for
#        end%function
        pass

    def visualize(self, **kwargs):
        """visualize

        Allows for 3D presentation of unit cell by allow for a & b
        coordinate of atoms.
        Also add magnetization per atom.

        Todo: use the avogadro project as plugin
        Todo: create unit cell from CIF file e.g. by xrayutilities
        plugin.

        """
        import matplotlib.pyplot as plt
        import matplotlib.cm as cmx

        strains = kwargs.get('strains', 0)
        if not isinstance(strains, np.ndarray):
            strains = np.array([strains])

        colors = [cmx.Dark2(x) for x in np.linspace(0, 1, self.num_atoms)]
        atom_ids = self.get_atom_ids()

        for strain in strains:
            plt.figure()
            atoms_plotted = np.zeros_like(atom_ids)
            for j in range(self.num_atoms):
                if not atoms_plotted[atom_ids.index(self.atoms[j][0].id)]:
                    label = self.atoms[j][0].id
                    atoms_plotted[atom_ids.index(self.atoms[j][0].id)] = True
                    plt.plot(1+j, self.atoms[j][1](strain), 'o',
                             MarkerSize=10,
                             markeredgecolor=[0, 0, 0],
                             markerfaceColor=colors[atom_ids.index(self.atoms[j][0].id)],
                             label=label)
                else:
                    label = '_nolegend_'
                    plt.plot(1+j, self.atoms[j][1](strain), 'o',
                             MarkerSize=10,
                             markeredgecolor=[0, 0, 0],
                             markerfaceColor=colors[atom_ids.index(self.atoms[j][0].id)],
                             label=label)

            plt.axis([0.1, self.num_atoms+0.9, -0.1, (1.1+np.max(strains))])
            plt.grid(True)

            plt.title('Strain: {:0.2f}%'.format(strain))
            plt.ylabel('relative Position')
            plt.xlabel('# Atoms')
            plt.legend()
            plt.show()

    def get_property_dict(self, **kwargs):
        """get_property_dict

        Returns a dictionary with all parameters. objects or dicts and
        objects are converted to strings. if a type is given, only these
        properties are returned.

        """
        # initialize input parser and define defaults and validators
        types = ['all', 'heat', 'phonon', 'XRD', 'optical']
        properties_by_types = {'heat': ['_c_axis', '_area', '_volume', '_opt_pen_depth',
                                        'therm_cond_str', 'heat_capacity_str',
                                        'int_heat_capacity_str', 'sub_system_coupling_str',
                                        'num_sub_systems'],
                               'phonon': ['num_sub_systems', 'int_lin_therm_exp_str', '_c_axis',
                                          '_mass', '_spring_const', '_phonon_damping'],
                               'XRD': ['num_atoms', 'atoms', '_area', '_deb_wal_fac', '_c_axis'],
                               'optical': ['_c_axis', '_opt_pen_depth', 'opt_ref_index',
                                           'opt_ref_index_per_strain'],
                               }

        types = kwargs.get('types', 'all')
        attrs = vars(self)
        # define the property names by the given type
        if types == 'all':
            S = attrs
        else:
            S = dict((key, value) for key, value in attrs.items()
                     if key in properties_by_types[types])

        return S

    def check_cell_array_input(self, inputs):
        """ check_cell_array_input

        Checks the input for inputs which are cell arrays of function
        handles, such as the heat capacity which is a cell array of N
        function handles.

        """
        output = []
        outputStrs = []
        # if the input is not a list, we convert it to one
        if not isinstance(inputs, list):
            inputs = [inputs]
        # traverse each list element and convert it to a function handle
        for input in inputs:
            if isfunction(input):
                raise ValueError('Please use string representation of function!')
                output.append(input)
                outputStrs.append('no str representation available')
            elif isinstance(input, str):
                try:
                    output.append(eval(input))
                    outputStrs.append(input)
                except Exception as e:
                    print('String input for unit cell property ' + input + ' \
                        cannot be converted to function handle!')
                    print(e)
            elif isinstance(input, (int, float)):
                output.append(eval('lambda T: {:f}'.format(input)))
                outputStrs.append('lambda T: {:f}'.format(input))
            elif isinstance(input, object):
                output.append(eval('lambda T: {:f}'.format(input.to_base_units().magnitude)))
                outputStrs.append('lambda T: {:f}'.format(input.to_base_units().magnitude))
            else:
                raise ValueError('Unit cell property input has to be a single or '
                                 'cell array of numerics, function handles or strings which can be'
                                 'converted into a function handle!')

        return(output, outputStrs)

    @property
    def int_heat_capacity(self):
        """get int_heat_capacity

        Returns the anti-derrivative of the temperature-dependent heat
        :math:`c(T)` capacity function. If the int_heat_capacity
        property is not set, the symbolic integration is performed.

        """
        if hasattr(self, '_int_heat_capacity') and isinstance(self._int_heat_capacity, list):
            return self._int_heat_capacity
        else:
            self._int_heat_capacity = []
            self.int_heat_capacity_str = []
            try:
                T = Symbol('T')
                for i, hcs in enumerate(self.heat_capacity_str):
                    integral = integrate(hcs.split(':')[1], T)
                    self._int_heat_capacity.append(lambdify(T, integral))
                    self.int_heat_capacity_str.append('lambda T : ' + str(integral))

            except Exception as e:
                print('The sympy integration did not work. You can set the'
                      'analytical anti-derivative of the heat capacity'
                      'of your unit cells as lambda function of the temperature'
                      'T by typing UC.int_heat_capacity = lambda T: c(T)'
                      'where UC is the name of the unit cell object.')
                print(e)

        return self._int_heat_capacity

    @int_heat_capacity.setter
    def int_heat_capacity(self, int_heat_capacity):
        """set int_heat_capacity

        Set the integrated heat capacity manually when no sympy is
        installed.

        """
        self._int_heat_capacity, self.int_heat_capacity_str = self.check_cell_array_input(
                int_heat_capacity)

    @property
    def int_lin_therm_exp(self):
        """get int_lin_therm_exp

        Returns the anti-derrivative of the integrated
        temperature-dependent linear thermal expansion function. If the
        int_lin_therm_exp property is not set, the symbolic integration
        is performed.

        """

        if hasattr(self, '_int_lin_therm_exp') and isinstance(self._int_lin_therm_exp, list):
            return self._int_lin_therm_exp
        else:
            self._int_lin_therm_exp = []
            self.int_lin_therm_exp_str = []
            try:
                T = Symbol('T')
                for i, ltes in enumerate(self.lin_therm_exp_str):
                    integral = integrate(ltes.split(':')[1], T)
                    self._int_lin_therm_exp.append(lambdify(T, integral))
                    self.int_lin_therm_exp_str.append('lambda T : ' + str(integral))

            except Exception as e:
                print('The sympy integration did not work. You can set the'
                      'the analytical anti-derivative of the heat capacity'
                      'of your unit cells as lambda function of the temperature'
                      'T by typing UC.int_heat_capacity = lambda T: c(T)'
                      'where UC is the name of the unit cell object.')
                print(e)

        return self._int_lin_therm_exp

    @int_lin_therm_exp.setter
    def int_lin_therm_exp(self, int_lin_therm_exp):
        """set int_lin_therm_exp

        Set the integrated linear thermal expansion coefficient manually
        when no sympy installed.

        """
        self._int_lin_therm_exp, self.int_lin_therm_exp_str = self.check_cell_array_input(
                int_lin_therm_exp)

    def add_atom(self, atom, position):
        """ add_atom

        Adds an atomBase/atomMixed at a relative position of the unit
        cell.

        Update the mass, density and spring constant of the unit cell
        automatically:

        .. math:: \kappa = m \cdot (v_s / c)^2

        """
        position_str = ''
        # test the input type of the position
        if isfunction(position):
            raise ValueError('Please use string representation of function!')
            pass
        elif isinstance(position, str):
            try:
                position_str = position
                position = eval(position)
            except Exception as e:
                print('String input for unit cell property ' + position + ' \
                    cannot be converted to function handle!')
                print(e)
        elif isinstance(position, (int, float)):
            position_str = 'lambda strain: {:e}*(strain+1)'.format(position)
            position = eval(position_str)
        else:
            raise ValueError('Atom position input has to be a scalar, or string'
                             'which can be converted into a lambda function!')

        # add the atom at the end of the array
        self.atoms.append([atom, position, position_str])
        # increase the number of atoms
        self.num_atoms = self.num_atoms + 1

        self.mass = 0*u.kg
        for i in range(self.num_atoms):
            self.mass = self.mass + self.atoms[i][0].mass

        self.density = self.mass / self.volume
        # set mass per unit area (do not know if necessary)
        self.mass = self.mass * 1*u.angstrom**2 / self.area
        self.calc_spring_const()

    def add_multiple_atoms(self, atom, position, Nb):
        """add_multiple_atoms

        Adds multiple atomBase/atomMixed at a relative position of the
        unit cell.

        """
        for i in range(Nb):
            self.addAtom(atom, position)

    def calc_spring_const(self):
        """calc_spring_const

        Calculates the spring constant of the unit cell from the mass,
        sound velocity and c-axis

        .. math:: k = m \, \left(\\frac{v}{c}\\right)^2

        """
        self.spring_const[0] = (self._mass * (self._sound_vel/self._c_axis)**2)

    def get_acoustic_impedance(self):
        """get_acoustic_impedance"""
        Z = np.sqrt(self.spring_const[0] * self.mass/self.area**2)
        return Z

    def set_ho_spring_constants(self, HO):
        """set_ho_spring_constants

        Set the higher orders of the spring constant for anharmonic
        phonon simulations.

        """
        # reset old higher order spring constants
        self.spring_const = np.delete(self.spring_const, np.r_[1:len(self.spring_const)])
        self.spring_const = np.hstack((self.spring_const, HO))

    def get_atom_ids(self):
        """get_atom_ids

        Returns a cell array of all atom ids in the unit cell.

        """
        ids = []
        for i in range(self.num_atoms):
            if not self.atoms[i][0].id in ids:
                ids.append(self.atoms[i][0].id)

        return ids

    def get_atom_positions(self, *args):
        """get_atom_positions

        Returns a vector of all relative postion of the atoms in the
        unit cell.

        """
        if args:
            strain = args[0]
        else:
            strain = 0

        res = np.zeros([self.num_atoms])
        for i, atom in enumerate(self.atoms):
            res[i] = atom[1](strain)

        return res

    @property
    def a_axis(self):
        """float: in-plane a-axis [m]"""
        return Q_(self._a_axis, u.meter).to('angstrom')

    @a_axis.setter
    def a_axis(self, a_axis):
        """set.a_axis"""
        self._a_axis = a_axis.to_base_units().magnitude

    @property
    def b_axis(self):
        """float: in-plane b-axis [m]"""
        return Q_(self._b_axis, u.meter).to('angstrom')

    @b_axis.setter
    def b_axis(self, b_axis):
        """set.a_axis"""
        self._b_axis = b_axis.to_base_units().magnitude

    @property
    def c_axis(self):
        """float: out-of-plane c-axis [m]"""
        return Q_(self._c_axis, u.meter).to('angstrom')

    @c_axis.setter
    def c_axis(self, c_axis):
        """set.c_axis"""
        self._c_axis = c_axis.to_base_units().magnitude

    @property
    def mass(self):
        """float: mass of unit cell normalized to area of 1 Å² [kg]"""
        return Q_(self._mass, u.kg)

    @mass.setter
    def mass(self, mass):
        """set.mass"""
        self._mass = mass.to_base_units().magnitude

    @property
    def density(self):
        """float: density of the unitCell [kg/m³]"""
        return Q_(self._density, u.kg/u.m**3)

    @density.setter
    def density(self, density):
        """set.density"""
        self._density = density.to_base_units().magnitude

    @property
    def area(self):
        """
        float: area of epitaxial unit cells need for normation for
        correct intensities) [m²]

        """
        return Q_(self._area, u.m**2)

    @area.setter
    def area(self, area):
        """set.area"""
        self._area = area.to_base_units().magnitude

    @property
    def volume(self):
        """float: volume of unit cell [m³]"""
        return Q_(self._volume, u.m**3)

    @volume.setter
    def volume(self, volume):
        """set.volume"""
        self._volume = volume.to_base_units().magnitude

    @property
    def sound_vel(self):
        """float: sound velocity in the unit cell [m/s]"""
        return Q_(self._sound_vel, u.m/u.s)

    @sound_vel.setter
    def sound_vel(self, sound_vel):
        """set.sound_vel
        If the sound velocity is set, the spring constant is
        (re)calculated.
        """
        self._sound_vel = sound_vel.to_base_units().magnitude
        self.calc_spring_const()

    @property
    def phonon_damping(self):
        """float: damping constant of phonon propagation [kg/s]"""
        return Q_(self._phonon_damping, u.kg/u.s)

    @phonon_damping.setter
    def phonon_damping(self, phonon_damping):
        """set.phonon_damping"""
        self._phonon_damping = phonon_damping.to_base_units().magnitude

    @property
    def opt_pen_depth(self):
        """
        float: penetration depth for pump always for 1st subsystem
        light in the unit cell [m]

        """
        return Q_(self._opt_pen_depth, u.meter).to('nanometer')

    @opt_pen_depth.setter
    def opt_pen_depth(self, opt_pen_depth):
        """set.opt_pen_depth"""
        self._opt_pen_depth = opt_pen_depth.to_base_units().magnitude
