'''Helper functions.'''
from typing import Literal

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from scipy.linalg import sqrtm, fractional_matrix_power

class SmoothClampedReLU(nn.Module):
    def __init__(self, beta=50):
        super(SmoothClampedReLU, self).__init__()
        self.beta = beta
        
    def forward(self, x):
        # Smooth transition at x=0 (using softplus with high beta)
        activated = F.softplus(x, beta=self.beta)
        # Smooth transition at x=1 (using sigmoid scaled and shifted)
        # As x approaches infinity, this approaches 1
        clamped = activated - F.softplus(activated - 1, beta=self.beta)
        
        return clamped

def float_x(data):
    '''Set data array precision.'''
    return np.float32(data)

def matrix_power(M, power):
    """
    Compute the power of a matrix.
    :param M: Matrix to power.
    :param power: Power to raise the matrix to.
    :return: Matrix raised to the power - M^{power}.
    """
    return stable_matrix_power(M, power)
    # if len(M.shape) == 2:
    #     assert M.shape[0] == M.shape[1], "Matrix must be square"

    #     # gpu square root
    #     S, U = torch.linalg.eigh(M)
    #     S[S<0] = 0.
    #     return U @ torch.diag(S**power) @ U.T
    # elif len(M.shape) == 1:
    #     assert M.shape[0] > 0, "Vector must be non-empty"
    #     M[M<0] = 0.
    #     return M**power
    # else:
    #     raise ValueError(f"Invalid matrix shape for square root: {M.shape}")
    
def stable_matrix_power(M, power):
    """
    Compute the power of a matrix.
    :param M: Matrix to power.
    :param power: Power to raise the matrix to.
    :return: Matrix raised to the power - M^{power}.
    """
    if len(M.shape) == 2:
        assert M.shape[0] == M.shape[1], "Matrix must be square"

        # Handle NaNs
        if torch.isnan(M).all():
            print("All NaNs in matrix, returning identity")
            return torch.eye(M.shape[0], device=M.device, dtype=M.dtype)

        if torch.isnan(M).any():
            print("Some NaNs in matrix, replacing with 0")
            M = torch.nan_to_num(M, nan=0.0, posinf=1e12, neginf=-1e12)
            # Optional: scale to a reasonable magnitude
            scale = M.abs().max()
            if scale > 0:
                M = M / scale

        if M.shape[0] < 700:
            M_cpu = M.cpu().float()
            M_cpu.diagonal().add_(1e-8)
            U, S, _ = torch.linalg.svd(M_cpu)
            S[S<0] = 0.
            return (U @ torch.diag(S**power) @ U.T).to(device=M.device, dtype=M.dtype)
        else:
            M.diagonal().add_(1e-8)
            S, U = torch.linalg.eigh(M)
            S[S<0] = 0.
            return (U @ torch.diag(S**power) @ U.T).to(device=M.device, dtype=M.dtype)

    elif len(M.shape) == 1:
        # Handle NaNs
        if torch.isnan(M).all():
            print("All NaNs in vector, returning all ones")
            return torch.ones(M.shape[0], device=M.device, dtype=M.dtype)

        if torch.isnan(M).any():
            print("Some NaNs in vector, replacing with 0")
            M = torch.nan_to_num(M, nan=0.0, posinf=1e12, neginf=-1e12)
            # Optional: scale to a reasonable magnitude
            scale = M.abs().max()
            if scale > 0:
                M = M / scale

        assert M.shape[0] > 0, "Vector must be non-empty"
        M[M<0] = 0.
        return M**power
    else:
        raise ValueError(f"Invalid matrix shape for square root: {M.shape}")
