#
# 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 orjson
from collections.abc import Iterator
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import torch
from ctak_train.net import NET_DIMS
from torch import Tensor
from torch.nn import Module


@dataclass
class GameLog:
    iteration: np.ndarray  # [G] iteration the game was played in
    black: np.ndarray  # [G] candidate id playing black
    white: np.ndarray  # [G] candidate id playing white
    result: np.ndarray  # [G] outcome from black's perspective, {-1, 0, 1}

    @classmethod
    def empty(cls) -> GameLog:
        return cls(
            iteration=np.empty(0, dtype=np.int32),
            black=np.empty(0, dtype=np.int16),
            white=np.empty(0, dtype=np.int16),
            result=np.empty(0, dtype=np.float32),
        )

    def append(self, newer: GameLog) -> None:
        self.iteration = np.concatenate([self.iteration, newer.iteration])
        self.black = np.concatenate([self.black, newer.black])
        self.white = np.concatenate([self.white, newer.white])
        self.result = np.concatenate([self.result, newer.result])


@dataclass(frozen=True)
class Pool:
    head: Tensor  # [n, 2] fixed head inputs of the root state s_t
    cells: Tensor  # [n, n_cells] cell states 0..6 of s_t
    leaf_head: Tensor  # [n, 2] head inputs of the PV leaf pv_t
    leaf_cells: Tensor  # [n, n_cells] cell states of pv_t
    parity: Tensor  # [n] +1 when black is to move at s_t
    leaf_parity: Tensor  # [n] +1 when black is to move at pv_t
    terminal: Tensor  # [n] leaf is an exact terminal
    terminal_value: Tensor  # [n] exact terminal value, black's perspective
    chain_lens: Tensor  # [C] rows per chain
    chain_result: Tensor  # [C] game result, black's perspective
    chain_self_play: Tensor  # [C] chain comes from a self vs self game

    @classmethod
    def empty(cls, n_cells: int) -> Pool:
        def rows(dtype: torch.dtype, width: int = 0) -> Tensor:
            return torch.empty((0, width) if width else (0,), dtype=dtype)

        return cls(
            head=rows(torch.float32, 2),
            cells=rows(torch.uint8, n_cells),
            leaf_head=rows(torch.float32, 2),
            leaf_cells=rows(torch.uint8, n_cells),
            parity=rows(torch.float32),
            leaf_parity=rows(torch.float32),
            terminal=rows(torch.bool),
            terminal_value=rows(torch.float32),
            chain_lens=rows(torch.int64),
            chain_result=rows(torch.float32),
            chain_self_play=rows(torch.bool),
        )

    @property
    def n_rows(self) -> int:
        return self.head.shape[0]

    def select_chains(self, keep: Tensor) -> Pool:
        keep = keep.bool()
        rows = torch.repeat_interleave(
            torch.arange(self.chain_lens.shape[0], dtype=torch.int64), self.chain_lens
        )
        idx = torch.nonzero(keep[rows]).flatten()
        return Pool(
            head=self.head[idx],
            cells=self.cells[idx],
            leaf_head=self.leaf_head[idx],
            leaf_cells=self.leaf_cells[idx],
            parity=self.parity[idx],
            leaf_parity=self.leaf_parity[idx],
            terminal=self.terminal[idx],
            terminal_value=self.terminal_value[idx],
            chain_lens=self.chain_lens[keep],
            chain_result=self.chain_result[keep],
            chain_self_play=self.chain_self_play[keep],
        )


_ROW_FIELDS = (
    "head",
    "cells",
    "leaf_head",
    "leaf_cells",
    "parity",
    "leaf_parity",
    "terminal",
    "terminal_value",
)


def _chain_arrays(
    rows: list[dict], n_cells: int, source: Path
) -> dict[str, np.ndarray]:

    def column(part: str, key: str, dtype: np.dtype) -> np.ndarray:
        return np.array([r[part][key] for r in rows], dtype=dtype)

    out = {
        "head": column("root", "head", np.float32),
        "cells": column("root", "cells", np.uint8),
        "leaf_head": column("leaf", "head", np.float32),
        "leaf_cells": column("leaf", "cells", np.uint8),
        "parity": column("root", "parity", np.float32),
        "leaf_parity": column("leaf", "parity", np.float32),
        "terminal": column("leaf", "terminal", np.bool_),
        "terminal_value": column("leaf", "terminal_value", np.float32),
    }
    for name in ("cells", "leaf_cells"):
        if out[name].shape[1] != n_cells:
            raise ValueError(
                f"{source}: {name} width {out[name].shape[1]}, expected {n_cells}"
            )
    return out


def _game_colours(game: dict) -> tuple[int, int]:
    a, b = int(game["cand_A"]), int(game["cand_B"])
    for r in game["rows"]:
        if r["root"]["parity"] > 0:
            cand = int(r["cand"])
            return cand, b if cand == a else a
    recorded = {int(r["cand"]) for r in game["rows"]}
    return (a, b) if a not in recorded else (b, a)


