# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_polytope.ipynb.

# %% auto 0
__all__ = ['Polytope', 'NFDPolytope']

# %% ../nbs/01_polytope.ipynb 2
import torch
import numpy as np

from typing import  Union
import gc
import warnings

from fastcore.basics import patch

# %% ../nbs/01_polytope.ipynb 3
from . import warningss

# %% ../nbs/01_polytope.ipynb 4
class Polytope:
    
    def __init__(self, 
                A_in:Union[torch.Tensor, np.ndarray], # Inequality Restriction Matrix 
                b_in:Union[torch.Tensor, np.ndarray], # Inequality Vector
                dtype=torch.float16, # torch dtype
                device = None, # device used cpu or cuda
                copy:bool=False, # bool for creating a copy of the restrictions
                requires_grad:bool=False # If we want to keep the gradient
                 ) -> None:
        
        self.dtype = dtype
        if device is None:
            self.device = 'cpu' if A_in.get_device()==-1 else f'cuda:{A_in.get_device()}'
        else:
            self.device = device
        self._assert_valid_device_(device=device)
        
        self._check_dtype_()
        self._check_intra_constraint_dimensions_(A_in,b_in,'Inequality')
        
        self.copy = copy
        if copy:
            warnings.warn('The object will create a copy of the tensors, so memory usage will increase')
        else:
            warnings.warn('The object will not create a copy of the tensors, so modifications will be reflected in the object')

        if requires_grad:
            warnings.warn('The tensors will accumalate the gradient of the operations')
        self.requires_grad = requires_grad
        
        self.mI = A_in.shape[0]
        self.n = A_in.shape[1]        
        self.A_in = self._process_tensor_or_array_(A_in,'A_in')
        self.b_in = self._process_tensor_or_array_(b_in,'b_in')
        gc.collect()
    
    
    def _check_intra_constraint_dimensions_(self, 
                        A:Union[torch.Tensor, np.ndarray], 
                        b:Union[torch.Tensor, np.ndarray], 
                        constraint:str=None):
    
        assert(A.shape[0] == b.shape[0]), f"{constraint} has dimension mismatch: A has shape {A.shape} and b {b.shape}."
    
    
    def _process_tensor_or_array_(self, 
                                 A:Union[torch.Tensor, np.ndarray], 
                                 restriction:str=None):
        gc.collect()
        dtype = self.dtype
        if isinstance(A, torch.Tensor): # Torch Tensor
            # A is already a PyTorch tensor
            if self.copy:
                tensor_A = A.clone()
            else:
                tensor_A = A
            if dtype is not None and A.dtype != dtype:
                tensor_A = tensor_A.to(dtype)
            return tensor_A.requires_grad_(self.requires_grad)
        elif isinstance(A, np.ndarray): # Numpy array
            # A is a NumPy array, convert it to a PyTorch tensor
            if self.copy:
                tensor_A = torch.tensor(A, dtype=dtype)
            else:
                tensor_A = torch.from_numpy(A, dtype=dtype)
            return tensor_A.requires_grad_(self.requires_grad)
        else:
            raise ValueError(f"Input {restriction} must be a NumPy array or PyTorch tensor")

    def _check_dtype_(self):
        valid_dtypes = (torch.float16, torch.float32, torch.float64)
        assert self.dtype in valid_dtypes, f'{self.dtype} is not a valid PyTorch float data type. Valid dtypes: {valid_dtypes}'
        if '16' in str(self.dtype): 
            long_message = f'The dtype {self.dtype} is typically used with GPU architectures. If you are using CPU, consider using 32 or 64-bit dtypes. \
Certain operations may be casted to 32 or 64 bits to enhance numerical stability.'

            warnings.warn(long_message)
            
    def _assert_valid_device_(self, device):
        """
        Asserts that the given device is valid in PyTorch.

        Args:
        device (str): The device to check, e.g., 'cuda:0', 'cpu'.
        """
        # Check if the device is 'cpu' or 'cuda'. If 'cuda', check if it's available
        if device == 'cpu' or (device.startswith('cuda') and torch.cuda.is_available()):
            # If device is cuda, additionally check if the specified index is valid
            if device.startswith('cuda'):
                device_index = int(device.split(':')[1]) if ':' in device else 0
                assert device_index < torch.cuda.device_count(), "CUDA device index is out of range."
        else:
            raise AssertionError(f"Device '{device}' is not valid. Please choose 'cpu' or a valid 'cuda' device.")


# %% ../nbs/01_polytope.ipynb 5
@patch
def send_to_device(self:Polytope, device:str=None):
    self.A_in = self.A_in.to(device)
    self.b_in = self.b_in.to(device)
    self.device = device
    
@patch
def cast_precision(self:Polytope, dtype=None):
    if dtype is None:
        dtype = self.dtype
    self.A_in = self.A_in.to(dtype)
    self.b_in = self.b_in.to(dtype)
    self.dtype = dtype 

# %% ../nbs/01_polytope.ipynb 6
@patch
def __str__(self:Polytope):
    string = f'Numeric Precision (dtype) {self.dtype}\n'
    string = string + f'Device: {self.device}\n'
    string = string + f'A_in: {self.A_in.shape} \n'
    string = string + f'b_in: {self.b_in.shape}'
    return string

@patch
def __repr__(self:Polytope):
    return self.__str__()

