import torch
import thop
from typing import Iterable, Type
from torch import nn, optim, Tensor
from abc4d.core.ab import Model, ModelState, MetricSpec
from abc4d.adapters.pytorch.ckpt_loader import pack_state_dict, load_packed_state_dict
from abc4d.utils.opt_chain import OptChain
from abc4d.utils.pprint import pprint


class PytorchModel(nn.Module, Model):

    def __init__(
        self,
        optimizer_cls: Type[optim.Optimizer] = None,  # type: ignore
        optimizer_defaults: dict = None,  # type: ignore
        device: str = "cpu",  # type: ignore
        path: str = "",
    ):
        nn.Module.__init__(self)
        self.optimizer_cls = optimizer_cls
        self.optimizer_defaults = optimizer_defaults or {}
        self.device = torch.device(device)
        Model.__init__(self, path=path)
        self.to(self.device)
        
    # Below are the methods that can / are required to be overridden.
    
    def build_model(self) -> None:
        raise NotImplementedError()
        
    def build_optimizer(self) -> None:
        if self.optimizer_cls is None:
            self.optimizer = None
        else:
            self.optimizer = self.optimizer_cls(self.parameters(), **self.optimizer_defaults)
    
    def baseline_parameters(self) -> Iterable[nn.Parameter]:
        """
        For control group, return self.parameters()
        For treatment group, return a subset of self.parameters()
        """
        raise NotImplementedError()
    
    def migrate_baseline_params(self, checkpoint):
        return checkpoint
            
    def forward(self, batch):
        raise NotImplementedError()

    def compute_loss(self, batch, output):
        raise NotImplementedError()

    def compute_metrics(self, batch, output) -> dict[str, MetricSpec]:
        raise NotImplementedError()
        
    # Below are the methods that are not required to be overridden.
            
    def save_checkpoint(self):
        return pack_state_dict(self, self.optimizer)  # type: ignore
        
    def load_checkpoint(self, checkpoint):
        self.optimizer = load_packed_state_dict(self, checkpoint, self.optimizer_cls, self.optimizer_defaults, self.device)  # type: ignore

    def freeze_baseline(self):
        for param in self.baseline_parameters():
            param.requires_grad = False

    def unfreeze_baseline(self):
        for param in self.baseline_parameters():
            param.requires_grad = True
        
    def update_parameters(self, batch, output):
        if self.optimizer is not None:
            loss = self.compute_loss(batch, output)
            if loss.requires_grad:
                loss.backward()
                self.optimizer.step()
                self.optimizer.zero_grad()

    def train_on_batch(self, batch):
        output = self(batch)
        self.update_parameters(batch, output)
        metrics = self.compute_metrics(batch, output)
        return metrics

    def update_summary(self, batch):
        flops, params = thop.profile(self, inputs=(batch,))  # type: ignore
        self.summary = f"{self.__class__.__name__}(MACs={int(flops/2):,}, Params={int(params):,})"  # type: ignore