import time
import torch
from dataclasses import dataclass
from typing import Protocol
from tqdm import tqdm


class Preconditioner(Protocol):
    def apply(self, r: torch.Tensor) -> torch.Tensor:
        ...


@dataclass
class SolverResult:
    solution: torch.Tensor
    iterations: int
    converged: bool
    final_residual: float
    history_abs_residual: list[float]
    history_rel_residual: list[float]
    history_time: list[float]


class FGMRES:

    def __init__(self, restart: int = 10, max_iters: int = 100,
                 rtol: float = 1e-8, timeout: float | None = None):
        self.restart = restart
        self.max_iters = max_iters
        self.rtol = rtol
        self.timeout = timeout

    def solve(self, A: torch.Tensor, b: torch.Tensor,
              M: Preconditioner | None = None,
              x0: torch.Tensor | None = None,
              progress_bar: bool = True) -> SolverResult:

        if x0 is None:
            x = torch.zeros_like(b)
        else:
            x = x0.clone()

        norm_b = torch.linalg.norm(b)
        if norm_b < 1e-15:
            return SolverResult(
                solution=x,
                iterations=0,
                converged=True,
                final_residual=0.0,
                history_abs_residual=[0.0],
                history_rel_residual=[0.0],
                history_time=[0.0]
            )

        hist_abs_res = []
        hist_rel_res = []
        hist_time = []

        n = len(b)
        V = torch.zeros(n, self.restart + 1, dtype=b.dtype, device=b.device)
        Z = torch.zeros(n, self.restart, dtype=b.dtype, device=b.device)
        H = torch.zeros(self.restart + 1, self.restart, dtype=b.dtype, device=b.device)
        g = torch.zeros(self.restart + 1, dtype=b.dtype, device=b.device)
        c = torch.zeros(self.restart, dtype=b.dtype, device=b.device)
        s = torch.zeros(self.restart, dtype=b.dtype, device=b.device)

        pbar = None
        if progress_bar:
            if self.timeout is None:
                pbar = tqdm(total=self.max_iters, desc="FGMRES")
            else:
                pbar = tqdm(desc="FGMRES")
            pbar.update()

        tic = time.time()

        r = b - A @ x
        beta = torch.linalg.norm(r)
        abs_res = beta.item()
        rel_res = abs_res / norm_b.item()
        hist_abs_res.append(abs_res)
        hist_rel_res.append(rel_res)
        hist_time.append(time.time() - tic)

        iters = 0
        converged = False

        while True:
            V[:, 0] = r / beta
            g.zero_()
            g[0] = beta

            j_final = 0
            for j in range(self.restart):
                if M is not None:
                    Z[:, j] = M.apply(V[:, j])
                else:
                    Z[:, j] = V[:, j]

                w = A @ Z[:, j]

                for k in range(j + 1):
                    H[k, j] = torch.dot(V[:, k], w)
                    w = w - H[k, j] * V[:, k]

                H[j + 1, j] = torch.linalg.norm(w)
                breakdown = bool(H[j + 1, j].abs() < 1e-14)

                if not breakdown:
                    V[:, j + 1] = w / H[j + 1, j]

                for k in range(j):
                    tmp = c[k] * H[k, j] + s[k] * H[k + 1, j]
                    H[k + 1, j] = -s[k] * H[k, j] + c[k] * H[k + 1, j]
                    H[k, j] = tmp

                t = torch.sqrt(H[j, j] ** 2 + H[j + 1, j] ** 2)
                if bool(t.abs() < 1e-14):
                    c[j] = 1.0
                    s[j] = 0.0
                else:
                    c[j] = H[j, j] / t
                    s[j] = H[j + 1, j] / t
                H[j, j] = c[j] * H[j, j] + s[j] * H[j + 1, j]
                H[j + 1, j] = 0
                g[j + 1] = -s[j] * g[j]
                g[j] = c[j] * g[j]

                abs_res = torch.abs(g[j + 1]).item()
                rel_res = abs_res / norm_b.item()
                hist_abs_res.append(abs_res)
                hist_rel_res.append(rel_res)
                hist_time.append(time.time() - tic)

                iters += 1
                j_final = j

                if pbar is not None:
                    pbar.update()

                should_stop = (
                    breakdown or
                    rel_res < self.rtol or
                    (self.timeout is None and iters >= self.max_iters) or
                    (self.timeout is not None and hist_time[-1] >= self.timeout)
                )

                if should_stop:
                    converged = rel_res < self.rtol or breakdown
                    break

            y = torch.linalg.solve_triangular(
                H[:j_final + 1, :j_final + 1],
                g[:j_final + 1].view(j_final + 1, 1),
                upper=True
            ).view(j_final + 1)

            x = x + Z[:, :j_final + 1] @ y
            r = b - A @ x
            beta = torch.linalg.norm(r)

            if converged or iters >= self.max_iters:
                break

            if self.timeout is not None and hist_time[-1] >= self.timeout:
                break

            g.zero_()

        if pbar is not None:
            pbar.close()

        return SolverResult(
            solution=x,
            iterations=iters,
            converged=converged,
            final_residual=hist_rel_res[-1],
            history_abs_residual=hist_abs_res,
            history_rel_residual=hist_rel_res,
            history_time=hist_time
        )


class Arnoldi:

    def build(self, A: torch.Tensor, m: int = 100,
              v0: torch.Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor]:
        n = A.shape[0]

        if v0 is None:
            v0 = torch.randn(n, dtype=A.dtype, device=A.device)

        beta = torch.linalg.norm(v0)

        V = torch.zeros(n, m + 1, dtype=A.dtype, device=A.device)
        H = torch.zeros(m + 1, m, dtype=A.dtype, device=A.device)

        V[:, 0] = v0 / beta

        for j in range(m):
            w = A @ V[:, j]
            for k in range(j + 1):
                H[k, j] = torch.dot(V[:, k], w)
                w = w - H[k, j] * V[:, k]
            H[j + 1, j] = torch.linalg.norm(w)

            if bool(H[j + 1, j].abs() < 1e-14):
                return V[:, :j + 2], H[:j + 2, :j + 1]

            V[:, j + 1] = w / H[j + 1, j]

        return V, H
