#
# Copyright IBM Corp. 2024 - 2024
# SPDX-License-Identifier: MIT
#
import torch


def model_info(model, verbose=False):
    # Plots a line-by-line description of a PyTorch model
    n_p = sum(x.numel() for x in model.parameters())  # number parameters
    n_g = sum(
        x.numel() for x in model.parameters() if x.requires_grad
    )  # number gradients
    if verbose:
        print(
            "%5s %40s %9s %12s %20s %10s %10s"
            % ("layer", "name", "gradient", "parameters", "shape", "mu", "sigma")
        )
        for i, (name, p) in enumerate(model.named_parameters()):
            name = name.replace("module_list.", "")
            print(
                "%5g %40s %9s %12g %20s %10.3g %10.3g"
                % (
                    i,
                    name,
                    p.requires_grad,
                    p.numel(),
                    list(p.shape),
                    p.mean(),
                    p.std(),
                )
            )

    try:  # FLOPS
        from thop import profile

        macs, _ = profile(model, inputs=(torch.zeros(1, 3, 480, 640),), verbose=False)
        fs = ", %.1f GFLOPS" % (macs / 1e9 * 2)
    except Exception:
        fs = ""

    print(
        "Model Summary: %g layers, %g parameters, %g gradients%s"
        % (len(list(model.parameters())), n_p, n_g, fs)
    )


