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 = {
    "aliases": ["g", "gromacs"],
    "application": "[Gromacs](http://www.gromacs.org)",
    "extension": ".gro",
    "water": "Atomic positions",
    "solute": "Atomic positions",
    "hb": "none",
    "remarks": "Default format.",
    "suboptions": "water_model: 3site, 4site, 6site, tip4p, etc. (water model name).",
}


def _to_gro(
    cellmat: np.ndarray,
    waters: Dict[int, Molecule],
    guests: List[Molecule],
    ions: Dict[int, Molecule],
    command_line: str = None,
) -> str:
    """Output in Gromacs format."""
    logger = getLogger("_to_gro")
    logger.info("Generating .gro...")

    if not (cellmat[0, 1] == 0 and cellmat[0, 2] == 0 and cellmat[1, 2] == 0):
        logger.info(
            "  The specified reshaping matrix does not obey the requirements for Gromacs' 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)

    atoms = []
    residue_count = 1
    top_append = ""
    for mol_name, mols in molecules.items():
        for mol in mols:
            for name, position in zip(mol.labels, mol.sites):
                atoms.append([residue_count, mol.name, name, position])
            residue_count += 1
        top_append += f"{mol_name:5s}{len(mols):>8d}\n"

    logger.info(
        f"Append the following lines to the .top file:\n[ molecules ]\n{top_append}"
    )
    logger.info("  Total number of atoms: {0}".format(len(atoms)))
    if len(atoms) > 99999:
        logger.warn(
            "  Fixed-digit format of Gromacs cannot deal with atoms more than 99999. Residue number and atom number are set appropriately."
        )
    s = ""
    if command_line:
        s += f"{command_line} # "
    s += "Generated by GenIce https://github.com/vitroid/GenIce \n"
    s += f"{len(atoms)}\n"
    for i, atom in enumerate(atoms):
        resno, resname, atomname, position = atom
        position = position @ rotmat
        if len(atoms) > 99999:
            s += f"{99999:5d}{resname:5s}{atomname:>5s}{99999:5d}{position[0]:8.3f}{position[1]:8.3f}{position[2]:8.3f}\n"
        else:
            s += f"{resno:5d}{resname:5s}{atomname:>5s}{i + 1:5d}{position[0]:8.3f}{position[1]:8.3f}{position[2]:8.3f}\n"
    cellmat = cellmat @ rotmat
    # if all close to zero (not exactly zero,)
    if (
        abs(cellmat[1, 0]) < 1e-8
        and abs(cellmat[2, 0]) < 1e-8
        and abs(cellmat[2, 1]) < 1e-8
    ):
        s += "    {0:.8f} {1:.8f} {2:.8f}\n".format(
            cellmat[0, 0], cellmat[1, 1], cellmat[2, 2]
        )
    else:
        s += "    {0:.8f} {1:.8f} {2:.8f} {3:.8f} {4:.8f} {5:.8f} {6:.8f} {7:.8f} {8:.8f}\n".format(
            cellmat[0, 0],
            cellmat[1, 1],
            cellmat[2, 2],
            cellmat[0, 1],
            cellmat[0, 2],
            cellmat[1, 0],
            cellmat[1, 2],
            cellmat[2, 0],
            cellmat[2, 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]]:
    """
    gromacs プラグインのオプションを処理する。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: Any,
):
    """
    Gromacs形式で出力

    Args:
        genice: GenIce3インスタンス
        file: 出力先ファイル
        options: プラグインオプション
            - safe_importにより自動的に辞書にパースされる
            - プラグイン側では既に辞書形式になっている
    """
    logger = getLogger("gromacs.dump")
    if name not in ("gromacs", ""):
        raise ValueError(f"name must be 'gromacs' 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_gro(
        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))
