"""
Generate a hydrogen-disordered honeycomb bilayer ice.

CLI example: genice3 bilayer --size 6,10
             genice3 bilayer --size 6,10 --sw 0.2
"""

import re
from typing import Any, Dict, Tuple

from cif2ice import cellvectors
from genice3.util import density_in_g_cm3
import genice3.unitcell

import networkx as nx
import numpy as np

import itertools as it
from cycless.cycles import cycles_iter
from logging import getLogger


desc = {
    "ref": {"bilayer": "Koga 1997"},
    "brief": "A Bilayer Honeycomb Ice Phase in Hydrophobic Nanopores.",
    "options": [
        {
            "name": "size",
            "help": "Size of the bilayer (x,y). x must be a multiple of 3.",
            "required": True,
            "example": "6,10",
        },
        {
            "name": "sw",
            "help": "Ratio of Stone-Wales defects to the number of sites.",
            "required": False,
            "example": "0.2",
        },
    ],
    "test": (
        {"options": "--size 6,10"},
        {"options": "--size 15,18 --sw 0.2", "seed": 1000},
    ),
}


def _scalar(v: Any) -> Any:
    if isinstance(v, (list, tuple)) and len(v) == 1:
        return v[0]
    return v


def _parse_size(raw: str) -> Tuple[int, int]:
    parts = re.split(r"[\s,]+", str(raw).strip())
    if len(parts) != 2:
        raise ValueError(
            "bilayer --size must be two integers (e.g. 6,10 or 6 10)"
        )
    try:
        nx_val, ny_val = int(parts[0]), int(parts[1])
    except ValueError as e:
        raise ValueError(
            "bilayer --size must be two integers (e.g. 6,10)"
        ) from e
    if nx_val % 3 != 0:
        raise ValueError("bilayer --size: x must be a multiple of 3.")
    return nx_val, ny_val


def parse_options(options: Dict[str, Any]) -> Tuple[Dict[str, Any], Dict[str, Any]]:
    """
    bilayer 固有: size, sw。
    残りは基底 UnitCell.parse_options に渡す。
    """
    processed: Dict[str, Any] = {}
    unprocessed = dict(options)
    if "size" in unprocessed:
        raw = _scalar(unprocessed.pop("size"))
        processed["size"] = _parse_size(raw)
    if "sw" in unprocessed:
        raw = _scalar(unprocessed.pop("sw"))
        try:
            sw_val = float(raw)
        except (TypeError, ValueError) as e:
            raise ValueError(
                "bilayer --sw must be a number (ratio of Stone-Wales defects)"
            ) from e
        if not 0 <= sw_val <= 1:
            raise ValueError("bilayer --sw must be between 0 and 1.")
        processed["sw"] = sw_val
    base_processed, base_unprocessed = genice3.unitcell.UnitCell.parse_options(
        unprocessed
    )
    processed.update(base_processed)
    return processed, base_unprocessed


# Unnecessary when pairs are given.
# bondlen = 2.0 * 3.0 / 8.0**0.5 * 1.1

# import matplotlib.pyplot as plt


def remove_node(cycle, v):
    cycle.remove(v)


def insert_node(cycle, v, pair):
    a, b = pair
    apos = cycle.index(a)
    if cycle[apos - 1] == b:
        cycle.insert(apos, v)
    else:
        cycle.insert(apos + 1, v)


def stone_wales(g, cycles):
    while True:
        i = np.random.randint(len(g))
        ineis = set(g[i])
        j = np.random.choice(list(ineis))
        jneis = set(g[j])

        owners = g.edges[i, j]["owners"]

        ineis -= {j}
        jneis -= {i}

        j1, j2 = list(jneis)
        i1, i2 = list(ineis)
        for o1 in owners:
            cycle = cycles[o1]
            assert i in cycle
            assert j in cycle
            if i1 not in cycle:
                continue
            if j2 in cycle:
                j1, j2 = j2, j1
            break

        # i1 and j1 is in the same side
        # cycle o1 includes i1 and j1
        if owners[0] == o1:
            o2 = owners[1]
        else:
            o2 = owners[0]

        owners = g.edges[i, i1]["owners"]
        if owners[0] == o1:
            o3 = owners[1]
        else:
            o3 = owners[0]

        owners = g.edges[j, j1]["owners"]
        if owners[0] == o1:
            o4 = owners[1]
        else:
            o4 = owners[0]
        if (
            len(cycles[o1]) < 6
            or len(cycles[o2]) < 6
            or len(cycles[o3]) > 6
            or len(cycles[o4]) > 6
        ):
            continue
        break
    g.add_edge(i1, j)
    g.edges[i1, j]["owners"] = [o1, o3]
    g.add_edge(i, j2)
    g.edges[i, j2]["owners"] = [o2, o4]
    g.edges[i, j]["owners"] = [o3, o4]

    g.remove_edge(i, i1)
    g.remove_edge(j, j2)
    # print(cycles[o1])
    # print(cycles[o2])
    # print(cycles[o3])
    # print(cycles[o4])

    remove_node(cycles[o1], i)
    remove_node(cycles[o2], j)
    insert_node(cycles[o3], j, (i, i1))
    insert_node(cycles[o4], i, (j, j2))
    # print(cycles[o1])
    # print(cycles[o2])
    # print(cycles[o3])
    # print(cycles[o4])


