from dataclasses import dataclass, field
from torch import nn
from typing import *
from utils import parse_structure

import torch.nn.functional as F
import torch

@dataclass
class OptConfig:
    name:str = 'Adam'
    args:Dict = field(default_factory=dict) 

@dataclass
class LossConfig:
    name:str = 'BCELoss'
    args:Dict = field(default_factory=dict)

def parse_loss(cfg: Dict)->nn.Module:
    cfg:LossConfig = parse_structure(OptConfig, cfg)
    loss = getattr(torch.nn, cfg.name)(**cfg.args)
    return loss

def parse_optimizer(cfg: Dict, model:nn.Module)->torch.optim.Optimizer:
    cfg:OptConfig = parse_structure(OptConfig, cfg)
    params = model.parameters()
    optim = getattr(torch.optim, cfg.name)(params, **cfg.args)
    return optim

def parse_scheduler(cfg: Dict, optimizer: torch.optim.Optimizer):
    lr_scheduler = getattr(torch.optim.lr_scheduler, cfg.name)(optimizer, **cfg.args)
    return lr_scheduler