import itertools
import numpy as np
from ase.lattice import bravais_lattices, UnconventionalLattice, bravais_names
from ase.cell import Cell

"""This module implements a crude method to recognize most Bravais lattices.

Suppose we use the ase.lattice module to generate many
lattices of some particular type, say, BCT(a, c), and then we
Niggli-reduce all of them.  The Niggli-reduced forms are not
immediately recognizable, but we know the mapping from each reduced
form back to the original form.  As it turns out, there are apparently
5 such mappings (the proof is left as an exercise for the reader).

Hence, presumably, every BCT lattice (whether generated by BCT(a, c)
or in some other form) Niggli-reduces to a form which, through the
inverse of one of those five operations, is mapped back to the
recognizable one.  Knowing all five operations (or equivalence
classes), we can characterize any BCT lattice.  Same goes for the
other lattices of sufficiently low dimension.

For MCL, MCLC, and TRI, we may not recognize all forms correctly,
but we aspire that this will work for all common inputs."""


niggli_op_table = {  # Generated by generate_niggli_op_table()
 'BCC': [(1, 0, 0, 0, 1, 0, 0, 0, 1)],
 'BCT': [(1, 0, 0, 0, 1, 0, 0, 0, 1),
         (0, 1, 0, 0, 0, 1, 1, 0, 0),
         (0, 1, 0, 1, 0, 0, 1, 1, -1),
         (-1, 0, 1, 0, 1, 0, -1, 1, 0),
         (1, 1, 0, 1, 0, 0, 0, 0, -1)],
 'CUB': [(1, 0, 0, 0, 1, 0, 0, 0, 1)],
 'FCC': [(1, 0, 0, 0, 1, 0, 0, 0, 1)],
 'HEX': [(1, 0, 0, 0, 1, 0, 0, 0, 1), (0, 1, 0, 0, 0, 1, 1, 0, 0)],
 'ORC': [(1, 0, 0, 0, 1, 0, 0, 0, 1)],
 'ORCC': [(1, 0, 0, 0, 1, 0, 0, 0, 1),
          (1, 0, -1, 1, 0, 0, 0, -1, 0),
          (-1, 1, 0, -1, 0, 0, 0, 0, 1),
          (0, 1, 0, 0, 0, 1, 1, 0, 0),
          (0, -1, 1, 0, -1, 0, 1, 0, 0)],
 'ORCF': [(0, -1, 0, 0, 1, -1, 1, 0, 0), (-1, 0, 0, 1, 0, 1, 1, 1, 0)],
 'ORCI': [(0, 0, -1, 0, -1, 0, -1, 0, 0),
          (0, 0, 1, -1, 0, 0, -1, -1, 0),
          (0, 1, 0, 1, 0, 0, 1, 1, -1),
          (0, -1, 0, 1, 0, -1, 1, -1, 0)],
 'RHL': [(0, -1, 0, 1, 1, 1, -1, 0, 0),
         (1, 0, 0, 0, 1, 0, 0, 0, 1),
         (1, -1, 0, 1, 0, -1, 1, 0, 0)],
 'TET': [(1, 0, 0, 0, 1, 0, 0, 0, 1), (0, 1, 0, 0, 0, 1, 1, 0, 0)],
 'MCL': [(0, 0, 1, -1, -1, 0, 1, 0, 0),
         (-1, 0, 0, 0, 1, 0, 0, 0, -1),
         (0, 0, -1, 1, 1, 0, 0, -1, 0),
         (0, -1, 0, 1, 0, 1, -1, 0, 0),
         (0, 1, 0, -1, 0, -1, 0, 0, 1),
         (-1, 0, 0, 0, 1, 1, 0, 0, -1),
         (0, 1, 0, 1, 0, -1, -1, 0, 0),
         (0, 0, 1, 1, -1, 0, 0, 1, 0),
         (0, 1, 0, -1, 0, 0, 0, 0, 1),
         (0, 0, -1, -1, 1, 0, 1, 0, 0),
         (1, 0, 0, 0, 1, -1, 0, 0, 1),
         (0, -1, 0, -1, 0, 1, 0, 0, -1),
         (-1, 0, 0, 0, -1, 1, 0, 1, 0),
         (1, 0, 0, 0, -1, -1, 0, 1, 0),
         (0, 0, -1, 1, 0, 0, 0, -1, 0)],
 'MCLC': [(1, 1, 1, 1, 0, 1, 0, 0, -1),
          (1, 1, 1, 1, 1, 0, -1, 0, 0),
          (1, -1, 1, -1, 0, 1, 0, 0, -1),
          (-1, 1, 0, 1, 0, 0, 0, 0, -1),
          (1, 0, 0, 0, 1, 0, 0, 0, 1),
          (-1, 0, -1, 1, -1, -1, 0, 0, 1),
          (1, -1, -1, 1, -1, 0, -1, 0, 0),
          (-1, -1, 0, -1, 0, -1, 1, 0, 0),
          (1, 0, 1, 1, 0, 0, 0, 1, 0),
          (-1, 1, 0, -1, 0, 1, 1, 0, 0),
          (0, -1, 1, -1, 0, 1, 0, 0, -1),
          (-1, -1, 0, -1, 0, 0, 0, 0, -1),
          (-1, -1, 1, -1, 0, 1, 0, 0, -1),
          (1, 0, 0, 0, -1, 1, 0, 0, -1),
          (-1, 0, -1, 0, -1, -1, 0, 0, 1),
          (1, 0, -1, -1, 1, -1, 0, 0, 1),
          (1, -1, 1, 1, -1, 0, 0, 1, 0),
          (0, -1, 0, 1, 0, -1, 0, 0, 1),
          (-1, 0, 0, 1, 1, 1, 0, 0, -1),
          (1, 0, -1, 0, 1, -1, 0, 0, 1),
          (-1, 1, 0, 1, 1, -1, 0, -1, 0),
          (1, 1, -1, 1, -1, 0, -1, 0, 0),
          (-1, -1, -1, -1, -1, 0, 0, 1, 0),
          (-1, 1, 1, 1, 0, 1, 0, 0, -1),
          (-1, 0, 0, 0, -1, 0, 0, 0, 1),
          (-1, -1, 1, 1, -1, 0, 0, 1, 0),
          (1, 1, 0, -1, 0, -1, 0, 0, 1)],
 'TRI': [(0, -1, 0, -1, 0, 0, 0, 0, -1),
         (0, 1, 0, 0, 0, 1, 1, 0, 0),
         (0, 0, -1, 0, -1, 0, -1, 1, 0),
         (0, 0, 1, 0, 1, 0, -1, 0, 0),
         (0, -1, 0, 0, 0, -1, 1, 1, 1),
         (0, 1, 0, 0, 0, 1, 1, -1, 0),
         (0, 0, -1, 0, -1, 0, -1, 0, 0),
         (-1, 1, 0, 0, 0, -1, 0, -1, 0),
         (0, 0, 1, 1, -1, 0, 0, 1, 0),
         (0, 0, -1, 1, 1, 1, 0, -1, 0),
         (-1, 0, 0, 0, 1, 0, 0, -1, -1),
         (0, 0, 1, 1, 0, 0, 0, 1, 0),
         (0, 0, 1, 0, 1, 0, -1, -1, -1),
         (-1, 0, 0, 0, 0, -1, 0, -1, 0),
         (0, -1, 0, 0, 0, -1, 1, 0, 0),
         (1, 0, 0, 0, 1, 0, 0, 0, 1),
         (0, 0, -1, -1, 0, 0, 1, 1, 1),
         (0, 0, -1, -1, 0, 0, 0, 1, 0),
         (-1, -1, -1, 0, 0, 1, 0, 1, 0)]
}

