from collections.abc import Iterable
import re

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, ContextResMPNN
from ..solver.fgmres import FGMRES, SolverResult

_NET_CLASSES = {
    "ContextResGCN": ContextResGCN,
    "BaselineResGCN": BaselineResGCN,
    "ContextResMPNN": ContextResMPNN,
}

_DTYPE_MAP = {
    "float32": torch.float32,
    "float64": torch.float64,
    "float16": torch.float16,
    "bfloat16": torch.bfloat16,
}

_HF_REPO_PATTERN = re.compile(r"^[\w.-]+/[\w.-]+$")


@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 = 1e-4
    grad_clip_norm: float = 1.0
    grad_accumulation_steps: int = 1


class MatrixPFN:

    def __init__(self, net: ContextResGCN | BaselineResGCN | ContextResMPNN, 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
        self.metadata: dict | None = 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_target = x.to(self.model_device).to(self.dtype)

        return F.l1_loss(x_pred, x_target)

    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:
                torch.nn.utils.clip_grad_norm_(self.net.parameters(), config.grad_clip_norm)
                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 self.uses_context:
            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 _net_config(self) -> dict:
        net = self.net
        config = {
            "net_type": type(net).__name__,
            "num_layers": net.num_layers,
            "embed": net.embed,
            "hidden": net.hidden,
            "drop_rate": net.dropout.p,
            "scale_input": net.scale_input,
            "dtype": str(net.dtype).removeprefix("torch."),
        }
        if hasattr(net, "num_context_pairs"):
            config["num_context_pairs"] = net.num_context_pairs
        return config

    def save_pretrained(self, path: str, *, metadata: dict | None = None) -> None:
        checkpoint = {
            "config": self._net_config(),
            "state_dict": self.net.state_dict(),
        }
        if metadata is not None:
            checkpoint["metadata"] = metadata
        torch.save(checkpoint, path)

    @classmethod
    def from_pretrained(cls, path: str, *,
                        revision: str = "main",
                        device: torch.device | None = None) -> "MatrixPFN":
        if device is None:
            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        if path.startswith(("http://", "https://")):
            checkpoint = torch.hub.load_state_dict_from_url(path, map_location="cpu")
        elif _HF_REPO_PATTERN.match(path):
            url = f"https://huggingface.co/{path}/resolve/{revision}/model.pt"
            checkpoint = torch.hub.load_state_dict_from_url(url, map_location="cpu")
        else:
            checkpoint = torch.load(path, map_location="cpu", weights_only=True)

        config = checkpoint["config"]
        net_cls = _NET_CLASSES[config["net_type"]]

        kwargs = {
            "num_layers": config["num_layers"],
            "embed": config["embed"],
            "hidden": config["hidden"],
            "drop_rate": config["drop_rate"],
            "scale_input": config.get("scale_input", True),
            "dtype": _DTYPE_MAP[config.get("dtype", "float32")],
        }
        if "num_context_pairs" in config:
            kwargs["num_context_pairs"] = config["num_context_pairs"]

        net = net_cls(**kwargs)
        net.load_state_dict(checkpoint["state_dict"])
        net = net.to(device)

        instance = cls(net, device)
        instance.metadata = checkpoint.get("metadata")
        return instance

    def solve(self, A, b, *,
              rtol: float = 1e-8, restart: int = 30, max_iters: int = 300,
              timeout: float | None = None,
              num_context_pairs: int | None = None,
              x0=None,
              progress_bar: bool = True) -> SolverResult:
        import numpy as np
        from scipy.sparse import issparse
        from .sparse_convert import scipy_to_torch_csc

        if issparse(A):
            A = scipy_to_torch_csc(A)
        if isinstance(b, np.ndarray):
            b = torch.from_numpy(np.ascontiguousarray(b)).to(dtype=torch.float64)
        if x0 is not None and isinstance(x0, np.ndarray):
            x0 = torch.from_numpy(np.ascontiguousarray(x0)).to(dtype=torch.float64)

        if num_context_pairs is None:
            num_context_pairs = getattr(self.net, "num_context_pairs", 5)

        self.prepare_for_solve(A, num_context_pairs=num_context_pairs)

        solver = FGMRES(restart=restart, max_iters=max_iters, rtol=rtol, timeout=timeout)
        return solver.solve(A, b, M=self, x0=x0, progress_bar=progress_bar)

    def save_to_hub(self, repo_id: str, *,
                    metadata: dict | None = None,
                    commit_message: str = "Upload model") -> str:
        from huggingface_hub import HfApi
        import tempfile
        import os

        api = HfApi()
        api.create_repo(repo_id, repo_type="model", exist_ok=True)

        with tempfile.TemporaryDirectory() as tmp:
            model_path = os.path.join(tmp, "model.pt")
            self.save_pretrained(model_path, metadata=metadata)

            card_path = os.path.join(tmp, "README.md")
            with open(card_path, "w") as f:
                f.write(self._generate_model_card(repo_id, metadata))

            api.upload_folder(
                folder_path=tmp,
                repo_id=repo_id,
                commit_message=commit_message,
            )

        return f"https://huggingface.co/{repo_id}"

    def _generate_model_card(self, repo_id: str, metadata: dict | None = None) -> str:
        config = self._net_config()
        total_params = sum(p.numel() for p in self.net.parameters())
        model_name = repo_id.split("/")[-1]

        lines = [
            "---",
            "library_name: matrixpfn",
            "tags:",
            "  - preconditioner",
            "  - sparse-linear-systems",
            "  - graph-neural-network",
            "  - pytorch",
            "---",
            "",
            f"# {model_name}",
            "",
            "GNN-based learned preconditioner for sparse linear systems.",
            "",
        ]

        if metadata and "version" in metadata:
            lines.append(f"**Version**: {metadata['version']}")
            lines.append("")

        lines.extend([
            "## Usage",
            "",
            "```python",
            "import numpy as np",
            "from scipy.io import mmread",
            "from matrixpfn import MatrixPFN",
            "",
            f'pfn = MatrixPFN.from_pretrained("{repo_id}")',
            "",
            "A = mmread(\"matrix.mtx\")  # any scipy sparse matrix",
            "b = A @ np.random.randn(A.shape[0])",
            "",
            "result = pfn.solve(A, b)  # accepts scipy sparse directly",
            "print(f\"Converged: {result.converged} in {result.iterations} iterations\")",
            "```",
            "",
            "## Architecture",
            "",
            f"| Parameter | Value |",
            f"|-----------|-------|",
            f"| Network | {config['net_type']} |",
            f"| Layers | {config['num_layers']} |",
            f"| Embed | {config['embed']} |",
            f"| Hidden | {config['hidden']} |",
            f"| Context pairs | {config.get('num_context_pairs', 'N/A')} |",
            f"| Parameters | {total_params:,} |",
            f"| dtype | {config['dtype']} |",
        ])

        if metadata and "training" in metadata:
            lines.extend(["", "## Training", ""])
            for key, value in metadata["training"].items():
                lines.append(f"- **{key}**: {value}")

        if metadata and "benchmark" in metadata:
            benchmark = metadata["benchmark"]
            if isinstance(benchmark, list) and benchmark:
                columns = list(benchmark[0].keys())
                lines.extend(["", "## Benchmark", ""])
                lines.append("| " + " | ".join(columns) + " |")
                lines.append("| " + " | ".join("---" for _ in columns) + " |")
                for row in benchmark:
                    lines.append("| " + " | ".join(str(row.get(c, "")) for c in columns) + " |")

        lines.append("")
        return "\n".join(lines)

    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]

        from ..generator.base import make_context_pairs
        context_pairs = make_context_pairs(indices, values_tensor, n, num_context_pairs, sparse_device)

        self.set_matrix(A, diag_A, context_pairs)
