import sys
from logging import getLogger
from io import TextIOWrapper
from typing import Dict, List, Any, Tuple
from collections import defaultdict

import numpy as np

from cif2ice import cellshape, cellvectors
from genice3.molecule import Molecule
from genice3.genice import GenIce3
from genice3.exporter import parse_water_model_option
format_desc = {
    "suboptions": "water_model: 3site, 4site, 6site, tip4p, etc. (water model name).",
    "aliases": ["lammps", "lmp"],
    "application": "[LAMMPS](https://www.lammps.org/)",
    "extension": ".lammps",
    "water": "Atomic positions",
    "solute": "Atomic positions",
    "hb": "none",
    "remarks": "Yet to be verified.",
}


def _to_lammps_data(
    cellmat: np.ndarray,
    waters: Dict[int, Molecule],
    guests: List[Molecule],
    ions: Dict[int, Molecule],
    command_line: str = None,
) -> str:
    """Output in LAMMPS data format.

    parametersには、hooksで指定したstageの結果が含まれる。
    """
    logger = getLogger("_to_lammps_data")
    logger.info("Generating LAMMPS data file... (Yet to be verified)")

    # セル行列をLAMMPS形式に変換
    # LAMMPSは直交座標系を前提とするため、必要に応じて変換
    if cellmat[0, 1] != 0 or cellmat[0, 2] != 0 or cellmat[1, 2] != 0:
        logger.info(
            "  The specified reshaping matrix does not obey the requirements for LAMMPS' unit cell convention."
        )
        a, b, c, A, B, C = cellshape(cellmat)
        rotmat = np.linalg.inv(cellmat) @ cellvectors(a, b, c, A, B, C)
        logger.info("  The reshape matrix is reoriented.")
    else:
        rotmat = np.eye(3)

    # 分子種ごとにソートする必要がある。
    molecules = defaultdict(list)
    for water in waters.values():
        molecules[water.name].append(water)
    for guest in guests:
        molecules[guest.name].append(guest)
    for ion in ions.values():
        molecules[ion.name].append(ion)

    # 原子タイプのマッピングを作成
    atom_types = {}
    type_counter = 1
    for mol_name in sorted(molecules.keys()):
        mols = molecules[mol_name]
        if len(mols) > 0:
            # 各分子の原子ラベルを取得してタイプを割り当て
            sample_mol = mols[0]
            for label in sample_mol.labels:
                type_key = f"{mol_name}_{label}"
                if type_key not in atom_types:
                    atom_types[type_key] = type_counter
                    type_counter += 1

    # 原子リストを作成
    atoms = []
    molecule_id = 1
    for mol_name in sorted(molecules.keys()):
        mols = molecules[mol_name]
        for mol in mols:
            for name, position in zip(mol.labels, mol.sites):
                type_key = f"{mol_name}_{name}"
                atom_type = atom_types[type_key]
                atoms.append([molecule_id, atom_type, name, position])
            molecule_id += 1

    # セル行列を回転
    cellmat = cellmat @ rotmat

    # LAMMPSデータファイル形式で出力
    s = ""
    if command_line:
        s += f"# {command_line}\n"
    s += "# Generated by GenIce https://github.com/vitroid/GenIce\n"
    s += f"\n{len(atoms)} atoms\n"
    s += f"{len(atom_types)} atom types\n"

    # ボックスサイズを計算
    # LAMMPSは直交座標系を前提とする
    # セルの原点を(0,0,0)として、各軸の範囲を計算
    xlo, xhi = 0.0, cellmat[0, 0]
    ylo, yhi = 0.0, cellmat[1, 1]
    zlo, zhi = 0.0, cellmat[2, 2]
    s += f"\n0.0 {cellmat[0, 0]:.8f} xlo xhi\n"
    s += f"0.0 {cellmat[1, 1]:.8f} ylo yhi\n"
    s += f"0.0 {cellmat[2, 2]:.8f} zlo zhi\n"
    if not (
        abs(cellmat[1, 0]) < 1e-8
        and abs(cellmat[2, 0]) < 1e-8
        and abs(cellmat[2, 1]) < 1e-8
    ):
        # 非直交セル（tilt形式）
        xy = cellmat[1, 0]
        xz = cellmat[2, 0]
        yz = cellmat[2, 1]
        s += f"{xy:.8f} {xz:.8f} {yz:.8f} xy xz yz\n"

    s += "\nAtoms\n\n"

    # 原子座標をnmからAngstromに変換（LAMMPSは通常Angstromを使用）
    # ただし、unitsコマンドで指定できるので、nmのままでも可
    # ここではnmのまま出力（units realまたはunits metalの場合）
    # 電荷は0.0をデフォルトとする（必要に応じて拡張可能）
    charge = 0.0
    for i, atom in enumerate(atoms):
        mol_id, atom_type, atomname, position = atom
        position = position @ rotmat
        s += f"{i + 1:8d} {mol_id:8d} {atom_type:4d} {charge:10.6f} {position[0]:15.8f} {position[1]:15.8f} {position[2]:15.8f}\n"

    logger.info(f"  Total number of atoms: {len(atoms)}")
    logger.info(f"  Total number of atom types: {len(atom_types)}")
    logger.info(f"  Total number of molecules: {molecule_id - 1}")

    return s


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


def parse_options(options: Dict[str, Any]) -> Tuple[Dict[str, Any], Dict[str, Any]]:
    """
    lammps プラグインのオプションを処理する。water は water_model のエイリアス。
    """
    options = dict(options)
    if "water" in options and "water_model" not in options:
        options["water_model"] = options["water"]
    processed: Dict[str, Any] = {}
    unprocessed = dict(options)
    if "water_model" in options:
        processed["water_model"] = _scalar(options["water_model"])
        unprocessed.pop("water_model", None)
    return processed, unprocessed


def dumps(
    genice: GenIce3,
    water_model: str = "3site",
    command_line: str = "",
    name: str = "",  # dummy
    **kwargs,
):
    """
    LAMMPS形式で出力

    Args:
        genice: GenIce3インスタンス
        water_model: 水分子モデル名
        command_line: コマンドライン文字列
        name: プラグイン名（"lammps"または""）
    """
    logger = getLogger("lammps.dump")
    if name not in ("lammps", ""):
        raise ValueError(f"name must be 'lammps' or '', got: {name}")
    water_model = parse_water_model_option(water_model)
    waters = genice.water_molecules(water_model=water_model)
    guests = genice.guest_molecules()
    ions = genice.substitutional_ions()

    return _to_lammps_data(
        cellmat=genice.cell,
        waters=waters,
        guests=guests,
        ions=ions,
        command_line=command_line,
    )


def dump(genice: GenIce3, file: TextIOWrapper = sys.stdout, **options):
    file.write(dumps(genice, **options))
