from collections.abc import Iterable

import torch
import torch.nn.functional as F
from torch.optim import Optimizer
from torch.optim.lr_scheduler import LRScheduler
from dataclasses import dataclass
from tqdm import tqdm

from ..generator.base import BatchMatrixData
from ..nn.context_resgcn import ContextResGCN, BaselineResGCN


@dataclass
class TrainingConfig:
    batch_size: int = 16
    epochs: int = 2000
    matrices_per_epoch: int = 1
    num_context_pairs: int = 5
    learning_rate: float = 1e-3
    weight_decay: float = 0.0
    grad_accumulation_steps: int = 1


class MatrixPFN:

    def __init__(self, net: ContextResGCN | BaselineResGCN, model_device: torch.device):
        self.net = net
        self.model_device = model_device
        self.dtype = net.dtype
        self.uses_context = hasattr(net, 'context_aggregator')
        self.A = None
        self.diag_A = None
        self.context_pairs = None

    def set_matrix(self, A: torch.Tensor, diag_A: torch.Tensor,
                   context_pairs: torch.Tensor | None = None):
        self.A = A
        self.diag_A = diag_A.to(self.model_device).to(self.dtype)
        if context_pairs is not None:
            self.context_pairs = context_pairs.to(self.model_device).to(self.dtype)
        else:
            self.context_pairs = None
        self.net.set_matrix(A)

    def _batch_data_to_sparse(self, data: BatchMatrixData) -> tuple[
        torch.Tensor, torch.Tensor, torch.Tensor
    ]:
        A = torch.sparse_coo_tensor(
            data.indices, data.values[0], (data.n, data.n)
        ).coalesce().to_sparse_csc()
        return A, data.diagonals[0], data.context_pairs[0]

    def _train_step(self, A: torch.Tensor, diag_A: torch.Tensor,
                    context_pairs: torch.Tensor, training_batch_size: int) -> torch.Tensor:
        n = A.shape[0]
        x = torch.randn(n, training_batch_size, dtype=torch.float64, device=A.device)
        b = A @ x

        self.net.set_matrix(A)

        b_input = b.to(self.model_device).to(self.dtype)
        diag_input = diag_A.to(self.model_device).to(self.dtype)
        context_input = context_pairs.to(self.model_device).to(self.dtype)

        if self.uses_context:
            x_pred = self.net(b_input, diag_input, context_input)
        else:
            x_pred = self.net(b_input)

        x_pred_cpu = x_pred.to(A.device).to(torch.float64)
        b_pred = (A @ x_pred_cpu).to(self.model_device).to(self.dtype)

        return F.l1_loss(b_pred, b_input)

    def train(
        self,
        data_source: Iterable[BatchMatrixData],
        config: TrainingConfig,
        optimizer: Optimizer,
        scheduler: LRScheduler | None = None,
        checkpoint_path: str | None = None,
        progress_bar: bool = True,
    ) -> dict:
        self.net.train()
        optimizer.zero_grad()

        history = {"loss": [], "best_loss": float("inf"), "best_epoch": -1}

        pbar = tqdm(total=config.epochs, desc="Training") if progress_bar else None
        data_iter = iter(data_source)

        for epoch in range(config.epochs):
            epoch_loss = 0.0

            for _ in range(config.matrices_per_epoch):
                data = next(data_iter)
                A, diag_A, context_pairs = self._batch_data_to_sparse(data)

                loss = self._train_step(A, diag_A, context_pairs, config.batch_size)
                epoch_loss += loss.item()

                scaled_loss = loss / (config.matrices_per_epoch * config.grad_accumulation_steps)
                scaled_loss.backward()

            if (epoch + 1) % config.grad_accumulation_steps == 0:
                optimizer.step()
                optimizer.zero_grad()
                if scheduler is not None:
                    scheduler.step()

            avg_loss = epoch_loss / config.matrices_per_epoch
            history["loss"].append(avg_loss)

            if avg_loss < history["best_loss"]:
                history["best_loss"] = avg_loss
                history["best_epoch"] = epoch
                if checkpoint_path is not None:
                    torch.save(self.net.state_dict(), checkpoint_path)

            if pbar is not None:
                pbar.set_description(f"Loss: {avg_loss:.4e}")
                pbar.update()

        if pbar is not None:
            pbar.close()

        return history

    @torch.no_grad()
    def apply(self, r: torch.Tensor) -> torch.Tensor:
        self.net.eval()

        original_device = r.device
        r_model = r.to(self.model_device).to(self.dtype)
        r_model = r_model.view(-1, 1)

        if isinstance(self.net, (ContextResGCN, ContextResMPNN)):
            z = self.net(r_model, self.diag_A, self.context_pairs)
        else:
            z = self.net(r_model)

        z = z.view(-1)
        z = z.to(original_device).double()

        return z

    def prepare_for_solve(self, A: torch.Tensor, num_context_pairs: int = 5):
        n = A.shape[0]
        sparse_device = A.device

        if A.layout == torch.sparse_csc:
            A_coo = A.to_sparse_coo().coalesce()
        else:
            A_coo = A.coalesce() if A.is_sparse else A
        indices = A_coo.indices()
        values_tensor = A_coo.values()
        diag_mask = indices[0] == indices[1]
        diag_A = torch.zeros(n, dtype=torch.float64, device=sparse_device)
        diag_A[indices[0, diag_mask]] = values_tensor[diag_mask]

        x_context = torch.randn(n, num_context_pairs, dtype=torch.float64, device=sparse_device)
        context_pairs = torch.stack([x_context, A @ x_context], dim=-1)

        self.set_matrix(A, diag_A, context_pairs)