# relaxation


def relax(g, rpos, cell, cycles=None, k=1.0):
    for loop in range(100):
        forces = np.zeros_like(rpos)
        for i, j in g.edges():
            d = rpos[i] - rpos[j]
            d -= np.floor(d + 0.5)
            d *= cell
            L = np.linalg.norm(d)
            f = k * (L - 1.0)
            f *= d / L
            forces[i] -= f
            forces[j] += f
        if cycles is not None:
            for cycle in cycles:
                for i, j in zip(cycle, cycle[-2:] + cycle[:-2]):
                    d = rpos[i] - rpos[j]
                    d -= np.floor(d + 0.5)
                    d *= cell
                    L = np.linalg.norm(d)
                    f = k * (L - 3**0.5)
                    f *= d / L
                    forces[i] -= f
                    forces[j] += f
        rpos += forces / cell


# def draw(g, rpos, cell):
#     nodes = rpos * cell
#     plt.axis('equal')
#     plt.scatter(nodes[:,0], nodes[:,1])
#     for i,j in g.edges():
#         d = rpos[i] - rpos[j]
#         s = np.floor(d+0.5)
#         if np.allclose(s, 0):
#             seg = np.vstack([nodes[j], nodes[j]+d*cell])
#             plt.plot(seg[:,0], seg[:,1], "-")
#         else:
#             d -= s
#             seg = np.vstack([nodes[j], nodes[j]+d*cell])
#             plt.plot(seg[:,0], seg[:,1], "-")
#             seg = np.vstack([nodes[i]-d*cell, nodes[i]])
#             plt.plot(seg[:,0], seg[:,1], "-")


class UnitCell(genice3.unitcell.UnitCell):
    def __init__(self, **kwargs):
        logger = getLogger()

        size = kwargs.get("size")
        if size is None:
            raise ValueError(
                "bilayer unitcell requires --size x,y (e.g. genice3 bilayer --size 6,10)"
            )
        if isinstance(size, (list, tuple)) and len(size) == 2:
            NX, NY = int(size[0]), int(size[1])
        else:
            NX, NY = _parse_size(_scalar(size))
        if NX % 3 != 0:
            raise ValueError("bilayer --size: x must be a multiple of 3.")
        sw = kwargs.get("sw", 0.0)
        try:
            sw = float(sw)
        except (TypeError, ValueError):
            sw = 0.0
        if not 0 <= sw <= 1:
            raise ValueError("bilayer --sw must be between 0 and 1.")

        logger.info(f"Bilayer ice of size ({NX}x{NY})")

        y1 = 3**0.5 / 2
        cell = np.array([NX, y1 * NY])

        nodes1 = [
            (i, j * y1 * 2) for i in range(NX) for j in range(NY // 2) if i % 3 != 2
        ]
        nodes2 = [
            (i + 0.5, (j + 0.5) * y1 * 2)
            for i in range(NX)
            for j in range(NY // 2)
            if i % 3 != 0
        ]
        nodes = np.array(nodes1 + nodes2)

        rpos = nodes / cell

        g = nx.Graph()
        for i, j in it.combinations(range(nodes.shape[0]), 2):
            d = rpos[i] - rpos[j]
            d -= np.floor(d + 0.5)
            d *= cell
            if d @ d < 1.01**2:
                g.add_edge(i, j)

        cycles = [list(cycle) for cycle in cycles_iter(g, 6)]

        for edge in g.edges():
            g.edges[edge]["owners"] = []
        for i, cycle in enumerate(cycles):
            for j, k in zip(cycle, cycle[-1:] + cycle[:-1]):
                g.edges[j, k]["owners"].append(i)

        # fig = plt.figure(figsize=(10,10))
        # draw(g, rpos, cell)
        # plt.show()

        # introduce defects
        Nnode = len(rpos)
        Nshuffle = int(Nnode * sw)
        # Nshuffle = 10
        for i in range(Nshuffle):
            stone_wales(g, cycles)

        relax(g, rpos, cell, cycles, k=0.1)

        cell = cellvectors(a=cell[0], b=cell[1], c=10.0) * 0.276

        density = density_in_g_cm3(Nnode * 2, cell)

        coord = "relative"
        waters = np.zeros([Nnode * 2, 3])
        waters[:Nnode, 0:2] = rpos
        waters[Nnode:, 0:2] = rpos
        waters[Nnode:, 2] = 0.276 / cell[2, 2]
        # self.bondlen = 1.1*0.276

        pairs = []
        for i, j in g.edges():
            pairs.append([i, j])
            pairs.append([i + Nnode, j + Nnode])
        for i in range(Nnode):
            pairs.append([i, i + Nnode])

        super().__init__(
            cell=cell,
            lattice_sites=waters,
            coord=coord,
            density=density,
            graph=nx.Graph(pairs),
            # bondlen=0.276 * 1.2,
            # **kwargs,
        )
