# -*- coding: utf-8 -*-
"""
top_writer.py
-------------
Writes the GROMACS .top topology file from a SystemModel.

All Gromacs-specific string formatting (column widths, section headers)
is contained here.
"""
from __future__ import annotations
from ....core.system_model import SystemModel, MoleculeTopology, DihedralRecord


# ---------------------------------------------------------------------------
# Low-level formatting helpers  (ported from original udf2gro.py)
# ---------------------------------------------------------------------------

def _strl(total: str, data: str, n: int) -> str:
    """Append left-justified data padded/truncated to n chars, then a space."""
    return total + data.ljust(n)[:n] + " "


def _strr(total: str, data: str, n: int) -> str:
    """Append right-justified data padded/truncated to n chars, then a space."""
    return total + data.rjust(n)[:n] + " "


def _f2s(value: float, ndigit: int) -> str:
    """Fixed-precision float string (same as float2str in original)."""
    return "{:.15f}".format(value)[:ndigit]


# ---------------------------------------------------------------------------
# TopWriter
# ---------------------------------------------------------------------------

class TopWriter:
    """Generates the content of a .top file and writes it."""

    def write(self, model: SystemModel, filepath: str) -> None:
        """SystemModelから .top トポロジファイルを書き出す。

        Args:
            model: 中間表現のシステムモデル。
            filepath: 出力先ファイルパス。
        """
        with open(filepath, "w") as f:
            f.write(self._build_header(model))
            for topo in model.mol_topologies:
                self._write_moleculetype(f, model, topo)
            f.write(self._build_system_molecules(model))

    # ------------------------------------------------------------------
    # Header: defaults + atomtypes
    # ------------------------------------------------------------------

    def _build_header(self, model: SystemModel) -> str:
        s = "; Generated by using ABMPTools\n; input: " + model.udf_path + "\n\n"

        # [ defaults ]
        s += "[ defaults ]\n"
        s = _strl(s, ";nbfunc",    10)
        s = _strl(s, "comb-rule",  10)
        s = _strl(s, "gen-pairs",  10)
        s = _strl(s, "fudgeLJ",    10)
        s = _strl(s, "fudgeQQ",    10)
        s += "\n"

        s = _strl(s, str(1),                  10)
        s = _strl(s, str(model.comb_rule),    10)
        s = _strl(s, "yes",                   10)
        if model.flags14 == 0:
            s = _strl(s, str(0.0), 10)
            s = _strl(s, str(0.0), 10)
        else:
            s = _strl(s, str(model.fudgeLJ), 10)
            s = _strl(s, str(model.fudgeQQ), 10)
        s += "\n\n"

        # [ atomtypes ]
        s += "[ atomtypes ]\n"
        s += ";"
        s = _strr(s, "atom type",         11)
        s = _strr(s, "mass",              12)
        s = _strr(s, "q",                 12)
        s = _strr(s, "particle_type",     13)
        s = _strr(s, "sigma[nm]",         12)
        s = _strr(s, "epsion[kJ mol^-1]", 17)
        s += "\n"

        for at in model.atom_types:
            s = _strr(s, at.name,         12)
            s = _strr(s, str(at.mass),    12)
            s = _strr(s, "0.0",           12)
            s = _strr(s, "A",             12)
            if at.sigma != 0.0 or at.epsilon != 0.0:
                s = _strr(s, _f2s(at.sigma,   12), 12)
                s = _strr(s, _f2s(at.epsilon, 12), 12)
            else:
                s = _strr(s, "0.0", 12)
                s = _strr(s, "0.0", 12)
            s += "\n"
        s += "\n"

        # [ nonbond_params ] is disabled (flg_nbp = False in original)
        # [ bondtypes ] etc. are also disabled (flg_bondpots = False)

        s += "; And here is the topology.\n"
        return s

    # ------------------------------------------------------------------
    # Per-molecule sections
    # ------------------------------------------------------------------

    def _write_moleculetype(
        self, f, model: SystemModel, topo: MoleculeTopology
    ) -> None:

        # [ moleculetype ]
        s = "[ moleculetype ]\n"
        s += ";"
        s = _strl(s, "Name",   11)
        s = _strl(s, "nrexcl", 6)
        s += "\n"
        s = _strl(s, topo.gro_name, 12)
        s = _strl(s, str(topo.nrexcl), 5)
        s += "\n\n"

        # [ atoms ]
        s += "[ atoms ]\n"
        s += ";"
        s = _strr(s, "nr",     11)
        s = _strr(s, "type",   12)
        s = _strr(s, "resnr",  12)
        s = _strr(s, "residu", 12)
        s = _strr(s, "atom",   12)
        s = _strr(s, "cgnr",   12)
        s = _strr(s, "charge", 12)
        s += "\n"
        f.write(s)

        for atom in topo.atoms:
            s = ""
            s = _strr(s, str(atom.index),     11)
            s = _strr(s, atom.type_name,      12)
            s = _strr(s, "1",                 12)
            s = _strr(s, topo.gro_name,       12)
            s = _strr(s, atom.gro_name,       12)
            s = _strr(s, str(atom.index),     12)
            s = _strr(s, _f2s(atom.charge, 12), 12)
            s += "\n"
            f.write(s)

        f.write("\n\n")

        # [ bonds ]
        s = "[ bonds ]\n"
        s += ";"
        s = _strr(s, "ai",                  11)
        s = _strr(s, "aj",                  12)
        s = _strr(s, "funct ",              12)
        s = _strr(s, "b0[nm]",             12)
        s = _strr(s, "kb[kJ mol^-1 nm^-2]", 19)
        s += "\n"
        f.write(s)

        for bond in topo.bonds:
            s = ""
            s = _strr(s, str(bond.atom1), 12)
            s = _strr(s, str(bond.atom2), 12)
            s = _strr(s, bond.funct,      12)
            s = _strr(s, _f2s(bond.r0, 12), 12)
            s = _strr(s, _f2s(bond.kb, 12), 12)
            s += "\n"
            f.write(s)

        f.write("\n\n")

        # [ pairs ]
        s = "[ pairs ]\n"
        s += ";"
        s = _strr(s, "ai",               11)
        s = _strr(s, "aj",               12)
        s = _strr(s, "funct ",           12)
        s = _strr(s, "sigma[nm]",        12)
        s = _strr(s, "epsion[kJ mol^-1]", 17)
        s += "\n"
        f.write(s)

        for pair in topo.pairs:
            s = ""
            s = _strr(s, str(pair.atom1), 12)
            s = _strr(s, str(pair.atom2), 12)
            s = _strr(s, "1",             12)
            s += "\n"
            f.write(s)

        f.write("\n\n")

        # [ angles ]
        s = "[ angles ]\n"
        s += ";"
        s = _strr(s, "ai",                       11)
        s = _strr(s, "aj",                       12)
        s = _strr(s, "ak",                       12)
        s = _strr(s, "funct ",                   12)
        s = _strr(s, "th0[deg]",                12)
        s = _strr(s, "cth[kJ mol^-1 rad^-2]",   21)
        s += "\n"
        f.write(s)

        for ang in topo.angles:
            s = ""
            s = _strr(s, str(ang.atom1),         11)
            s = _strr(s, str(ang.atom2),         11)
            s = _strr(s, str(ang.atom3),         11)
            s = _strr(s, "1",                    11)
            s = _strr(s, _f2s(ang.theta0, 12),  12)
            s = _strr(s, _f2s(ang.k, 12),       12)
            s += "\n"
            f.write(s)

        f.write("\n\n")

        # [ dihedrals ]
        s = "[ dihedrals ]\n"
        s += ";"
        s = _strr(s, "ai",  11)
        s = _strr(s, "aj",  12)
        s = _strr(s, "ak",  12)
        s = _strr(s, "al",  12)

        # determine header style (Amber vs Cosine_Polynomial)
        use_amber_header = any(
            d.funct in ("1", "9") and len(d.params) == 3
            for d in topo.dihedrals
            if d.funct not in ("3",)
        )
        # Actually match original: if any Amber type present -> phi/k/mult header
        # We check by looking at whether there is any funct-1/9 with 3 params
        # (which is how Amber dihedrals are stored)
        has_amber = any(d.funct in ("1", "9") for d in topo.dihedrals)
        has_rb    = any(d.funct == "3"         for d in topo.dihedrals)

        # Original code: if "Amber" in tors_pot_type_list -> use phi_s header
        # We stored this implicitly: funct=1 with 3 params could be Amber or RB-derived
        # We need to use the same header the original chose.
        # Original logic: check tors_pot_type_list (all torsion potential types in UDF)
        # We'll use the presence of funct "3" (Ryckaert-Bell) to decide:
        # if no "3" funct -> Amber header
        if not has_rb:
            # Amber header
            s = _strr(s, "phi_s[deg]", 12)
            s = _strr(s, "k[kJ mol^-1]", 12)
            s = _strr(s, "mult", 12)
        else:
            s = _strr(s, "C0", 12)
            s = _strr(s, "C1", 12)
            s = _strr(s, "C2", 12)
            s = _strr(s, "C3", 12)
            s = _strr(s, "C4", 12)
            s = _strr(s, "C5", 12)
        s += "\n"
        f.write(s)

        for dih in topo.dihedrals:
            s = ""
            s = _strr(s, str(dih.atom1), 12)
            s = _strr(s, str(dih.atom2), 12)
            s = _strr(s, str(dih.atom3), 12)
            s = _strr(s, str(dih.atom4), 12)
            s = _strr(s, dih.funct,      12)

            if dih.funct in ("1", "4", "9"):
                # params = [phi_s, k, mult]
                s = _strr(s, _f2s(dih.params[0], 12), 12)
                s = _strr(s, _f2s(dih.params[1], 12), 12)
                s = _strr(s, str(int(dih.params[2])), 12)
            elif dih.funct == "3":
                # params = [C0, C1, C2, C3, C4, C5]
                for c in dih.params:
                    s = _strr(s, _f2s(c, 12), 12)
            s += "\n"
            f.write(s)

        f.write("\n\n")

    # ------------------------------------------------------------------
    # [ system ] and [ molecules ]
    # ------------------------------------------------------------------

    def _build_system_molecules(self, model: SystemModel) -> str:
        s = "\n\n"
        s += "[ system ]\n"
        s += ";"
        s = _strl(s, "Name", 11)
        s += "\n"
        s = _strl(s, model.title, len(model.title))
        s += "\n\n"

        s += "[ molecules ]\n"
        for gro_name, count in model.mol_sequence:
            s = _strr(s, gro_name, 11)
            s += str(count)
            s += "\n"
        s += "\n\n"
        return s
