"""
Q128 — 130-Bit Precision GPU Runtime Python Wrapper
====================================================

130-bit precision (exceeds IEEE FP128) matrix multiplication on GPU.
Achieves 10^15-10^16x improvement over standard FP64.

API versions:
- 130-bit: q128_gemm_host(A, B, C_hi, C_lo, C_res, M, N, K) — full precision
- 106-bit: q128_gemm_host_106(A, B, C_hi, C_lo, M, N, K) — legacy (used here)

Usage:
    from q128 import Q128

    q = Q128()  # Auto-loads DLL
    C = q.gemm(A, B)  # Returns high-precision result

    # Or get the (hi, lo) pair directly:
    C_hi, C_lo = q.gemm_pair(A, B)

Precision Verification:
    from q128 import verify_precision
    verify_precision(K=256)  # Shows true precision vs ground truth

Author: Clouthier Simulation Labs
Date: February 2026
"""

import ctypes
import numpy as np
from decimal import Decimal, getcontext
import os
import platform

__version__ = "2.0.0"
__author__ = "Clouthier Simulation Labs"


def D(x):
    """
    Convert float to EXACT Decimal representation.

    CRITICAL: Use from_float(), NOT str()!
    - Decimal(str(1.1)) = 1.1 (loses bits!)
    - Decimal.from_float(1.1) = 1.100000000000000088817... (exact!)
    """
    return Decimal.from_float(float(x))


class Q128:
    """
    Q128 Quad-Precision GEMM Runtime

    Provides true ~10^-31 precision matrix multiplication using
    Error-Free Transformations (EFT) on GPU.
    """

    def __init__(self, dll_path=None):
        """
        Initialize Q128 runtime.

        Args:
            dll_path: Path to Q128 DLL/SO. Auto-detected if None.
        """
        if dll_path is None:
            dll_path = self._find_dll()

        self._load_dll(dll_path)
        self._setup_functions()

    def _find_dll(self):
        """Find Q128 library in standard locations."""
        if platform.system() == "Windows":
            candidates = [
                os.path.join(os.path.dirname(__file__), "build", "libq128.dll"),
                os.path.join(os.path.dirname(__file__), "libq128.dll"),
                r"C:\Q128\build\libq128.dll",
                "libq128.dll",
            ]
            ext = ".dll"
        else:
            candidates = [
                os.path.join(os.path.dirname(__file__), "libq128.so"),
                "./libq128.so",
                "/usr/local/lib/libq128.so",
            ]
            ext = ".so"

        for path in candidates:
            if os.path.exists(path):
                return path

        raise FileNotFoundError(
            f"Q128 library not found. Build it first or specify dll_path.\n"
            f"Searched: {candidates}"
        )

    def _load_dll(self, path):
        """Load the Q128 DLL."""
        self._lib = ctypes.CDLL(path)
        self._path = path

    def _setup_functions(self):
        """Setup ctypes function signatures."""
        P_D = ctypes.POINTER(ctypes.c_double)
        P_F = ctypes.POINTER(ctypes.c_float)
        I = ctypes.c_int

        # 130-bit API: q128_gemm_host(A, B, C_hi, C_lo, C_res, M, N, K)
        self._lib.q128_gemm_host.argtypes = [P_D, P_D, P_D, P_D, P_F, I, I, I]
        self._lib.q128_gemm_host.restype = I

        self._lib.q128_version.restype = ctypes.c_char_p

        self._P_D = P_D
        self._P_F = P_F
        self._I = I

    @property
    def version(self):
        """Get Q128 version string."""
        return self._lib.q128_version().decode()

    def gemm_pair(self, A, B):
        """
        130-bit precision matrix multiply returning (hi, lo) pair.

        The full precision result is: C = C_hi + C_lo (+ C_res for extra bits)

        Args:
            A: Matrix of shape (M, K), float64
            B: Matrix of shape (K, N), float64

        Returns:
            (C_hi, C_lo): Tuple of matrices, each (M, N) float64
                         Full result: C_hi + C_lo with ~10^-31 precision
        """
        A = np.ascontiguousarray(A, dtype=np.float64)
        B = np.ascontiguousarray(B, dtype=np.float64)

        M, K = A.shape
        K2, N = B.shape

        if K != K2:
            raise ValueError(f"Matrix dimension mismatch: A is {M}x{K}, B is {K2}x{N}")

        C_hi = np.zeros((M, N), dtype=np.float64)
        C_lo = np.zeros((M, N), dtype=np.float64)
        C_res = np.zeros((M, N), dtype=np.float32)

        ret = self._lib.q128_gemm_host(
            A.ctypes.data_as(self._P_D),
            B.ctypes.data_as(self._P_D),
            C_hi.ctypes.data_as(self._P_D),
            C_lo.ctypes.data_as(self._P_D),
            C_res.ctypes.data_as(self._P_F),
            self._I(M), self._I(N), self._I(K)
        )

        if ret != 0:
            raise RuntimeError(f"Q128 error: {ret}")

        return C_hi, C_lo

    def gemm(self, A, B):
        """
        Quad-precision matrix multiply returning single result.

        Note: The returned float64 matrix only holds ~16 digits.
        For full ~32 digit precision, use gemm_pair() and combine
        using Decimal arithmetic.

        Args:
            A: Matrix of shape (M, K), float64
            B: Matrix of shape (K, N), float64

        Returns:
            C: Matrix of shape (M, N), float64 (high part only)
        """
        C_hi, C_lo = self.gemm_pair(A, B)
        return C_hi + C_lo

    def gemm_decimal(self, A, B):
        """
        Quad-precision matrix multiply with Decimal output.

        Returns the full ~32 digit precision result as Decimal objects.
        Slower but preserves all computed precision.

        Args:
            A: Matrix of shape (M, K), float64
            B: Matrix of shape (K, N), float64

        Returns:
            C: 2D list of Decimal objects with full precision
        """
        C_hi, C_lo = self.gemm_pair(A, B)
        M, N = C_hi.shape

        result = []
        for i in range(M):
            row = []
            for j in range(N):
                val = D(C_hi[i, j]) + D(C_lo[i, j])
                row.append(val)
            result.append(row)

        return result


