from dataclasses import dataclass
from typing import TypeVar, Generic, Type, Optional
from functools import wraps
import time
import random

import torch as T
import torch.nn as nn

# @TODO: remove si_module from codebase
# we use this in our research codebase to make modules from callable configs
si_module_TpV = TypeVar('si_module_TpV')
def si_module(cls: Type[si_module_TpV]) -> Type[si_module_TpV]:
    if not hasattr(cls, 'Config') or not isinstance(cls.Config, type):
        class Config:
            pass
        cls.Config = Config
    
    cls.Config = dataclass(cls.Config)
    
    class ConfigWrapper(cls.Config, Generic[si_module_TpV]):
        def __call__(self, *args, **kwargs) -> si_module_TpV:
            if len(kwargs) > 0:
                config_dict = {field.name: getattr(self, field.name) for field in self.__dataclass_fields__.values()}
                config_dict.update(kwargs)
                new_config = type(self)(**config_dict)
                return cls(new_config)
            else:
                return cls(self, *args)
    
    ConfigWrapper.__module__ = cls.__module__
    ConfigWrapper.__name__ = f"{cls.__name__}Config"
    ConfigWrapper.__qualname__ = f"{cls.__qualname__}.Config"
    
    cls.Config = ConfigWrapper
    
    original_init = cls.__init__
    def new_init(self, *args, **kwargs):
        self.c = next((arg for arg in args if isinstance(arg, cls.Config)), None) or next((arg for arg in kwargs.values() if isinstance(arg, cls.Config)), None)
        original_init(self, *args, **kwargs)
        self.register_buffer('_device_tracker', T.Tensor(), persistent=False)
    
    cls.__init__ = new_init
    
    @property
    def device(self):
        return self._device_tracker.device
    
    @property
    def dtype(self):
        return self._device_tracker.dtype
    
    cls.device = device
    cls.dtype = dtype
    
    return cls


def get_activation(nonlinear_activation, nonlinear_activation_params={}):
    if hasattr(nn, nonlinear_activation):
        return getattr(nn, nonlinear_activation)(**nonlinear_activation_params)
    else:
        raise NotImplementedError(f"Activation {nonlinear_activation} not found in torch.nn")


def exists(v):
    return v is not None

def isnt(v):
    return not exists(v)

def truthyexists(v):
    return exists(v) and v is not False

def truthyattr(obj, attr):
    return hasattr(obj, attr) and truthyexists(getattr(obj, attr))

defaultT = TypeVar('defaultT')

def default(*args: Optional[defaultT]) -> Optional[defaultT]:
    for arg in args:
        if exists(arg):
            return arg
    return None

def maybe(fn):
    @wraps(fn)
    def inner(x, *args, **kwargs):
        if not exists(x):
            return x
        return fn(x, *args, **kwargs)
    return inner