# XXX The TRI list was generated by looping over all TRI structures in
# the COD (Crystallography Open Database) and seeing what operations
# were necessary to map all those to standard form.  Hence if the
# data does not cover all possible inputs, we could miss something.
#
# Looping over all possible TRI lattices in general would generate
# 100+ operations, we don't want to tabulate that.


def lattice_loop(latcls, length_grid, angle_grid):
    """Yield all lattices defined by the length and angle grids."""
    param_grids = []
    for varname in latcls.parameters:
        # Actually we could choose one parameter, a, to always be 1,
        # reducing the dimension of the problem by 1.  The lattice
        # recognition code should do something like that as well, but
        # it doesn't.  This could affect the impact of the eps value
        # on lattice determination, so we just loop over the whole
        # thing in order not to worry.
        if latcls.name in ['MCL', 'MCLC']:
            special_var = 'c'
        else:
            special_var = 'a'
        if varname == special_var:
            values = np.ones(1)
        elif varname in 'abc':
            values = length_grid
        elif varname in ['alpha', 'beta', 'gamma']:
            values = angle_grid
        else:
            raise ValueError(varname)
        param_grids.append(values)

    for latpars in itertools.product(*param_grids):
        kwargs = dict(zip(latcls.parameters, latpars))
        try:
            lat = latcls(**kwargs)
        except (UnconventionalLattice, AssertionError):
            # XXX assertion error can happen because cellpar_to_cell
            # makes certain assumptions.  Should be investigated.
            # {'b': 0.1, 'gamma': 60.0, 'c': 0.1, 'a': 1.0,
            #   'alpha': 30.0, 'beta': 30.0} <-- this won't work
            pass
        else:
            yield lat