def load_window(
    files: list[Path], size: int, n_candidates: int, iteration: int
) -> tuple[list[Pool], GameLog]:
    """Training pools for every file in the window, plus the game records of
    the newest file only - the caller's GameLog already holds everything older."""
    n_cells = NET_DIMS[size][0] - 2
    blocks = [{f: [] for f in _ROW_FIELDS} for _ in range(n_candidates)]
    chain_lens: list[list[int]] = [[] for _ in range(n_candidates)]
    chain_result: list[list[float]] = [[] for _ in range(n_candidates)]
    chain_self_play: list[list[bool]] = [[] for _ in range(n_candidates)]
    black, white, result = [], [], []

    for file_index, path in enumerate(files):
        is_newest = file_index == len(files) - 1
        with open(path, "rb") as f:
            for line in f:
                if not line.strip():
                    continue
                game = orjson.loads(line)
                game_result = float(game["result"])
                self_play = int(game["cand_A"]) == int(game["cand_B"])
                if is_newest:
                    b, w = _game_colours(game)
                    black.append(b)
                    white.append(w)
                    result.append(game_result)

                by_cand: dict[int, list[dict]] = {}
                for r in game["rows"]:
                    by_cand.setdefault(int(r["cand"]), []).append(r)

                for cand, rows in by_cand.items():
                    parity = [r["root"]["parity"] for r in rows]
                    for sign in (1, -1):
                        chain = [
                            r for r, p in zip(rows, parity) if (p > 0) == (sign > 0)
                        ]
                        if not chain:
                            continue
                        for name, array in _chain_arrays(chain, n_cells, path).items():
                            blocks[cand][name].append(array)
                        chain_lens[cand].append(len(chain))
                        chain_result[cand].append(game_result)
                        chain_self_play[cand].append(self_play)

    def pool_of(cand: int) -> Pool:
        if not chain_lens[cand]:
            return Pool.empty(n_cells)
        arrays = {
            name: torch.from_numpy(np.concatenate(blocks[cand][name]))
            for name in _ROW_FIELDS
        }
        return Pool(
            **arrays,
            chain_lens=torch.tensor(chain_lens[cand], dtype=torch.int64),
            chain_result=torch.tensor(chain_result[cand], dtype=torch.float32),
            chain_self_play=torch.tensor(chain_self_play[cand], dtype=torch.bool),
        )

    games = GameLog(
        iteration=np.full(len(black), iteration, dtype=np.int32),
        black=np.array(black, dtype=np.int16),
        white=np.array(white, dtype=np.int16),
        result=np.array(result, dtype=np.float32),
    )
    return [pool_of(c) for c in range(n_candidates)], games


@dataclass
class DataSplit:
    pool: Pool
    train_idx: Tensor
    val_idx: Tensor
    targets: Tensor  # [n] aligned with pool rows, refreshed per epoch
    chain_end: Tensor  # [C] last row index of each chain
    wave_order: Tensor  # [n] rows ordered by distance from their chain's end
    wave_bounds: Tensor  # [d] start of distance d in wave_order

    @property
    def n(self) -> int:
        return self.pool.n_rows

    @classmethod
    def build(cls, pool: Pool, val_frac: float) -> DataSplit:
        n = pool.n_rows
        if n == 0:
            empty = torch.empty(0, dtype=torch.int64)
            return cls(pool, empty, empty, torch.empty(0), empty, empty, empty)

        # Smoothing is G[t] = (1-λ)V[t] + λG[t+1] with G at the chain end.
        # Chains are contiguous, so every row that is d plies from its chain's
        # end can be updated in one vectorised wavefront step per d.
        lens = pool.chain_lens
        chain_end = torch.cumsum(lens, dim=0) - 1
        chain_of_row = torch.repeat_interleave(
            torch.arange(lens.shape[0], dtype=torch.int64), lens
        )
        dist = chain_end[chain_of_row] - torch.arange(n, dtype=torch.int64)
        wave_order = torch.argsort(dist, stable=True)
        wave_bounds = torch.searchsorted(
            dist[wave_order], torch.arange(int(dist.max()) + 2, dtype=torch.int64)
        )

        n_val = max(1, int(n * val_frac))
        perm = torch.randperm(n)
        return cls(
            pool=pool,
            train_idx=perm[n_val:],
            val_idx=perm[:n_val],
            targets=torch.zeros(n),
            chain_end=chain_end,
            wave_order=wave_order,
            wave_bounds=wave_bounds,
        )

    def refresh_targets(self, net: Module, lam: float) -> None:
        pool = self.pool
        net.eval()
        with torch.no_grad():
            # Net(pv_t), black's perspective; exact terminals override the net
            leaf = net(pool.leaf_head, pool.leaf_cells) * pool.leaf_parity
            leaf = torch.where(pool.terminal, pool.terminal_value, leaf)

        targets = torch.zeros_like(leaf)
        targets[self.chain_end] = pool.chain_result
        for lo, hi in zip(self.wave_bounds[1:-1], self.wave_bounds[2:]):
            idx = self.wave_order[lo:hi]
            targets[idx] = (1.0 - lam) * leaf[idx] + lam * targets[idx + 1]
        self.targets = targets

    def batches(
        self, batch_size: int
    ) -> Iterator[tuple[Tensor, Tensor, Tensor, Tensor]]:
        order = torch.randperm(self.train_idx.shape[0])
        train_targets = self.targets[self.train_idx]
        for start in range(0, order.shape[0], batch_size):
            idx = order[start : start + batch_size]
            rows = self.train_idx[idx]
            yield (
                self.pool.head[rows],
                self.pool.cells[rows],
                self.pool.parity[rows],
                train_targets[idx],
            )