from dataclasses import dataclass, field
from typing import *
from .ops import parse_optimizer, parse_scheduler, parse_loss
from utils import parse_structure
from metrics import BaseMetrics

import lightning.pytorch as pl
import torch
import os

@dataclass
class BaseSystemConfig:
    model_type:str = 'BaseModel'
    model:Dict = field(default_factory=dict)
    optimizer:Dict = field(default_factory=dict)
    scheduler:Dict = field(default_factory=dict)
    loss:Dict = field(default_factory=dict)
    metrics:Dict = field(default_factory=dict)
    log_on_step: bool = False
    log_on_epoch: bool = True
    log_prog_bar: bool = True
    log_logger: bool = True

class BaseSystem(pl.LightningModule):
    cfg: BaseSystemConfig

    def __init__(self, cfg: Dict, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)

        self.cfg = parse_structure(BaseSystemConfig, cfg)
        self.criterion = parse_loss(self.cfg.loss)
        self.metrics_func = BaseMetrics(self.cfg.metrics)
    
    def log_metrics(self, metrics:Dict[str, float]):
        for name, value in metrics.items():
            self.log(
                name, 
                value, 
                on_step=self.cfg.log_on_step, 
                on_epoch=self.cfg.log_on_epoch, 
                prog_bar=self.cfg.log_prog_bar, 
                logger=self.cfg.log_logger
            )
    
    def __str__(self):
        return self.__class__.__name__
    
    def set_save_dir(self, path:str):
        self.trial_dir = path
        os.makedirs(self.trial_dir, exist_ok=True)
    
    def configure_optimizers(self):
        if self.model is None:
            raise ValueError(f'self.model is not initialized')
        optimizer = parse_optimizer(self.cfg.optimizer, self.model)
        return {
            "optimizer": optimizer,
            "lr_scheduler": {"scheduler": parse_scheduler(self.cfg.scheduler, optimizer)}
        }

    def on_fit_start(self) -> None:
        super().on_fit_start()
        print('[INFO]: Experiment Started')
    
    def on_fit_end(self) -> None:
        super().on_fit_end()
        print('[INFO]: Experiment Ended')
        with open(os.path.join(self.trial_dir, 'done.txt'), 'w') as file:
            file.close()