#!/usr/bin/python
"""
Usage: genice3 c1te
"""
import genice3.unitcell
from genice3.util import (
    atomdic,
    symmetry_operators,
    waters_and_pairs,
    density_in_g_cm3,
)
from logging import getLogger
import numpy as np
import networkx as nx
from cif2ice import cellvectors

desc = {
    "ref": {"C1": "Teeratchanan 2015"},
    "brief": "Hydrogen-ordered hydrogen hydrate C1 by Teeratchanan. (Positions of guests are supplied.)",
}


def pick_atoms(atoms, names, repeat=(1, 1, 1)):
    nrep = np.array(repeat)
    for atomname, fracpos in atoms:
        if atomname in names:
            for x in range(repeat[0]):
                for y in range(repeat[1]):
                    for z in range(repeat[2]):
                        yield atomname, (fracpos + np.array([x, y, z])) / nrep


class UnitCell(genice3.unitcell.UnitCell):
    SUPPORTS_ION_DOPING = False  # 水素秩序氷

    def __init__(self, **kwargs):
        logger = getLogger()

        # Ref. C1
        atoms = """
        O1 0.2228 0.1966 0.0454
        O2 0.5236 0.8974  0.1466
        H1 0.8184 0.5344  0.3198
        H2 0.2244 0.2188 0.2045
        H3 0.4384 0.8815  0.1496
        H4 0.5729 -0.0253 0.2297
        Ne1 0.0000 0.0000 0.7361
        """

        # Ref. C1
        # space group: R-3 No. 148
        # in a rhombus cell
        symops = """
         x,            y,            z
         -y,           x-y,           z
        -x+y,          -x,            z

         -x,           -y,           -z
          y,          -x+y,          -z
         x-y,           x,           -z

          x+2/3,            y+1/3,           z+1/3
         -y+2/3,          x-y+1/3,           z+1/3
       -x+y+2/3,           -x+1/3,           z+1/3

         -x+2/3,           -y+1/3,          -z+1/3
          y+2/3,         -x+y+1/3,          -z+1/3
        x-y+2/3,            x+1/3,          -z+1/3

          x+1/3,            y+2/3,           z+2/3
         -y+1/3,          x-y+2/3,           z+2/3
       -x+y+1/3,           -x+2/3,           z+2/3

         -x+1/3,           -y+2/3,          -z+2/3
          y+1/3,         -x+y+2/3,          -z+2/3
        x-y+1/3,            x+2/3,          -z+2/3
        """.replace(
            ",", " "
        )

        # Ref. C1
        a = 12.673 / 10.0  # nm
        c = 6.017 / 10.0  # nm
        C = 120

        cell = cellvectors(a, a, c, C=C)

        atomd = atomdic(atoms)
        sops = symmetry_operators(symops)
        waters, fixed = waters_and_pairs(cell, atomd, sops)

        # self.cagetype = []
        # self.cagepos = []
        # for name, pos in pick_atoms(atoms, ("Ne1",)):
        #     self.cagetype.append(name)
        #     self.cagepos.append(pos)

        density = density_in_g_cm3(len(waters), cell)
        coord = "relative"
        super().__init__(
            cell=cell,
            lattice_sites=waters,
            density=density,
            coord=coord,
            graph=nx.Graph(fixed),
            fixed=nx.DiGraph(fixed),
            **kwargs,
        )
