#
# 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 dataclasses import dataclass
from itertools import product
from pathlib import Path

import torch
from torch.nn import MSELoss
from tqdm import tqdm

from ctak_train.data import DataSplit, Pool
from ctak_train.net import NET_DIMS, Net
from ctak_train.selfplay import cand_copy_for


@dataclass(frozen=True)
class DataGenConfig:
    size: int
    depth: int
    min_ply: int
    games: int
    procs: int
    replay_iters: int
    root: Path
    train_only_self_play: bool


@dataclass(frozen=True)
class TrainingGridSearchConfig:
    smooth_factor: list[float]
    epochs: list[int]
    batch: list[int]
    lr: list[float]
    val_frac: list[float]

    def make_configs(self, data_gen_config: DataGenConfig) -> list["Candidate"]:
        inp_dim, hidden = NET_DIMS[data_gen_config.size]
        flds = self.__dataclass_fields__.keys()
        almost_cands = [
            dict(zip(flds, vals)) for vals in product(*(getattr(self, k) for k in flds))
        ]
        for ident, almost_cand in enumerate(almost_cands):
            net = Net(inp_dim, hidden)
            net.to_c(
                size=data_gen_config.size,
                ident=ident,
                path=data_gen_config.root
                / "include"
                / f"cand_{ident}_weights_{data_gen_config.size}.c",
                note=f"{ident=}, v000, random init",
            )
            cand_copy_for(
                6 if data_gen_config.size == 5 else 5, ident, data_gen_config.root
            )
            almost_cand["net"] = net
            almost_cand["train_only_self_play"] = data_gen_config.train_only_self_play
        return [Candidate(**d) for d in almost_cands]


@dataclass
class Candidate:
    train_only_self_play: bool
    smooth_factor: float
    epochs: int
    batch: int
    lr: float
    val_frac: float
    net: "Net"

    def __str__(self) -> str:
        return f"(λ={self.smooth_factor}, E={self.epochs}, BS={self.batch}, η={self.lr:.2e}, V%={self.val_frac * 100.0:.1f})"

    def training_pool(self, pool: Pool) -> Pool:
        if self.train_only_self_play:
            return pool.select_chains(pool.chain_self_play)
        return pool

    def fit(self, pool: Pool) -> tuple[float, float]:
        # The train/val split is fixed but targets are refreshed
        split = DataSplit.build(self.training_pool(pool), self.val_frac)
        if split.n == 0:
            return math.nan, math.nan

        optimizer = torch.optim.AdamW(self.net.parameters(), lr=self.lr)
        loss_fn = MSELoss()

        pbar = tqdm(range(self.epochs), desc="train", unit="ep")
        for _ in pbar:
            split.refresh_targets(self.net, self.smooth_factor)

            self.net.train()
            for head, cells, parity, target in split.batches(self.batch):
                optimizer.zero_grad()
                predictions = self.net(head, cells) * parity
                loss = loss_fn(predictions, target)
                loss.backward()
                optimizer.step()

            with torch.no_grad():
                self.net.eval()
                val = split.val_idx
                predictions = self.net(split.pool.head[val], split.pool.cells[val])
                predictions = predictions * split.pool.parity[val]
                targets = split.targets[val]
                mse = loss_fn(predictions, targets).item()
                corr = torch.corrcoef(torch.stack([predictions, targets]))[0, 1].item()

            pbar.set_postfix(val_corr=f"{corr:.4f}", val_mse=f"{mse:.4f}")

        return corr, mse