# def init_seeds(seed=0):
#     torch.manual_seed(seed)
#
#     # Reduce randomness (may be slower on Tesla GPUs)
#     # https://pytorch.org/docs/stable/notes/randomness.html
#     if seed == 0:
#         cudnn.deterministic = False
#         cudnn.benchmark = True
#
#
# def select_device(device='', apex=False, batch_size=None):
#     # device = 'cpu' or '0' or '0,1,2,3'
#     cpu_request = device.lower() == 'cpu'
#     if device and not cpu_request:  # if device requested other than 'cpu'
#         os.environ['CUDA_VISIBLE_DEVICES'] = device  # set environment variable
#         # check availablity
#         assert torch.cuda.is_available(), 'CUDA unavailable, invalid device %s requested' % device
#
#     cuda = False if cpu_request else torch.cuda.is_available()
#     if cuda:
#         c = 1024 ** 2  # bytes to MB
#         ng = torch.cuda.device_count()
#         if ng > 1 and batch_size:  # check that batch_size is compatible with device_count
#             assert batch_size % ng == 0, 'batch-size %g not multiple of GPU count %g' % \
#                 (batch_size, ng)
#         x = [torch.cuda.get_device_properties(i) for i in range(ng)]
#         # apex for mixed precision https://github.com/NVIDIA/apex
#         s = 'Using CUDA ' + ('Apex ' if apex else '')
#         for i in range(0, ng):
#             if i == 1:
#                 s = ' ' * len(s)
#             print("%sdevice%g _CudaDeviceProperties(name='%s', total_memory=%dMB)" %
#                   (s, i, x[i].name, x[i].total_memory / c))
#     else:
#         print('Using CPU')
#
#     print('')  # skip a line
#     return torch.device('cuda:0' if cuda else 'cpu')
#
#
# def time_synchronized():
#     torch.cuda.synchronize() if torch.cuda.is_available() else None
#     return time.time()
#
#
# def initialize_weights(model):
#     for m in model.modules():
#         t = type(m)
#         if t is nn.Conv2d:
#             pass  # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
#         elif t is nn.BatchNorm2d:
#             m.eps = 1e-4
#             m.momentum = 0.03
#         elif t in [nn.LeakyReLU, nn.ReLU, nn.ReLU6]:
#             m.inplace = True
#
#
# def find_modules(model, mclass=nn.Conv2d):
#     # finds layer indices matching module class 'mclass'
#     return [i for i, m in enumerate(model.module_list) if isinstance(m, mclass)]
#
#
# def fuse_conv_and_bn(conv, bn):
#     # https://tehnokv.com/posts/fusing-batchnorm-and-conv/
#     with torch.no_grad():
#         # init
#         fusedconv = torch.nn.Conv2d(conv.in_channels,
#                                     conv.out_channels,
#                                     kernel_size=conv.kernel_size,
#                                     stride=conv.stride,
#                                     padding=conv.padding,
#                                     bias=True)
#
#         # prepare filters
#         w_conv = conv.weight.clone().view(conv.out_channels, -1)
#         w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var)))
#         fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.size()))
#
#         # prepare spatial bias
#         if conv.bias is not None:
#             b_conv = conv.bias
#         else:
#             b_conv = torch.zeros(conv.weight.size(0))
#         b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps))
#         fusedconv.bias.copy_(torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn)
#
#         return fusedconv
#
#
# def load_classifier(name='resnet101', n=2):
#     # Loads a pretrained model reshaped to n-class output
#     import pretrainedmodels  # https://github.com/Cadene/pretrained-models.pytorch#torchvision
#     model = pretrainedmodels.__dict__[name](num_classes=1000, pretrained='imagenet')
#
#     # Display model properties
#     for x in ['model.input_size', 'model.input_space', 'model.input_range', 'model.mean',
#               'model.std']:
#         print(x + ' =', eval(x))
#
#     # Reshape output to n classes
#     filters = model.last_linear.weight.shape[1]
#     model.last_linear.bias = torch.nn.Parameter(torch.zeros(n))
#     model.last_linear.weight = torch.nn.Parameter(torch.zeros(n, filters))
#     model.last_linear.out_features = n
#     return model
#
#
# def scale_img(img, ratio=1.0, same_shape=True):  # img(16,3,256,416), r=ratio
#     # scales img(bs,3,y,x) by ratio
#     h, w = img.shape[2:]
#     s = (int(h * ratio), int(w * ratio))  # new size
#     img = F.interpolate(img, size=s, mode='bilinear', align_corners=False)  # resize
#     if not same_shape:  # pad/crop img
#         gs = 64  # (pixels) grid size
#         h, w = [math.ceil(x * ratio / gs) * gs for x in (h, w)]
#     return F.pad(img, [0, w - s[1], 0, h - s[0]], value=0.447)  # value = imagenet mean
#
#
# class ModelEMA:
#     """ Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models
#     Keep a moving average of everything in the model state_dict (parameters and buffers).
#     This is intended to allow functionality like
#     https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage
#     A smoothed version of the weights is necessary for some training schemes to perform well.
#     E.g. Google's hyper-params for training MNASNet, MobileNet-V3, EfficientNet, etc that use
#     RMSprop with a short 2.4-3 epoch decay period and slow LR decay rate of .96-.99 requires EMA
#     smoothing of weights to match results. Pay attention to the decay constant you are using
#     relative to your update count per epoch.
#     To keep EMA from using GPU resources, set device='cpu'. This will save a bit of memory but
#     disable validation of the EMA weights. Validation will have to be done manually in a separate
#     process, or after the training stops converging.
#     This class is sensitive where it is initialized in the sequence of model init,
#     GPU assignment and distributed training wrappers.
#     I've tested with the sequence in my own train.py for torch.DataParallel, apex.DDP, and
#     single-GPU.
#     """
#
#     def __init__(self, model, decay=0.9999, device=''):
#         # make a copy of the model for accumulating moving average of weights
#         self.ema = deepcopy(model)
#         self.ema.eval()
#         self.updates = 0  # number of EMA updates
#         # decay exponential ramp (to help early epochs)
#         self.decay = lambda x: decay * (1 - math.exp(-x / 2000))
#         self.device = device  # perform ema on different device from model if set
#         if device:
#             self.ema.to(device=device)
#         for p in self.ema.parameters():
#             p.requires_grad_(False)
#
#     def update(self, model):
#         self.updates += 1
#         d = self.decay(self.updates)
#         with torch.no_grad():
#             if type(model) in (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel):
#                 msd, esd = model.module.state_dict(), self.ema.module.state_dict()
#             else:
#                 msd, esd = model.state_dict(), self.ema.state_dict()
#
#             for k, v in esd.items():
#                 if v.dtype.is_floating_point:
#                     v *= d
#                     v += (1. - d) * msd[k].detach()
#
#     def update_attr(self, model):
#         # Assign attributes (which may change during training)
#         for k in model.__dict__.keys():
#             if not k.startswith('_'):
#                 setattr(self.ema, k, getattr(model, k))
