# coding: utf-8
# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.
"""
Low-level objects providing an abstraction for the objects involved in the calculation.
"""
from __future__ import unicode_literals, division, print_function

import collections
import abc
import six
import numpy as np
import pymatgen.core.units as units

from pprint import pformat
from monty.design_patterns import singleton
from monty.collections import AttrDict
from enum import Enum
from monty.json import MSONable
from pymatgen.core.units import ArrayWithUnit
from pymatgen.core.lattice import Lattice
from pymatgen.core.structure import Structure
from pymatgen.util.serialization import pmg_serialize
from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
from monty.json import MontyEncoder, MontyDecoder


def lattice_from_abivars(cls=None, *args, **kwargs):
    """
    Returns a `Lattice` object from a dictionary
    with the Abinit variables `acell` and either `rprim` in Bohr or `angdeg`
    If acell is not given, the Abinit default is used i.e. [1,1,1] Bohr

    Args:
        cls: Lattice class to be instantiated. pymatgen.core.lattice.Lattice if `cls` is None

    Example:

        lattice_from_abivars(acell=3*[10], rprim=np.eye(3))
    """
    cls = Lattice if cls is None else cls
    kwargs.update(dict(*args))
    d = kwargs

    rprim = d.get("rprim", None)
    angdeg = d.get("angdeg", None)
    acell = d["acell"]

    if rprim is not None:
        if angdeg is not None:
            raise ValueError("angdeg and rprimd are mutually exclusive")
        rprim = np.reshape(rprim, (3,3))
        rprimd = [float(acell[i]) * rprim[i] for i in range(3)]
        # Call pymatgen constructors (note that pymatgen uses Angstrom instead of Bohr).
        return cls(ArrayWithUnit(rprimd, "bohr").to("ang"))

    elif angdeg is not None:
        angdeg = np.reshape(angdeg, 3)

        if np.any(angdeg <= 0.):
            raise ValueError("Angles must be > 0 but got %s" % str(angdeg))
        if angdeg.sum() >= 360.:
            raise ValueError("The sum of angdeg must be lower that 360, angdeg %s" % str(angdeg))

        # This code follows the implementation in ingeo.F90
        # See also http://www.abinit.org/doc/helpfiles/for-v7.8/input_variables/varbas.html#angdeg
        tol12 = 1e-12
        pi, sin, cos, sqrt = np.pi, np.sin, np.cos, np.sqrt
        rprim = np.zeros((3,3))
        if (abs(angdeg[0] -angdeg[1]) < tol12 and abs(angdeg[1] - angdeg[2]) < tol12 and
            abs(angdeg[0]-90.) + abs(angdeg[1]-90.) + abs(angdeg[2] -90) > tol12):
            # Treat the case of equal angles (except all right angles):
            # generates trigonal symmetry wrt third axis
            cosang = cos(pi * angdeg[0]/180.0)
            a2 = 2.0/3.0*(1.0 - cosang)
            aa = sqrt(a2)
            cc = sqrt(1.0-a2)
            rprim[0,0] = aa     ; rprim[0,1] = 0.0              ; rprim[0,2] = cc
            rprim[1,0] = -0.5*aa; rprim[1,1] = sqrt(3.0)*0.5*aa ; rprim[1,2] = cc
            rprim[2,0] = -0.5*aa; rprim[2,1] = -sqrt(3.0)*0.5*aa; rprim[2,2] = cc
        else:
            # Treat all the other cases
            rprim[0,0] = 1.0
            rprim[1,0] = cos(pi*angdeg[2]/180.)
            rprim[1,1] = sin(pi*angdeg[2]/180.)
            rprim[2,0] = cos(pi*angdeg[1]/180.)
            rprim[2,1] = (cos(pi*angdeg[0]/180.0)-rprim[1,0]*rprim[2,0])/rprim[1,1]
            rprim[2,2] = sqrt(1.0-rprim[2,0]**2-rprim[2,1]**2)

        # Call pymatgen constructors (note that pymatgen uses Angstrom instead of Bohr).
        rprimd = [float(acell[i]) * rprim[i] for i in range(3)]
        return cls(ArrayWithUnit(rprimd, "bohr").to("ang"))

    raise ValueError("Don't know how to construct a Lattice from dict:\n%s" % pformat(d))


def structure_from_abivars(cls=None, *args, **kwargs):
    """
    Build a :class:`Structure` object from a dictionary with ABINIT variables.

    Args:
        cls: Structure class to be instantiated. pymatgen.core.structure.Structure if cls is None

    example:

        al_structure = structure_from_abivars(
            acell=3*[7.5],
            rprim=[0.0, 0.5, 0.5,
                   0.5, 0.0, 0.5,
                   0.5, 0.5, 0.0],
            typat=1,
            xred=[0.0, 0.0, 0.0],
            ntypat=1,
            znucl=13,
        )

    `xred` can be replaced with `xcart` or `xangst`.