#
# 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/>.

import math
from pathlib import Path

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from ctak_train.candidate import Candidate
from ctak_train.data import GameLog
from scipy.optimize import minimize

matplotlib.use("Agg")

# Elo points per unit of logit strength
ELO_PER_LOGIT = 400.0 / math.log(10.0)

# Spread of the zero-mean prior on the ratings, in logit units. Head-to-head
# results alone leave the likelihood unbounded while the data is thin (a pair
# that has never split draws runs off to +-inf).
PRIOR_SD_LOGIT = 1.0


def fit_strengths(log: GameLog, n_candidates: int) -> tuple[np.ndarray, np.ndarray]:
    """Bradley-Terry logistic model (draws scored 0.5) fit in natural-logit
    units, mean anchored to 0, reported as Elo points around 1500. Returns
    (strengths, +-1 sd).
    """
    rated = log.black != log.white
    black = log.black[rated].astype(np.int64)
    white = log.white[rated].astype(np.int64)
    if n_candidates < 2 or black.size == 0:
        return np.full(n_candidates, 1500.0), np.zeros(n_candidates)
    score = (log.result[rated].astype(np.float64) + 1.0) / 2.0

    # n_candidates - 1 free parameters; the last candidate anchors the mean to 0
    design = np.eye(n_candidates, n_candidates - 1)
    design[n_candidates - 1] = -1.0

    def nll_grad(theta: np.ndarray) -> tuple[float, np.ndarray]:
        rating = design @ theta
        x = np.clip(rating[black] - rating[white], -350.0, 350.0)
        p = 1.0 / (1.0 + np.exp(-x))
        sp = np.maximum(x, 0.0) + np.log1p(np.exp(-np.abs(x)))  # softplus(x)
        sp_neg = np.maximum(-x, 0.0) + np.log1p(np.exp(-np.abs(x)))
        # nll = score*softplus(-x) + (1-score)*softplus(x): finite for all x
        nll = float(np.sum(score * sp_neg + (1.0 - score) * sp))
        nll += 0.5 * float(rating @ rating) / PRIOR_SD_LOGIT**2

        grad_rating = rating / PRIOR_SD_LOGIT**2
        np.add.at(grad_rating, black, p - score)
        np.add.at(grad_rating, white, score - p)
        return nll, design.T @ grad_rating

    def hessian(theta: np.ndarray) -> np.ndarray:
        rating = design @ theta
        x = np.clip(rating[black] - rating[white], -350.0, 350.0)
        p = 1.0 / (1.0 + np.exp(-x))
        w = p * (1.0 - p)

        diag = np.bincount(black, weights=w, minlength=n_candidates) + np.bincount(
            white, weights=w, minlength=n_candidates
        )
        off = np.bincount(
            black * n_candidates + white, weights=-w, minlength=n_candidates**2
        ).reshape(n_candidates, n_candidates)
        hess_rating = (
            off
            + np.diag(diag - np.diag(off))
            + np.eye(n_candidates) / PRIOR_SD_LOGIT**2
        )
        return design.T @ hess_rating @ design

    res = minimize(nll_grad, np.zeros(n_candidates - 1), jac=True, method="L-BFGS-B")
    if not res.success:
        print(f"strength fit did not converge: {res.message}")
    cov = design @ np.linalg.pinv(hessian(res.x)) @ design.T
    se = np.sqrt(np.maximum(np.diag(cov), 0.0)) * ELO_PER_LOGIT
    return 1500.0 + (design @ res.x) * ELO_PER_LOGIT, se


def plot_strengths(
    history: list[np.ndarray], candidates: list[Candidate], path: Path
) -> None:
    fig, ax = plt.subplots(figsize=(10, 6))
    iterations = range(1, len(history) + 1)
    for idx, cand in enumerate(candidates):
        ax.plot(
            iterations,
            [row[idx] for row in history],
            marker=".",
            ms=3,
            lw=1.2,
            label=f"cand {idx:03d} {cand}",
        )
    ax.axhline(1500.0, color="gray", lw=0.5)
    ax.set_xscale("log")
    ax.set_xlabel("Iteration")
    ax.set_ylabel("Elo")
    ax.set_title("Candidate strength")
    ax.legend(fontsize=8)
    fig.tight_layout()
    _ = fig.savefig(path, dpi=150)
    plt.close(fig)