# %% ../nbs/01_polytope.ipynb 7
class NFDPolytope(Polytope):
    
    def __init__(self, 
                A_in:Union[torch.Tensor, np.ndarray], 
                b_in:Union[torch.Tensor, np.ndarray], 
                A_eq:Union[torch.Tensor, np.ndarray], 
                b_eq:Union[torch.Tensor, np.ndarray], 
                dtype=torch.float32, # Default dtype for Non-Fully-Dimensioonal Polytopes is 32 bits
                device = None,
                copy:bool=False,
                requires_grad:bool=False
                 ) -> None:
        
        self._check_intra_constraint_dimensions_(A_in,b_in,'Inequality')
        self._check_intra_constraint_dimensions_(A_eq,b_eq,'Equality')
        self._check_inter_constraint_dimensions_(A_in, A_eq)
        
        if '16' in str(dtype):
            warnings.warn('Float16 precision was selected, since there are some equality restrictions a Projection matrix must be computed. Be sure to evaluate the numerical stability of the algorithm. We suggest using float32 or bigger.')

        super().__init__(A_in, b_in, dtype, device,copy, requires_grad)
        self.mE = A_eq.shape[0]
        self.A_eq = self._process_tensor_or_array_(A_eq,'A_eq')
        self.b_eq = self._process_tensor_or_array_(b_eq,'b_eq')
        self.projection_matrix = None
        gc.collect()
        
        
    def _check_inter_constraint_dimensions_(self,A_in,A_eq):
        
        assert(A_in.shape[1] == A_eq.shape[1]), f"Constraints have dimension mismatch: A_in has shape {A_in.shape} and  A_eq {A_eq.shape}."
        
        

# %% ../nbs/01_polytope.ipynb 8
@patch
def send_to_device(self:NFDPolytope, device:str):
    Polytope.send_to_device(self,device)
    self.device = None
    self.A_eq = self.A_eq.to(device)
    self.b_eq = self.b_eq.to(device)
    if self.projection_matrix is not None:
        self.projection_matrix = self.projection_matrix.to(device)
    self.device = device

@patch
def cast_precision(self:NFDPolytope, dtype=None):
    Polytope.cast_precision(self, dtype)
    self.dtype = None
    self.A_eq = self.A_eq.to(dtype)
    self.b_eq = self.b_eq.to(dtype)
    if self.projection_matrix is not None:
        self.projection_matrix = self.projection_matrix.to(dtype)
    self.dtype = dtype 

# %% ../nbs/01_polytope.ipynb 9
@patch
def compute_projection_matrix(self:NFDPolytope, device:str, solver_precision=torch.float64):

    ## Set Device
    original_device = self.device
    self.send_to_device(device)
    
    ## Set Precision
    if '64' in str(solver_precision):
        solver_precision_tensor = torch.float64
    elif '32' in str(solver_precision):
        solver_precision_tensor = torch.float32
    elif '16' in str(solver_precision):
        solver_precision_tensor = torch.float16
    original_precision = self.dtype
    error_precision = self.dtype
        
    if ('cuda' not in device) & ('16' in str(original_precision)):
        warnings.warn('Float16 precision was chosen for the polytope, but the "device=cpu" option is selected. Tensors will be temporarily cast to float32 for computing the Projection Matrix and stability evaluation. If you wish to use float16 precision, please select "device=cuda".')
        solver_precision_tensor = torch.float32
        error_precision = torch.float32
    self.cast_precision(solver_precision_tensor)
    
    # Compute (A A')^(-1)
    A_eq_t = torch.transpose(self.A_eq, 0, 1)
    A_eq_mm_A_eq_t = torch.matmul(self.A_eq, A_eq_t)
    
    # Compute (A A')^(-1)A
    la = torch.linalg.solve(A_eq_mm_A_eq_t,self.A_eq)

    # Check numerical stability of (A A')^(-1) (AA') - I
    est = torch.mm(la.to(error_precision), A_eq_t.to(error_precision))
    est = torch.max(torch.abs(est - torch.eye(est.shape[0], device=device).to(error_precision)))
    print(f"Max non zero error for term (A A')^(-1)A at precision {error_precision}: ", est)
    del est

    # Compute I - A'(A A')^(-1)A
    projection_matrix = torch.matmul(A_eq_t,la)
    projection_matrix = torch.eye(projection_matrix.shape[0], device=device) - projection_matrix

    # Free Memory
    del A_eq_mm_A_eq_t
    del A_eq_t
    del la
    gc.collect()
    self.projection_matrix =  projection_matrix
        
    self.send_to_device(original_device)    
    self.cast_precision(original_precision)
              
        

# %% ../nbs/01_polytope.ipynb 10
@patch
def project_h(self:NFDPolytope, x):
    assert(self.projection_matrix is not None), f'Projection Matrix has not been computed.'
        
    x_projected = torch.matmul(self.projection_matrix,x)
    return x_projected

# %% ../nbs/01_polytope.ipynb 11
@patch
def __str__(self:NFDPolytope):
    string = Polytope.__str__(self)
    string = string + f'\nA_eq: {self.A_eq.shape} \nb_eq: {self.b_eq.shape}'
    if self.projection_matrix is not None:
        string = string + f'\nProjection Matrix: {self.projection_matrix.shape}'
    return string

@patch
def __repr__(self:NFDPolytope):
    return self.__str__()
