import os
from abc import ABC, abstractmethod

import torch
import torch.nn as nn
from torchsummary import summary

class BaseModel(nn.Module, ABC):
  def __init__(self):
    super().__init__()
    self.best_loss = 1000000

  @abstractmethod
  def forward(self, x):
    pass

  @abstractmethod
  def test(self):
    pass

  @property
  def device(self):
    return next(self.parameters()).device

  def restore_checkpoint(self, ckpt_file, optimizer=None, affect_weights=True):
    """
    Restores checkpoint from a pth file and restores optimizer state.
    Args:
      ckpt_file (str): A PyTorch pth file containing model weights.
      optimizer (Optimizer): A vanilla optimizer to have its state restored from.
    Returns:
      int: Global step variable where the model was last checkpointed.
    """
    if not ckpt_file:
      raise ValueError("No checkpoint file to be restored.")

    try:
      ckpt_dict = torch.load(ckpt_file)

    except RuntimeError:
      ckpt_dict = torch.load(ckpt_file, map_location=lambda storage, loc: storage)
    # Restore model weights if needed
    if affect_weights:
      self.load_state_dict(ckpt_dict['model_state_dict'])
    # Restore optimizer status if existing. Evaluation doesn't need this
    if optimizer:
      optimizer.load_state_dict(ckpt_dict['optimizer_state_dict'])
    # Return global step
    return ckpt_dict, optimizer

  def count_params(self):
    """
    Computes the number of parameters in this model.
    Args: None
    Returns:
        int: Total number of weight parameters for this model.
        int: Total number of trainable parameters for this model.
    """
    num_total_params = sum(p.numel() for p in self.parameters())
    num_trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad)
    return num_total_params, num_trainable_params

  def inference(self, input_tensor):
    self.eval()
    with torch.no_grad():
      output = self.forward(input_tensor)
      if isinstance(output, tuple):
        output = output[0]
      return output.cpu().detach()