import os
from pathlib import Path

import pandas as pd
import lightning as L

import models
from datasets.loader.unpad import unpad_batch
from metrics import check_metric_is_better, get_all_metrics


class MlPipeline(L.LightningModule):
    def __init__(self, config):
        super().__init__()
        self.save_hyperparameters()
        self.task = config["task"]
        self.los_info = config["los_info"]
        self.model_name = config["model"]
        self.main_metric = config["main_metric"]
        self.cur_best_performance = {}

        model_class = getattr(models, self.model_name)
        self.model = model_class(**config)

        self.test_performance = {}
        self.test_outputs = {}
        checkpoint_folder = f'logs/train/{config["dataset"]}/{config["task"]}/{config["model"]}-fold{config["fold"]}-seed{config["seed"]}/checkpoints/'
        Path(checkpoint_folder).mkdir(parents=True, exist_ok=True)
        self.checkpoint_path = os.path.join(checkpoint_folder, 'best.ckpt')

    def forward(self, x):
        pass
    def training_step(self, batch, batch_idx):
        # the batch is large enough to contain the whole training set
        x, y, lens, pid = batch
        x, y = unpad_batch(x, y, lens)
        self.model.fit(x, y) # y contains both [outcome, los]
    def validation_step(self, batch, batch_idx):
        x, y, lens, pid = batch
        x, y = unpad_batch(x, y, lens)
        y_hat = self.model.predict(x) # y_hat is the prediction results, outcome or los
        metrics = get_all_metrics(y_hat, y, self.task, self.los_info)
        # for k, v in metrics.items(): self.log(k, v)
        main_score = metrics[self.main_metric]
        if check_metric_is_better(self.cur_best_performance, self.main_metric, main_score, self.task):
            self.cur_best_performance = metrics
            for k, v in metrics.items(): self.log("best_"+k, v)
            pd.to_pickle(self.model, self.checkpoint_path)
        return main_score
    def test_step(self, batch, batch_idx):
        x, y, lens, pid = batch
        x, y = unpad_batch(x, y, lens)
        self.model = pd.read_pickle(self.checkpoint_path)
        y_hat = self.model.predict(x)
        self.test_performance = get_all_metrics(y_hat, y, self.task, self.los_info)
        self.test_outputs = {'preds': y_hat, 'labels': y}
        return self.test_performance
    def configure_optimizers(self):
        pass