def find_niggli_ops(latcls, length_grid, angle_grid):
    niggli_ops = {}


    for lat in lattice_loop(latcls, length_grid, angle_grid):
        cell = lat.tocell()

        try:
            rcell, op = cell.niggli_reduce()
        except RuntimeError:
            print('Niggli reduce did not converge')
            continue
        assert op.dtype == int
        op_key = tuple(op.ravel())

        if op_key in niggli_ops:
            niggli_ops[op_key] += 1
        else:
            niggli_ops[op_key] = 1

        rcell_test = Cell(op.T @ cell)
        rcellpar_test = rcell_test.cellpar()
        rcellpar = rcell.cellpar()
        err = np.abs(rcellpar_test - rcellpar).max()
        assert err < 1e-7, err

    return niggli_ops


def find_all_niggli_ops(length_grid, angle_grid, lattices=None):
    all_niggli_ops = {}
    if lattices is None:
        lattices = [name for name in bravais_names
                    if name not in ['MCL', 'MCLC', 'TRI']]

    for latname in lattices:
        latcls = bravais_lattices[latname]
        if latcls.ndim < 3:
            continue

        print('Working on {}...'.format(latname))
        niggli_ops = find_niggli_ops(latcls, length_grid, angle_grid)
        print('Found {} ops for {}'.format(len(niggli_ops), latname))
        for key, count in niggli_ops.items():
            print('  {:>40}: {}'.format(str(np.array(key)), count))
        print()
        all_niggli_ops[latname] = niggli_ops
    return all_niggli_ops


def generate_niggli_op_table(lattices=None,
                             length_grid=None,
                             angle_grid=None):

    if length_grid is None:
        length_grid = np.logspace(-0.5, 1.5, 50).round(3)
    if angle_grid is None:
        angle_grid = np.linspace(10, 179, 50).round()
    all_niggli_ops_and_counts = find_all_niggli_ops(length_grid, angle_grid,
                                                    lattices=lattices)

    niggli_op_table = {}
    for latname, ops in all_niggli_ops_and_counts.items():
        ops = [op for op in ops if np.abs(op).max() < 2]
        niggli_op_table[latname] = ops

    import pprint
    print(pprint.pformat(niggli_op_table))
    return niggli_op_table

def test():
    length_grid = np.logspace(-0.5, 1.5, 11).round(3)
    angle_grid = np.linspace(10, 179, 11).round()
    #all_ops = find_all_niggli_ops(length_grid, angle_grid)
    #niggli_op_table.clear()
    #niggli_op_table.update(all_ops)

    for latname in bravais_names:
        if latname in ['MCL', 'MCLC', 'TRI']:
            continue
        latcls = bravais_lattices[latname]
        if latcls.ndim != 3:
            continue

        print('Check', latname)
        maxerr = 0.0

        for lat in lattice_loop(latcls, length_grid, angle_grid):
            cell = lat.tocell()
            from ase.lattice import identify_lattice
            out_lat, op = identify_lattice(cell, eps=2e-4)

            # Some lattices represent simpler lattices,
            # e.g. TET(a, a) is cubic.  What we need to check is that
            # the cell parameters are the same.
            cellpar = cell.cellpar()
            outcellpar = out_lat.tocell().cellpar()
            err = np.abs(outcellpar - cellpar).max()
            maxerr = max(err, maxerr)
            if lat.name != out_lat.name:
                print(repr(lat), '-->', repr(out_lat))
            assert err < 1e-8, (err, repr(lat), repr(out_lat))

        print('    OK.  Maxerr={}'.format(maxerr))

if __name__ == '__main__':
    import sys
    lattices = sys.argv[1:]
    if not lattices:
        lattices = None
    length_grid = np.logspace(-1, 1, 60)
    angle_grid = np.linspace(30, 90, 60 + 59)
    table = generate_niggli_op_table(lattices=lattices,
                                     angle_grid=angle_grid,
                                     length_grid=length_grid)
    for key in table:
        print('{}: {}'.format(key, len(table[key])))