def compute_ground_truth(A, B):
    """
    Compute exact GEMM result using Decimal arithmetic.

    CRITICAL: Uses Decimal.from_float() for exact conversion.

    Args:
        A, B: Input matrices (numpy float64)

    Returns:
        2D list of Decimal objects with exact result
    """
    getcontext().prec = 150

    M, K = A.shape
    K2, N = B.shape

    result = []
    for i in range(M):
        row = []
        for j in range(N):
            acc = Decimal(0)
            for k in range(K):
                acc += D(A[i, k]) * D(B[k, j])
            row.append(acc)
        result.append(row)

    return result


def verify_precision(K=256, M=8, N=8, seed=42):
    """
    Verify Q128 achieves true quad precision.

    Compares Q128 output against exact Decimal ground truth.

    Args:
        K: Inner dimension (larger = more accumulation)
        M, N: Output matrix dimensions
        seed: Random seed for reproducibility

    Returns:
        dict with precision metrics
    """
    getcontext().prec = 150

    np.random.seed(seed)
    A = np.random.randn(M, K).astype(np.float64)
    B = np.random.randn(K, N).astype(np.float64)

    # Q128 result
    q = Q128()
    C_hi, C_lo = q.gemm_pair(A, B)

    # Ground truth
    C_truth = compute_ground_truth(A, B)

    # Compute errors
    fp64_errors = []
    q128_errors = []

    for i in range(M):
        for j in range(N):
            truth = C_truth[i][j]

            # FP64 error
            fp64_val = D(np.dot(A[i, :], B[:, j]))
            fp64_err = abs(fp64_val - truth) / abs(truth)
            fp64_errors.append(float(fp64_err))

            # Q128 error
            q128_val = D(C_hi[i, j]) + D(C_lo[i, j])
            q128_err = abs(q128_val - truth) / abs(truth)
            q128_errors.append(float(q128_err))

    result = {
        'K': K,
        'M': M,
        'N': N,
        'fp64_max_error': max(fp64_errors),
        'fp64_mean_error': sum(fp64_errors) / len(fp64_errors),
        'q128_max_error': max(q128_errors),
        'q128_mean_error': sum(q128_errors) / len(q128_errors),
        'improvement': max(fp64_errors) / max(q128_errors),
    }

    print(f"Q128 Precision Verification (K={K}):")
    print(f"  FP64 max error:  {result['fp64_max_error']:.3e}")
    print(f"  Q128 max error:  {result['q128_max_error']:.3e}")
    print(f"  Improvement:     {result['improvement']:.2e}x")

    return result


# Convenience function for quick testing
def test():
    """Quick test of Q128 functionality."""
    print("Q128 Quick Test")
    print("=" * 50)

    q = Q128()
    print(f"Version: {q.version}")

    # Small test
    A = np.random.randn(4, 4).astype(np.float64)
    B = np.random.randn(4, 4).astype(np.float64)

    C_hi, C_lo = q.gemm_pair(A, B)

    print(f"Test GEMM: 4x4 @ 4x4")
    print(f"  C_hi[0,0] = {C_hi[0,0]:.16e}")
    print(f"  C_lo[0,0] = {C_lo[0,0]:.16e}")
    print(f"  |C_lo/C_hi| = {abs(C_lo[0,0]/C_hi[0,0]):.2e}")
    print()

    # Precision verification
    verify_precision(K=64)

    print()
    print("Q128 working correctly!")


if __name__ == "__main__":
    test()
