#
# This file is part of ctak.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with ctak. If not, see <https://www.gnu.org/licenses/>.

from pathlib import Path

import numpy as np
import torch
from torch import Tensor
from torch.nn import Embedding, Linear, Module

# (input dim, hidden dim) per board size; must match include/weights_<size>.h
NET_DIMS: dict[int, tuple[int, int]] = {
    5: (27, 64),  # FIVE_INP_NUM, FIVE_DENSE_NUM
    6: (38, 64),  # SIX_INP_NUM, SIX_DENSE_NUM
}


def _fmt_flat(a: np.ndarray, per_line: int = 8) -> str:
    a = np.ravel(np.asarray(a, dtype=np.float32))
    lines: list[str] = []
    for i in range(0, len(a), per_line):
        chunk = ", ".join(f"{v:.9g}" for v in a[i : i + per_line])
        if i + per_line < len(a):
            chunk += ","
        lines.append("    " + chunk)
    return "\n".join(lines)


def _fmt_matrix(mat: np.ndarray, per_line: int = 8) -> str:
    mat = np.asarray(mat, dtype=np.float32)
    rows: list[str] = []
    for r in mat:
        chunks = [
            ", ".join(f"{v:.9g}" for v in r[i : i + per_line])
            for i in range(0, len(r), per_line)
        ]
        rows.append("    {" + ",\n     ".join(chunks) + "},")
    return "\n".join(rows)


def fast_tanh(x: Tensor) -> Tensor:
    # Must match fast_tanh in include/nn.c: degree-11 fit of tanh on [-3, 3],
    # saturating to +-1 outside.
    x2 = x * x
    poly = x * (
        0.997136032
        + x2
        * (
            -0.305215124
            + x2
            * (
                0.0832466371
                + x2 * (-0.0139947924 + x2 * (0.00123106314 + x2 * -4.29748644e-05))
            )
        )
    )
    return torch.where(x > 3.0, 1.0, torch.where(x < -3.0, -1.0, poly))


def fast_silu(x: Tensor) -> Tensor:
    # Must match fast_silu in include/nn.c: [4,2] Pade approximant with cutoff.
    # The cutoff below is evaluated at x = 0 (instead of x) because torch
    # evaluates both branches of where: 1 + 0.1 * x is 0 for x = -10.
    xs = torch.where(x < -4.6, 0.0, x)
    poly = xs * (0.5 + xs * (0.25 + xs * (0.05 + 0.004167 * xs))) / (1 + 0.1 * xs)
    return torch.where(x < -4.6, 0.0, torch.where(x > 4.6, x, poly))


class Net(Module):
    def __init__(self, inp_dim: int, hidden: int) -> None:
        super().__init__()
        # learned cell-state embedding, see nn.c
        self.embedding = Embedding(7, 1)
        self.hidden = Linear(inp_dim, hidden, bias=True)
        self.output = Linear(hidden, 1, bias=True)

    def forward(self, head: Tensor, cells: Tensor) -> Tensor:
        # head: [n, 2] fixed head inputs; cells: [n, n_cells] raw cell states
        # (0..6) - the only place they become anything other than integers
        inp = torch.cat([head, self.embedding(cells.long()).flatten(1)], dim=1)
        hidden = fast_silu(self.hidden(inp))
        return fast_tanh(self.output(hidden)).squeeze(-1)

    def to_c(
        self,
        *,
        size: int,
        ident: int | None,
        path: Path,
        note: str,
    ) -> None:
        upper = "FIVE" if size == 5 else "SIX"
        prefix = (f"cand_{ident}_" if ident is not None else "") + upper
        guard = f"WEIGHTS_{upper}" + (f"_{ident}" if ident is not None else "")

        E = self.embedding.weight.detach().cpu().numpy()
        W = self.hidden.weight.detach().cpu().numpy()
        B = self.hidden.bias.detach().cpu().numpy()
        LW = self.output.weight.detach().cpu().numpy().ravel()
        LB = float(self.output.bias.detach().cpu().item())

        with open(path, "w") as f:
            lines = [
                "/* generated by nn_train.py, do not edit by hand */\n",
                f"/* {note} */\n",
                f"#ifndef {guard}\n#define {guard}\n",
                f'#include "weights_{size}.h"\n\n',
                f"const float {prefix}_embedding[{len(E)}] =\n{{\n{_fmt_flat(E, per_line=7)}\n}};\n\n",
                f"const float {prefix}_dense_weights[{upper}_INP_NUM][{upper}_DENSE_NUM] NN_ALIGNED =\n{{\n{_fmt_matrix(W.T)}\n}};\n\n",
                f"const float {prefix}_dense_biases[{upper}_DENSE_NUM] =\n{{\n{_fmt_flat(B)}\n}};\n\n",
                f"const float {prefix}_logit_weights[{upper}_DENSE_NUM] NN_ALIGNED =\n{{\n{_fmt_flat(LW)}\n}};\n\n",
                f"const float {prefix}_logit_bias = {LB:.9g};\n",
                "#endif",
            ]
            for line in lines:
                _ = f.write(line)