import torch
from custom.any4d.a4d_model import VidarModel
from custom.any4d.a4d_logistics import dict_to_cuda_recursive

from cosmos_predict2.utils.optim_instantiate import get_base_scheduler
from imaginaire.lazy_config import instantiate


class ModelRegression(VidarModel):
    def __init__(self, config):
        super().__init__(config)

    def run(self, data_batch):
        if 'vggt' in self.pipe.extra_nets:
            output_batch, loss = self.pipe.extra_nets['vggt'](data_batch)
        else:
            raise ValueError('No network available for regression')
        return output_batch, loss

    # New function, added for i4 adaption
    def init_optimizer_scheduler(self, optimizer_config, scheduler_config):
        model = []
        if self.net is not None:
            model.append(self.net) # Add DiT parameters
        model += [val for val in self.pipe.extra_nets.values()] # Add extra networks parameters
        optimizer = instantiate(optimizer_config, model=model)
        scheduler = get_base_scheduler(optimizer, self, scheduler_config)
        return optimizer, scheduler

    def run_modules(self, data_batch, phase):
        # NOTE: broken -- DataTransforms.prepare_batch takes no `phase` kwarg and now returns
        # (batch, fps); left as-is since ModelRegression is not actively maintained right now.
        if self.transforms is not None:
            data_batch['anydata'] = self.transforms.prepare_batch(data_batch['anydata'], phase=phase)
        return dict_to_cuda_recursive(data_batch)

    def training_step(
        self,
        data_batch: dict[str, torch.Tensor],
        iteration: int,
        local_path: str = None,
        directives: dict = None,
    ):
        self.pipe.device = self.device
        data_batch = self.run_modules(data_batch, phase='train')    # Run data through modules
        return self.run(data_batch)                                 # Run and return model

    @torch.no_grad()
    def validation_step(
        self,
        data_batch: dict[str, torch.Tensor],
        iteration: int,
        dataloader_key: str = None,
        local_path: str = None,
        directives: dict = None,
        val_iter: int = 0,
    ):
        self.pipe.device = self.device
        data_batch = self.run_modules(data_batch, phase='val')  # Run data through modules
        return self.run(data_batch)                             # Run and return model