import collections
import itertools
import json

import click
import numpy

from .lattice import Atom
from .lattice import Vertex
from .lattice import Lattice
from .lattice import NanoParticle
from .lattice import DepletedLattice
from .material import Material


Geometry = collections.namedtuple(
    'Geometry', ['sites', 'links', 'axes', 'patch']
)


def compute_geometry(lattice, material):
    locator = material.locator()

    lsites = []
    linteractions = []
    axes = []
    # Patch the ids according to the removed material
    new_ids = {}

    for idx, site in enumerate(lattice.sites()):
        lsites.append(site)
        new_ids[lattice.index(site)] = idx
        axis = numpy.array([0.0, 0.0, 0.0])
        for interaction in lattice.interactions_for(site):
            linteractions.append(interaction)
            axis += locator.locate(interaction.target) - locator.locate(site)
        axes.append(axis)

    return Geometry(lsites, linteractions, axes, new_ids)


def echo_header(material, sites, links):
    atom_kinds = material.atom_kinds()
    click.echo('{}\t{}\t{}\t{}'.format(
        len(sites), len(links), len(atom_kinds), 0))
    click.echo('\n'.join(atom_kinds))


def echo_sites(material, lattice, sites, axes, patch=None):
    patch = patch or {}
    locator = material.locator()
    anisotropy = material.anisotropy
    for site, axis in zip(sites, axes):
        px, py, pz = tuple(locator.locate(site))
        index = lattice.index(site)
        val, axis = anisotropy.from_axis(axis)
        ax, ay, az = axis
        click.echo(
            '{uuid}\t{px}\t{py}\t{pz}\t{spin}\t'
            '{ax}\t{ay}\t{az}\t{k}\t{kind}'.format(
                uuid=patch.get(index, index),