# coding: utf-8
"""
Exporter that outputs a unitcell plugin (Python module) whose unit cell
is the current supercell (拡大胞). The generated plugin can be used as
a GenIce3 unitcell with rep=1x1x1 to reproduce the same structure.
"""

import sys
from io import TextIOWrapper
from typing import Any, Dict

import numpy as np
import networkx as nx

from genice3.genice import GenIce3

format_desc = {
    "aliases": ["py", "python"],
    "application": "Python unitcell plugin",
    "extension": ".py",
    "water": "none",
    "solute": "none",
    "hb": "graph",
    "remarks": "Outputs a unitcell plugin where the supercell becomes the new unit cell.",
    "suboptions": "name: output unitcell module name (default: exported).",
}


def _format_array(arr: np.ndarray, indent: str = "    ") -> str:
    """Format numpy array as Python list of lists for source code."""
    lst = arr.tolist()
    lines = [
        "np.array([",
        *[f"{indent}[{', '.join(repr(x) for x in row)}]," for row in lst],
        f"{indent}])",
    ]
    return "\n".join(lines)


def _format_graph_edges(G: nx.Graph) -> str:
    """Format graph edges as Python list of tuples."""
    edges = list(G.edges())
    if not edges:
        return "[]"
    # 1 line per edge if many, else one line
    if len(edges) <= 4:
        return repr(edges)
    parts = [f"        {repr(e)}," for e in edges]
    return "[\n" + "\n".join(parts) + "\n    ]"


def dumps(genice: GenIce3, name: str = "exported", **kwargs: Any) -> str:
    """
    Generate Python source code for a unitcell plugin that defines
    the current supercell as the new unit cell.

    Args:
        genice: GenIce3 instance (after structure is built).
        name: Base name for the generated class / module (used in docstring).
        **kwargs: Ignored (for API compatibility).

    Returns:
        Python source code as a string.

    Raises:
        ValueError: If anion, cation, or guest is defined (this exporter does not output them).
    """
    anions = getattr(genice, "anions", {}) or {}
    cations = getattr(genice, "cations", {}) or {}
    guests = getattr(genice, "guests", {}) or {}
    spot_guests = getattr(genice, "spot_guests", {}) or {}

    if anions or cations:
        raise ValueError(
            "python exporter does not support anion/cation. "
            "Export without spot_anion/spot_cation and without unitcell anion/cation."
        )
    if guests or spot_guests:
        raise ValueError(
            "python exporter does not support guest. "
            "Export without guest/spot_guest."
        )

    # Trigger reactive computation of extended-cell data
    cell = genice.cell
    lattice_sites = genice.lattice_sites
    graph = genice.graph
    fixed = genice.fixed_edges

    cell_str = _format_array(cell)
    sites_str = _format_array(lattice_sites)
    graph_edges_str = _format_graph_edges(graph)
    fixed_edges = list(fixed.edges())
    fixed_str = _format_graph_edges(nx.DiGraph(fixed_edges)) if fixed_edges else "[]"

    lines = [
        "# coding: utf-8",
        "# Auto-generated unitcell plugin: supercell as new unit cell",
        "# Generated by GenIce3 exporter python.py",
        "",
        "import genice3.unitcell",
        "import numpy as np",
        "import networkx as nx",
        "",
        "desc = {",
        '    "ref": {},',
        f'    "usage": "Unit cell exported from supercell ({name}).",',
        f'    "brief": "Exported supercell as unit cell ({name}).",',
        "}",
        "",
        "",
        "class UnitCell(genice3.unitcell.UnitCell):",
        "    \"\"\"Unit cell identical to the exported supercell (rep=1x1x1).\"\"\"",
        "",
        "    def __init__(self, **kwargs):",
        "        cell = " + cell_str.replace("\n", "\n        "),
        "",
        "        lattice_sites = " + sites_str.replace("\n", "\n        "),
        "",
        "        graph = nx.Graph()",
        "        graph.add_edges_from(" + graph_edges_str + ")",
        "",
        "        fixed = nx.DiGraph()",
        "        fixed.add_edges_from(" + fixed_str + ")",
        "",
        "        super().__init__(",
        "            cell=cell,",
        "            lattice_sites=lattice_sites,",
        '            coord="relative",',
        "            graph=graph,",
        "            fixed=fixed,",
        "            **kwargs,",
        "        )",
        "",
    ]
    return "\n".join(lines)


def dump(genice: GenIce3, file: TextIOWrapper = sys.stdout, **options: Any) -> None:
    """Write the unitcell plugin source to file."""
    name = options.get("name", "exported")
    file.write(dumps(genice, name=name, **options))
