import pytorch_lightning as pl

import models
from systems.utils import parse_optimizer, parse_scheduler, update_module_step
from utils.mixins import SaverMixin
from utils.misc import config_to_primitive, get_rank

import sys, os, io

class BaseSystem(pl.LightningModule, SaverMixin):
    """
    Two ways to print to console:
    1. self.print: correctly handle progress bar
    2. rank_zero_info: use the logging module
    """
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.rank = get_rank()
        self.model = models.make(self.config.model.name, self.config.model)
        self.prepare()
        self.debugging = False

        self.planner = None
        self.last_nbv = 100000000

        self.test_psnr = -1.0
        self.test_ssim = -1.0

        self.new_view_id = 0

        self.stage = None
        self.local_iter = 0
        self.end_training_flag = False
        self.make_model = False

        class TeeFile(io.TextIOBase):
            def __init__(self, log_file, stdout):
                self.log_file = log_file
                self.stdout = stdout

            def write(self, s):
                self.log_file.write(s)
                self.stdout.write(s)
                return len(s)

            def flush(self):
                self.log_file.flush()
                self.stdout.flush()

        log_file_path = os.path.join(self.config.exp_dir, self.config.trial_name, "output.log")
        os.makedirs(os.path.dirname(log_file_path), exist_ok=True)

        log_file = open(log_file_path, "w", encoding='utf-8')
        tee = TeeFile(log_file, sys.stdout)

        sys.stdout = tee 
        sys.stderr = tee 

    def prepare(self):
        pass

    def forward(self, batch):
        raise NotImplementedError
    
    def C(self, value):
        if isinstance(value, int) or isinstance(value, float):
            pass
        else:
            value = config_to_primitive(value)
            if not isinstance(value, list):
                raise TypeError('Scalar specification only supports list, got', type(value))
            if len(value) == 3:
                value = [0] + value
            assert len(value) == 4
            start_step, start_value, end_value, end_step = value
            if isinstance(end_step, int):
                current_step = self.global_step
                value = start_value + (end_value - start_value) * max(min(1.0, (current_step - start_step) / (end_step - start_step)), 0.0)
            elif isinstance(end_step, float):
                current_step = self.current_epoch
                value = start_value + (end_value - start_value) * max(min(1.0, (current_step - start_step) / (end_step - start_step)), 0.0)
        return value
    
    def preprocess_data(self, batch, stage):
        pass

    """
    Implementing on_after_batch_transfer of DataModule does the same.
    But on_after_batch_transfer does not support DP.
    """
    def on_train_batch_start(self, batch, batch_idx, unused=0):
        self.stage = "train"
        self.dataset = self.trainer.datamodule.train_dataloader().dataset
        self.preprocess_data(batch, 'train')
        update_module_step(self.model, self.current_epoch, self.local_iter)
    
        if hasattr(self, "end_training_flag"): 
            if self.end_training_flag: 
                return -1

    def on_validation_batch_start(self, batch, batch_idx, dataloader_idx):
        self.stage = "val"

        self.dataset = self.trainer.datamodule.val_dataloader().dataset
        self.preprocess_data(batch, 'validation')
        #update_module_step(self.model, self.current_epoch, self.local_iter)
    
    def on_test_batch_start(self, batch, batch_idx, dataloader_idx):
        self.stage = "test"

        self.dataset = self.trainer.datamodule.test_dataloader().dataset
        self.preprocess_data(batch, 'test')
        update_module_step(self.model, self.current_epoch, self.local_iter)

    def on_predict_batch_start(self, batch, batch_idx, dataloader_idx):
        self.stage = "pred"

        self.dataset = self.trainer.datamodule.predict_dataloader().dataset
        self.preprocess_data(batch, 'predict')
        update_module_step(self.model, self.current_epoch, self.local_iter)
    
    def training_step(self, batch, batch_idx):
        raise NotImplementedError
    
    """
    # aggregate outputs from different devices (DP)
    def training_step_end(self, out):
        pass
    """
    
    """
    # aggregate outputs from different iterations
    def training_epoch_end(self, out):
        pass
    """
    
    def validation_step(self, batch, batch_idx):
        raise NotImplementedError
    
    """
    # aggregate outputs from different devices when using DP
    def validation_step_end(self, out):
        pass
    """
    
    def validation_epoch_end(self, out):
        """
        Gather metrics from all devices, compute mean.
        Purge repeated results using data index.
        """
        raise NotImplementedError

    def test_step(self, batch, batch_idx):        
        raise NotImplementedError
    
    def test_epoch_end(self, out):
        """
        Gather metrics from all devices, compute mean.
        Purge repeated results using data index.
        """
        raise NotImplementedError

    def export(self):
        raise NotImplementedError

    def configure_optimizers(self):
        optim = parse_optimizer(self.config.system.optimizer, self.model)
        ret = {
            'optimizer': optim,
        }
        if 'scheduler' in self.config.system:
            ret.update({
                'lr_scheduler': parse_scheduler(self.config.system.scheduler, optim),
            })
        return ret

